From 673f0369550cd027ea54f7ec8fa858948482a720 Mon Sep 17 00:00:00 2001
From: Vikhyath Mondreti <vikhyath@simstudio.ai>
Date: Wed, 10 Jun 2026 14:44:16 -0700
Subject: [PATCH] fix(attribution): workspace id attr should be best-effort for
 self hosted users

---
 .../app/api/billing/update-cost/route.test.ts | 42 ++++++++++++++--
 apps/sim/app/api/billing/update-cost/route.ts | 48 +++++++++++++++++--
 apps/sim/lib/api/contracts/subscription.ts    | 14 +++---
 3 files changed, 91 insertions(+), 13 deletions(-)

diff --git a/apps/sim/app/api/billing/update-cost/route.test.ts b/apps/sim/app/api/billing/update-cost/route.test.ts
index 42769756897..1fb8f8c294d 100644
--- a/apps/sim/app/api/billing/update-cost/route.test.ts
+++ b/apps/sim/app/api/billing/update-cost/route.test.ts
@@ -1,7 +1,7 @@
 /**
  * @vitest-environment node
  */
-import { createMockRequest } from '@sim/testing'
+import { createMockRequest, dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 const {
@@ -16,6 +16,8 @@ const {
   mockCheckAndBillOverageThreshold: vi.fn(),
 }))
 
+vi.mock('@sim/db', () => dbChainMock)
+
 vi.mock('@/lib/copilot/request/http', () => ({
   checkInternalApiKey: mockCheckInternalApiKey,
 }))
@@ -47,10 +49,12 @@ 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 () => {
@@ -120,7 +124,7 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
     expect(mockCheckAndBillOverageThreshold).not.toHaveBeenCalled()
   })
 
-  it('rejects with 400 when workspaceId is omitted (contract-required, fail loud)', async () => {
+  it('records unattributed when workspaceId is omitted (headless client)', async () => {
     const res = await POST(
       createMockRequest(
         'POST',
@@ -128,7 +132,37 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
         { 'x-api-key': 'internal' }
       )
     )
-    expect(res.status).toBe(400)
-    expect(mockRecordUsage).not.toHaveBeenCalled()
+    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',
+    })
   })
 })
diff --git a/apps/sim/app/api/billing/update-cost/route.ts b/apps/sim/app/api/billing/update-cost/route.ts
index 42b615b4d25..92ccce1e8a8 100644
--- a/apps/sim/app/api/billing/update-cost/route.ts
+++ b/apps/sim/app/api/billing/update-cost/route.ts
@@ -1,6 +1,9 @@
 import type { Span } from '@opentelemetry/api'
+import { db } from '@sim/db'
+import { workspace } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
-import { toError } from '@sim/utils/errors'
+import { getPostgresConstraintName, getPostgresErrorCode, toError } from '@sim/utils/errors'
+import { eq } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
 import { billingUpdateCostContract } from '@/lib/api/contracts/subscription'
 import { parseRequest } from '@/lib/api/server'
@@ -17,6 +20,35 @@ import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 
 const logger = createLogger('BillingUpdateCostAPI')
 
+/**
+ * Resolves the request-supplied workspace to one that exists in this
+ * deployment. Workspace attribution on the usage ledger is best-effort:
+ * self-hosted and headless clients bill through this endpoint with workspace
+ * IDs from their own databases, and `usage_log.workspace_id` carries an FK to
+ * `workspace`, so stamping a foreign ID would fail the entire flush with an
+ * FK violation and strand real cost in the caller's dead-letter queue.
+ * Unknown workspaces are recorded unattributed instead — billing is keyed on
+ * the user's billing entity and never depends on the workspace.
+ */
+async function resolveAttributableWorkspaceId(
+  requestId: string,
+  workspaceId: string | undefined
+): Promise<string | undefined> {
+  if (!workspaceId) return undefined
+
+  const [row] = await db
+    .select({ id: workspace.id })
+    .from(workspace)
+    .where(eq(workspace.id, workspaceId))
+    .limit(1)
+  if (row) return row.id
+
+  logger.warn(`[${requestId}] Workspace not found in this deployment; recording unattributed`, {
+    workspaceId,
+  })
+  return undefined
+}
+
 /**
  * POST /api/billing/update-cost
  * Update user cost with a pre-calculated cost value (internal API key auth required)
@@ -129,6 +161,8 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
       source,
     })
 
+    const attributedWorkspaceId = await resolveAttributableWorkspaceId(requestId, workspaceId)
+
     // 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
@@ -141,7 +175,7 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
     if (idempotencyKey) {
       const result = await recordCumulativeUsage({
         userId,
-        workspaceId,
+        workspaceId: attributedWorkspaceId,
         source,
         model,
         cost,
@@ -160,7 +194,7 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
     } else {
       await recordUsage({
         userId,
-        workspaceId,
+        workspaceId: attributedWorkspaceId,
         entries: [
           {
             category: 'model',
@@ -229,8 +263,16 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
   } catch (error) {
     const duration = Date.now() - startTime
 
+    // Surface the underlying Postgres failure (e.g. 23503 FK violation vs a
+    // lock timeout) — Drizzle's "Failed query" wrapper alone cannot
+    // distinguish them, which made the dead-workspace incident undiagnosable
+    // from logs.
+    const pgCode = getPostgresErrorCode(error)
+    const pgConstraint = getPostgresConstraintName(error)
     logger.error(`[${requestId}] Cost update failed`, {
       error: toError(error).message,
+      ...(pgCode && { pgCode }),
+      ...(pgConstraint && { pgConstraint }),
       stack: error instanceof Error ? error.stack : undefined,
       duration,
     })
diff --git a/apps/sim/lib/api/contracts/subscription.ts b/apps/sim/lib/api/contracts/subscription.ts
index c2c3e3af402..b4960f4dedd 100644
--- a/apps/sim/lib/api/contracts/subscription.ts
+++ b/apps/sim/lib/api/contracts/subscription.ts
@@ -21,13 +21,15 @@ export const billingUpdateCostBodySchema = z.object({
     .default('copilot'),
   idempotencyKey: z.string().min(1).optional(),
   /**
-   * Originating workspace. Stamped onto `usage_log.workspaceId` so mothership/
-   * copilot cost is attributable to org-owned workspaces (per-member usage).
-   * Required: the Go mothership always resolves a workspace for a billed request,
-   * so a missing value is a bug to surface (fail loud) rather than silently drop
-   * the cost from the per-member meter.
+   * Originating workspace, used for org-workspace cost attribution on hosted
+   * Sim. Best-effort by design: self-hosted and headless clients bill through
+   * this endpoint with workspace IDs that exist only in their own deployment
+   * (or with none at all — the Go client omits the field when empty), so the
+   * value is optional and the route only stamps it onto the ledger when it
+   * resolves to a workspace in this deployment. Billing is keyed on the
+   * user's billing entity and must never fail over attribution metadata.
    */
-  workspaceId: z.string().min(1),
+  workspaceId: z.string().min(1).optional(),
 })
 export type BillingUpdateCostBody = z.input<typeof billingUpdateCostBodySchema>