Compare commits

...

8 Commits

Author SHA1 Message Date
Theodore Li
36e6464992 Record usage to user stats table 2026-02-13 11:41:32 -08:00
Theodore Li
2a36143f46 Add warning comment if default calculation is used 2026-02-13 11:16:17 -08:00
Theodore Li
c12e92c807 Consolidate byok type definitions 2026-02-13 10:18:37 -08:00
Theodore Li
d174a6a3fb Add telemetry 2026-02-13 09:53:18 -08:00
Theodore Li
8a78f8047a Add custom pricing, switch to exa as first hosted key 2026-02-13 09:40:06 -08:00
Theodore Li
e5c8aec07d Add rate limiting (3 tries, exponential backoff) 2026-02-12 19:16:28 -08:00
Theodore Li
3e6527a540 Handle required fields correctly for hosted keys 2026-02-12 19:08:08 -08:00
Theodore Li
2cdb89681b feat(hosted keys): Implement serper hosted key 2026-02-12 18:32:00 -08:00
22 changed files with 781 additions and 17 deletions

View File

@@ -12,7 +12,7 @@ import { getUserEntityPermissions, getWorkspaceById } from '@/lib/workspaces/per
const logger = createLogger('WorkspaceBYOKKeysAPI') 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({ const UpsertKeySchema = z.object({
providerId: z.enum(VALID_PROVIDERS), providerId: z.enum(VALID_PROVIDERS),

View File

@@ -3,6 +3,7 @@ import {
buildCanonicalIndex, buildCanonicalIndex,
evaluateSubBlockCondition, evaluateSubBlockCondition,
isSubBlockFeatureEnabled, isSubBlockFeatureEnabled,
isSubBlockHiddenByHostedKey,
isSubBlockVisibleForMode, isSubBlockVisibleForMode,
} from '@/lib/workflows/subblocks/visibility' } from '@/lib/workflows/subblocks/visibility'
import type { BlockConfig, SubBlockConfig, SubBlockType } from '@/blocks/types' import type { BlockConfig, SubBlockConfig, SubBlockType } from '@/blocks/types'
@@ -108,6 +109,9 @@ export function useEditorSubblockLayout(
// Check required feature if specified - declarative feature gating // Check required feature if specified - declarative feature gating
if (!isSubBlockFeatureEnabled(block)) return false 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) // Special handling for trigger-config type (legacy trigger configuration UI)
if (block.type === ('trigger-config' as SubBlockType)) { if (block.type === ('trigger-config' as SubBlockType)) {
const isPureTriggerBlock = config?.triggers?.enabled && config.category === 'triggers' const isPureTriggerBlock = config?.triggers?.enabled && config.category === 'triggers'

View File

@@ -15,6 +15,7 @@ import {
evaluateSubBlockCondition, evaluateSubBlockCondition,
hasAdvancedValues, hasAdvancedValues,
isSubBlockFeatureEnabled, isSubBlockFeatureEnabled,
isSubBlockHiddenByHostedKey,
isSubBlockVisibleForMode, isSubBlockVisibleForMode,
resolveDependencyValue, resolveDependencyValue,
} from '@/lib/workflows/subblocks/visibility' } from '@/lib/workflows/subblocks/visibility'
@@ -828,6 +829,7 @@ export const WorkflowBlock = memo(function WorkflowBlock({
if (block.hidden) return false if (block.hidden) return false
if (block.hideFromPreview) return false if (block.hideFromPreview) return false
if (!isSubBlockFeatureEnabled(block)) return false if (!isSubBlockFeatureEnabled(block)) return false
if (isSubBlockHiddenByHostedKey(block)) return false
const isPureTriggerBlock = config?.triggers?.enabled && config.category === 'triggers' const isPureTriggerBlock = config?.triggers?.enabled && config.category === 'triggers'

View File

@@ -13,15 +13,15 @@ import {
ModalFooter, ModalFooter,
ModalHeader, ModalHeader,
} from '@/components/emcn' } 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 { Skeleton } from '@/components/ui'
import { import {
type BYOKKey, type BYOKKey,
type BYOKProviderId,
useBYOKKeys, useBYOKKeys,
useDeleteBYOKKey, useDeleteBYOKKey,
useUpsertBYOKKey, useUpsertBYOKKey,
} from '@/hooks/queries/byok-keys' } from '@/hooks/queries/byok-keys'
import type { BYOKProviderId } from '@/tools/types'
const logger = createLogger('BYOKSettings') const logger = createLogger('BYOKSettings')
@@ -60,6 +60,13 @@ const PROVIDERS: {
description: 'LLM calls and Knowledge Base OCR', description: 'LLM calls and Knowledge Base OCR',
placeholder: 'Enter your API key', 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() { function BYOKKeySkeleton() {

View File

@@ -297,6 +297,7 @@ export const ExaBlock: BlockConfig<ExaResponse> = {
placeholder: 'Enter your Exa API key', placeholder: 'Enter your Exa API key',
password: true, password: true,
required: true, required: true,
hideWhenHosted: true,
}, },
], ],
tools: { tools: {

View File

@@ -243,6 +243,7 @@ export interface SubBlockConfig {
hidden?: boolean hidden?: boolean
hideFromPreview?: boolean // Hide this subblock from the workflow block preview 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 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 description?: string
tooltip?: string // Tooltip text displayed via info icon next to the title tooltip?: string // Tooltip text displayed via info icon next to the title
value?: (params: Record<string, any>) => string value?: (params: Record<string, any>) => string

View File

@@ -98,10 +98,21 @@ export class GenericBlockHandler implements BlockHandler {
} }
const output = result.output const output = result.output
let cost = null
if (output?.cost) { // Merge costs from output (e.g., AI model costs) and result (e.g., hosted key costs)
cost = output.cost // TODO: migrate model usage to output cost.
const outputCost = output?.cost
const resultCost = result.cost
let cost = null
if (outputCost || resultCost) {
cost = {
input: (outputCost?.input || 0) + (resultCost?.input || 0),
output: (outputCost?.output || 0) + (resultCost?.output || 0),
total: (outputCost?.total || 0) + (resultCost?.total || 0),
tokens: outputCost?.tokens,
model: outputCost?.model,
}
} }
if (cost) { if (cost) {

View File

@@ -1,11 +1,10 @@
import { createLogger } from '@sim/logger' import { createLogger } from '@sim/logger'
import { keepPreviousData, useMutation, useQuery, useQueryClient } from '@tanstack/react-query' import { keepPreviousData, useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
import { API_ENDPOINTS } from '@/stores/constants' import { API_ENDPOINTS } from '@/stores/constants'
import type { BYOKProviderId } from '@/tools/types'
const logger = createLogger('BYOKKeysQueries') const logger = createLogger('BYOKKeysQueries')
export type BYOKProviderId = 'openai' | 'anthropic' | 'google' | 'mistral'
export interface BYOKKey { export interface BYOKKey {
id: string id: string
providerId: BYOKProviderId providerId: BYOKProviderId

View File

@@ -7,11 +7,10 @@ import { isHosted } from '@/lib/core/config/feature-flags'
import { decryptSecret } from '@/lib/core/security/encryption' import { decryptSecret } from '@/lib/core/security/encryption'
import { getHostedModels } from '@/providers/models' import { getHostedModels } from '@/providers/models'
import { useProvidersStore } from '@/stores/providers/store' import { useProvidersStore } from '@/stores/providers/store'
import type { BYOKProviderId } from '@/tools/types'
const logger = createLogger('BYOKKeys') const logger = createLogger('BYOKKeys')
export type BYOKProviderId = 'openai' | 'anthropic' | 'google' | 'mistral'
export interface BYOKKeyResult { export interface BYOKKeyResult {
apiKey: string apiKey: string
isBYOK: true isBYOK: true

View File

@@ -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<string, never> export type FixedUsageMetadata = Record<string, unknown>
/** /**
* Union type for all metadata types * Union type for all metadata types
@@ -60,6 +60,8 @@ export interface LogFixedUsageParams {
workspaceId?: string workspaceId?: string
workflowId?: string workflowId?: string
executionId?: 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<void>
category: 'fixed', category: 'fixed',
source: params.source, source: params.source,
description: params.description, description: params.description,
metadata: null, metadata: params.metadata ?? null,
cost: params.cost.toString(), cost: params.cost.toString(),
workspaceId: params.workspaceId ?? null, workspaceId: params.workspaceId ?? null,
workflowId: params.workflowId ?? null, workflowId: params.workflowId ?? null,

View File

@@ -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) * Track chat deployed (workflow deployed as chat interface)
*/ */

View File

@@ -1,4 +1,5 @@
import { getEnv, isTruthy } from '@/lib/core/config/env' import { getEnv, isTruthy } from '@/lib/core/config/env'
import { isHosted } from '@/lib/core/config/feature-flags'
import type { SubBlockConfig } from '@/blocks/types' import type { SubBlockConfig } from '@/blocks/types'
export type CanonicalMode = 'basic' | 'advanced' export type CanonicalMode = 'basic' | 'advanced'
@@ -270,3 +271,12 @@ export function isSubBlockFeatureEnabled(subBlock: SubBlockConfig): boolean {
if (!subBlock.requiresFeature) return true if (!subBlock.requiresFeature) return true
return isTruthy(getEnv(subBlock.requiresFeature)) 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
}

View File

@@ -10,6 +10,7 @@ import {
isCanonicalPair, isCanonicalPair,
isNonEmptyValue, isNonEmptyValue,
isSubBlockFeatureEnabled, isSubBlockFeatureEnabled,
isSubBlockHiddenByHostedKey,
resolveCanonicalMode, resolveCanonicalMode,
} from '@/lib/workflows/subblocks/visibility' } from '@/lib/workflows/subblocks/visibility'
import { getBlock } from '@/blocks' import { getBlock } from '@/blocks'
@@ -49,6 +50,7 @@ function shouldSerializeSubBlock(
canonicalModeOverrides?: CanonicalModeOverrides canonicalModeOverrides?: CanonicalModeOverrides
): boolean { ): boolean {
if (!isSubBlockFeatureEnabled(subBlockConfig)) return false if (!isSubBlockFeatureEnabled(subBlockConfig)) return false
if (isSubBlockHiddenByHostedKey(subBlockConfig)) return false
if (subBlockConfig.mode === 'trigger') { if (subBlockConfig.mode === 'trigger') {
if (!isTriggerContext && !isTriggerCategory) return false if (!isTriggerContext && !isTriggerCategory) return false

View File

@@ -1,6 +1,9 @@
import { createLogger } from '@sim/logger'
import type { ExaAnswerParams, ExaAnswerResponse } from '@/tools/exa/types' import type { ExaAnswerParams, ExaAnswerResponse } from '@/tools/exa/types'
import type { ToolConfig } from '@/tools/types' import type { ToolConfig } from '@/tools/types'
const logger = createLogger('ExaAnswerTool')
export const answerTool: ToolConfig<ExaAnswerParams, ExaAnswerResponse> = { export const answerTool: ToolConfig<ExaAnswerParams, ExaAnswerResponse> = {
id: 'exa_answer', id: 'exa_answer',
name: 'Exa Answer', name: 'Exa Answer',
@@ -27,6 +30,23 @@ export const answerTool: ToolConfig<ExaAnswerParams, ExaAnswerResponse> = {
description: 'Exa AI API Key', description: 'Exa AI API Key',
}, },
}, },
hosting: {
envKeys: ['EXA_API_KEY'],
apiKeyParam: 'apiKey',
byokProviderId: 'exa',
pricing: {
type: 'custom',
getCost: (_params, response) => {
// Use costDollars from Exa API response
if (response.costDollars?.total) {
return { cost: response.costDollars.total, metadata: { costDollars: response.costDollars } }
}
// Fallback: $5/1000 requests
logger.warn('Exa answer response missing costDollars, using fallback pricing')
return 0.005
},
},
},
request: { request: {
url: 'https://api.exa.ai/answer', url: 'https://api.exa.ai/answer',
@@ -61,6 +81,7 @@ export const answerTool: ToolConfig<ExaAnswerParams, ExaAnswerResponse> = {
url: citation.url, url: citation.url,
text: citation.text || '', text: citation.text || '',
})) || [], })) || [],
costDollars: data.costDollars,
}, },
} }
}, },

View File

@@ -1,6 +1,9 @@
import { createLogger } from '@sim/logger'
import type { ExaFindSimilarLinksParams, ExaFindSimilarLinksResponse } from '@/tools/exa/types' import type { ExaFindSimilarLinksParams, ExaFindSimilarLinksResponse } from '@/tools/exa/types'
import type { ToolConfig } from '@/tools/types' import type { ToolConfig } from '@/tools/types'
const logger = createLogger('ExaFindSimilarLinksTool')
export const findSimilarLinksTool: ToolConfig< export const findSimilarLinksTool: ToolConfig<
ExaFindSimilarLinksParams, ExaFindSimilarLinksParams,
ExaFindSimilarLinksResponse ExaFindSimilarLinksResponse
@@ -76,6 +79,24 @@ export const findSimilarLinksTool: ToolConfig<
description: 'Exa AI API Key', description: 'Exa AI API Key',
}, },
}, },
hosting: {
envKeys: ['EXA_API_KEY'],
apiKeyParam: 'apiKey',
byokProviderId: 'exa',
pricing: {
type: 'custom',
getCost: (_params, response) => {
// Use costDollars from Exa API response
if (response.costDollars?.total) {
return { cost: response.costDollars.total, metadata: { costDollars: response.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 resultCount = response.similarLinks?.length || 0
return resultCount <= 25 ? 0.005 : 0.025
},
},
},
request: { request: {
url: 'https://api.exa.ai/findSimilar', url: 'https://api.exa.ai/findSimilar',
@@ -140,6 +161,7 @@ export const findSimilarLinksTool: ToolConfig<
highlights: result.highlights, highlights: result.highlights,
score: result.score || 0, score: result.score || 0,
})), })),
costDollars: data.costDollars,
}, },
} }
}, },

