diff --git a/apps/sim/app/api/providers/route.test.ts b/apps/sim/app/api/providers/route.test.ts new file mode 100644 index 00000000000..97161bc7972 --- /dev/null +++ b/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() + }) +}) diff --git a/apps/sim/app/api/providers/route.ts b/apps/sim/app/api/providers/route.ts index 49b90074943..dbb87f071c9 100644 --- a/apps/sim/app/api/providers/route.ts +++ b/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, }) diff --git a/apps/sim/executor/handlers/agent/agent-handler.test.ts b/apps/sim/executor/handlers/agent/agent-handler.test.ts index 6990a719bf2..0c0eafb57bb 100644 --- a/apps/sim/executor/handlers/agent/agent-handler.test.ts +++ b/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[] = [] diff --git a/apps/sim/executor/handlers/agent/agent-handler.ts b/apps/sim/executor/handlers/agent/agent-handler.ts index 2373b8e6527..81018c20bf1 100644 --- a/apps/sim/executor/handlers/agent/agent-handler.ts +++ b/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, diff --git a/apps/sim/providers/types.ts b/apps/sim/providers/types.ts index 01752a9c1b8..289957a5296 100644 --- a/apps/sim/providers/types.ts +++ b/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 diff --git a/apps/sim/providers/utils.test.ts b/apps/sim/providers/utils.test.ts index 1697c12f80d..87d576b79c5 100644 --- a/apps/sim/providers/utils.test.ts +++ b/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 = { diff --git a/apps/sim/providers/utils.ts b/apps/sim/providers/utils.ts index 8ca3017354a..1c6904c5c30 100644 --- a/apps/sim/providers/utils.ts +++ b/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 isDeployedContext?: boolean callChain?: string[] + billingAttribution?: BillingAttributionSnapshot } ): { toolParams: Record @@ -1310,10 +1312,10 @@ 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 } : {}), @@ -1321,6 +1323,9 @@ export function prepareToolExecution( ? { isDeployedContext: request.isDeployedContext } : {}), ...(request.callChain ? { callChain: request.callChain } : {}), + ...(request.billingAttribution + ? { billingAttribution: request.billingAttribution } + : {}), }, } : {}), diff --git a/bun.lock b/bun.lock index 454f3857417..eaee1d6f50f 100644 --- a/bun.lock +++ b/bun.lock @@ -1,6 +1,5 @@ { "lockfileVersion": 1, - "configVersion": 0, "workspaces": { "": { "name": "simstudio", @@ -3063,7 +3062,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=="], @@ -4421,6 +4420,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=="], @@ -4609,14 +4612,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=="], "gcp-metadata/gaxios": ["gaxios@7.1.3", "", { "dependencies": { "extend": "^3.0.2", "https-proxy-agent": "^7.0.1", "node-fetch": "^3.3.2", "rimraf": "^5.0.1" } }, "sha512-YGGyuEdVIjqxkxVH1pUTMY/XtmmsApXrCVv5EU25iX6inEPbV+VakJfLealkBtJN69AQmh1eGOdCl9Sm1UP6XQ=="], @@ -4763,6 +4762,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=="],