update

Author: icecrasher321

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

@@ -6,7 +6,6 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 const {
   mockCheckInternalApiKey,
-  mockRecordUsage,
   mockRecordCumulativeUsage,
   mockCheckAndBillOverageThreshold,
   mockCheckAndBillPayerOverageThreshold,
@@ -19,7 +18,6 @@ const {
   billingState,
 } = vi.hoisted(() => ({
   mockCheckInternalApiKey: vi.fn(),
-  mockRecordUsage: vi.fn(),
   mockRecordCumulativeUsage: vi.fn(),
   mockCheckAndBillOverageThreshold: vi.fn(),
   mockCheckAndBillPayerOverageThreshold: vi.fn(),
@@ -59,7 +57,6 @@ vi.mock('@/lib/copilot/request/otel', () => ({
 
 vi.mock('@/lib/billing/core/usage-log', () => ({
   CumulativeUsageContextMismatchError: MockCumulativeUsageContextMismatchError,
-  recordUsage: mockRecordUsage,
   recordCumulativeUsage: mockRecordCumulativeUsage,
 }))
 
@@ -150,13 +147,21 @@ const OLD_GO_OPAQUE_WORKSPACE_UPDATE_COST_BODY = {
   workspaceId: 'local-self-hosted-workspace',
 } as const
 
+const KEYLESS_UPDATE_COST_BODY = {
+  userId: 'user-1',
+  cost: 0.5,
+  model: 'gpt',
+  inputTokens: 1,
+  outputTokens: 2,
+  source: 'copilot',
+} as const
+
 describe('POST /api/billing/update-cost — workspaceId attribution', () => {
   beforeEach(() => {
     vi.clearAllMocks()
     billingState.isBillingEnabled = true
     billingState.isCopilotBillingProtocolRequired = false
     mockCheckInternalApiKey.mockReturnValue({ success: true })
-    mockRecordUsage.mockResolvedValue(undefined)
     mockRecordCumulativeUsage.mockResolvedValue({ billed: true, delta: 0.5, total: 0.5 })
     mockCheckAndBillOverageThreshold.mockResolvedValue(undefined)
     mockCheckAndBillPayerOverageThreshold.mockResolvedValue(undefined)
@@ -191,7 +196,6 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
       error: 'Invalid internal API key',
     })
     expect(mockCheckInternalApiKey).toHaveBeenCalledTimes(1)
-    expect(mockRecordUsage).not.toHaveBeenCalled()
     expect(mockRecordCumulativeUsage).not.toHaveBeenCalled()
   })
 
@@ -218,7 +222,6 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
       data: { billingEnabled: false },
     })
     expect(mockCheckInternalApiKey).toHaveBeenCalledTimes(1)
-    expect(mockRecordUsage).not.toHaveBeenCalled()
     expect(mockRecordCumulativeUsage).not.toHaveBeenCalled()
   })
 
@@ -232,6 +235,17 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
     ).toBe(true)
   })
 
