fix(agent): pass through billing attribution to tools

apps/sim/app/api/providers/route.test.ts

@@ -0,0 +1,148 @@
+/**
+ * @vitest-environment node
+ */
+import { createMockRequest, hybridAuthMockFns } from '@sim/testing'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const {
+  mockExecuteProviderRequest,
+  mockRequireBillingAttributionHeader,
+  mockCheckWorkspaceAccess,
+  mockAuthorizeCredentialUse,
+} = vi.hoisted(() => ({
+  mockExecuteProviderRequest: vi.fn(),
+  mockRequireBillingAttributionHeader: vi.fn(),
+  mockCheckWorkspaceAccess: vi.fn(),
+  mockAuthorizeCredentialUse: vi.fn(),
+}))
+
+vi.mock('@/providers', () => ({
+  executeProviderRequest: mockExecuteProviderRequest,
+}))
+
+vi.mock('@/lib/billing/core/billing-attribution', () => ({
+  BILLING_ATTRIBUTION_HEADER: 'x-sim-billing-attribution',
+  requireBillingAttributionHeader: mockRequireBillingAttributionHeader,
+}))
+
+vi.mock('@/lib/workspaces/permissions/utils', () => ({
+  checkWorkspaceAccess: mockCheckWorkspaceAccess,
+}))
+
+vi.mock('@/lib/auth/credential-access', () => ({
+  authorizeCredentialUse: mockAuthorizeCredentialUse,
+}))
+
+vi.mock('@/app/api/auth/oauth/utils', () => ({
+  getServiceAccountToken: vi.fn(),
+  refreshTokenIfNeeded: vi.fn(),
+  resolveOAuthAccountId: vi.fn(),
+}))
+
+vi.mock('@/ee/access-control/utils/permission-check', () => ({
+  assertPermissionsAllowed: vi.fn(),
+  IntegrationNotAllowedError: class IntegrationNotAllowedError extends Error {},
+  ModelNotAllowedError: class ModelNotAllowedError extends Error {},
+  ProviderNotAllowedError: class ProviderNotAllowedError extends Error {},
+}))
+
+import { POST } from '@/app/api/providers/route'
+
+const BILLING_ATTRIBUTION = {
+  actorUserId: 'user-1',
+  workspaceId: 'ws-1',
+  organizationId: 'org-1',
+  billedAccountUserId: 'owner-1',
+  billingEntity: { type: 'organization', id: 'org-1' },
+  billingPeriod: {
+    start: '2026-07-01T00:00:00.000Z',
+    end: '2026-08-01T00:00:00.000Z',
+  },
+  payerSubscription: null,
+}
+
+describe('POST /api/providers', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    hybridAuthMockFns.mockCheckInternalAuth.mockResolvedValue({
+      success: true,
+      userId: 'user-1',
+      authType: 'internal_jwt',
+    })
+    mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true })
+    mockRequireBillingAttributionHeader.mockReturnValue(BILLING_ATTRIBUTION)
+    mockExecuteProviderRequest.mockResolvedValue({
+      content: 'hello',
+      model: 'gpt-4o',
+      tokens: { input: 1, output: 1, total: 2 },
+    })
+  })
+
+  it('validates the attribution header and forwards it to executeProviderRequest', async () => {
+    const res = await POST(
+      createMockRequest(
+        'POST',
+        { provider: 'openai', model: 'gpt-4o', workspaceId: 'ws-1' },
+        { 'x-sim-billing-attribution': 'encoded-attribution' }
+      )
+    )
+
+    expect(res.status).toBe(200)
+    expect(mockRequireBillingAttributionHeader).toHaveBeenCalledWith(expect.anything(), {
+      actorUserId: 'user-1',
+      workspaceId: 'ws-1',
+    })
+    expect(mockExecuteProviderRequest).toHaveBeenCalledWith(
+      'openai',
+      expect.objectContaining({ billingAttribution: BILLING_ATTRIBUTION })
+    )
+  })
+
+  it('executes without attribution when the header is absent', async () => {
+    const res = await POST(
+      createMockRequest('POST', { provider: 'openai', model: 'gpt-4o', workspaceId: 'ws-1' })
+    )
+
+    expect(res.status).toBe(200)
+    expect(mockRequireBillingAttributionHeader).not.toHaveBeenCalled()
+    expect(mockExecuteProviderRequest).toHaveBeenCalledWith(
+      'openai',
+      expect.objectContaining({ billingAttribution: undefined })
+    )
+  })
+
+  it('rejects an attribution header when the body has no workspaceId to validate against', async () => {
+    const res = await POST(
+      createMockRequest(
+        'POST',
+        { provider: 'openai', model: 'gpt-4o' },
+        { 'x-sim-billing-attribution': 'encoded-attribution' }
+      )
+    )
+
+    expect(res.status).toBe(400)
+    expect(mockRequireBillingAttributionHeader).not.toHaveBeenCalled()
+    expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
+  })
+
+  it('rejects with 400 when the attribution header does not match the authenticated scope', async () => {
+    mockRequireBillingAttributionHeader.mockImplementation(() => {
+      throw new Error('Billing attribution header does not match the authenticated request scope')
+    })
+
+    const res = await POST(
+      createMockRequest(
+        'POST',
+        { provider: 'openai', model: 'gpt-4o', workspaceId: 'ws-1' },
+        { 'x-sim-billing-attribution': 'encoded-attribution' }
+      )
+    )
+
+    expect(res.status).toBe(400)
+    const body = await res.json()
+    expect(body.error).toBe(
+      'Billing attribution header does not match the authenticated request scope'
+    )
+    expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
+  })
+})

