11import { createLogger } from '@sim/logger'
22import { generateInternalToken } from '@/lib/auth/internal'
3+ import { getBYOKKey } from '@/lib/api-key/byok'
4+ import { logFixedUsage } from '@/lib/billing/core/usage-log'
5+ import { env } from '@/lib/core/config/env'
36import { DEFAULT_EXECUTION_TIMEOUT_MS } from '@/lib/core/execution-limits'
47import {
58 secureFetchWithPinnedIP ,
@@ -13,7 +16,12 @@ import { resolveSkillContent } from '@/executor/handlers/agent/skills-resolver'
1316import type { ExecutionContext } from '@/executor/types'
1417import type { ErrorInfo } from '@/tools/error-extractors'
1518import { extractErrorMessage } from '@/tools/error-extractors'
16- import type { OAuthTokenPayload , ToolConfig , ToolResponse } from '@/tools/types'
19+ import type {
20+ OAuthTokenPayload ,
21+ ToolConfig ,
22+ ToolHostingPricing ,
23+ ToolResponse ,
24+ } from '@/tools/types'
1725import {
1826 formatRequestParams ,
1927 getTool ,
@@ -23,6 +31,150 @@ import {
2331
2432const logger = createLogger ( 'Tools' )
2533
34+ /**
35+ * Get a hosted API key from environment variables
36+ * Supports rotation when multiple keys are configured
37+ */
38+ function getHostedKeyFromEnv ( envKeys : string [ ] ) : string | null {
39+ const keys = envKeys
40+ . map ( ( key ) => env [ key as keyof typeof env ] )
41+ . filter ( ( value ) : value is string => Boolean ( value ) )
42+
43+ if ( keys . length === 0 ) return null
44+
45+ // Round-robin rotation based on current minute
46+ const currentMinute = Math . floor ( Date . now ( ) / 60000 )
47+ const keyIndex = currentMinute % keys . length
48+
49+ return keys [ keyIndex ]
50+ }
51+
52+ /**
53+ * Inject hosted API key if tool supports it and user didn't provide one.
54+ * Checks BYOK workspace keys first, then falls back to hosted env keys.
55+ * Returns whether a hosted (billable) key was injected.
56+ */
57+ async function injectHostedKeyIfNeeded (
58+ tool : ToolConfig ,
59+ params : Record < string , unknown > ,
60+ executionContext : ExecutionContext | undefined ,
61+ requestId : string
62+ ) : Promise < boolean > {
63+ if ( ! tool . hosting ) return false
64+
65+ const { envKeys, apiKeyParam, byokProviderId } = tool . hosting
66+ const userProvidedKey = params [ apiKeyParam ]
67+
68+ if ( userProvidedKey ) {
69+ logger . debug ( `[${ requestId } ${ tool . id } )
70+ return false
71+ }
72+
73+ // Check BYOK workspace key first
74+ if ( byokProviderId && executionContext ?. workspaceId ) {
75+ try {
76+ const byokResult = await getBYOKKey (
77+ executionContext . workspaceId ,
78+ byokProviderId as 'openai' | 'anthropic' | 'google' | 'mistral' | 'serper'
79+ )
80+ if ( byokResult ) {
81+ params [ apiKeyParam ] = byokResult . apiKey
82+ logger . info ( `[${ requestId } ${ tool . id } )
83+ return false // Don't bill - user's own key
84+ }
85+ } catch ( error ) {
86+ logger . error ( `[${ requestId } ${ tool . id } , error )
87+ // Fall through to hosted key
88+ }
89+ }
90+
91+ // Fall back to hosted env key
92+ const hostedKey = getHostedKeyFromEnv ( envKeys )
93+ if ( ! hostedKey ) {
94+ logger . debug ( `[${ requestId } ${ tool . id } )
95+ return false
96+ }
97+
98+ params [ apiKeyParam ] = hostedKey
99+ logger . info ( `[${ requestId } ${ tool . id } )
100+ return true // Bill the user
101+ }
102+
103+ /**
104+ * Calculate cost based on pricing model
105+ */
106+ function calculateToolCost (
107+ pricing : ToolHostingPricing ,
108+ params : Record < string , unknown > ,
109+ response : Record < string , unknown >
110+ ) : number {
111+ switch ( pricing . type ) {
112+ case 'per_request' :
113+ return pricing . cost
114+
115+ case 'per_unit' : {
116+ const usage = pricing . getUsage ( params , response )
117+ return usage * pricing . costPerUnit
118+ }
119+
120+ case 'per_result' : {
121+ const resultCount = pricing . getResultCount ( response )
122+ const billableResults = pricing . maxResults
123+ ? Math . min ( resultCount , pricing . maxResults )
124+ : resultCount
125+ return billableResults * pricing . costPerResult
126+ }
127+
128+ case 'per_second' : {
129+ const duration = pricing . getDuration ( response )
130+ const billableDuration = pricing . minimumSeconds
131+ ? Math . max ( duration , pricing . minimumSeconds )
132+ : duration
133+ return billableDuration * pricing . costPerSecond
134+ }
135+
136+ default : {
137+ const exhaustiveCheck : never = pricing
138+ throw new Error ( `Unknown pricing type: ${ ( exhaustiveCheck as ToolHostingPricing ) . type } )
139+ }
140+ }
141+ }
142+
143+ /**
144+ * Log usage for a tool that used a hosted API key
145+ */
146+ async function logHostedToolUsage (
147+ tool : ToolConfig ,
148+ params : Record < string , unknown > ,
149+ response : Record < string , unknown > ,
150+ executionContext : ExecutionContext | undefined ,
151+ requestId : string
152+ ) : Promise < void > {
153+ if ( ! tool . hosting ?. pricing || ! executionContext ?. userId ) {
154+ return
155+ }
156+
157+ const cost = calculateToolCost ( tool . hosting . pricing , params , response )
158+
159+ if ( cost <= 0 ) return
160+
161+ try {
162+ await logFixedUsage ( {
163+ userId : executionContext . userId ,
164+ source : 'workflow' ,
165+ description : `tool:${ tool . id } ,
166+ cost,
167+ workspaceId : executionContext . workspaceId ,
168+ workflowId : executionContext . workflowId ,
169+ executionId : executionContext . executionId ,
170+ } )
171+ logger . debug ( `[${ requestId } ${ tool . id } ${ cost } )
172+ } catch ( error ) {
173+ logger . error ( `[${ requestId } ${ tool . id } , error )
174+ // Don't throw - usage logging should not break the main flow
175+ }
176+ }
177+
26178/**
27179 * Normalizes a tool ID by stripping resource ID suffix (UUID).
28180 * Workflow tools: 'workflow_executor_<uuid>' -> 'workflow_executor'
@@ -279,6 +431,14 @@ export async function executeTool(
279431 throw new Error ( `Tool not found: ${ toolId } )
280432 }
281433
434+ // Inject hosted API key if tool supports it and user didn't provide one
435+ const isUsingHostedKey = await injectHostedKeyIfNeeded (
436+ tool ,
437+ contextParams ,
438+ executionContext ,
439+ requestId
440+ )
441+
282442 // If we have a credential parameter, fetch the access token
283443 if ( contextParams . credential ) {
284444 logger . info (
@@ -387,6 +547,11 @@ export async function executeTool(
387547 // Process file outputs if execution context is available
388548 finalResult = await processFileOutputs ( finalResult , tool , executionContext )
389549
550+ // Log usage for hosted key if execution was successful
551+ if ( isUsingHostedKey && finalResult . success ) {
552+ await logHostedToolUsage ( tool , contextParams , finalResult . output , executionContext , requestId )
553+ }
554+
390555 // Add timing data to the result
391556 const endTime = new Date ( )
392557 const endTimeISO = endTime . toISOString ( )
@@ -420,6 +585,11 @@ export async function executeTool(
420585 // Process file outputs if execution context is available
421586 finalResult = await processFileOutputs ( finalResult , tool , executionContext )
422587
588+ // Log usage for hosted key if execution was successful
589+ if ( isUsingHostedKey && finalResult . success ) {
590+ await logHostedToolUsage ( tool , contextParams , finalResult . output , executionContext , requestId )
591+ }
592+
423593 // Add timing data to the result
424594 const endTime = new Date ( )
425595 const endTimeISO = endTime . toISOString ( )
0 commit comments