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

Author: icecrasher321

apps/sim/lib/copilot/tool-executor/executor.test.ts

@@ -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)

apps/sim/lib/copilot/tool-executor/executor.ts

@@ -161,6 +161,7 @@ function buildAppToolParams(
     requestMode: context.requestMode,
     currentAgentId: context.currentAgentId,
     enforceCredentialAccess: true,
+    ...(context.billingAttribution ? { billingAttribution: context.billingAttribution } : {}),
   }
 
   return result

apps/sim/tools/index.test.ts

@@ -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,

apps/sim/tools/index.ts

@@ -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 } } : {}),
     }
   }
 }

apps/sim/tools/llm/chat.ts

@@ -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 } : {}),
       }
     },
   },