diff --git a/apps/sim/app/api/workspaces/[id]/byok-keys/route.ts b/apps/sim/app/api/workspaces/[id]/byok-keys/route.ts index 3078555350..5a8eb86f2d 100644 --- a/apps/sim/app/api/workspaces/[id]/byok-keys/route.ts +++ b/apps/sim/app/api/workspaces/[id]/byok-keys/route.ts @@ -12,7 +12,7 @@ import { getUserEntityPermissions, getWorkspaceById } from '@/lib/workspaces/per const logger = createLogger('WorkspaceBYOKKeysAPI') -const VALID_PROVIDERS = ['openai', 'anthropic', 'google', 'mistral'] as const +const VALID_PROVIDERS = ['openai', 'anthropic', 'google', 'mistral', 'exa'] as const const UpsertKeySchema = z.object({ providerId: z.enum(VALID_PROVIDERS), diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/hooks/use-editor-subblock-layout.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/hooks/use-editor-subblock-layout.ts index 50d3f416e4..9f81bb3955 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/hooks/use-editor-subblock-layout.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/hooks/use-editor-subblock-layout.ts @@ -3,6 +3,7 @@ import { buildCanonicalIndex, evaluateSubBlockCondition, isSubBlockFeatureEnabled, + isSubBlockHiddenByHostedKey, isSubBlockVisibleForMode, } from '@/lib/workflows/subblocks/visibility' import type { BlockConfig, SubBlockConfig, SubBlockType } from '@/blocks/types' @@ -108,6 +109,9 @@ export function useEditorSubblockLayout( // Check required feature if specified - declarative feature gating if (!isSubBlockFeatureEnabled(block)) return false + // Hide tool API key fields when hosted key is available + if (isSubBlockHiddenByHostedKey(block)) return false + // Special handling for trigger-config type (legacy trigger configuration UI) if (block.type === ('trigger-config' as SubBlockType)) { const isPureTriggerBlock = config?.triggers?.enabled && config.category === 'triggers' diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/workflow-block/workflow-block.tsx b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/workflow-block/workflow-block.tsx index c0f89e2b3e..339a535e98 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/workflow-block/workflow-block.tsx +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/workflow-block/workflow-block.tsx @@ -15,6 +15,7 @@ import { evaluateSubBlockCondition, hasAdvancedValues, isSubBlockFeatureEnabled, + isSubBlockHiddenByHostedKey, isSubBlockVisibleForMode, resolveDependencyValue, } from '@/lib/workflows/subblocks/visibility' @@ -828,6 +829,7 @@ export const WorkflowBlock = memo(function WorkflowBlock({ if (block.hidden) return false if (block.hideFromPreview) return false if (!isSubBlockFeatureEnabled(block)) return false + if (isSubBlockHiddenByHostedKey(block)) return false const isPureTriggerBlock = config?.triggers?.enabled && config.category === 'triggers' diff --git a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/byok/byok.tsx b/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/byok/byok.tsx index b8304402b3..39f308d9e8 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/byok/byok.tsx +++ b/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/byok/byok.tsx @@ -13,15 +13,15 @@ import { ModalFooter, ModalHeader, } from '@/components/emcn' -import { AnthropicIcon, GeminiIcon, MistralIcon, OpenAIIcon } from '@/components/icons' +import { AnthropicIcon, ExaAIIcon, GeminiIcon, MistralIcon, OpenAIIcon } from '@/components/icons' import { Skeleton } from '@/components/ui' import { type BYOKKey, - type BYOKProviderId, useBYOKKeys, useDeleteBYOKKey, useUpsertBYOKKey, } from '@/hooks/queries/byok-keys' +import type { BYOKProviderId } from '@/tools/types' const logger = createLogger('BYOKSettings') @@ -60,6 +60,13 @@ const PROVIDERS: { description: 'LLM calls and Knowledge Base OCR', placeholder: 'Enter your API key', }, + { + id: 'exa', + name: 'Exa', + icon: ExaAIIcon, + description: 'AI-powered search and research', + placeholder: 'Enter your Exa API key', + }, ] function BYOKKeySkeleton() { diff --git a/apps/sim/blocks/blocks/exa.ts b/apps/sim/blocks/blocks/exa.ts index 43a7c88383..481fbdf1c3 100644 --- a/apps/sim/blocks/blocks/exa.ts +++ b/apps/sim/blocks/blocks/exa.ts @@ -297,6 +297,7 @@ export const ExaBlock: BlockConfig = { placeholder: 'Enter your Exa API key', password: true, required: true, + hideWhenHosted: true, }, ], tools: { diff --git a/apps/sim/blocks/types.ts b/apps/sim/blocks/types.ts index 08a716925f..def037eeec 100644 --- a/apps/sim/blocks/types.ts +++ b/apps/sim/blocks/types.ts @@ -243,6 +243,7 @@ export interface SubBlockConfig { hidden?: boolean hideFromPreview?: boolean // Hide this subblock from the workflow block preview requiresFeature?: string // Environment variable name that must be truthy for this subblock to be visible + hideWhenHosted?: boolean // Hide this subblock when running on hosted sim description?: string tooltip?: string // Tooltip text displayed via info icon next to the title value?: (params: Record) => string diff --git a/apps/sim/executor/handlers/generic/generic-handler.test.ts b/apps/sim/executor/handlers/generic/generic-handler.test.ts index 3a107df40a..cf18f8a254 100644 --- a/apps/sim/executor/handlers/generic/generic-handler.test.ts +++ b/apps/sim/executor/handlers/generic/generic-handler.test.ts @@ -147,219 +147,4 @@ describe('GenericBlockHandler', () => { 'Block execution of Some Custom Tool failed with no error message' ) }) - - describe('Knowledge block cost tracking', () => { - beforeEach(() => { - // Set up knowledge block mock - mockBlock = { - ...mockBlock, - config: { tool: 'knowledge_search', params: {} }, - } - - mockTool = { - ...mockTool, - id: 'knowledge_search', - name: 'Knowledge Search', - } - - mockGetTool.mockImplementation((toolId) => { - if (toolId === 'knowledge_search') { - return mockTool - } - return undefined - }) - }) - - it.concurrent( - 'should extract and restructure cost information from knowledge tools', - async () => { - const inputs = { query: 'test query' } - const mockToolResponse = { - success: true, - output: { - results: [], - query: 'test query', - totalResults: 0, - cost: { - input: 0.00001042, - output: 0, - total: 0.00001042, - tokens: { - input: 521, - output: 0, - total: 521, - }, - model: 'text-embedding-3-small', - pricing: { - input: 0.02, - output: 0, - updatedAt: '2025-07-10', - }, - }, - }, - } - - mockExecuteTool.mockResolvedValue(mockToolResponse) - - const result = await handler.execute(mockContext, mockBlock, inputs) - - // Verify cost information is restructured correctly for enhanced logging - expect(result).toEqual({ - results: [], - query: 'test query', - totalResults: 0, - cost: { - input: 0.00001042, - output: 0, - total: 0.00001042, - }, - tokens: { - input: 521, - output: 0, - total: 521, - }, - model: 'text-embedding-3-small', - }) - } - ) - - it.concurrent('should handle knowledge_upload_chunk cost information', async () => { - // Update to upload_chunk tool - mockBlock.config.tool = 'knowledge_upload_chunk' - mockTool.id = 'knowledge_upload_chunk' - mockTool.name = 'Knowledge Upload Chunk' - - mockGetTool.mockImplementation((toolId) => { - if (toolId === 'knowledge_upload_chunk') { - return mockTool - } - return undefined - }) - - const inputs = { content: 'test content' } - const mockToolResponse = { - success: true, - output: { - data: { - id: 'chunk-123', - content: 'test content', - chunkIndex: 0, - }, - message: 'Successfully uploaded chunk', - documentId: 'doc-123', - cost: { - input: 0.00000521, - output: 0, - total: 0.00000521, - tokens: { - input: 260, - output: 0, - total: 260, - }, - model: 'text-embedding-3-small', - pricing: { - input: 0.02, - output: 0, - updatedAt: '2025-07-10', - }, - }, - }, - } - - mockExecuteTool.mockResolvedValue(mockToolResponse) - - const result = await handler.execute(mockContext, mockBlock, inputs) - - // Verify cost information is restructured correctly - expect(result).toEqual({ - data: { - id: 'chunk-123', - content: 'test content', - chunkIndex: 0, - }, - message: 'Successfully uploaded chunk', - documentId: 'doc-123', - cost: { - input: 0.00000521, - output: 0, - total: 0.00000521, - }, - tokens: { - input: 260, - output: 0, - total: 260, - }, - model: 'text-embedding-3-small', - }) - }) - - it('should pass through output unchanged for knowledge tools without cost info', async () => { - const inputs = { query: 'test query' } - const mockToolResponse = { - success: true, - output: { - results: [], - query: 'test query', - totalResults: 0, - // No cost information - }, - } - - mockExecuteTool.mockResolvedValue(mockToolResponse) - - const result = await handler.execute(mockContext, mockBlock, inputs) - - // Should return original output without cost transformation - expect(result).toEqual({ - results: [], - query: 'test query', - totalResults: 0, - }) - }) - - it.concurrent( - 'should process cost info for all tools (universal cost extraction)', - async () => { - mockBlock.config.tool = 'some_other_tool' - mockTool.id = 'some_other_tool' - - mockGetTool.mockImplementation((toolId) => { - if (toolId === 'some_other_tool') { - return mockTool - } - return undefined - }) - - const inputs = { param: 'value' } - const mockToolResponse = { - success: true, - output: { - result: 'success', - cost: { - input: 0.001, - output: 0.002, - total: 0.003, - tokens: { input: 100, output: 50, total: 150 }, - model: 'some-model', - }, - }, - } - - mockExecuteTool.mockResolvedValue(mockToolResponse) - - const result = await handler.execute(mockContext, mockBlock, inputs) - - expect(result).toEqual({ - result: 'success', - cost: { - input: 0.001, - output: 0.002, - total: 0.003, - }, - tokens: { input: 100, output: 50, total: 150 }, - model: 'some-model', - }) - } - ) - }) }) diff --git a/apps/sim/executor/handlers/generic/generic-handler.ts b/apps/sim/executor/handlers/generic/generic-handler.ts index c6a6b7e9f3..9a9cec6e61 100644 --- a/apps/sim/executor/handlers/generic/generic-handler.ts +++ b/apps/sim/executor/handlers/generic/generic-handler.ts @@ -97,27 +97,7 @@ export class GenericBlockHandler implements BlockHandler { throw error } - const output = result.output - let cost = null - - if (output?.cost) { - cost = output.cost - } - - if (cost) { - return { - ...output, - cost: { - input: cost.input, - output: cost.output, - total: cost.total, - }, - tokens: cost.tokens, - model: cost.model, - } - } - - return output + return result.output } catch (error: any) { if (!error.message || error.message === 'undefined (undefined)') { let errorMessage = `Block execution of ${tool?.name || block.config.tool} failed` diff --git a/apps/sim/hooks/queries/byok-keys.ts b/apps/sim/hooks/queries/byok-keys.ts index 26d348d5a7..167238f4a1 100644 --- a/apps/sim/hooks/queries/byok-keys.ts +++ b/apps/sim/hooks/queries/byok-keys.ts @@ -1,11 +1,10 @@ import { createLogger } from '@sim/logger' import { keepPreviousData, useMutation, useQuery, useQueryClient } from '@tanstack/react-query' import { API_ENDPOINTS } from '@/stores/constants' +import type { BYOKProviderId } from '@/tools/types' const logger = createLogger('BYOKKeysQueries') -export type BYOKProviderId = 'openai' | 'anthropic' | 'google' | 'mistral' - export interface BYOKKey { id: string providerId: BYOKProviderId diff --git a/apps/sim/lib/api-key/byok.ts b/apps/sim/lib/api-key/byok.ts index 04a35adb42..127feb9af3 100644 --- a/apps/sim/lib/api-key/byok.ts +++ b/apps/sim/lib/api-key/byok.ts @@ -7,11 +7,10 @@ import { isHosted } from '@/lib/core/config/feature-flags' import { decryptSecret } from '@/lib/core/security/encryption' import { getHostedModels } from '@/providers/models' import { useProvidersStore } from '@/stores/providers/store' +import type { BYOKProviderId } from '@/tools/types' const logger = createLogger('BYOKKeys') -export type BYOKProviderId = 'openai' | 'anthropic' | 'google' | 'mistral' - export interface BYOKKeyResult { apiKey: string isBYOK: true diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index b21fb552f7..50883c5fc4 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -25,9 +25,9 @@ export interface ModelUsageMetadata { } /** - * Metadata for 'fixed' category charges (currently empty, extensible) + * Metadata for 'fixed' category charges (e.g., tool cost breakdown) */ -export type FixedUsageMetadata = Record +export type FixedUsageMetadata = Record /** * Union type for all metadata types @@ -60,6 +60,8 @@ export interface LogFixedUsageParams { workspaceId?: string workflowId?: string executionId?: string + /** Optional metadata (e.g., tool cost breakdown from API) */ + metadata?: FixedUsageMetadata } /** @@ -119,7 +121,7 @@ export async function logFixedUsage(params: LogFixedUsageParams): Promise category: 'fixed', source: params.source, description: params.description, - metadata: null, + metadata: params.metadata ?? null, cost: params.cost.toString(), workspaceId: params.workspaceId ?? null, workflowId: params.workflowId ?? null, diff --git a/apps/sim/lib/core/telemetry.ts b/apps/sim/lib/core/telemetry.ts index c12fe1303a..47d46c7cf4 100644 --- a/apps/sim/lib/core/telemetry.ts +++ b/apps/sim/lib/core/telemetry.ts @@ -934,6 +934,31 @@ export const PlatformEvents = { }) }, + /** + * Track hosted key throttled (rate limited) + */ + hostedKeyThrottled: (attrs: { + toolId: string + envVarName: string + attempt: number + maxRetries: number + delayMs: number + userId?: string + workspaceId?: string + workflowId?: string + }) => { + trackPlatformEvent('platform.hosted_key.throttled', { + 'tool.id': attrs.toolId, + 'hosted_key.env_var': attrs.envVarName, + 'throttle.attempt': attrs.attempt, + 'throttle.max_retries': attrs.maxRetries, + 'throttle.delay_ms': attrs.delayMs, + ...(attrs.userId && { 'user.id': attrs.userId }), + ...(attrs.workspaceId && { 'workspace.id': attrs.workspaceId }), + ...(attrs.workflowId && { 'workflow.id': attrs.workflowId }), + }) + }, + /** * Track chat deployed (workflow deployed as chat interface) */ diff --git a/apps/sim/lib/workflows/subblocks/visibility.ts b/apps/sim/lib/workflows/subblocks/visibility.ts index 1ce0076b44..ac39244c28 100644 --- a/apps/sim/lib/workflows/subblocks/visibility.ts +++ b/apps/sim/lib/workflows/subblocks/visibility.ts @@ -1,4 +1,5 @@ import { getEnv, isTruthy } from '@/lib/core/config/env' +import { isHosted } from '@/lib/core/config/feature-flags' import type { SubBlockConfig } from '@/blocks/types' export type CanonicalMode = 'basic' | 'advanced' @@ -270,3 +271,12 @@ export function isSubBlockFeatureEnabled(subBlock: SubBlockConfig): boolean { if (!subBlock.requiresFeature) return true return isTruthy(getEnv(subBlock.requiresFeature)) } + +/** + * Check if a subblock should be hidden because we're running on hosted Sim. + * Used for tool API key fields that should be hidden when Sim provides hosted keys. + */ +export function isSubBlockHiddenByHostedKey(subBlock: SubBlockConfig): boolean { + if (!subBlock.hideWhenHosted) return false + return isHosted +} diff --git a/apps/sim/serializer/index.ts b/apps/sim/serializer/index.ts index 622667d9fc..53b72e17ca 100644 --- a/apps/sim/serializer/index.ts +++ b/apps/sim/serializer/index.ts @@ -10,6 +10,7 @@ import { isCanonicalPair, isNonEmptyValue, isSubBlockFeatureEnabled, + isSubBlockHiddenByHostedKey, resolveCanonicalMode, } from '@/lib/workflows/subblocks/visibility' import { getBlock } from '@/blocks' @@ -49,6 +50,7 @@ function shouldSerializeSubBlock( canonicalModeOverrides?: CanonicalModeOverrides ): boolean { if (!isSubBlockFeatureEnabled(subBlockConfig)) return false + if (isSubBlockHiddenByHostedKey(subBlockConfig)) return false if (subBlockConfig.mode === 'trigger') { if (!isTriggerContext && !isTriggerCategory) return false diff --git a/apps/sim/tools/exa/answer.ts b/apps/sim/tools/exa/answer.ts index 95c29e0e68..8e43e135f5 100644 --- a/apps/sim/tools/exa/answer.ts +++ b/apps/sim/tools/exa/answer.ts @@ -1,6 +1,9 @@ +import { createLogger } from '@sim/logger' import type { ExaAnswerParams, ExaAnswerResponse } from '@/tools/exa/types' import type { ToolConfig } from '@/tools/types' +const logger = createLogger('ExaAnswerTool') + export const answerTool: ToolConfig = { id: 'exa_answer', name: 'Exa Answer', @@ -27,6 +30,24 @@ export const answerTool: ToolConfig = { description: 'Exa AI API Key', }, }, + hosting: { + envKeys: ['EXA_API_KEY_1', 'EXA_API_KEY_2', 'EXA_API_KEY_3'], + apiKeyParam: 'apiKey', + byokProviderId: 'exa', + pricing: { + type: 'custom', + getCost: (_params, output) => { + // Use _costDollars from Exa API response (internal field, stripped from final output) + const costDollars = output._costDollars as { total?: number } | undefined + if (costDollars?.total) { + return { cost: costDollars.total, metadata: { costDollars } } + } + // Fallback: $5/1000 requests + logger.warn('Exa answer response missing costDollars, using fallback pricing') + return 0.005 + }, + }, + }, request: { url: 'https://api.exa.ai/answer', @@ -61,6 +82,7 @@ export const answerTool: ToolConfig = { url: citation.url, text: citation.text || '', })) || [], + _costDollars: data.costDollars, }, } }, diff --git a/apps/sim/tools/exa/find_similar_links.ts b/apps/sim/tools/exa/find_similar_links.ts index 0996061a3d..6e693789dc 100644 --- a/apps/sim/tools/exa/find_similar_links.ts +++ b/apps/sim/tools/exa/find_similar_links.ts @@ -1,6 +1,9 @@ +import { createLogger } from '@sim/logger' import type { ExaFindSimilarLinksParams, ExaFindSimilarLinksResponse } from '@/tools/exa/types' import type { ToolConfig } from '@/tools/types' +const logger = createLogger('ExaFindSimilarLinksTool') + export const findSimilarLinksTool: ToolConfig< ExaFindSimilarLinksParams, ExaFindSimilarLinksResponse @@ -76,6 +79,26 @@ export const findSimilarLinksTool: ToolConfig< description: 'Exa AI API Key', }, }, + hosting: { + envKeys: ['EXA_API_KEY_1', 'EXA_API_KEY_2', 'EXA_API_KEY_3'], + apiKeyParam: 'apiKey', + byokProviderId: 'exa', + pricing: { + type: 'custom', + getCost: (_params, output) => { + // Use _costDollars from Exa API response (internal field, stripped from final output) + const costDollars = output._costDollars as { total?: number } | undefined + if (costDollars?.total) { + return { cost: costDollars.total, metadata: { costDollars } } + } + // Fallback: $5/1000 (1-25 results) or $25/1000 (26-100 results) + logger.warn('Exa find_similar_links response missing costDollars, using fallback pricing') + const similarLinks = output.similarLinks as unknown[] | undefined + const resultCount = similarLinks?.length || 0 + return resultCount <= 25 ? 0.005 : 0.025 + }, + }, + }, request: { url: 'https://api.exa.ai/findSimilar', @@ -140,6 +163,7 @@ export const findSimilarLinksTool: ToolConfig< highlights: result.highlights, score: result.score || 0, })), + _costDollars: data.costDollars, }, } }, diff --git a/apps/sim/tools/exa/get_contents.ts b/apps/sim/tools/exa/get_contents.ts index be44b70222..449f7a5959 100644 --- a/apps/sim/tools/exa/get_contents.ts +++ b/apps/sim/tools/exa/get_contents.ts @@ -1,6 +1,9 @@ +import { createLogger } from '@sim/logger' import type { ExaGetContentsParams, ExaGetContentsResponse } from '@/tools/exa/types' import type { ToolConfig } from '@/tools/types' +const logger = createLogger('ExaGetContentsTool') + export const getContentsTool: ToolConfig = { id: 'exa_get_contents', name: 'Exa Get Contents', @@ -61,6 +64,25 @@ export const getContentsTool: ToolConfig { + // Use _costDollars from Exa API response (internal field, stripped from final output) + const costDollars = output._costDollars as { total?: number } | undefined + if (costDollars?.total) { + return { cost: costDollars.total, metadata: { costDollars } } + } + // Fallback: $1/1000 pages + logger.warn('Exa get_contents response missing costDollars, using fallback pricing') + const results = output.results as unknown[] | undefined + return (results?.length || 0) * 0.001 + }, + }, + }, request: { url: 'https://api.exa.ai/contents', @@ -132,6 +154,7 @@ export const getContentsTool: ToolConfig = description: 'Exa AI API Key', }, }, + hosting: { + envKeys: ['EXA_API_KEY_1', 'EXA_API_KEY_2', 'EXA_API_KEY_3'], + apiKeyParam: 'apiKey', + byokProviderId: 'exa', + pricing: { + type: 'custom', + getCost: (params, output) => { + // Use _costDollars from Exa API response (internal field, stripped from final output) + const costDollars = output._costDollars as { total?: number } | undefined + if (costDollars?.total) { + return { cost: costDollars.total, metadata: { costDollars } } + } + + // Fallback to estimate if cost not available + logger.warn('Exa research response missing costDollars, using fallback pricing') + const model = params.model || 'exa-research' + return model === 'exa-research-pro' ? 0.055 : 0.03 + }, + }, + }, request: { url: 'https://api.exa.ai/research/v1', @@ -111,6 +131,8 @@ export const researchTool: ToolConfig = score: 1.0, }, ], + // Include cost breakdown for pricing calculation (internal field, stripped from final output) + _costDollars: taskData.costDollars, } return result } diff --git a/apps/sim/tools/exa/search.ts b/apps/sim/tools/exa/search.ts index a4099dfeec..f0cc7afd00 100644 --- a/apps/sim/tools/exa/search.ts +++ b/apps/sim/tools/exa/search.ts @@ -1,6 +1,9 @@ +import { createLogger } from '@sim/logger' import type { ExaSearchParams, ExaSearchResponse } from '@/tools/exa/types' import type { ToolConfig } from '@/tools/types' +const logger = createLogger('ExaSearchTool') + export const searchTool: ToolConfig = { id: 'exa_search', name: 'Exa Search', @@ -86,6 +89,31 @@ export const searchTool: ToolConfig = { description: 'Exa AI API Key', }, }, + hosting: { + envKeys: ['EXA_API_KEY_1', 'EXA_API_KEY_2', 'EXA_API_KEY_3'], + apiKeyParam: 'apiKey', + byokProviderId: 'exa', + pricing: { + type: 'custom', + getCost: (params, output) => { + // Use _costDollars from Exa API response (internal field, stripped from final output) + const costDollars = output._costDollars as { total?: number } | undefined + if (costDollars?.total) { + return { cost: costDollars.total, metadata: { costDollars } } + } + + // Fallback: estimate based on search type and result count + logger.warn('Exa search response missing costDollars, using fallback pricing') + const isDeepSearch = params.type === 'neural' + if (isDeepSearch) { + return 0.015 + } + const results = output.results as unknown[] | undefined + const resultCount = results?.length || 0 + return resultCount <= 25 ? 0.005 : 0.025 + }, + }, + }, request: { url: 'https://api.exa.ai/search', @@ -167,6 +195,7 @@ export const searchTool: ToolConfig = { highlights: result.highlights, score: result.score, })), + _costDollars: data.costDollars, }, } }, diff --git a/apps/sim/tools/exa/types.ts b/apps/sim/tools/exa/types.ts index bcdf63d1a2..e3b1dc7319 100644 --- a/apps/sim/tools/exa/types.ts +++ b/apps/sim/tools/exa/types.ts @@ -6,6 +6,11 @@ export interface ExaBaseParams { apiKey: string } +/** Cost breakdown returned by Exa API responses */ +export interface ExaCostDollars { + total: number +} + // Search tool types export interface ExaSearchParams extends ExaBaseParams { query: string @@ -50,6 +55,7 @@ export interface ExaSearchResult { export interface ExaSearchResponse extends ToolResponse { output: { results: ExaSearchResult[] + costDollars?: ExaCostDollars } } @@ -78,6 +84,7 @@ export interface ExaGetContentsResult { export interface ExaGetContentsResponse extends ToolResponse { output: { results: ExaGetContentsResult[] + costDollars?: ExaCostDollars } } @@ -120,6 +127,7 @@ export interface ExaSimilarLink { export interface ExaFindSimilarLinksResponse extends ToolResponse { output: { similarLinks: ExaSimilarLink[] + costDollars?: ExaCostDollars } } @@ -137,6 +145,7 @@ export interface ExaAnswerResponse extends ToolResponse { url: string text: string }[] + costDollars?: ExaCostDollars } } @@ -158,6 +167,7 @@ export interface ExaResearchResponse extends ToolResponse { author?: string score: number }[] + costDollars?: ExaCostDollars } } diff --git a/apps/sim/tools/index.test.ts b/apps/sim/tools/index.test.ts index 9a20977ae8..b5e61583fe 100644 --- a/apps/sim/tools/index.test.ts +++ b/apps/sim/tools/index.test.ts @@ -15,52 +15,74 @@ import { } from '@sim/testing' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -// Mock custom tools query - must be hoisted before imports -vi.mock('@/hooks/queries/custom-tools', () => ({ - getCustomTool: (toolId: string) => { - if (toolId === 'custom-tool-123') { - return { - id: 'custom-tool-123', - title: 'Custom Weather Tool', - code: 'return { result: "Weather data" }', - schema: { - function: { - description: 'Get weather information', - parameters: { - type: 'object', - properties: { - location: { type: 'string', description: 'City name' }, - unit: { type: 'string', description: 'Unit (metric/imperial)' }, - }, - required: ['location'], - }, - }, - }, - } - } - return undefined +// Hoisted mock state - these are available to vi.mock factories +const { mockIsHosted, mockEnv, mockGetBYOKKey, mockLogFixedUsage } = vi.hoisted(() => ({ + mockIsHosted: { value: false }, + mockEnv: { NEXT_PUBLIC_APP_URL: 'http://localhost:3000' } as Record, + mockGetBYOKKey: vi.fn(), + mockLogFixedUsage: vi.fn(), +})) + +// Mock feature flags +vi.mock('@/lib/core/config/feature-flags', () => ({ + get isHosted() { + return mockIsHosted.value }, - getCustomTools: () => [ - { - id: 'custom-tool-123', - title: 'Custom Weather Tool', - code: 'return { result: "Weather data" }', - schema: { - function: { - description: 'Get weather information', - parameters: { - type: 'object', - properties: { - location: { type: 'string', description: 'City name' }, - unit: { type: 'string', description: 'Unit (metric/imperial)' }, - }, - required: ['location'], + isProd: false, + isDev: true, + isTest: true, +})) + +// Mock env config to control hosted key availability +vi.mock('@/lib/core/config/env', () => ({ + env: new Proxy({} as Record, { + get: (_target, prop: string) => mockEnv[prop], + }), + getEnv: (key: string) => mockEnv[key], + isTruthy: (val: unknown) => val === true || val === 'true' || val === '1', + isFalsy: (val: unknown) => val === false || val === 'false' || val === '0', +})) + +// Mock getBYOKKey +vi.mock('@/lib/api-key/byok', () => ({ + getBYOKKey: (...args: unknown[]) => mockGetBYOKKey(...args), +})) + +// Mock logFixedUsage for billing +vi.mock('@/lib/billing/core/usage-log', () => ({ + logFixedUsage: (...args: unknown[]) => mockLogFixedUsage(...args), +})) + +// Mock custom tools - define mock data inside factory function +vi.mock('@/hooks/queries/custom-tools', () => { + const mockCustomTool = { + id: 'custom-tool-123', + title: 'Custom Weather Tool', + code: 'return { result: "Weather data" }', + schema: { + function: { + description: 'Get weather information', + parameters: { + type: 'object', + properties: { + location: { type: 'string', description: 'City name' }, + unit: { type: 'string', description: 'Unit (metric/imperial)' }, }, + required: ['location'], }, }, }, - ], -})) + } + return { + getCustomTool: (toolId: string) => { + if (toolId === 'custom-tool-123') { + return mockCustomTool + } + return undefined + }, + getCustomTools: () => [mockCustomTool], + } +}) import { executeTool } from '@/tools/index' import { tools } from '@/tools/registry' @@ -959,3 +981,654 @@ describe('MCP Tool Execution', () => { expect(result.timing).toBeDefined() }) }) + +describe('Hosted Key Injection', () => { + let cleanupEnvVars: () => void + + beforeEach(() => { + process.env.NEXT_PUBLIC_APP_URL = 'http://localhost:3000' + cleanupEnvVars = setupEnvVars({ NEXT_PUBLIC_APP_URL: 'http://localhost:3000' }) + vi.clearAllMocks() + mockGetBYOKKey.mockReset() + mockLogFixedUsage.mockReset() + }) + + afterEach(() => { + vi.resetAllMocks() + cleanupEnvVars() + }) + + it('should not inject hosted key when tool has no hosting config', async () => { + const mockTool = { + id: 'test_no_hosting', + name: 'Test No Hosting', + description: 'A test tool without hosting config', + version: '1.0.0', + params: {}, + request: { + url: '/api/test/endpoint', + method: 'POST' as const, + headers: () => ({ 'Content-Type': 'application/json' }), + }, + transformResponse: vi.fn().mockResolvedValue({ + success: true, + output: { result: 'success' }, + }), + } + + const originalTools = { ...tools } + ;(tools as any).test_no_hosting = mockTool + + global.fetch = Object.assign( + vi.fn().mockImplementation(async () => ({ + ok: true, + status: 200, + headers: new Headers(), + json: () => Promise.resolve({ success: true }), + })), + { preconnect: vi.fn() } + ) as typeof fetch + + const mockContext = createToolExecutionContext() + await executeTool('test_no_hosting', {}, false, mockContext) + + // BYOK should not be called since there's no hosting config + expect(mockGetBYOKKey).not.toHaveBeenCalled() + + Object.assign(tools, originalTools) + }) + + it('should check BYOK key first when tool has hosting config', async () => { + // Note: isHosted is mocked to false by default, so hosted key injection won't happen + // This test verifies the flow when isHosted would be true + const mockTool = { + id: 'test_with_hosting', + name: 'Test With Hosting', + description: 'A test tool with hosting config', + version: '1.0.0', + params: { + apiKey: { type: 'string', required: true }, + }, + hosting: { + envKeys: ['TEST_API_KEY'], + apiKeyParam: 'apiKey', + byokProviderId: 'exa', + pricing: { + type: 'per_request' as const, + cost: 0.005, + }, + }, + request: { + url: '/api/test/endpoint', + method: 'POST' as const, + headers: (params: any) => ({ + 'Content-Type': 'application/json', + 'x-api-key': params.apiKey, + }), + }, + transformResponse: vi.fn().mockResolvedValue({ + success: true, + output: { result: 'success' }, + }), + } + + const originalTools = { ...tools } + ;(tools as any).test_with_hosting = mockTool + + // Mock BYOK returning a key + mockGetBYOKKey.mockResolvedValue({ apiKey: 'byok-test-key', isBYOK: true }) + + global.fetch = Object.assign( + vi.fn().mockImplementation(async () => ({ + ok: true, + status: 200, + headers: new Headers(), + json: () => Promise.resolve({ success: true }), + })), + { preconnect: vi.fn() } + ) as typeof fetch + + const mockContext = createToolExecutionContext() + await executeTool('test_with_hosting', {}, false, mockContext) + + // With isHosted=false, BYOK won't be called - this is expected behavior + // The test documents the current behavior + Object.assign(tools, originalTools) + }) + + it('should use per_request pricing model correctly', async () => { + const mockTool = { + id: 'test_per_request_pricing', + name: 'Test Per Request Pricing', + description: 'A test tool with per_request pricing', + version: '1.0.0', + params: { + apiKey: { type: 'string', required: true }, + }, + hosting: { + envKeys: ['TEST_API_KEY'], + apiKeyParam: 'apiKey', + byokProviderId: 'exa', + pricing: { + type: 'per_request' as const, + cost: 0.005, + }, + }, + request: { + url: '/api/test/endpoint', + method: 'POST' as const, + headers: (params: any) => ({ + 'Content-Type': 'application/json', + 'x-api-key': params.apiKey, + }), + }, + transformResponse: vi.fn().mockResolvedValue({ + success: true, + output: { result: 'success' }, + }), + } + + // Verify pricing config structure + expect(mockTool.hosting.pricing.type).toBe('per_request') + expect(mockTool.hosting.pricing.cost).toBe(0.005) + }) + + it('should use custom pricing model correctly', async () => { + const mockGetCost = vi.fn().mockReturnValue({ cost: 0.01, metadata: { breakdown: 'test' } }) + + const mockTool = { + id: 'test_custom_pricing', + name: 'Test Custom Pricing', + description: 'A test tool with custom pricing', + version: '1.0.0', + params: { + apiKey: { type: 'string', required: true }, + }, + hosting: { + envKeys: ['TEST_API_KEY'], + apiKeyParam: 'apiKey', + byokProviderId: 'exa', + pricing: { + type: 'custom' as const, + getCost: mockGetCost, + }, + }, + request: { + url: '/api/test/endpoint', + method: 'POST' as const, + headers: (params: any) => ({ + 'Content-Type': 'application/json', + 'x-api-key': params.apiKey, + }), + }, + transformResponse: vi.fn().mockResolvedValue({ + success: true, + output: { result: 'success', costDollars: { total: 0.01 } }, + }), + } + + // Verify pricing config structure + expect(mockTool.hosting.pricing.type).toBe('custom') + expect(typeof mockTool.hosting.pricing.getCost).toBe('function') + + // Test getCost returns expected value + const result = mockTool.hosting.pricing.getCost({}, { costDollars: { total: 0.01 } }) + expect(result).toEqual({ cost: 0.01, metadata: { breakdown: 'test' } }) + }) + + it('should handle custom pricing returning a number', async () => { + const mockGetCost = vi.fn().mockReturnValue(0.005) + + const mockTool = { + id: 'test_custom_pricing_number', + name: 'Test Custom Pricing Number', + description: 'A test tool with custom pricing returning number', + version: '1.0.0', + params: { + apiKey: { type: 'string', required: true }, + }, + hosting: { + envKeys: ['TEST_API_KEY'], + apiKeyParam: 'apiKey', + byokProviderId: 'exa', + pricing: { + type: 'custom' as const, + getCost: mockGetCost, + }, + }, + request: { + url: '/api/test/endpoint', + method: 'POST' as const, + headers: (params: any) => ({ + 'Content-Type': 'application/json', + 'x-api-key': params.apiKey, + }), + }, + } + + // Test getCost returns a number + const result = mockTool.hosting.pricing.getCost({}, {}) + expect(result).toBe(0.005) + }) +}) + +describe('Rate Limiting and Retry Logic', () => { + let cleanupEnvVars: () => void + + beforeEach(() => { + process.env.NEXT_PUBLIC_APP_URL = 'http://localhost:3000' + cleanupEnvVars = setupEnvVars({ + NEXT_PUBLIC_APP_URL: 'http://localhost:3000', + }) + vi.clearAllMocks() + mockIsHosted.value = true + mockEnv.TEST_HOSTED_KEY = 'test-hosted-api-key' + mockGetBYOKKey.mockResolvedValue(null) + }) + + afterEach(() => { + vi.resetAllMocks() + cleanupEnvVars() + mockIsHosted.value = false + mockEnv.TEST_HOSTED_KEY = undefined + }) + + it('should retry on 429 rate limit errors with exponential backoff', async () => { + let attemptCount = 0 + + const mockTool = { + id: 'test_rate_limit', + name: 'Test Rate Limit', + description: 'A test tool for rate limiting', + version: '1.0.0', + params: { + apiKey: { type: 'string', required: false }, + }, + hosting: { + envKeys: ['TEST_HOSTED_KEY'], + apiKeyParam: 'apiKey', + pricing: { + type: 'per_request' as const, + cost: 0.001, + }, + }, + request: { + url: '/api/test/rate-limit', + method: 'POST' as const, + headers: () => ({ 'Content-Type': 'application/json' }), + }, + transformResponse: vi.fn().mockResolvedValue({ + success: true, + output: { result: 'success' }, + }), + } + + const originalTools = { ...tools } + ;(tools as any).test_rate_limit = mockTool + + global.fetch = Object.assign( + vi.fn().mockImplementation(async () => { + attemptCount++ + if (attemptCount < 3) { + // Return a proper 429 response - the code extracts error, attaches status, and throws + return { + ok: false, + status: 429, + statusText: 'Too Many Requests', + headers: new Headers(), + json: () => Promise.resolve({ error: 'Rate limited' }), + text: () => Promise.resolve('Rate limited'), + } + } + return { + ok: true, + status: 200, + headers: new Headers(), + json: () => Promise.resolve({ success: true }), + } + }), + { preconnect: vi.fn() } + ) as typeof fetch + + const mockContext = createToolExecutionContext() + const result = await executeTool('test_rate_limit', {}, false, mockContext) + + // Should succeed after retries + expect(result.success).toBe(true) + // Should have made 3 attempts (2 failures + 1 success) + expect(attemptCount).toBe(3) + + Object.assign(tools, originalTools) + }) + + it('should fail after max retries on persistent rate limiting', async () => { + const mockTool = { + id: 'test_persistent_rate_limit', + name: 'Test Persistent Rate Limit', + description: 'A test tool for persistent rate limiting', + version: '1.0.0', + params: { + apiKey: { type: 'string', required: false }, + }, + hosting: { + envKeys: ['TEST_HOSTED_KEY'], + apiKeyParam: 'apiKey', + pricing: { + type: 'per_request' as const, + cost: 0.001, + }, + }, + request: { + url: '/api/test/persistent-rate-limit', + method: 'POST' as const, + headers: () => ({ 'Content-Type': 'application/json' }), + }, + } + + const originalTools = { ...tools } + ;(tools as any).test_persistent_rate_limit = mockTool + + global.fetch = Object.assign( + vi.fn().mockImplementation(async () => { + // Always return 429 to test max retries exhaustion + return { + ok: false, + status: 429, + statusText: 'Too Many Requests', + headers: new Headers(), + json: () => Promise.resolve({ error: 'Rate limited' }), + text: () => Promise.resolve('Rate limited'), + } + }), + { preconnect: vi.fn() } + ) as typeof fetch + + const mockContext = createToolExecutionContext() + const result = await executeTool('test_persistent_rate_limit', {}, false, mockContext) + + // Should fail after all retries exhausted + expect(result.success).toBe(false) + expect(result.error).toContain('Rate limited') + + Object.assign(tools, originalTools) + }) + + it('should not retry on non-rate-limit errors', async () => { + let attemptCount = 0 + + const mockTool = { + id: 'test_no_retry', + name: 'Test No Retry', + description: 'A test tool that should not retry', + version: '1.0.0', + params: { + apiKey: { type: 'string', required: false }, + }, + hosting: { + envKeys: ['TEST_HOSTED_KEY'], + apiKeyParam: 'apiKey', + pricing: { + type: 'per_request' as const, + cost: 0.001, + }, + }, + request: { + url: '/api/test/no-retry', + method: 'POST' as const, + headers: () => ({ 'Content-Type': 'application/json' }), + }, + } + + const originalTools = { ...tools } + ;(tools as any).test_no_retry = mockTool + + global.fetch = Object.assign( + vi.fn().mockImplementation(async () => { + attemptCount++ + // Return a 400 response - should not trigger retry logic + return { + ok: false, + status: 400, + statusText: 'Bad Request', + headers: new Headers(), + json: () => Promise.resolve({ error: 'Bad request' }), + text: () => Promise.resolve('Bad request'), + } + }), + { preconnect: vi.fn() } + ) as typeof fetch + + const mockContext = createToolExecutionContext() + const result = await executeTool('test_no_retry', {}, false, mockContext) + + // Should fail immediately without retries + expect(result.success).toBe(false) + expect(attemptCount).toBe(1) + + Object.assign(tools, originalTools) + }) +}) + +describe('Cost Field Handling', () => { + let cleanupEnvVars: () => void + + beforeEach(() => { + process.env.NEXT_PUBLIC_APP_URL = 'http://localhost:3000' + cleanupEnvVars = setupEnvVars({ + NEXT_PUBLIC_APP_URL: 'http://localhost:3000', + }) + vi.clearAllMocks() + mockIsHosted.value = true + mockEnv.TEST_HOSTED_KEY = 'test-hosted-api-key' + mockGetBYOKKey.mockResolvedValue(null) + mockLogFixedUsage.mockResolvedValue(undefined) + }) + + afterEach(() => { + vi.resetAllMocks() + cleanupEnvVars() + mockIsHosted.value = false + mockEnv.TEST_HOSTED_KEY = undefined + }) + + it('should add cost to output when using hosted key with per_request pricing', async () => { + const mockTool = { + id: 'test_cost_per_request', + name: 'Test Cost Per Request', + description: 'A test tool with per_request pricing', + version: '1.0.0', + params: { + apiKey: { type: 'string', required: false }, + }, + hosting: { + envKeys: ['TEST_HOSTED_KEY'], + apiKeyParam: 'apiKey', + pricing: { + type: 'per_request' as const, + cost: 0.005, + }, + }, + request: { + url: '/api/test/cost', + method: 'POST' as const, + headers: () => ({ 'Content-Type': 'application/json' }), + }, + transformResponse: vi.fn().mockResolvedValue({ + success: true, + output: { result: 'success' }, + }), + } + + const originalTools = { ...tools } + ;(tools as any).test_cost_per_request = mockTool + + global.fetch = Object.assign( + vi.fn().mockImplementation(async () => ({ + ok: true, + status: 200, + headers: new Headers(), + json: () => Promise.resolve({ success: true }), + })), + { preconnect: vi.fn() } + ) as typeof fetch + + const mockContext = createToolExecutionContext({ + userId: 'user-123', + } as any) + const result = await executeTool('test_cost_per_request', {}, false, mockContext) + + expect(result.success).toBe(true) + // Note: In test environment, hosted key injection may not work due to env mocking complexity. + // The cost calculation logic is tested via the pricing model tests above. + // This test verifies the tool execution flow when hosted key IS available (by checking output structure). + if (result.output.cost) { + expect(result.output.cost.total).toBe(0.005) + // Should have logged usage + expect(mockLogFixedUsage).toHaveBeenCalledWith( + expect.objectContaining({ + userId: 'user-123', + cost: 0.005, + description: 'tool:test_cost_per_request', + }) + ) + } + + Object.assign(tools, originalTools) + }) + + it('should not add cost when not using hosted key', async () => { + mockIsHosted.value = false + + const mockTool = { + id: 'test_no_hosted_cost', + name: 'Test No Hosted Cost', + description: 'A test tool without hosted key', + version: '1.0.0', + params: { + apiKey: { type: 'string', required: true }, + }, + hosting: { + envKeys: ['TEST_HOSTED_KEY'], + apiKeyParam: 'apiKey', + pricing: { + type: 'per_request' as const, + cost: 0.005, + }, + }, + request: { + url: '/api/test/no-hosted', + method: 'POST' as const, + headers: () => ({ 'Content-Type': 'application/json' }), + }, + transformResponse: vi.fn().mockResolvedValue({ + success: true, + output: { result: 'success' }, + }), + } + + const originalTools = { ...tools } + ;(tools as any).test_no_hosted_cost = mockTool + + global.fetch = Object.assign( + vi.fn().mockImplementation(async () => ({ + ok: true, + status: 200, + headers: new Headers(), + json: () => Promise.resolve({ success: true }), + })), + { preconnect: vi.fn() } + ) as typeof fetch + + const mockContext = createToolExecutionContext() + // Pass user's own API key + const result = await executeTool( + 'test_no_hosted_cost', + { apiKey: 'user-api-key' }, + false, + mockContext + ) + + expect(result.success).toBe(true) + // Should not have cost since user provided their own key + expect(result.output.cost).toBeUndefined() + // Should not have logged usage + expect(mockLogFixedUsage).not.toHaveBeenCalled() + + Object.assign(tools, originalTools) + }) + + it('should use custom pricing getCost function', async () => { + const mockGetCost = vi.fn().mockReturnValue({ + cost: 0.015, + metadata: { mode: 'advanced', results: 10 }, + }) + + const mockTool = { + id: 'test_custom_pricing_cost', + name: 'Test Custom Pricing Cost', + description: 'A test tool with custom pricing', + version: '1.0.0', + params: { + apiKey: { type: 'string', required: false }, + mode: { type: 'string', required: false }, + }, + hosting: { + envKeys: ['TEST_HOSTED_KEY'], + apiKeyParam: 'apiKey', + pricing: { + type: 'custom' as const, + getCost: mockGetCost, + }, + }, + request: { + url: '/api/test/custom-pricing', + method: 'POST' as const, + headers: () => ({ 'Content-Type': 'application/json' }), + }, + transformResponse: vi.fn().mockResolvedValue({ + success: true, + output: { result: 'success', results: 10 }, + }), + } + + const originalTools = { ...tools } + ;(tools as any).test_custom_pricing_cost = mockTool + + global.fetch = Object.assign( + vi.fn().mockImplementation(async () => ({ + ok: true, + status: 200, + headers: new Headers(), + json: () => Promise.resolve({ success: true }), + })), + { preconnect: vi.fn() } + ) as typeof fetch + + const mockContext = createToolExecutionContext({ + userId: 'user-123', + } as any) + const result = await executeTool( + 'test_custom_pricing_cost', + { mode: 'advanced' }, + false, + mockContext + ) + + expect(result.success).toBe(true) + expect(result.output.cost).toBeDefined() + expect(result.output.cost.total).toBe(0.015) + + // getCost should have been called with params and output + expect(mockGetCost).toHaveBeenCalled() + + // Should have logged usage with metadata + expect(mockLogFixedUsage).toHaveBeenCalledWith( + expect.objectContaining({ + cost: 0.015, + metadata: { mode: 'advanced', results: 10 }, + }) + ) + + Object.assign(tools, originalTools) + }) +}) diff --git a/apps/sim/tools/index.ts b/apps/sim/tools/index.ts index 040a40a272..46dfd7fc05 100644 --- a/apps/sim/tools/index.ts +++ b/apps/sim/tools/index.ts @@ -1,10 +1,15 @@ import { createLogger } from '@sim/logger' +import { getBYOKKey } from '@/lib/api-key/byok' import { generateInternalToken } from '@/lib/auth/internal' +import { logFixedUsage } from '@/lib/billing/core/usage-log' +import { env } from '@/lib/core/config/env' +import { isHosted } from '@/lib/core/config/feature-flags' import { DEFAULT_EXECUTION_TIMEOUT_MS } from '@/lib/core/execution-limits' import { secureFetchWithPinnedIP, validateUrlWithDNS, } from '@/lib/core/security/input-validation.server' +import { PlatformEvents } from '@/lib/core/telemetry' import { generateRequestId } from '@/lib/core/utils/request' import { getBaseUrl } from '@/lib/core/utils/urls' import { parseMcpToolId } from '@/lib/mcp/utils' @@ -13,7 +18,13 @@ import { resolveSkillContent } from '@/executor/handlers/agent/skills-resolver' import type { ExecutionContext } from '@/executor/types' import type { ErrorInfo } from '@/tools/error-extractors' import { extractErrorMessage } from '@/tools/error-extractors' -import type { OAuthTokenPayload, ToolConfig, ToolResponse } from '@/tools/types' +import type { + BYOKProviderId, + OAuthTokenPayload, + ToolConfig, + ToolHostingPricing, + ToolResponse, +} from '@/tools/types' import { formatRequestParams, getTool, @@ -23,6 +34,246 @@ import { const logger = createLogger('Tools') +/** Result from hosted key lookup */ +interface HostedKeyResult { + key: string + envVarName: string +} + +/** + * Get a hosted API key from environment variables + * Supports rotation when multiple keys are configured + * Returns both the key and which env var it came from + */ +function getHostedKeyFromEnv(envKeys: string[]): HostedKeyResult | null { + const keysWithNames = envKeys + .map((envVarName) => ({ envVarName, key: env[envVarName as keyof typeof env] })) + .filter((item): item is { envVarName: string; key: string } => Boolean(item.key)) + + if (keysWithNames.length === 0) return null + + // Round-robin rotation based on current minute + const currentMinute = Math.floor(Date.now() / 60000) + const keyIndex = currentMinute % keysWithNames.length + + return keysWithNames[keyIndex] +} + +/** Result from hosted key injection */ +interface HostedKeyInjectionResult { + isUsingHostedKey: boolean + envVarName?: string +} + +/** + * Inject hosted API key if tool supports it and user didn't provide one. + * Checks BYOK workspace keys first, then falls back to hosted env keys. + * Returns whether a hosted (billable) key was injected and which env var it came from. + */ +async function injectHostedKeyIfNeeded( + tool: ToolConfig, + params: Record, + executionContext: ExecutionContext | undefined, + requestId: string +): Promise { + if (!tool.hosting) return { isUsingHostedKey: false } + if (!isHosted) return { isUsingHostedKey: false } + + const { envKeys, apiKeyParam, byokProviderId } = tool.hosting + + // Check BYOK workspace key first + if (byokProviderId && executionContext?.workspaceId) { + try { + const byokResult = await getBYOKKey( + executionContext.workspaceId, + byokProviderId as BYOKProviderId + ) + if (byokResult) { + params[apiKeyParam] = byokResult.apiKey + logger.info(`[${requestId}] Using BYOK key for ${tool.id}`) + return { isUsingHostedKey: false } // Don't bill - user's own key + } + } catch (error) { + logger.error(`[${requestId}] Failed to get BYOK key for ${tool.id}:`, error) + // Fall through to hosted key + } + } + + // Fall back to hosted env key + const hostedKeyResult = getHostedKeyFromEnv(envKeys) + if (!hostedKeyResult) { + logger.debug(`[${requestId}] No hosted key available for ${tool.id}`) + return { isUsingHostedKey: false } + } + + params[apiKeyParam] = hostedKeyResult.key + logger.info(`[${requestId}] Using hosted key for ${tool.id} (${hostedKeyResult.envVarName})`) + return { isUsingHostedKey: true, envVarName: hostedKeyResult.envVarName } +} + +/** + * Check if an error is a rate limit (throttling) error + */ +function isRateLimitError(error: unknown): boolean { + if (error && typeof error === 'object') { + const status = (error as { status?: number }).status + // 429 = Too Many Requests, 503 = Service Unavailable (sometimes used for rate limiting) + if (status === 429 || status === 503) return true + } + return false +} + +/** Context for retry with throttle tracking */ +interface RetryContext { + requestId: string + toolId: string + envVarName: string + executionContext?: ExecutionContext +} + +/** + * Execute a function with exponential backoff retry for rate limiting errors. + * Only used for hosted key requests. Tracks throttling events via telemetry. + */ +async function executeWithRetry( + fn: () => Promise, + context: RetryContext, + maxRetries = 3, + baseDelayMs = 1000 +): Promise { + const { requestId, toolId, envVarName, executionContext } = context + let lastError: unknown + + for (let attempt = 0; attempt <= maxRetries; attempt++) { + try { + return await fn() + } catch (error) { + lastError = error + + if (!isRateLimitError(error) || attempt === maxRetries) { + throw error + } + + const delayMs = baseDelayMs * 2 ** attempt + + // Track throttling event via telemetry + PlatformEvents.hostedKeyThrottled({ + toolId, + envVarName, + attempt: attempt + 1, + maxRetries, + delayMs, + userId: executionContext?.userId, + workspaceId: executionContext?.workspaceId, + workflowId: executionContext?.workflowId, + }) + + logger.warn( + `[${requestId}] Rate limited for ${toolId} (${envVarName}), retrying in ${delayMs}ms (attempt ${attempt + 1}/${maxRetries})` + ) + await new Promise((resolve) => setTimeout(resolve, delayMs)) + } + } + + throw lastError +} + +/** Result from cost calculation */ +interface ToolCostResult { + cost: number + metadata?: Record +} + +/** + * Calculate cost based on pricing model + */ +function calculateToolCost( + pricing: ToolHostingPricing, + params: Record, + response: Record +): ToolCostResult { + switch (pricing.type) { + case 'per_request': + return { cost: pricing.cost } + + case 'custom': { + const result = pricing.getCost(params, response) + if (typeof result === 'number') { + return { cost: result } + } + return result + } + + default: { + const exhaustiveCheck: never = pricing + throw new Error(`Unknown pricing type: ${(exhaustiveCheck as ToolHostingPricing).type}`) + } + } +} + +interface HostedKeyCostResult { + cost: number + metadata?: Record +} + +/** + * Calculate and log hosted key cost for a tool execution. + * Logs to usageLog for audit trail and returns cost + metadata for output. + */ +async function processHostedKeyCost( + tool: ToolConfig, + params: Record, + response: Record, + executionContext: ExecutionContext | undefined, + requestId: string +): Promise { + if (!tool.hosting?.pricing) { + return { cost: 0 } + } + + const { cost, metadata } = calculateToolCost(tool.hosting.pricing, params, response) + + if (cost <= 0) return { cost: 0 } + + // Log to usageLog table for audit trail + if (executionContext?.userId) { + try { + await logFixedUsage({ + userId: executionContext.userId, + source: 'workflow', + description: `tool:${tool.id}`, + cost, + workspaceId: executionContext.workspaceId, + workflowId: executionContext.workflowId, + executionId: executionContext.executionId, + metadata, + }) + logger.debug( + `[${requestId}] Logged hosted key cost for ${tool.id}: $${cost}`, + metadata ? { metadata } : {} + ) + } catch (error) { + logger.error(`[${requestId}] Failed to log hosted key usage for ${tool.id}:`, error) + } + } + + return { cost, metadata } +} + +/** + * Strips internal fields (keys starting with underscore) from output. + * Used to hide internal data (e.g., _costDollars) from end users. + */ +function stripInternalFields(output: Record): Record { + const result: Record = {} + for (const [key, value] of Object.entries(output)) { + if (!key.startsWith('_')) { + result[key] = value + } + } + return result +} + /** * Normalizes a tool ID by stripping resource ID suffix (UUID). * Workflow tools: 'workflow_executor_' -> 'workflow_executor' @@ -279,6 +530,14 @@ export async function executeTool( throw new Error(`Tool not found: ${toolId}`) } + // Inject hosted API key if tool supports it and user didn't provide one + const hostedKeyInfo = await injectHostedKeyIfNeeded( + tool, + contextParams, + executionContext, + requestId + ) + // If we have a credential parameter, fetch the access token if (contextParams.credential) { logger.info( @@ -391,8 +650,33 @@ export async function executeTool( const endTime = new Date() const endTimeISO = endTime.toISOString() const duration = endTime.getTime() - startTime.getTime() + + // Calculate hosted key cost and merge into output.cost + if (hostedKeyInfo.isUsingHostedKey && finalResult.success) { + const { cost: hostedKeyCost, metadata } = await processHostedKeyCost( + tool, + contextParams, + finalResult.output, + executionContext, + requestId + ) + if (hostedKeyCost > 0) { + finalResult.output = { + ...finalResult.output, + cost: { + total: hostedKeyCost, + ...metadata, + }, + } + } + } + + // Strip internal fields (keys starting with _) from output before returning + const strippedOutput = stripInternalFields(finalResult.output || {}) + return { ...finalResult, + output: strippedOutput, timing: { startTime: startTimeISO, endTime: endTimeISO, @@ -402,7 +686,15 @@ export async function executeTool( } // Execute the tool request directly (internal routes use regular fetch, external use SSRF-protected fetch) - const result = await executeToolRequest(toolId, tool, contextParams) + // Wrap with retry logic for hosted keys to handle rate limiting due to higher usage + const result = hostedKeyInfo.isUsingHostedKey + ? await executeWithRetry(() => executeToolRequest(toolId, tool, contextParams), { + requestId, + toolId, + envVarName: hostedKeyInfo.envVarName!, + executionContext, + }) + : await executeToolRequest(toolId, tool, contextParams) // Apply post-processing if available and not skipped let finalResult = result @@ -424,8 +716,33 @@ export async function executeTool( const endTime = new Date() const endTimeISO = endTime.toISOString() const duration = endTime.getTime() - startTime.getTime() + + // Calculate hosted key cost and merge into output.cost + if (hostedKeyInfo.isUsingHostedKey && finalResult.success) { + const { cost: hostedKeyCost, metadata } = await processHostedKeyCost( + tool, + contextParams, + finalResult.output, + executionContext, + requestId + ) + if (hostedKeyCost > 0) { + finalResult.output = { + ...finalResult.output, + cost: { + total: hostedKeyCost, + ...metadata, + }, + } + } + } + + // Strip internal fields (keys starting with _) from output before returning + const strippedOutput = stripInternalFields(finalResult.output || {}) + return { ...finalResult, + output: strippedOutput, timing: { startTime: startTimeISO, endTime: endTimeISO, diff --git a/apps/sim/tools/knowledge/knowledge.test.ts b/apps/sim/tools/knowledge/knowledge.test.ts new file mode 100644 index 0000000000..1dd0f28771 --- /dev/null +++ b/apps/sim/tools/knowledge/knowledge.test.ts @@ -0,0 +1,202 @@ +/** + * @vitest-environment node + * + * Knowledge Tools Unit Tests + * + * Tests for knowledge_search and knowledge_upload_chunk tools, + * specifically the cost restructuring in transformResponse. + */ + +import { describe, expect, it } from 'vitest' +import { knowledgeSearchTool } from '@/tools/knowledge/search' +import { knowledgeUploadChunkTool } from '@/tools/knowledge/upload_chunk' + +/** + * Creates a mock Response object for testing transformResponse + */ +function createMockResponse(data: unknown): Response { + return { + json: async () => data, + ok: true, + status: 200, + } as Response +} + +describe('Knowledge Tools', () => { + describe('knowledgeSearchTool', () => { + describe('transformResponse', () => { + it('should restructure cost information for logging', async () => { + const apiResponse = { + data: { + results: [{ content: 'test result', similarity: 0.95 }], + query: 'test query', + totalResults: 1, + cost: { + input: 0.00001042, + output: 0, + total: 0.00001042, + tokens: { + prompt: 521, + completion: 0, + total: 521, + }, + model: 'text-embedding-3-small', + pricing: { + input: 0.02, + output: 0, + updatedAt: '2025-07-10', + }, + }, + }, + } + + const result = await knowledgeSearchTool.transformResponse!(createMockResponse(apiResponse)) + + expect(result.success).toBe(true) + expect(result.output).toEqual({ + results: [{ content: 'test result', similarity: 0.95 }], + query: 'test query', + totalResults: 1, + cost: { + input: 0.00001042, + output: 0, + total: 0.00001042, + }, + tokens: { + prompt: 521, + completion: 0, + total: 521, + }, + model: 'text-embedding-3-small', + }) + }) + + it('should handle response without cost information', async () => { + const apiResponse = { + data: { + results: [], + query: 'test query', + totalResults: 0, + }, + } + + const result = await knowledgeSearchTool.transformResponse!(createMockResponse(apiResponse)) + + expect(result.success).toBe(true) + expect(result.output).toEqual({ + results: [], + query: 'test query', + totalResults: 0, + }) + expect(result.output.cost).toBeUndefined() + expect(result.output.tokens).toBeUndefined() + expect(result.output.model).toBeUndefined() + }) + + it('should handle response with partial cost information', async () => { + const apiResponse = { + data: { + results: [], + query: 'test query', + totalResults: 0, + cost: { + input: 0.001, + output: 0, + total: 0.001, + // No tokens or model + }, + }, + } + + const result = await knowledgeSearchTool.transformResponse!(createMockResponse(apiResponse)) + + expect(result.success).toBe(true) + expect(result.output.cost).toEqual({ + input: 0.001, + output: 0, + total: 0.001, + }) + expect(result.output.tokens).toBeUndefined() + expect(result.output.model).toBeUndefined() + }) + }) + }) + + describe('knowledgeUploadChunkTool', () => { + describe('transformResponse', () => { + it('should restructure cost information for logging', async () => { + const apiResponse = { + data: { + id: 'chunk-123', + chunkIndex: 0, + content: 'test content', + contentLength: 12, + tokenCount: 3, + enabled: true, + documentId: 'doc-456', + documentName: 'Test Document', + createdAt: '2025-01-01T00:00:00Z', + updatedAt: '2025-01-01T00:00:00Z', + cost: { + input: 0.00000521, + output: 0, + total: 0.00000521, + tokens: { + prompt: 260, + completion: 0, + total: 260, + }, + model: 'text-embedding-3-small', + pricing: { + input: 0.02, + output: 0, + updatedAt: '2025-07-10', + }, + }, + }, + } + + const result = await knowledgeUploadChunkTool.transformResponse!( + createMockResponse(apiResponse) + ) + + expect(result.success).toBe(true) + expect(result.output.cost).toEqual({ + input: 0.00000521, + output: 0, + total: 0.00000521, + }) + expect(result.output.tokens).toEqual({ + prompt: 260, + completion: 0, + total: 260, + }) + expect(result.output.model).toBe('text-embedding-3-small') + expect(result.output.data.chunkId).toBe('chunk-123') + expect(result.output.documentId).toBe('doc-456') + }) + + it('should handle response without cost information', async () => { + const apiResponse = { + data: { + id: 'chunk-123', + chunkIndex: 0, + content: 'test content', + documentId: 'doc-456', + documentName: 'Test Document', + }, + } + + const result = await knowledgeUploadChunkTool.transformResponse!( + createMockResponse(apiResponse) + ) + + expect(result.success).toBe(true) + expect(result.output.cost).toBeUndefined() + expect(result.output.tokens).toBeUndefined() + expect(result.output.model).toBeUndefined() + expect(result.output.data.chunkId).toBe('chunk-123') + }) + }) + }) +}) diff --git a/apps/sim/tools/knowledge/search.ts b/apps/sim/tools/knowledge/search.ts index 574017d083..af82111adc 100644 --- a/apps/sim/tools/knowledge/search.ts +++ b/apps/sim/tools/knowledge/search.ts @@ -80,13 +80,24 @@ export const knowledgeSearchTool: ToolConfig = { const result = await response.json() const data = result.data || result + // Restructure cost: extract tokens/model to top level for logging + let costFields: Record = {} + if (data.cost && typeof data.cost === 'object') { + const { tokens, model, input, output: outputCost, total } = data.cost + costFields = { + cost: { input, output: outputCost, total }, + ...(tokens && { tokens }), + ...(model && { model }), + } + } + return { success: true, output: { results: data.results || [], query: data.query, totalResults: data.totalResults || 0, - cost: data.cost, + ...costFields, }, } }, diff --git a/apps/sim/tools/knowledge/upload_chunk.ts b/apps/sim/tools/knowledge/upload_chunk.ts index 24e07ee24a..d7ad0fd93b 100644 --- a/apps/sim/tools/knowledge/upload_chunk.ts +++ b/apps/sim/tools/knowledge/upload_chunk.ts @@ -52,6 +52,17 @@ export const knowledgeUploadChunkTool: ToolConfig = {} + if (data.cost && typeof data.cost === 'object') { + const { tokens, model, input, output: outputCost, total } = data.cost + costFields = { + cost: { input, output: outputCost, total }, + ...(tokens && { tokens }), + ...(model && { model }), + } + } + return { success: true, output: { @@ -68,7 +79,7 @@ export const knowledgeUploadChunkTool: ToolConfig { * Maps param IDs to their enrichment configuration. */ schemaEnrichment?: Record + + /** + * Hosted API key configuration for this tool. + * When configured, the tool can use Sim's hosted API keys if user doesn't provide their own. + * Usage is billed according to the pricing config. + */ + hosting?: ToolHostingConfig

} export interface TableRow { @@ -170,3 +179,46 @@ export interface SchemaEnrichmentConfig { required?: string[] } | null> } + +/** + * Pricing models for hosted API key usage + */ +/** Flat fee per API call (e.g., Serper search) */ +export interface PerRequestPricing { + type: 'per_request' + /** Cost per request in dollars */ + cost: number +} + +/** Result from custom pricing calculation */ +export interface CustomPricingResult { + /** Cost in dollars */ + cost: number + /** Optional metadata about the cost calculation (e.g., breakdown from API) */ + metadata?: Record +} + +/** Custom pricing calculated from params and response (e.g., Exa with different modes/result counts) */ +export interface CustomPricing

> { + type: 'custom' + /** Calculate cost based on request params and response output. Fields starting with _ are internal. */ + getCost: (params: P, output: Record) => number | CustomPricingResult +} + +/** Union of all pricing models */ +export type ToolHostingPricing

> = PerRequestPricing | CustomPricing

+ +/** + * Configuration for hosted API key support + * When configured, the tool can use Sim's hosted API keys if user doesn't provide their own + */ +export interface ToolHostingConfig

> { + /** Environment variable names to check for hosted keys (supports rotation with multiple keys) */ + envKeys: string[] + /** The parameter name that receives the API key */ + apiKeyParam: string + /** BYOK provider ID for workspace key lookup */ + byokProviderId?: BYOKProviderId + /** Pricing when using hosted key */ + pricing: ToolHostingPricing

+}