View File

@@ -1,6 +1,9 @@
import { createLogger } from '@sim/logger'
import type { ExaGetContentsParams, ExaGetContentsResponse } from '@/tools/exa/types' import type { ExaGetContentsParams, ExaGetContentsResponse } from '@/tools/exa/types'
import type { ToolConfig } from '@/tools/types' import type { ToolConfig } from '@/tools/types'
const logger = createLogger('ExaGetContentsTool')
export const getContentsTool: ToolConfig<ExaGetContentsParams, ExaGetContentsResponse> = { export const getContentsTool: ToolConfig<ExaGetContentsParams, ExaGetContentsResponse> = {
id: 'exa_get_contents', id: 'exa_get_contents',
name: 'Exa Get Contents', name: 'Exa Get Contents',
@@ -61,6 +64,23 @@ export const getContentsTool: ToolConfig<ExaGetContentsParams, ExaGetContentsRes
description: 'Exa AI API Key', description: 'Exa AI API Key',
}, },
}, },
hosting: {
envKeys: ['EXA_API_KEY'],
apiKeyParam: 'apiKey',
byokProviderId: 'exa',
pricing: {
type: 'custom',
getCost: (_params, response) => {
// Use costDollars from Exa API response
if (response.costDollars?.total) {
return { cost: response.costDollars.total, metadata: { costDollars: response.costDollars } }
}
// Fallback: $1/1000 pages
logger.warn('Exa get_contents response missing costDollars, using fallback pricing')
return (response.results?.length || 0) * 0.001
},
},
},
request: { request: {
url: 'https://api.exa.ai/contents', url: 'https://api.exa.ai/contents',
@@ -132,6 +152,7 @@ export const getContentsTool: ToolConfig<ExaGetContentsParams, ExaGetContentsRes
summary: result.summary || '', summary: result.summary || '',
highlights: result.highlights, highlights: result.highlights,
})), })),
costDollars: data.costDollars,
}, },
} }
}, },

