fix(mothership): share and verify billing callback contracts

Author: Sg312

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

@@ -1,6 +1,7 @@
 /**
  * @vitest-environment node
  */
+import { readFileSync } from 'node:fs'
 import { createMockRequest, resetEnvFlagsMock, setEnvFlags } from '@sim/testing'
 import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
@@ -80,6 +81,16 @@ vi.mock('@/lib/billing/threshold-billing', () => ({
 }))
 
 import { billingUpdateCostBodySchema } from '@/lib/api/contracts/subscription'
+import {
+  BillingCallbackBody,
+  BillingCallbackHeaders,
+  BillingProtocol,
+  BillingProtocolHeaders,
+} from '@/lib/mothership/generated/billing'
+import {
+  BILLING_PROTOCOL_HEADERS,
+  COPILOT_BILLING_PROTOCOL,
+} from '@/lib/mothership/generated/billing-protocol-v1'
 import { POST } from '@/app/api/billing/update-cost/route'
 
 afterAll(resetEnvFlagsMock)
@@ -167,6 +178,45 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
     })
   })
 
+  it('keeps the worker callback protocol aligned with the Go-produced wire constants', () => {
+    expect(Object.values(BillingProtocol).sort()).toEqual(
+      Object.values(COPILOT_BILLING_PROTOCOL).sort()
+    )
+    expect(BillingProtocolHeaders).toEqual(BILLING_PROTOCOL_HEADERS)
+  })
+
+  it.skipIf(!process.env.BILLING_WIRE_FIXTURE)(
+    'accepts the actual worker HTTP callbacks through the Sim handler',
+    async () => {
+      const path = process.env.BILLING_WIRE_FIXTURE
+      if (!path) throw new Error('Missing worker callback fixture')
+      const receipts: { body: unknown; headers: Record<string, string> }[] = JSON.parse(
+        readFileSync(path, 'utf8')
+      )
+      expect(receipts.length).toBeGreaterThan(5)
+      setEnvFlags({ isBillingEnabled: true, isHosted: true })
+      for (const receipt of receipts) {
+        const body = BillingCallbackBody.parse(receipt.body)
+        BillingCallbackHeaders.parse(receipt.headers)
+        mockRequireBillingAttributionHeader.mockReturnValue({
+          ...ATTRIBUTION,
+          workspaceId: body.workspaceId,
+        })
+        const result = await POST(createMockRequest('POST', body, receipt.headers))
+        expect(result.status, JSON.stringify(await result.clone().json())).toBe(200)
+        expect(mockRecordCumulativeUsage).toHaveBeenLastCalledWith(
+          expect.objectContaining({
+            userId: body.userId,
+            cost: body.cost,
+            model: body.model,
+            eventKey: `update-cost:${body.idempotencyKey}`,
+            metadata: { inputTokens: body.inputTokens, outputTokens: body.outputTokens },
+          })
+        )
+      }
+    }
+  )
+
   it('returns 401 for a billing-disabled request without valid internal auth', async () => {
     setEnvFlags({ isBillingEnabled: false })
     mockCheckInternalApiKey.mockReturnValue({ success: false, error: 'Invalid internal API key' })

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

@@ -1,16 +1,7 @@
 import { z } from 'zod'
 import { workspaceIdSchema } from '@/lib/api/contracts/primitives'
 import { defineRouteContract } from '@/lib/api/contracts/types'
-import { INTERNAL_CHAT_BILLING_SOURCES } from '@/lib/billing/usage-sources'
-import {
-  BILLING_ACCOUNT_DECISION_HEADER,
-  BILLING_ACCOUNT_DECISION_HEADER_MAX_BYTES,
-  BILLING_ATTRIBUTION_HEADER,
-  BILLING_ATTRIBUTION_HEADER_MAX_BYTES,
-  BILLING_REQUEST_ID_HEADER,
-  COPILOT_BILLING_PROTOCOL_HEADER,
-  COPILOT_BILLING_PROTOCOL_VALUES,
-} from '@/lib/mothership/generated/billing-protocol-v1'
+import { BillingCallbackBody, BillingCallbackHeaders } from '@/lib/mothership/generated/billing'
 
 const booleanQueryParamSchema = z
   .preprocess((value) => {
@@ -21,35 +12,9 @@ const booleanQueryParamSchema = z
   .optional()
   .default(false)
 
-export const billingUpdateCostBodySchema = z.object({
-  userId: z.string().min(1, 'User ID is required'),
-  cost: z.number().min(0, 'Cost must be a non-negative number'),
-  model: z.string().min(1, 'Model is required'),
-  inputTokens: z.number().min(0).default(0),
-  outputTokens: z.number().min(0).default(0),
-  source: z.enum(INTERNAL_CHAT_BILLING_SOURCES).default('copilot'),
-  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
-   * supply an ID from another deployment or omit it. Modern protocols bind a
-   * locally known workspace to their immutable envelope. Markerless local
-   * self-hosted callbacks re-resolve current workspace payer state; unknown
-   * workspaces remain account-only.
-   */
-  workspaceId: z.string().min(1).optional(),
-})
+export const billingUpdateCostBodySchema = BillingCallbackBody
 export type BillingUpdateCostBody = z.input<typeof billingUpdateCostBodySchema>
-
-export const billingUpdateCostHeadersSchema = z.object({
-  [COPILOT_BILLING_PROTOCOL_HEADER]: z.enum(COPILOT_BILLING_PROTOCOL_VALUES).optional(),
-  [BILLING_REQUEST_ID_HEADER]: z.string().uuid().optional(),
-  [BILLING_ATTRIBUTION_HEADER]: z.string().max(BILLING_ATTRIBUTION_HEADER_MAX_BYTES).optional(),
-  [BILLING_ACCOUNT_DECISION_HEADER]: z
-    .string()
-    .max(BILLING_ACCOUNT_DECISION_HEADER_MAX_BYTES)
-    .optional(),
-})
+export const billingUpdateCostHeadersSchema = BillingCallbackHeaders
 export type BillingUpdateCostHeaders = z.input<typeof billingUpdateCostHeadersSchema>
 
 export const billingSwitchPlanBodySchema = z.object({

apps/sim/lib/mothership/generated/billing.ts

@@ -0,0 +1,42 @@
+// GENERATED — do not edit. Source of truth: mothership worker packages/contracts/src/billing.ts
+// Regenerate with `bun run contracts:sync` in the worker.
+
+import { z } from "zod";
+
+export const BillingProtocol = {
+  attributed: "attribution-v1",
+  direct: "direct-v1",
+  previous: "legacy-v0",
+} as const;
+
+export const BillingProtocolHeaders = {
+  protocol: "x-sim-billing-protocol",
+  requestId: "x-sim-billing-request-id",
+  attribution: "x-sim-billing-attribution",
+  accountDecision: "x-sim-billing-account-decision",
+} as const;
+
+export const BillingCallbackBody = z.object({
+  userId: z.string().min(1, "User ID is required"),
+  cost: z.number().min(0, "Cost must be a non-negative number"),
+  model: z.string().min(1, "Model is required"),
+  inputTokens: z.number().min(0).default(0),
+  outputTokens: z.number().min(0).default(0),
+  source: z.enum(["copilot", "workspace-chat", "mcp_copilot", "mothership_block"]).default("copilot"),
+  idempotencyKey: z.string().min(1, "Idempotency key is required"),
+  workspaceId: z.string().min(1).optional(),
+});
+export type BillingCallbackBody = z.infer<typeof BillingCallbackBody>;
+
+export const BillingCallbackHeaders = z.object({
+  [BillingProtocolHeaders.protocol]: z.enum(Object.values(BillingProtocol)).optional(),
+  [BillingProtocolHeaders.requestId]: z.string().uuid().optional(),
+  [BillingProtocolHeaders.attribution]: z.string().max(8192).optional(),
+  [BillingProtocolHeaders.accountDecision]: z.string().max(2048).optional(),
+});
+
+export const BillingCallbackResult = z.object({
+  success: z.boolean(),
+  code: z.string().optional(),
+});
+export const BillingDuplicateCode = "DUPLICATE_BILLING_EVENT";