apps/sim/app/api/providers/route.ts

@@ -8,6 +8,11 @@ import { executeProviderContract } from '@/lib/api/contracts/providers'
 import { parseRequest } from '@/lib/api/server'
 import { authorizeCredentialUse } from '@/lib/auth/credential-access'
 import { checkInternalAuth } from '@/lib/auth/hybrid'
+import {
+  BILLING_ATTRIBUTION_HEADER,
+  type BillingAttributionSnapshot,
+  requireBillingAttributionHeader,
+} from '@/lib/billing/core/billing-attribution'
 import { generateRequestId } from '@/lib/core/utils/request'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { checkWorkspaceAccess } from '@/lib/workspaces/permissions/utils'
@@ -177,11 +182,41 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
       )
     }
 
+    /**
+     * Nested tool calls made by the LLM (e.g. knowledge search) hit internal
+     * routes that require the upstream billing decision. This route is
+     * internal-JWT-only, so the caller's attribution arrives as a header;
+     * validate it against the authenticated scope and thread it through. A
+     * header the route cannot validate is a caller protocol error (400), never
+     * silently dropped.
+     */
+    let billingAttribution: BillingAttributionSnapshot | undefined
+    if (request.headers.get(BILLING_ATTRIBUTION_HEADER)) {
+      if (!workspaceId) {
+        return NextResponse.json(
+          { error: 'workspaceId is required when billing attribution is supplied' },
+          { status: 400 }
+        )
+      }
+      try {
+        billingAttribution = requireBillingAttributionHeader(request.headers, {
+          actorUserId: auth.userId,
+          workspaceId,
+        })
+      } catch (error) {
+        return NextResponse.json(
+          { error: getErrorMessage(error, 'Invalid billing attribution header') },
+          { status: 400 }
+        )
+      }
+    }
+
     logger.info(`[${requestId}] Executing provider request`, {
       provider,
       model,
       workflowId,
       hasApiKey: !!finalApiKey,
+      hasBillingAttribution: !!billingAttribution,
     })
 
     const response = await executeProviderRequest(provider, {
@@ -209,6 +244,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
       workflowVariables,
       blockData,
       blockNameMapping,
+      billingAttribution,
       reasoningEffort,
       verbosity,
     })

apps/sim/executor/handlers/agent/agent-handler.test.ts

@@ -1828,6 +1828,42 @@ describe('AgentBlockHandler', () => {
       expect(providerCallArgs.callChain).toEqual(['wf-parent', 'test-workflow-456'])
     })
 