View File

@@ -34,6 +34,25 @@ export const researchTool: ToolConfig<ExaResearchParams, ExaResearchResponse> =
description: 'Exa AI API Key', description: 'Exa AI API Key',
}, },
}, },
hosting: {
envKeys: ['EXA_API_KEY'],
apiKeyParam: 'apiKey',
byokProviderId: 'exa',
pricing: {
type: 'custom',
getCost: (params, response) => {
// Use costDollars from Exa API response
if (response.costDollars?.total) {
return { cost: response.costDollars.total, metadata: { costDollars: response.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: { request: {
url: 'https://api.exa.ai/research/v1', url: 'https://api.exa.ai/research/v1',
@@ -111,6 +130,8 @@ export const researchTool: ToolConfig<ExaResearchParams, ExaResearchResponse> =
score: 1.0, score: 1.0,
}, },
], ],
// Include cost breakdown for pricing calculation
costDollars: taskData.costDollars,
} }
return result return result
} }

View File

@@ -1,6 +1,9 @@
import { createLogger } from '@sim/logger'
import type { ExaSearchParams, ExaSearchResponse } from '@/tools/exa/types' import type { ExaSearchParams, ExaSearchResponse } from '@/tools/exa/types'
import type { ToolConfig } from '@/tools/types' import type { ToolConfig } from '@/tools/types'
const logger = createLogger('ExaSearchTool')
export const searchTool: ToolConfig<ExaSearchParams, ExaSearchResponse> = { export const searchTool: ToolConfig<ExaSearchParams, ExaSearchResponse> = {
id: 'exa_search', id: 'exa_search',
name: 'Exa Search', name: 'Exa Search',
@@ -86,6 +89,29 @@ export const searchTool: ToolConfig<ExaSearchParams, ExaSearchResponse> = {
description: 'Exa AI API Key', description: 'Exa AI API Key',
}, },
}, },
hosting: {
envKeys: ['EXA_API_KEY'],
apiKeyParam: 'apiKey',
byokProviderId: 'exa',
pricing: {
type: 'custom',
getCost: (params, response) => {
// Use costDollars from Exa API response
if (response.costDollars?.total) {
return { cost: response.costDollars.total, metadata: { costDollars: response.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 resultCount = response.results?.length || 0
return resultCount <= 25 ? 0.005 : 0.025
},
},
},
request: { request: {
url: 'https://api.exa.ai/search', url: 'https://api.exa.ai/search',
@@ -167,6 +193,7 @@ export const searchTool: ToolConfig<ExaSearchParams, ExaSearchResponse> = {
highlights: result.highlights, highlights: result.highlights,
score: result.score, score: result.score,
})), })),
costDollars: data.costDollars,
}, },
} }
}, },

