Skip to content

Commit 8eaf401

Browse files
author
Theodore Li
committed
Handle hosted agent tool calls
1 parent 5d04ae5 commit 8eaf401

File tree

22 files changed

+159
-59
lines changed

22 files changed

+159
-59
lines changed

apps/sim/lib/billing/core/usage-log.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ export type UsageLogSource = 'workflow' | 'wand' | 'copilot' | 'mcp_copilot'
2222
export interface ModelUsageMetadata {
2323
inputTokens: number
2424
outputTokens: number
25+
toolCost?: number
2526
}
2627

2728
/**
@@ -44,6 +45,7 @@ export interface LogModelUsageParams {
4445
inputTokens: number
4546
outputTokens: number
4647
cost: number
48+
toolCost?: number
4749
workspaceId?: string
4850
workflowId?: string
4951
executionId?: string
@@ -76,6 +78,7 @@ export async function logModelUsage(params: LogModelUsageParams): Promise<void>
7678
const metadata: ModelUsageMetadata = {
7779
inputTokens: params.inputTokens,
7880
outputTokens: params.outputTokens,
81+
...(params.toolCost != null && params.toolCost > 0 && { toolCost: params.toolCost }),
7982
}
8083

8184
await db.insert(usageLog).values({
@@ -157,6 +160,7 @@ export interface LogWorkflowUsageBatchParams {
157160
{
158161
total: number
159162
tokens: { input: number; output: number }
163+
toolCost?: number
160164
}
161165
>
162166
}
@@ -209,6 +213,8 @@ export async function logWorkflowUsageBatch(params: LogWorkflowUsageBatchParams)
209213
metadata: {
210214
inputTokens: modelData.tokens.input,
211215
outputTokens: modelData.tokens.output,
216+
...(modelData.toolCost != null &&
217+
modelData.toolCost > 0 && { toolCost: modelData.toolCost }),
212218
},
213219
cost: modelData.total.toString(),
214220
workspaceId: params.workspaceId ?? null,

apps/sim/lib/copilot/orchestrator/tool-executor/integration-tools.ts

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,19 +6,18 @@ import type {
66
ToolCallResult,
77
ToolCallState,
88
} from '@/lib/copilot/orchestrator/types'
9+
import { isHosted } from '@/lib/core/config/feature-flags'
910
import { generateRequestId } from '@/lib/core/utils/request'
1011
import { getEffectiveDecryptedEnv } from '@/lib/environment/utils'
1112
import { refreshTokenIfNeeded } from '@/app/api/auth/oauth/utils'
1213
import { resolveEnvVarReferences } from '@/executor/utils/reference-validation'
1314
import { executeTool } from '@/tools'
15+
import type { ToolConfig } from '@/tools/types'
1416
import { resolveToolId } from '@/tools/utils'
1517

1618
export async function executeIntegrationToolDirect(
1719
toolCall: ToolCallState,
18-
toolConfig: {
19-
oauth?: { required?: boolean; provider?: string }
20-
params?: { apiKey?: { required?: boolean } }
21-
},
20+
toolConfig: ToolConfig,
2221
context: ExecutionContext
2322
): Promise<ToolCallResult> {
2423
const { userId, workflowId } = context
@@ -74,7 +73,8 @@ export async function executeIntegrationToolDirect(
7473
executionParams.accessToken = accessToken
7574
}
7675

77-
if (toolConfig.params?.apiKey?.required && !executionParams.apiKey) {
76+
const hasHostedKeySupport = isHosted && !!toolConfig.hosting
77+
if (toolConfig.params?.apiKey?.required && !executionParams.apiKey && !hasHostedKeySupport) {
7878
return {
7979
success: false,
8080
error: `API key not provided for ${toolName}. Use {{YOUR_API_KEY_ENV_VAR}} to reference your environment variable.`,
@@ -83,6 +83,7 @@ export async function executeIntegrationToolDirect(
8383

8484
executionParams._context = {
8585
workflowId,
86+
workspaceId,
8687
userId,
8788
}
8889

apps/sim/lib/logs/execution/logger.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,7 @@ export class ExecutionLogger implements IExecutionLoggerService {
181181
input: number
182182
output: number
183183
total: number
184+
toolCost?: number
184185
tokens: { input: number; output: number; total: number }
185186
}
186187
>
@@ -507,6 +508,7 @@ export class ExecutionLogger implements IExecutionLoggerService {
507508
input: number
508509
output: number
509510
total: number
511+
toolCost?: number
510512
tokens: { input: number; output: number; total: number }
511513
}
512514
>

apps/sim/lib/logs/execution/logging-factory.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ export function calculateCostSummary(traceSpans: any[]): {
9595
input: number
9696
output: number
9797
total: number
98+
toolCost?: number
9899
tokens: { input: number; output: number; total: number }
99100
}
100101
>
@@ -143,6 +144,7 @@ export function calculateCostSummary(traceSpans: any[]): {
143144
input: number
144145
output: number
145146
total: number
147+
toolCost?: number
146148
tokens: { input: number; output: number; total: number }
147149
}
148150
> = {}
@@ -171,6 +173,10 @@ export function calculateCostSummary(traceSpans: any[]): {
171173
models[model].tokens.input += span.tokens?.input ?? span.tokens?.prompt ?? 0
172174
models[model].tokens.output += span.tokens?.output ?? span.tokens?.completion ?? 0
173175
models[model].tokens.total += span.tokens?.total || 0
176+
177+
if (span.cost.toolCost) {
178+
models[model].toolCost = (models[model].toolCost || 0) + span.cost.toolCost
179+
}
174180
}
175181
}
176182

apps/sim/providers/anthropic/core.ts

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import {
1919
calculateCost,
2020
prepareToolExecution,
2121
prepareToolsWithUsageControl,
22+
sumToolCosts,
2223
} from '@/providers/utils'
2324
import { executeTool } from '@/tools'
2425

@@ -490,7 +491,7 @@ export async function executeAnthropicProviderRequest(
490491
}
491492

492493
const toolCalls = []
493-
const toolResults = []
494+
const toolResults: any[] = []
494495
const currentMessages = [...messages]
495496
let iterationCount = 0
496497
let hasUsedForcedTool = false
@@ -783,10 +784,12 @@ export async function executeAnthropicProviderRequest(
783784
}
784785

785786
const streamCost = calculateCost(request.model, usage.input_tokens, usage.output_tokens)
787+
const tc = sumToolCosts(toolResults)
786788
streamingResult.execution.output.cost = {
787789
input: accumulatedCost.input + streamCost.input,
788790
output: accumulatedCost.output + streamCost.output,
789-
total: accumulatedCost.total + streamCost.total,
791+
toolCost: tc || undefined,
792+
total: accumulatedCost.total + streamCost.total + tc,
790793
}
791794

792795
const streamEndTime = Date.now()
@@ -829,6 +832,7 @@ export async function executeAnthropicProviderRequest(
829832
cost: {
830833
input: accumulatedCost.input,
831834
output: accumulatedCost.output,
835+
toolCost: undefined as number | undefined,
832836
total: accumulatedCost.total,
833837
},
834838
},
@@ -901,7 +905,7 @@ export async function executeAnthropicProviderRequest(
901905
}
902906

903907
const toolCalls = []
904-
const toolResults = []
908+
const toolResults: any[] = []
905909
const currentMessages = [...messages]
906910
let iterationCount = 0
907911
let hasUsedForcedTool = false
@@ -1208,10 +1212,12 @@ export async function executeAnthropicProviderRequest(
12081212
}
12091213

12101214
const streamCost = calculateCost(request.model, usage.input_tokens, usage.output_tokens)
1215+
const tc2 = sumToolCosts(toolResults)
12111216
streamingResult.execution.output.cost = {
12121217
input: cost.input + streamCost.input,
12131218
output: cost.output + streamCost.output,
1214-
total: cost.total + streamCost.total,
1219+
toolCost: tc2 || undefined,
1220+
total: cost.total + streamCost.total + tc2,
12151221
}
12161222

12171223
const streamEndTime = Date.now()
@@ -1254,6 +1260,7 @@ export async function executeAnthropicProviderRequest(
12541260
cost: {
12551261
input: cost.input,
12561262
output: cost.output,
1263+
toolCost: undefined as number | undefined,
12571264
total: cost.total,
12581265
},
12591266
},

apps/sim/providers/azure-openai/index.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ import {
3535
calculateCost,
3636
prepareToolExecution,
3737
prepareToolsWithUsageControl,
38+
sumToolCosts,
3839
} from '@/providers/utils'
3940
import { executeTool } from '@/tools'
4041

@@ -499,10 +500,12 @@ async function executeChatCompletionsRequest(
499500
usage.prompt_tokens,
500501
usage.completion_tokens
501502
)
503+
const tc = sumToolCosts(toolResults)
502504
streamingResult.execution.output.cost = {
503505
input: accumulatedCost.input + streamCost.input,
504506
output: accumulatedCost.output + streamCost.output,
505-
total: accumulatedCost.total + streamCost.total,
507+
toolCost: tc || undefined,
508+
total: accumulatedCost.total + streamCost.total + tc,
506509
}
507510

508511
const streamEndTime = Date.now()

apps/sim/providers/bedrock/index.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import {
3333
calculateCost,
3434
prepareToolExecution,
3535
prepareToolsWithUsageControl,
36+
sumToolCosts,
3637
} from '@/providers/utils'
3738
import { executeTool } from '@/tools'
3839

@@ -815,10 +816,12 @@ export const bedrockProvider: ProviderConfig = {
815816
}
816817

817818
const streamCost = calculateCost(request.model, usage.inputTokens, usage.outputTokens)
819+
const tc = sumToolCosts(toolResults)
818820
streamingResult.execution.output.cost = {
819821
input: cost.input + streamCost.input,
820822
output: cost.output + streamCost.output,
821-
total: cost.total + streamCost.total,
823+
toolCost: tc || undefined,
824+
total: cost.total + streamCost.total + tc,
822825
}
823826

824827
const streamEndTime = Date.now()
@@ -861,6 +864,7 @@ export const bedrockProvider: ProviderConfig = {
861864
cost: {
862865
input: cost.input,
863866
output: cost.output,
867+
toolCost: undefined as number | undefined,
864868
total: cost.total,
865869
},
866870
},

apps/sim/providers/cerebras/index.ts

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import {
1616
calculateCost,
1717
prepareToolExecution,
1818
prepareToolsWithUsageControl,
19+
sumToolCosts,
1920
trackForcedToolUsage,
2021
} from '@/providers/utils'
2122
import { executeTool } from '@/tools'
@@ -195,7 +196,7 @@ export const cerebrasProvider: ProviderConfig = {
195196
total: currentResponse.usage?.total_tokens || 0,
196197
}
197198
const toolCalls = []
198-
const toolResults = []
199+
const toolResults: any[] = []
199200
const currentMessages = [...allMessages]
200201
let iterationCount = 0
201202

@@ -472,10 +473,12 @@ export const cerebrasProvider: ProviderConfig = {
472473
usage.prompt_tokens,
473474
usage.completion_tokens
474475
)
476+
const tc = sumToolCosts(toolResults)
475477
streamingResult.execution.output.cost = {
476478
input: accumulatedCost.input + streamCost.input,
477479
output: accumulatedCost.output + streamCost.output,
478-
total: accumulatedCost.total + streamCost.total,
480+
toolCost: tc || undefined,
481+
total: accumulatedCost.total + streamCost.total + tc,
479482
}
480483
}),
481484
execution: {
@@ -508,6 +511,7 @@ export const cerebrasProvider: ProviderConfig = {
508511
cost: {
509512
input: accumulatedCost.input,
510513
output: accumulatedCost.output,
514+
toolCost: undefined as number | undefined,
511515
total: accumulatedCost.total,
512516
},
513517
},

apps/sim/providers/deepseek/index.ts

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import {
1515
calculateCost,
1616
prepareToolExecution,
1717
prepareToolsWithUsageControl,
18+
sumToolCosts,
1819
trackForcedToolUsage,
1920
} from '@/providers/utils'
2021
import { executeTool } from '@/tools'
@@ -205,7 +206,7 @@ export const deepseekProvider: ProviderConfig = {
205206
total: currentResponse.usage?.total_tokens || 0,
206207
}
207208
const toolCalls = []
208-
const toolResults = []
209+
const toolResults: any[] = []
209210
const currentMessages = [...allMessages]
210211
let iterationCount = 0
211212
let hasUsedForcedTool = false
@@ -471,10 +472,12 @@ export const deepseekProvider: ProviderConfig = {
471472
usage.prompt_tokens,
472473
usage.completion_tokens
473474
)
475+
const tc = sumToolCosts(toolResults)
474476
streamingResult.execution.output.cost = {
475477
input: accumulatedCost.input + streamCost.input,
476478
output: accumulatedCost.output + streamCost.output,
477-
total: accumulatedCost.total + streamCost.total,
479+
toolCost: tc || undefined,
480+
total: accumulatedCost.total + streamCost.total + tc,
478481
}
479482
}
480483
),
@@ -508,6 +511,7 @@ export const deepseekProvider: ProviderConfig = {
508511
cost: {
509512
input: accumulatedCost.input,
510513
output: accumulatedCost.output,
514+
toolCost: undefined as number | undefined,
511515
total: accumulatedCost.total,
512516
},
513517
},

apps/sim/providers/gemini/core.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import {
3131
isDeepResearchModel,
3232
prepareToolExecution,
3333
prepareToolsWithUsageControl,
34+
sumToolCosts,
3435
} from '@/providers/utils'
3536
import { executeTool } from '@/tools'
3637
import type { ExecutionState, GeminiProviderType, GeminiUsage } from './types'
@@ -1163,10 +1164,12 @@ export async function executeGeminiRequest(
11631164
usage.promptTokenCount,
11641165
usage.candidatesTokenCount
11651166
)
1167+
const tc = sumToolCosts(state.toolResults)
11661168
streamingResult.execution.output.cost = {
11671169
input: accumulatedCost.input + streamCost.input,
11681170
output: accumulatedCost.output + streamCost.output,
1169-
total: accumulatedCost.total + streamCost.total,
1171+
toolCost: tc || undefined,
1172+
total: accumulatedCost.total + streamCost.total + tc,
11701173
pricing: streamCost.pricing,
11711174
}
11721175

0 commit comments

Comments
 (0)