+  it('rejects billing-enabled callbacks without a stable idempotency key', async () => {
+    const res = await POST(
+      createMockRequest('POST', KEYLESS_UPDATE_COST_BODY, { 'x-api-key': 'internal' })
+    )
+
+    expect(res.status).toBe(400)
+    expect(mockRecordCumulativeUsage).not.toHaveBeenCalled()
+    expect(mockCheckAndBillOverageThreshold).not.toHaveBeenCalled()
+    expect(mockCheckAndBillPayerOverageThreshold).not.toHaveBeenCalled()
+  })
+
   it('bills the routed workspace payer for the exact markerless hosted callback', async () => {
     const res = await POST(
       createMockRequest('POST', OLD_GO_HOSTED_UPDATE_COST_BODY, { 'x-api-key': 'internal' })
@@ -565,7 +579,6 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
       code: 'DUPLICATE_BILLING_EVENT',
     })
     expect(mockRecordCumulativeUsage).toHaveBeenCalledTimes(2)
-    expect(mockRecordUsage).not.toHaveBeenCalled()
     expect(mockResolveLegacyV0BillingAttribution).toHaveBeenCalledTimes(2)
     expect(mockCheckAndBillPayerOverageThreshold).toHaveBeenCalledTimes(2)
     expect(mockCheckAndBillPayerOverageThreshold).toHaveBeenNthCalledWith(

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

@@ -22,7 +22,6 @@ import {
   type CumulativeUsageContextField,
   CumulativeUsageContextMismatchError,
   recordCumulativeUsage,
-  recordUsage,
 } from '@/lib/billing/core/usage-log'
 import {
   checkAndBillOverageThreshold,
@@ -293,57 +292,30 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
       workspaceId: resolvedWorkspaceId,
     })
 
-    // Go sends the request's CUMULATIVE cost, possibly more than once (a
-    // mid-loop provider-error flush, then the recovered terminal flush, plus
-    // abort-race duplicates). Record it as a monotonic top-up: one ledger row
-    // per request holds the MAX cumulative and we bill only the delta, so
-    // partial + complete flushes converge to the true total exactly once — no
-    // under-billing on recovery, no over-billing on duplicates. When there is
-    // no idempotency key (shouldn't happen for real requests) we fall back to a
-    // plain append so cost is never silently dropped.
-    let billed = true
-    if (idempotencyKey) {
-      const result = await recordCumulativeUsage({
-        userId,
-        workspaceId: resolvedWorkspaceId,
-        ...billingContext,
-        source,
-        model,
-        cost,
-        eventKey: `update-cost:${idempotencyKey}`,
-        metadata: { inputTokens, outputTokens },
-      })
-      billed = result.billed
-      logger.info(`[${requestId}] Cumulative cost top-up`, {
-        userId,
-        source,
-        cumulativeCost: cost,
-        billedDelta: result.delta,
-        newTotal: result.total,
-        billed: result.billed,
-      })
-    } else {
-      await recordUsage({
-        userId,
-        workspaceId: resolvedWorkspaceId,
-        ...billingContext,
-        entries: [
-          {
-            category: 'model',
-            source,
-            description: model,
-            cost,
-            sourceReference: requestId,
-            metadata: { inputTokens, outputTokens },
-          },
-        ],
-      })
-      logger.info(`[${requestId}] Recorded usage (no idempotency key)`, {
-        userId,
-        addedCost: cost,
-        source,
-      })
-    }
+    /**
+     * Go sends cumulative cost across partial, terminal, and retry flushes.
+     * Every accepted callback has a stable key, so the maximum cumulative cost
+     * converges on one ledger event without underbilling or double-billing.
+     */
+    const result = await recordCumulativeUsage({
+      userId,
+      workspaceId: resolvedWorkspaceId,
+      ...billingContext,
+      source,
+      model,
+      cost,
+      eventKey: `update-cost:${idempotencyKey}`,
+      metadata: { inputTokens, outputTokens },
+    })
+    const billed = result.billed
+    logger.info(`[${requestId}] Cumulative cost top-up`, {
+      userId,
+      source,
+      cumulativeCost: cost,
+      billedDelta: result.delta,
+      newTotal: result.total,
+      billed: result.billed,
+    })
 
     // Reconcile the payer's ledger-backed threshold after every cumulative
     // callback, including duplicate retries after a prior settlement failure.

apps/sim/lib/api/contracts/subscription.ts

@@ -29,7 +29,7 @@ export const billingUpdateCostBodySchema = z.object({
   source: z
     .enum(['copilot', 'workspace-chat', 'mcp_copilot', 'mothership_block'])
     .default('copilot'),
-  idempotencyKey: z.string().min(1).optional(),
+  idempotencyKey: z.string().min(1, 'Idempotency key is required'),
   /**
    * Originating workspace, used for org-workspace cost attribution on hosted
    * Sim. The value remains optional because self-hosted/headless callers may