View File

@@ -6,6 +6,11 @@ export interface ExaBaseParams {
apiKey: string apiKey: string
} }
/** Cost breakdown returned by Exa API responses */
export interface ExaCostDollars {
total: number
}
// Search tool types // Search tool types
export interface ExaSearchParams extends ExaBaseParams { export interface ExaSearchParams extends ExaBaseParams {
query: string query: string
@@ -50,6 +55,7 @@ export interface ExaSearchResult {
export interface ExaSearchResponse extends ToolResponse { export interface ExaSearchResponse extends ToolResponse {
output: { output: {
results: ExaSearchResult[] results: ExaSearchResult[]
costDollars?: ExaCostDollars
} }
} }
@@ -78,6 +84,7 @@ export interface ExaGetContentsResult {
export interface ExaGetContentsResponse extends ToolResponse { export interface ExaGetContentsResponse extends ToolResponse {
output: { output: {
results: ExaGetContentsResult[] results: ExaGetContentsResult[]
costDollars?: ExaCostDollars
} }
} }
@@ -120,6 +127,7 @@ export interface ExaSimilarLink {
export interface ExaFindSimilarLinksResponse extends ToolResponse { export interface ExaFindSimilarLinksResponse extends ToolResponse {
output: { output: {
similarLinks: ExaSimilarLink[] similarLinks: ExaSimilarLink[]
costDollars?: ExaCostDollars
} }
} }
@@ -137,6 +145,7 @@ export interface ExaAnswerResponse extends ToolResponse {
url: string url: string
text: string text: string
}[] }[]
costDollars?: ExaCostDollars
} }
} }
@@ -158,6 +167,7 @@ export interface ExaResearchResponse extends ToolResponse {
author?: string author?: string
score: number score: number
}[] }[]
costDollars?: ExaCostDollars
} }
} }