+    it('should pass billingAttribution to executeProviderRequest so LLM tool calls carry it', async () => {
+      const billingAttribution = {
+        actorUserId: 'user-1',
+        workspaceId: 'test-workspace-123',
+        organizationId: 'organization-1',
+        billedAccountUserId: 'owner-1',
+        billingEntity: { type: 'organization', id: 'organization-1' },
+        billingPeriod: {
+          start: '2026-07-01T00:00:00.000Z',
+          end: '2026-08-01T00:00:00.000Z',
+        },
+        payerSubscription: null,
+      }
+
+      const inputs = {
+        model: 'gpt-4o',
+        userPrompt: 'Search the knowledge base',
+        apiKey: 'test-api-key',
+      }
+
+      const contextWithAttribution = {
+        ...mockContext,
+        workspaceId: 'test-workspace-123',
+        workflowId: 'test-workflow-456',
+        metadata: { ...mockContext.metadata, billingAttribution },
+      } as ExecutionContext
+
+      mockGetProviderFromModel.mockReturnValue('openai')
+
+      await handler.execute(contextWithAttribution, mockBlock, inputs)
+
+      expect(mockExecuteProviderRequest).toHaveBeenCalled()
+      const providerCallArgs = mockExecuteProviderRequest.mock.calls[0][1]
+      expect(providerCallArgs.billingAttribution).toEqual(billingAttribution)
+    })
+
     it('should handle multiple MCP tools from the same server efficiently', async () => {
       const fetchCalls: any[] = []
 

apps/sim/executor/handlers/agent/agent-handler.ts

@@ -1021,6 +1021,7 @@ export class AgentBlockHandler implements BlockHandler {
         blockNameMapping,
         isDeployedContext: ctx.isDeployedContext,
         callChain: ctx.callChain,
+        billingAttribution: ctx.metadata.billingAttribution,
         reasoningEffort: providerRequest.reasoningEffort,
         verbosity: providerRequest.verbosity,
         thinkingLevel: providerRequest.thinkingLevel,

apps/sim/providers/types.ts

@@ -1,3 +1,4 @@
+import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution'
 import type { ProviderTimingSegment, StreamingExecution, UserFile } from '@/executor/types'
 
 export type ProviderId =
@@ -185,6 +186,12 @@ export interface ProviderRequest {
   thinkingLevel?: string
   isDeployedContext?: boolean
   callChain?: string[]
+  /**
+   * Immutable actor/payer decision captured before execution. Propagated into
+   * the `_context` of every tool the LLM invokes so internal routes that
+   * require the billing attribution header (e.g. knowledge search) receive it.
+   */
+  billingAttribution?: BillingAttributionSnapshot
   /** Previous interaction ID for multi-turn Interactions API requests (deep research follow-ups) */
   previousInteractionId?: string
   abortSignal?: AbortSignal

apps/sim/providers/utils.test.ts

@@ -1290,6 +1290,67 @@ describe('prepareToolExecution', () => {
     })
   })
 
+  describe('_context propagation', () => {
+    const billingAttribution = {
+      actorUserId: 'user-1',
+      workspaceId: 'workspace-1',
+      organizationId: 'organization-1',
+      billedAccountUserId: 'owner-1',
+      billingEntity: { type: 'organization' as const, id: 'organization-1' },
+      billingPeriod: {
+        start: '2026-07-01T00:00:00.000Z',
+        end: '2026-08-01T00:00:00.000Z',
+      },
+      payerSubscription: null,
+    }
+
+    it.concurrent(
+      'should include billingAttribution in _context when the request carries it',
+      () => {
+        const tool = { params: {} }
+        const request = {
+          workflowId: 'wf-123',
+          workspaceId: 'workspace-1',
+          userId: 'user-1',
+          billingAttribution,
+        }
+
+        const { executionParams } = prepareToolExecution(tool, {}, request)
+
+        expect(executionParams._context.billingAttribution).toEqual(billingAttribution)
+      }
+    )
+
+    it.concurrent('should omit billingAttribution from _context when the request lacks it', () => {
+      const tool = { params: {} }
+      const request = { workflowId: 'wf-123', workspaceId: 'workspace-1' }
+
+      const { executionParams } = prepareToolExecution(tool, {}, request)
+
+      expect(executionParams._context).toBeDefined()
+      expect(executionParams._context).not.toHaveProperty('billingAttribution')
+    })
+
+    it.concurrent('should carry billingAttribution even when the request has no workflowId', () => {
+      const tool = { params: {} }
+      const request = { workspaceId: 'workspace-1', billingAttribution }
+
+      const { executionParams } = prepareToolExecution(tool, {}, request)
+
+      expect(executionParams._context.billingAttribution).toEqual(billingAttribution)
+      expect(executionParams._context.workspaceId).toBe('workspace-1')
+      expect(executionParams._context).not.toHaveProperty('workflowId')
+    })
+
+    it.concurrent('should not build _context when there is no workflowId or attribution', () => {
+      const tool = { params: {} }
+
+      const { executionParams } = prepareToolExecution(tool, {}, { workspaceId: 'workspace-1' })
+
+      expect(executionParams).not.toHaveProperty('_context')
+    })
+  })
+
   describe('inputMapping deep merge for workflow tools', () => {
     it.concurrent('should deep merge inputMapping when user provides empty object', () => {
       const tool = {

apps/sim/providers/utils.ts

@@ -4,6 +4,7 @@ import { omit } from '@sim/utils/object'
 import type OpenAI from 'openai'
 import type { ChatCompletionChunk } from 'openai/resources/chat/completions'
 import type { CompletionUsage } from 'openai/resources/completions'
+import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution'
 import { formatCreditCost } from '@/lib/billing/credits/conversion'
 import { env } from '@/lib/core/config/env'
 import { getBlacklistedProvidersFromEnv, isHosted } from '@/lib/core/config/env-flags'
@@ -1293,6 +1294,7 @@ export function prepareToolExecution(
     blockNameMapping?: Record<string, string>
     isDeployedContext?: boolean
     callChain?: string[]
+    billingAttribution?: BillingAttributionSnapshot
   }
 ): {
   toolParams: Record<string, any>
@@ -1310,17 +1312,20 @@ export function prepareToolExecution(
 
   const executionParams = {
     ...toolParams,
-    ...(request.workflowId
+    ...(request.workflowId || request.billingAttribution
       ? {
           _context: {
-            workflowId: request.workflowId,
+            ...(request.workflowId ? { workflowId: request.workflowId } : {}),
             ...(request.workspaceId ? { workspaceId: request.workspaceId } : {}),
             ...(request.chatId ? { chatId: request.chatId } : {}),
             ...(request.userId ? { userId: request.userId } : {}),
             ...(request.isDeployedContext !== undefined
               ? { isDeployedContext: request.isDeployedContext }
               : {}),
             ...(request.callChain ? { callChain: request.callChain } : {}),
+            ...(request.billingAttribution
+              ? { billingAttribution: request.billingAttribution }
+              : {}),
           },
         }
       : {}),