diff --git a/apps/sim/providers/utils.test.ts b/apps/sim/providers/utils.test.ts index 7c30ac3c3a..ea8d7d5133 100644 --- a/apps/sim/providers/utils.test.ts +++ b/apps/sim/providers/utils.test.ts @@ -24,6 +24,7 @@ import { MODELS_WITH_VERBOSITY, PROVIDERS_WITH_TOOL_USAGE_CONTROL, prepareToolsWithUsageControl, + shouldBillModelUsage, supportsTemperature, supportsToolUsageControl, transformCustomTool, @@ -40,6 +41,7 @@ describe('getApiKey', () => { beforeEach(() => { vi.clearAllMocks() + // @ts-expect-error - mocking boolean with different value isHostedSpy.mockReturnValue(false) module.require = vi.fn(() => ({ @@ -53,6 +55,7 @@ describe('getApiKey', () => { }) it('should return user-provided key when not in hosted environment', () => { + // @ts-expect-error - mocking boolean with different value isHostedSpy.mockReturnValue(false) // For OpenAI @@ -65,6 +68,7 @@ describe('getApiKey', () => { }) it('should throw error if no key provided in non-hosted environment', () => { + // @ts-expect-error - mocking boolean with different value isHostedSpy.mockReturnValue(false) expect(() => getApiKey('openai', 'gpt-4')).toThrow('API key is required for openai gpt-4') @@ -80,7 +84,8 @@ describe('getApiKey', () => { throw new Error('Rotation failed') }) - const key = getApiKey('openai', 'gpt-4', 'user-fallback-key') + // Use gpt-4o which IS in the hosted models list + const key = getApiKey('openai', 'gpt-4o', 'user-fallback-key') expect(key).toBe('user-fallback-key') }) @@ -91,7 +96,8 @@ describe('getApiKey', () => { throw new Error('Rotation failed') }) - expect(() => getApiKey('openai', 'gpt-4')).toThrow('No API key available for openai gpt-4') + // Use gpt-4o which IS in the hosted models list + expect(() => getApiKey('openai', 'gpt-4o')).toThrow('No API key available for openai gpt-4o') }) it('should require user key for non-OpenAI/Anthropic providers even in hosted environment', () => { @@ -104,6 +110,30 @@ describe('getApiKey', () => { 'API key is required for other-provider some-model' ) }) + + it('should require user key for models NOT in hosted list even if provider matches', () => { + isHostedSpy.mockReturnValue(true) + + // Models with version suffixes that are NOT in the hosted list should require user API key + // even though they're from anthropic/openai providers + + // User provides their own key - should work + const key1 = getApiKey('anthropic', 'claude-sonnet-4-20250514', 'user-key-anthropic') + expect(key1).toBe('user-key-anthropic') + + // No user key - should throw, NOT use server key + expect(() => getApiKey('anthropic', 'claude-sonnet-4-20250514')).toThrow( + 'API key is required for anthropic claude-sonnet-4-20250514' + ) + + // Same for OpenAI versioned models not in list + const key2 = getApiKey('openai', 'gpt-4o-2024-08-06', 'user-key-openai') + expect(key2).toBe('user-key-openai') + + expect(() => getApiKey('openai', 'gpt-4o-2024-08-06')).toThrow( + 'API key is required for openai gpt-4o-2024-08-06' + ) + }) }) describe('Model Capabilities', () => { @@ -476,6 +506,52 @@ describe('getHostedModels', () => { }) }) +describe('shouldBillModelUsage', () => { + it.concurrent('should return true for exact matches of hosted models', () => { + // OpenAI models + expect(shouldBillModelUsage('gpt-4o')).toBe(true) + expect(shouldBillModelUsage('o1')).toBe(true) + + // Anthropic models + expect(shouldBillModelUsage('claude-sonnet-4-0')).toBe(true) + expect(shouldBillModelUsage('claude-opus-4-0')).toBe(true) + + // Google models + expect(shouldBillModelUsage('gemini-2.5-pro')).toBe(true) + expect(shouldBillModelUsage('gemini-2.5-flash')).toBe(true) + }) + + it.concurrent('should return false for non-hosted models', () => { + // Other providers + expect(shouldBillModelUsage('deepseek-v3')).toBe(false) + expect(shouldBillModelUsage('grok-4-latest')).toBe(false) + + // Unknown models + expect(shouldBillModelUsage('unknown-model')).toBe(false) + }) + + it.concurrent('should return false for versioned model names not in hosted list', () => { + // Versioned model names that are NOT in the hosted list + // These should NOT be billed (user provides own API key) + expect(shouldBillModelUsage('claude-sonnet-4-20250514')).toBe(false) + expect(shouldBillModelUsage('gpt-4o-2024-08-06')).toBe(false) + expect(shouldBillModelUsage('claude-3-5-sonnet-20241022')).toBe(false) + }) + + it.concurrent('should be case insensitive', () => { + expect(shouldBillModelUsage('GPT-4O')).toBe(true) + expect(shouldBillModelUsage('Claude-Sonnet-4-0')).toBe(true) + expect(shouldBillModelUsage('GEMINI-2.5-PRO')).toBe(true) + }) + + it.concurrent('should not match partial model names', () => { + // Should not match partial/prefix models + expect(shouldBillModelUsage('gpt-4')).toBe(false) // gpt-4o is hosted, not gpt-4 + expect(shouldBillModelUsage('claude-sonnet')).toBe(false) + expect(shouldBillModelUsage('gemini')).toBe(false) + }) +}) + describe('Provider Management', () => { describe('getProviderFromModel', () => { it.concurrent('should return correct provider for known models', () => { diff --git a/apps/sim/providers/utils.ts b/apps/sim/providers/utils.ts index 85718667a5..2a755c3095 100644 --- a/apps/sim/providers/utils.ts +++ b/apps/sim/providers/utils.ts @@ -619,7 +619,7 @@ export function getHostedModels(): string[] { */ export function shouldBillModelUsage(model: string): boolean { const hostedModels = getHostedModels() - return hostedModels.includes(model) + return hostedModels.some((hostedModel) => model.toLowerCase() === hostedModel.toLowerCase()) } /** @@ -643,19 +643,22 @@ export function getApiKey(provider: string, model: string, userProvidedKey?: str const isGeminiModel = provider === 'google' if (isHosted && (isOpenAIModel || isClaudeModel || isGeminiModel)) { - try { - // Import the key rotation function - const { getRotatingApiKey } = require('@/lib/core/config/api-keys') - const serverKey = getRotatingApiKey(isGeminiModel ? 'gemini' : provider) - return serverKey - } catch (_error) { - // If server key fails and we have a user key, fallback to that - if (hasUserKey) { - return userProvidedKey! - } + // Only use server key if model is explicitly in our hosted list + const hostedModels = getHostedModels() + const isModelHosted = hostedModels.some((m) => m.toLowerCase() === model.toLowerCase()) + + if (isModelHosted) { + try { + const { getRotatingApiKey } = require('@/lib/core/config/api-keys') + const serverKey = getRotatingApiKey(isGeminiModel ? 'gemini' : provider) + return serverKey + } catch (_error) { + if (hasUserKey) { + return userProvidedKey! + } - // Otherwise, throw an error - throw new Error(`No API key available for ${provider} ${model}`) + throw new Error(`No API key available for ${provider} ${model}`) + } } }