View File

@@ -15,6 +15,27 @@ import {
} from '@sim/testing' } from '@sim/testing'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
// Mock isHosted flag - hoisted so we can control it per test
const mockIsHosted = vi.hoisted(() => ({ value: false }))
vi.mock('@/lib/core/config/feature-flags', () => ({
isHosted: mockIsHosted.value,
isProd: false,
isDev: true,
isTest: true,
}))
// Mock getBYOKKey - hoisted so we can control it per test
const mockGetBYOKKey = vi.hoisted(() => vi.fn())
vi.mock('@/lib/api-key/byok', () => ({
getBYOKKey: mockGetBYOKKey,
}))
// Mock logFixedUsage for billing
const mockLogFixedUsage = vi.hoisted(() => vi.fn())
vi.mock('@/lib/billing/core/usage-log', () => ({
logFixedUsage: mockLogFixedUsage,
}))
// Mock custom tools query - must be hoisted before imports // Mock custom tools query - must be hoisted before imports
vi.mock('@/hooks/queries/custom-tools', () => ({ vi.mock('@/hooks/queries/custom-tools', () => ({
getCustomTool: (toolId: string) => { getCustomTool: (toolId: string) => {
@@ -959,3 +980,233 @@ describe('MCP Tool Execution', () => {
expect(result.timing).toBeDefined() 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)
})
})

View File

@@ -1,5 +1,9 @@
import { createLogger } from '@sim/logger' import { createLogger } from '@sim/logger'
import { generateInternalToken } from '@/lib/auth/internal' import { generateInternalToken } from '@/lib/auth/internal'
import { getBYOKKey } from '@/lib/api-key/byok'
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 { DEFAULT_EXECUTION_TIMEOUT_MS } from '@/lib/core/execution-limits'
import { import {
secureFetchWithPinnedIP, secureFetchWithPinnedIP,
@@ -13,16 +17,239 @@ import { resolveSkillContent } from '@/executor/handlers/agent/skills-resolver'
import type { ExecutionContext } from '@/executor/types' import type { ExecutionContext } from '@/executor/types'
import type { ErrorInfo } from '@/tools/error-extractors' import type { ErrorInfo } from '@/tools/error-extractors'
import { extractErrorMessage } 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 { import {
formatRequestParams, formatRequestParams,
getTool, getTool,
getToolAsync, getToolAsync,
validateRequiredParametersAfterMerge, validateRequiredParametersAfterMerge,
} from '@/tools/utils' } from '@/tools/utils'
import { PlatformEvents } from '@/lib/core/telemetry'
const logger = createLogger('Tools') 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<string, unknown>,
executionContext: ExecutionContext | undefined,
requestId: string
): Promise<HostedKeyInjectionResult> {
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<T>(
fn: () => Promise<T>,
context: RetryContext,
maxRetries = 3,
baseDelayMs = 1000
): Promise<T> {
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 * Math.pow(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<string, unknown>
}
/**
* Calculate cost based on pricing model
*/
function calculateToolCost(
pricing: ToolHostingPricing,
params: Record<string, unknown>,
response: Record<string, unknown>
): 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}`)
}
}
}
/**
* Calculate and log hosted key cost for a tool execution.
* Logs to usageLog for audit trail and returns cost for accumulation in userStats.
*/
async function processHostedKeyCost(
tool: ToolConfig,
params: Record<string, unknown>,
response: Record<string, unknown>,
executionContext: ExecutionContext | undefined,
requestId: string
): Promise<number> {
if (!tool.hosting?.pricing) {
return 0
}
const { cost, metadata } = calculateToolCost(tool.hosting.pricing, params, response)
if (cost <= 0) return 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
}
/** /**
* Normalizes a tool ID by stripping resource ID suffix (UUID). * Normalizes a tool ID by stripping resource ID suffix (UUID).
* Workflow tools: 'workflow_executor_<uuid>' -> 'workflow_executor' * Workflow tools: 'workflow_executor_<uuid>' -> 'workflow_executor'
@@ -279,6 +506,14 @@ export async function executeTool(
throw new Error(`Tool not found: ${toolId}`) 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 we have a credential parameter, fetch the access token
if (contextParams.credential) { if (contextParams.credential) {
logger.info( logger.info(
@@ -391,7 +626,14 @@ export async function executeTool(
const endTime = new Date() const endTime = new Date()
const endTimeISO = endTime.toISOString() const endTimeISO = endTime.toISOString()
const duration = endTime.getTime() - startTime.getTime() const duration = endTime.getTime() - startTime.getTime()
return {
// Calculate and log hosted key cost if applicable
let hostedKeyCost = 0
if (hostedKeyInfo.isUsingHostedKey && finalResult.success) {
hostedKeyCost = await processHostedKeyCost(tool, contextParams, finalResult.output, executionContext, requestId)
}
const response: ToolResponse = {
...finalResult, ...finalResult,
timing: { timing: {
startTime: startTimeISO, startTime: startTimeISO,
@@ -399,10 +641,25 @@ export async function executeTool(
duration, duration,
}, },
} }
if (hostedKeyCost > 0) {
response.cost = { total: hostedKeyCost }
}
return response
} }
// Execute the tool request directly (internal routes use regular fetch, external use SSRF-protected fetch) // 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 // Apply post-processing if available and not skipped
let finalResult = result let finalResult = result
@@ -424,7 +681,14 @@ export async function executeTool(
const endTime = new Date() const endTime = new Date()
const endTimeISO = endTime.toISOString() const endTimeISO = endTime.toISOString()
const duration = endTime.getTime() - startTime.getTime() const duration = endTime.getTime() - startTime.getTime()
return {
// Calculate and log hosted key cost if applicable
let hostedKeyCost = 0
if (hostedKeyInfo.isUsingHostedKey && finalResult.success) {
hostedKeyCost = await processHostedKeyCost(tool, contextParams, finalResult.output, executionContext, requestId)
}
const response: ToolResponse = {
...finalResult, ...finalResult,
timing: { timing: {
startTime: startTimeISO, startTime: startTimeISO,
@@ -432,6 +696,10 @@ export async function executeTool(
duration, duration,
}, },
} }
if (hostedKeyCost > 0) {
response.cost = { total: hostedKeyCost }
}
return response
} catch (error: any) { } catch (error: any) {
logger.error(`[${requestId}] Error executing tool ${toolId}:`, { logger.error(`[${requestId}] Error executing tool ${toolId}:`, {
error: error instanceof Error ? error.message : String(error), error: error instanceof Error ? error.message : String(error),

View File

@@ -1,5 +1,7 @@
import type { OAuthService } from '@/lib/oauth' import type { OAuthService } from '@/lib/oauth'
export type BYOKProviderId = 'openai' | 'anthropic' | 'google' | 'mistral' | 'exa'
export type HttpMethod = 'GET' | 'POST' | 'PUT' | 'DELETE' | 'PATCH' | 'HEAD' export type HttpMethod = 'GET' | 'POST' | 'PUT' | 'DELETE' | 'PATCH' | 'HEAD'
export type OutputType = export type OutputType =
@@ -39,6 +41,12 @@ export interface ToolResponse {
endTime: string // ISO timestamp when the tool execution ended endTime: string // ISO timestamp when the tool execution ended
duration: number // Duration in milliseconds duration: number // Duration in milliseconds
} }
// Cost incurred by this tool execution (for billing)
cost?: {
total: number
input?: number
output?: number
}
} }
export interface OAuthConfig { export interface OAuthConfig {
@@ -127,6 +135,13 @@ export interface ToolConfig<P = any, R = any> {
* Maps param IDs to their enrichment configuration. * Maps param IDs to their enrichment configuration.
*/ */
schemaEnrichment?: Record<string, SchemaEnrichmentConfig> schemaEnrichment?: Record<string, SchemaEnrichmentConfig>
/**
* 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<P, R extends ToolResponse ? R : ToolResponse>
} }
export interface TableRow { export interface TableRow {
@@ -170,3 +185,48 @@ export interface SchemaEnrichmentConfig {
required?: string[] required?: string[]
} | null> } | 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<string, unknown>
}
/** Custom pricing calculated from params and response (e.g., Exa with different modes/result counts) */
export interface CustomPricing<P = Record<string, unknown>, R extends ToolResponse = ToolResponse> {
type: 'custom'
/** Calculate cost based on request params and response data. Returns cost or cost with metadata. */
getCost: (params: P, response: R['output']) => number | CustomPricingResult
}
/** Union of all pricing models */
export type ToolHostingPricing<P = Record<string, unknown>, R extends ToolResponse = ToolResponse> =
| PerRequestPricing
| CustomPricing<P, R>
/**
* 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<P = Record<string, unknown>, R extends ToolResponse = ToolResponse> {
/** 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<P, R>
}