mirror of
https://github.com/simstudioai/sim.git
synced 2026-09-24 15:45:35 +08:00
fix(mothership): billing for mcps, hosted key tools (#5740)
* fix(mothership): attribution for direct exec tools * fix hosted key charges for direct tool exec * fix llm chat tool
This commit is contained in:
Vikhyath Mondreti
@@ -62,6 +62,55 @@ describe('copilot tool executor fallback', () => {
|
||||
expect(result).toEqual({ success: true, output: { emails: [] } })
|
||||
})
|
||||
|
||||
it('threads billing attribution into _context for dynamic tools (MCP)', async () => {
|
||||
isKnownTool.mockReturnValue(false)
|
||||
isSimExecuted.mockReturnValue(false)
|
||||
executeAppTool.mockResolvedValue({ success: true, output: {} })
|
||||
|
||||
const billingAttribution = {
|
||||
actorUserId: 'user-1',
|
||||
workspaceId: 'ws-1',
|
||||
organizationId: null,
|
||||
billedAccountUserId: 'owner-1',
|
||||
billingEntity: { type: 'user', id: 'owner-1' },
|
||||
billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z' },
|
||||
payerSubscription: null,
|
||||
}
|
||||
|
||||
await executeTool(
|
||||
'mcp-server-1-web_search_exa',
|
||||
{ query: 'test' },
|
||||
{
|
||||
userId: 'user-1',
|
||||
workflowId: '',
|
||||
workspaceId: 'ws-1',
|
||||
billingAttribution: billingAttribution as never,
|
||||
}
|
||||
)
|
||||
|
||||
expect(executeAppTool).toHaveBeenCalledWith(
|
||||
'mcp-server-1-web_search_exa',
|
||||
expect.objectContaining({
|
||||
_context: expect.objectContaining({
|
||||
userId: 'user-1',
|
||||
workspaceId: 'ws-1',
|
||||
billingAttribution,
|
||||
}),
|
||||
})
|
||||
)
|
||||
})
|
||||
|
||||
it('omits billingAttribution from _context when the context has none', async () => {
|
||||
isKnownTool.mockReturnValue(false)
|
||||
isSimExecuted.mockReturnValue(false)
|
||||
executeAppTool.mockResolvedValue({ success: true, output: {} })
|
||||
|
||||
await executeTool('gmail_read', {}, { userId: 'user-1', workflowId: 'workflow-1' })
|
||||
|
||||
const appParams = executeAppTool.mock.calls[0][1] as Record<string, unknown>
|
||||
expect(appParams._context).not.toHaveProperty('billingAttribution')
|
||||
})
|
||||
|
||||
it('uses the registered handler for client-routed tools when running headless (Mothership block)', async () => {
|
||||
isKnownTool.mockReturnValue(true)
|
||||
isSimExecuted.mockReturnValue(false)
|
||||
|
||||
@@ -161,6 +161,7 @@ function buildAppToolParams(
|
||||
requestMode: context.requestMode,
|
||||
currentAgentId: context.currentAgentId,
|
||||
enforceCredentialAccess: true,
|
||||
...(context.billingAttribution ? { billingAttribution: context.billingAttribution } : {}),
|
||||
}
|
||||
|
||||
return result
|
||||
|
||||
@@ -2973,6 +2973,130 @@ describe('Cost Field Handling', () => {
|
||||
Object.assign(tools, originalTools)
|
||||
})
|
||||
|
||||
it('emits _serviceCost for copilot executions using a hosted key', async () => {
|
||||
const mockTool = {
|
||||
id: 'test_copilot_hosted_cost',
|
||||
name: 'Test Copilot Hosted Cost',
|
||||
description: 'A hosted-key tool invoked by the copilot',
|
||||
version: '1.0.0',
|
||||
params: {
|
||||
apiKey: { type: 'string', required: false },
|
||||
},
|
||||
hosting: {
|
||||
envKeyPrefix: 'TEST_HOSTED_KEY',
|
||||
apiKeyParam: 'apiKey',
|
||||
byokProviderId: 'exa',
|
||||
pricing: {
|
||||
type: 'per_request' as const,
|
||||
cost: 0.005,
|
||||
},
|
||||
rateLimit: {
|
||||
mode: 'per_request' as const,
|
||||
requestsPerMinute: 100,
|
||||
},
|
||||
},
|
||||
request: {
|
||||
url: '/api/test/copilot-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_copilot_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
|
||||
|
||||
// Copilot flow: no executionContext option; scope comes from _context,
|
||||
// which the copilot tool-executor stamps with copilotToolExecution.
|
||||
const result = await executeTool('test_copilot_hosted_cost', {
|
||||
_context: {
|
||||
userId: 'user-123',
|
||||
workspaceId: 'workspace-456',
|
||||
copilotToolExecution: true,
|
||||
},
|
||||
})
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
expect(result.output.cost).toEqual({ total: 0.005 })
|
||||
expect(result.output._serviceCost).toEqual({ service: 'exa', cost: 0.005 })
|
||||
|
||||
Object.assign(tools, originalTools)
|
||||
})
|
||||
|
||||
it('does not emit _serviceCost for workflow executions using a hosted key', async () => {
|
||||
const mockTool = {
|
||||
id: 'test_workflow_hosted_cost',
|
||||
name: 'Test Workflow Hosted Cost',
|
||||
description: 'A hosted-key tool invoked by a workflow',
|
||||
version: '1.0.0',
|
||||
params: {
|
||||
apiKey: { type: 'string', required: false },
|
||||
},
|
||||
hosting: {
|
||||
envKeyPrefix: 'TEST_HOSTED_KEY',
|
||||
apiKeyParam: 'apiKey',
|
||||
pricing: {
|
||||
type: 'per_request' as const,
|
||||
cost: 0.005,
|
||||
},
|
||||
rateLimit: {
|
||||
mode: 'per_request' as const,
|
||||
requestsPerMinute: 100,
|
||||
},
|
||||
},
|
||||
request: {
|
||||
url: '/api/test/workflow-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_workflow_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({ userId: 'user-123' } as any)
|
||||
const result = await executeTool(
|
||||
'test_workflow_hosted_cost',
|
||||
{},
|
||||
{ executionContext: mockContext }
|
||||
)
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
// Workflow executions bill through the execution ledger; emitting
|
||||
// _serviceCost here would double-bill via Go's service-charge path.
|
||||
expect(result.output.cost).toEqual({ total: 0.005 })
|
||||
expect(result.output._serviceCost).toBeUndefined()
|
||||
|
||||
Object.assign(tools, originalTools)
|
||||
})
|
||||
|
||||
it('should use custom pricing getCost function', async () => {
|
||||
const mockGetCost = vi.fn().mockReturnValue({
|
||||
cost: 0.015,
|
||||
|
||||
@@ -709,6 +709,15 @@ export function postProcessToolOutput(toolId: string, output: Record<string, unk
|
||||
/**
|
||||
* Apply post-execution hosted-key cost tracking to a successful tool result.
|
||||
* Reports custom dimension usage, calculates cost, and merges it into the output.
|
||||
*
|
||||
* Billing capture differs by caller:
|
||||
* - Workflow executions bill `output.cost.total` through trace spans and the
|
||||
* execution ledger (`recordUsage`), so the `cost` field alone suffices.
|
||||
* - Copilot tool executions have no execution ledger. Their only billing hook
|
||||
* is Go's `extractServiceCost`, which reads a top-level `_serviceCost` field
|
||||
* from the tool result and charges it through the per-round update-cost
|
||||
* callback (the same path the media tools use). Without it, hosted-key spend
|
||||
* from copilot-dispatched integration tools is never charged.
|
||||
*/
|
||||
async function applyHostedKeyCostToResult(
|
||||
finalResult: ToolResponse,
|
||||
@@ -734,12 +743,16 @@ async function applyHostedKeyCostToResult(
|
||||
hostedKeyMetrics.recordCostCharged(hostedKeyCost, { provider, tool: tool.id })
|
||||
|
||||
if (hostedKeyCost > 0) {
|
||||
const { copilotToolExecution } = resolveToolScope(params, executionContext)
|
||||
finalResult.output = {
|
||||
...finalResult.output,
|
||||
cost: {
|
||||
...metadata,
|
||||
total: hostedKeyCost,
|
||||
},
|
||||
// Copilot-only: workflow runs must not emit _serviceCost or the cost
|
||||
// would be billed twice (execution ledger + Go service charge).
|
||||
...(copilotToolExecution ? { _serviceCost: { service: provider, cost: hostedKeyCost } } : {}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,6 +19,10 @@ interface LLMChatParams {
|
||||
bedrockAccessKeyId?: string
|
||||
bedrockSecretKey?: string
|
||||
bedrockRegion?: string
|
||||
_context?: {
|
||||
workspaceId?: string
|
||||
workflowId?: string
|
||||
}
|
||||
}
|
||||
|
||||
interface LLMChatResponse extends ToolResponse {
|
||||
@@ -146,6 +150,11 @@ export const llmChatTool: ToolConfig<LLMChatParams, LLMChatResponse> = {
|
||||
bedrockAccessKeyId: params.bedrockAccessKeyId,
|
||||
bedrockSecretKey: params.bedrockSecretKey,
|
||||
bedrockRegion: params.bedrockRegion,
|
||||
// The executor attaches the billing-attribution header on internal
|
||||
// routes whenever the execution scope carries one, and /api/providers
|
||||
// rejects that header unless the body names the workspace it is
|
||||
// validated against.
|
||||
...(params._context?.workspaceId ? { workspaceId: params._context.workspaceId } : {}),
|
||||
}
|
||||
},
|
||||
},
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
{
|
||||
"lockfileVersion": 1,
|
||||
"configVersion": 0,
|
||||
"workspaces": {
|
||||
"": {
|
||||
"name": "simstudio",
|
||||
@@ -3086,7 +3085,7 @@
|
||||
|
||||
"lru.min": ["lru.min@1.1.4", "", {}, "sha512-DqC6n3QQ77zdFpCMASA1a3Jlb64Hv2N2DciFGkO/4L9+q/IpIAuRlKOvCXabtRW6cQf8usbmM6BE/TOPysCdIA=="],
|
||||
|
||||
"lucide-react": ["lucide-react@0.479.0", "", { "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-aBhNnveRhorBOK7uA4gDjgaf+YlHMdMhQ/3cupk6exM10hWlEU+2QtWYOfhXhjAsmdb6LeKR+NZnow4UxRRiTQ=="],
|
||||
"lucide-react": ["lucide-react@1.23.0", "", { "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-38BpJcD0JhFosxHApP/BYsBetLpQFRoTRzEzstM/XCc3jsAG7wqaY1lgVwxiUe3xqYE+lNxo2PkCmYwXWrwwIw=="],
|
||||
|
||||
"magic-string": ["magic-string@0.30.21", "", { "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.5" } }, "sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ=="],
|
||||
|
||||
@@ -4462,6 +4461,10 @@
|
||||
|
||||
"@sim/db/@types/node": ["@types/node@22.19.21", "", { "dependencies": { "undici-types": "~6.21.0" } }, "sha512-VMeFBSCKQKmm2swI2kW51SFusDqekC6q9trBCvJ/JliDchFSuoYYKN7yVNjPthP1HKZcx3U1gI/wTcEBjEFKTA=="],
|
||||
|
||||
"@sim/emcn/lucide-react": ["lucide-react@0.479.0", "", { "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-aBhNnveRhorBOK7uA4gDjgaf+YlHMdMhQ/3cupk6exM10hWlEU+2QtWYOfhXhjAsmdb6LeKR+NZnow4UxRRiTQ=="],
|
||||
|
||||
"@sim/workflow-renderer/lucide-react": ["lucide-react@0.479.0", "", { "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-aBhNnveRhorBOK7uA4gDjgaf+YlHMdMhQ/3cupk6exM10hWlEU+2QtWYOfhXhjAsmdb6LeKR+NZnow4UxRRiTQ=="],
|
||||
|
||||
"@smithy/middleware-compression/fflate": ["fflate@0.8.1", "", {}, "sha512-/exOvEuc+/iaUm105QIiOt4LpBdMTWsXxqR0HDF35vx3fmaKzw7354gTilCh5rkzEt8WYyG//ku3h3nRmd7CHQ=="],
|
||||
|
||||
"@socket.io/redis-adapter/debug": ["debug@4.3.7", "", { "dependencies": { "ms": "^2.1.3" } }, "sha512-Er2nc/H7RrMXZBFCEim6TCmMk02Z8vLC2Rbi1KEBggpo0fS6l0S1nnapwmIi3yW/+GOJap1Krg4w0Hg80oCqgQ=="],
|
||||
@@ -4660,14 +4663,10 @@
|
||||
|
||||
"fumadocs-openapi/ajv": ["ajv@8.20.0", "", { "dependencies": { "fast-deep-equal": "^3.1.3", "fast-uri": "^3.0.1", "json-schema-traverse": "^1.0.0", "require-from-string": "^2.0.2" } }, "sha512-Thbli+OlOj+iMPYFBVBfJ3OmCAnaSyNn4M1vz9T6Gka5Jt9ba/HIR56joy65tY6kx/FCF5VXNB819Y7/GUrBGA=="],
|
||||
|
||||
"fumadocs-openapi/lucide-react": ["lucide-react@1.23.0", "", { "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-38BpJcD0JhFosxHApP/BYsBetLpQFRoTRzEzstM/XCc3jsAG7wqaY1lgVwxiUe3xqYE+lNxo2PkCmYwXWrwwIw=="],
|
||||
|
||||
"fumadocs-openapi/tailwind-merge": ["tailwind-merge@3.6.0", "", {}, "sha512-uxL7qAVQriqRQPAyK3pj66VqskWqoZ37PW94jwOTwNfq/z9oyu1V+eqrZqtR2+fCiXdYOZe/Modt8GtvqNzu+w=="],
|
||||
|
||||
"fumadocs-ui/@radix-ui/react-slot": ["@radix-ui/react-slot@1.3.0", "", { "dependencies": { "@radix-ui/react-compose-refs": "1.1.3" }, "peerDependencies": { "@types/react": "*", "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react"] }, "sha512-MojKku4U/miO8Av4Dkb+ctMAQx7JmY96LmtDQlAarCRtd7rN52QCSzBF+XAvr5S6coSVj9HEPBgHAHKEJVk/WA=="],
|
||||
|
||||
"fumadocs-ui/lucide-react": ["lucide-react@1.23.0", "", { "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-38BpJcD0JhFosxHApP/BYsBetLpQFRoTRzEzstM/XCc3jsAG7wqaY1lgVwxiUe3xqYE+lNxo2PkCmYwXWrwwIw=="],
|
||||
|
||||
"fumadocs-ui/tailwind-merge": ["tailwind-merge@3.6.0", "", {}, "sha512-uxL7qAVQriqRQPAyK3pj66VqskWqoZ37PW94jwOTwNfq/z9oyu1V+eqrZqtR2+fCiXdYOZe/Modt8GtvqNzu+w=="],
|
||||
|
||||
"gaxios/node-fetch": ["node-fetch@2.7.0", "", { "dependencies": { "whatwg-url": "^5.0.0" }, "peerDependencies": { "encoding": "^0.1.0" }, "optionalPeers": ["encoding"] }, "sha512-c4FRfUm/dbcWZ7U+1Wq0AwCyFL+3nt2bEw05wfxSz+DWpWsitgmSgYmy2dQdWyKC1694ELPqMs/YzUSNozLt8A=="],
|
||||
@@ -4822,6 +4821,8 @@
|
||||
|
||||
"sharp/semver": ["semver@7.8.4", "", { "bin": { "semver": "bin/semver.js" } }, "sha512-rUCObTnP32Q08R2uuIrt7r9PlEonuTmtuXYcW6s5kjdlj3xbnwe+21yXptAUYcMAABLkYYTtnmzb3w3EDZfueA=="],
|
||||
|
||||
"sim/lucide-react": ["lucide-react@0.479.0", "", { "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-aBhNnveRhorBOK7uA4gDjgaf+YlHMdMhQ/3cupk6exM10hWlEU+2QtWYOfhXhjAsmdb6LeKR+NZnow4UxRRiTQ=="],
|
||||
|
||||
"sim/tailwindcss": ["tailwindcss@3.4.19", "", { "dependencies": { "@alloc/quick-lru": "^5.2.0", "arg": "^5.0.2", "chokidar": "^3.6.0", "didyoumean": "^1.2.2", "dlv": "^1.1.3", "fast-glob": "^3.3.2", "glob-parent": "^6.0.2", "is-glob": "^4.0.3", "jiti": "^1.21.7", "lilconfig": "^3.1.3", "micromatch": "^4.0.8", "normalize-path": "^3.0.0", "object-hash": "^3.0.0", "picocolors": "^1.1.1", "postcss": "^8.4.47", "postcss-import": "^15.1.0", "postcss-js": "^4.0.1", "postcss-load-config": "^4.0.2 || ^5.0 || ^6.0", "postcss-nested": "^6.2.0", "postcss-selector-parser": "^6.1.2", "resolve": "^1.22.8", "sucrase": "^3.35.0" }, "bin": { "tailwind": "lib/cli.js", "tailwindcss": "lib/cli.js" } }, "sha512-3ofp+LL8E+pK/JuPLPggVAIaEuhvIz4qNcf3nA1Xn2o/7fb7s/TYpHhwGDv1ZU3PkBluUVaF8PyCHcm48cKLWQ=="],
|
||||
|
||||
"simstudio/@types/node": ["@types/node@20.19.43", "", { "dependencies": { "undici-types": "~6.21.0" } }, "sha512-6oYBAi5ikg4Pl+kGsoYtawUMBT2zZMCvPNF7pVLnHZfd1zf38DRiWn/gT01RYCdUqkv7Fhr+C9ot4/tb+2sVvA=="],
|
||||
|
||||
Reference in New Issue
Block a user