tests

Author: icecrasher321

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

@@ -0,0 +1,147 @@
+/**
+ * @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('ignores the 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(200)
+    expect(mockRequireBillingAttributionHeader).not.toHaveBeenCalled()
+    expect(mockExecuteProviderRequest).toHaveBeenCalledWith(
+      'openai',
+      expect.objectContaining({ billingAttribution: undefined })
+    )
+  })
+
+  it('fails 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(500)
+    expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
+  })
+})

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/providers/utils.test.ts

@@ -1290,6 +1290,45 @@ 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')
+    })
+  })
+
   describe('inputMapping deep merge for workflow tools', () => {
     it.concurrent('should deep merge inputMapping when user provides empty object', () => {
       const tool = {