simstudioai--sim/apps/sim/app/api/billing/update-cost/route.test.ts

/**
 * @vitest-environment node
 */
import { createMockRequest, dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'

const {
  mockCheckInternalApiKey,
  mockRecordUsage,
  mockRecordCumulativeUsage,
  mockCheckAndBillOverageThreshold,
} = vi.hoisted(() => ({
  mockCheckInternalApiKey: vi.fn(),
  mockRecordUsage: vi.fn(),
  mockRecordCumulativeUsage: vi.fn(),
  mockCheckAndBillOverageThreshold: vi.fn(),
}))

vi.mock('@sim/db', () => dbChainMock)

vi.mock('@/lib/copilot/request/http', () => ({
  checkInternalApiKey: mockCheckInternalApiKey,
}))

vi.mock('@/lib/copilot/request/otel', () => ({
  withIncomingGoSpan: (
    _headers: unknown,
    _span: unknown,
    _attrs: unknown,
    fn: (span: { setAttribute: () => void; setAttributes: () => void }) => unknown
  ) => fn({ setAttribute: vi.fn(), setAttributes: vi.fn() }),
}))

vi.mock('@/lib/billing/core/usage-log', () => ({
  recordUsage: mockRecordUsage,
  recordCumulativeUsage: mockRecordCumulativeUsage,
}))

vi.mock('@/lib/billing/threshold-billing', () => ({
  checkAndBillOverageThreshold: mockCheckAndBillOverageThreshold,
}))

vi.mock('@/lib/core/config/env-flags', () => ({
  isBillingEnabled: true,
}))

import { POST } from '@/app/api/billing/update-cost/route'

describe('POST /api/billing/update-cost — workspaceId attribution', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    resetDbChainMock()
    mockCheckInternalApiKey.mockReturnValue({ success: true })
    mockRecordUsage.mockResolvedValue(undefined)
    mockRecordCumulativeUsage.mockResolvedValue({ billed: true, delta: 0.5, total: 0.5 })
    mockCheckAndBillOverageThreshold.mockResolvedValue(undefined)
    dbChainMockFns.limit.mockResolvedValue([{ id: 'ws-1' }])
  })

  it('stamps workspaceId onto recorded usage when provided (no idempotency key)', async () => {
    const res = await POST(
      createMockRequest(
        'POST',
        { userId: 'user-1', cost: 0.5, model: 'gpt', source: 'mcp_copilot', workspaceId: 'ws-1' },
        { 'x-api-key': 'internal' }
      )
    )
    expect(res.status).toBe(200)
    expect(mockRecordUsage).toHaveBeenCalledTimes(1)
    expect(mockRecordUsage.mock.calls[0][0]).toMatchObject({
      userId: 'user-1',
      workspaceId: 'ws-1',
    })
  })

  it('records cumulative cost via monotonic top-up when an idempotency key is present', async () => {
    const res = await POST(
      createMockRequest(
        'POST',
        {
          userId: 'user-1',
          cost: 0.4662453,
          model: 'claude-opus-4.8',
          source: 'workspace-chat',
          workspaceId: 'ws-1',
          idempotencyKey: 'msg-1-billing',
          inputTokens: 461371,
          outputTokens: 1686,
        },
        { 'x-api-key': 'internal' }
      )
    )
    expect(res.status).toBe(200)
    expect(mockRecordUsage).not.toHaveBeenCalled()
    expect(mockRecordCumulativeUsage).toHaveBeenCalledTimes(1)
    expect(mockRecordCumulativeUsage.mock.calls[0][0]).toMatchObject({
      userId: 'user-1',
      workspaceId: 'ws-1',
      source: 'workspace-chat',
      model: 'claude-opus-4.8',
      cost: 0.4662453,
      eventKey: 'update-cost:msg-1-billing',
    })
    expect(mockCheckAndBillOverageThreshold).toHaveBeenCalledWith('user-1')
  })

  it('returns 409 and skips overage when the cumulative is not higher (duplicate flush)', async () => {
    mockRecordCumulativeUsage.mockResolvedValue({ billed: false, delta: 0, total: 0.4662453 })
    const res = await POST(
      createMockRequest(
        'POST',
        {
          userId: 'user-1',
          cost: 0.4662453,
          model: 'claude-opus-4.8',
          source: 'workspace-chat',
          workspaceId: 'ws-1',
          idempotencyKey: 'msg-1-billing',
        },
        { 'x-api-key': 'internal' }
      )
    )
    expect(res.status).toBe(409)
    expect(mockCheckAndBillOverageThreshold).not.toHaveBeenCalled()
  })

  it('records unattributed when workspaceId is omitted (headless client)', async () => {
    const res = await POST(
      createMockRequest(
        'POST',
        { userId: 'user-1', cost: 0.5, model: 'gpt', source: 'copilot' },
        { 'x-api-key': 'internal' }
      )
    )
    expect(res.status).toBe(200)
    expect(dbChainMockFns.limit).not.toHaveBeenCalled()
    expect(mockRecordUsage).toHaveBeenCalledTimes(1)
    expect(mockRecordUsage.mock.calls[0][0]).toMatchObject({
      userId: 'user-1',
      workspaceId: undefined,
    })
  })

  it('records unattributed when the workspace does not exist in this deployment (self-hosted client)', async () => {
    dbChainMockFns.limit.mockResolvedValue([])
    const res = await POST(
      createMockRequest(
        'POST',
        {
          userId: 'user-1',
          cost: 0.5,
          model: 'claude-opus-4.8',
          source: 'workspace-chat',
          workspaceId: 'self-hosted-ws',
          idempotencyKey: 'msg-1-billing',
        },
        { 'x-api-key': 'internal' }
      )
    )
    expect(res.status).toBe(200)
    expect(mockRecordCumulativeUsage).toHaveBeenCalledTimes(1)
    expect(mockRecordCumulativeUsage.mock.calls[0][0]).toMatchObject({
      userId: 'user-1',
      workspaceId: undefined,
      eventKey: 'update-cost:msg-1-billing',
    })
  })
})