fix billing, tracing, retry lifecycle issues

Author: icecrasher321

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

@@ -4,13 +4,17 @@
 import { createMockRequest } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockCheckInternalApiKey, mockRecordUsage, mockCheckAndBillOverageThreshold } = vi.hoisted(
-  () => ({
-    mockCheckInternalApiKey: vi.fn(),
-    mockRecordUsage: vi.fn(),
-    mockCheckAndBillOverageThreshold: vi.fn(),
-  })
-)
+const {
+  mockCheckInternalApiKey,
+  mockRecordUsage,
+  mockRecordCumulativeUsage,
+  mockCheckAndBillOverageThreshold,
+} = vi.hoisted(() => ({
+  mockCheckInternalApiKey: vi.fn(),
+  mockRecordUsage: vi.fn(),
+  mockRecordCumulativeUsage: vi.fn(),
+  mockCheckAndBillOverageThreshold: vi.fn(),
+}))
 
 vi.mock('@/lib/copilot/request/http', () => ({
   checkInternalApiKey: mockCheckInternalApiKey,
@@ -27,19 +31,13 @@ vi.mock('@/lib/copilot/request/otel', () => ({
 
 vi.mock('@/lib/billing/core/usage-log', () => ({
   recordUsage: mockRecordUsage,
+  recordCumulativeUsage: mockRecordCumulativeUsage,
 }))
 
 vi.mock('@/lib/billing/threshold-billing', () => ({
   checkAndBillOverageThreshold: mockCheckAndBillOverageThreshold,
 }))
 
-vi.mock('@/lib/core/idempotency/service', () => ({
-  billingIdempotency: {
-    atomicallyClaim: vi.fn(),
-    release: vi.fn(),
-  },
-}))
-
 vi.mock('@/lib/core/config/feature-flags', () => ({
   isBillingEnabled: true,
 }))
@@ -51,10 +49,11 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
     vi.clearAllMocks()
     mockCheckInternalApiKey.mockReturnValue({ success: true })
     mockRecordUsage.mockResolvedValue(undefined)
+    mockRecordCumulativeUsage.mockResolvedValue({ billed: true, delta: 0.5, total: 0.5 })
     mockCheckAndBillOverageThreshold.mockResolvedValue(undefined)
   })
 
-  it('stamps workspaceId onto recorded usage when provided', async () => {
+  it('stamps workspaceId onto recorded usage when provided (no idempotency key)', async () => {
     const res = await POST(
       createMockRequest(
         'POST',
@@ -70,6 +69,57 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
     })
   })
 
+  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('rejects with 400 when workspaceId is omitted (contract-required, fail loud)', async () => {
     const res = await POST(
       createMockRequest(

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

@@ -4,15 +4,14 @@ import { toError } from '@sim/utils/errors'
 import { type NextRequest, NextResponse } from 'next/server'
 import { billingUpdateCostContract } from '@/lib/api/contracts/subscription'
 import { parseRequest } from '@/lib/api/server'
-import { recordUsage } from '@/lib/billing/core/usage-log'
+import { recordCumulativeUsage, recordUsage } from '@/lib/billing/core/usage-log'
 import { checkAndBillOverageThreshold } from '@/lib/billing/threshold-billing'
 import { BillingRouteOutcome } from '@/lib/copilot/generated/trace-attribute-values-v1'
 import { TraceAttr } from '@/lib/copilot/generated/trace-attributes-v1'
 import { TraceSpan } from '@/lib/copilot/generated/trace-spans-v1'
 import { checkInternalApiKey } from '@/lib/copilot/request/http'
 import { withIncomingGoSpan } from '@/lib/copilot/request/otel'
 import { isBillingEnabled } from '@/lib/core/config/feature-flags'
-import { type AtomicClaimResult, billingIdempotency } from '@/lib/core/idempotency/service'
 import { generateRequestId } from '@/lib/core/utils/request'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 
@@ -42,8 +41,6 @@ export const POST = withRouteHandler((req: NextRequest) =>
 async function updateCostInner(req: NextRequest, span: Span): Promise<NextResponse> {
   const requestId = generateRequestId()
   const startTime = Date.now()
-  let claim: AtomicClaimResult | null = null
-  let usageCommitted = false
 
   try {
     logger.info(`[${requestId}] Update cost request started`)
@@ -125,63 +122,92 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
       ...(idempotencyKey ? { [TraceAttr.BillingIdempotencyKey]: idempotencyKey } : {}),
     })
 
-    claim = idempotencyKey
-      ? await billingIdempotency.atomicallyClaim('update-cost', idempotencyKey)
-      : null
+    logger.info(`[${requestId}] Processing cost update`, {
+      userId,
+      cost,
+      model,
+      source,
+    })
 
-    if (claim && !claim.claimed) {
-      logger.warn(`[${requestId}] Duplicate billing update rejected`, {
-        idempotencyKey,
+    // 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,
+        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,
+        entries: [
+          {
+            category: 'model',
+            source,
+            description: model,
+            cost,
+            sourceReference: requestId,
+            metadata: { inputTokens, outputTokens },
+          },
+        ],
+      })
+      logger.info(`[${requestId}] Recorded usage (no idempotency key)`, {
+        userId,
+        addedCost: cost,
+        source,
+      })
+    }
+
+    const duration = Date.now() - startTime
+
+    // Same-or-lower cumulative than already recorded: nothing new to bill. Tell
+    // the caller via 409 (its existing "duplicate" outcome) without re-running
+    // overage billing.
+    if (!billed) {
+      logger.info(`[${requestId}] Duplicate/non-increasing cumulative cost; no new charge`, {
+        idempotencyKey,
+        userId,
+        cost,
       })
       span.setAttribute(TraceAttr.BillingOutcome, BillingRouteOutcome.DuplicateIdempotencyKey)
       span.setAttribute(TraceAttr.HttpStatusCode, 409)
+      span.setAttribute(TraceAttr.BillingDurationMs, duration)
       return NextResponse.json(
         {
           success: false,
-          error: 'Duplicate request: idempotency key already processed',
+          error: 'Duplicate request: cumulative cost already recorded',
           requestId,
         },
         { status: 409 }
       )
     }
 
-    logger.info(`[${requestId}] Processing cost update`, {
-      userId,
-      cost,
-      model,
-      source,
-    })
-
-    await recordUsage({
-      userId,
-      workspaceId,
-      entries: [
-        {
-          category: 'model',
-          source,
-          description: model,
-          cost,
-          eventKey: idempotencyKey ? `update-cost:${idempotencyKey}` : undefined,
-          sourceReference: idempotencyKey ? `update-cost:${idempotencyKey}` : requestId,
-          metadata: { inputTokens, outputTokens },
-        },
-      ],
-    })
-    usageCommitted = true
-
-    logger.info(`[${requestId}] Recorded usage`, {
-      userId,
-      addedCost: cost,
-      source,
-    })
-
-    // Check if user has hit overage threshold and bill incrementally
+    // Check if user has hit overage threshold and bill incrementally. Reads the
+    // (now topped-up) ledger total and is idempotent against billedOverage, so
+    // it is safe to run on every flush that records new cost.
     await checkAndBillOverageThreshold(userId)
 
-    const duration = Date.now() - startTime
-
     logger.info(`[${requestId}] Cost update completed successfully`, {
       userId,
       duration,
@@ -209,22 +235,9 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
       duration,
     })
 
-    if (claim?.claimed && !usageCommitted) {
-      await billingIdempotency
-        .release(claim.normalizedKey, claim.storageMethod)
-        .catch((releaseErr) => {
-          logger.warn(`[${requestId}] Failed to release idempotency claim`, {
-            error: toError(releaseErr).message,
-            normalizedKey: claim?.normalizedKey,
-          })
-        })
-    } else if (claim?.claimed && usageCommitted) {
-      logger.warn(
-        `[${requestId}] Error occurred after usage committed; retaining idempotency claim to prevent double-billing`,
-        { normalizedKey: claim.normalizedKey }
-      )
-    }
-
+    // The cumulative top-up runs in a single transaction (and a plain append is
+    // a single insert), so a failure here leaves nothing partially committed —
+    // a retry re-evaluates the max idempotently. No claim to release.
     span.setAttribute(TraceAttr.BillingOutcome, BillingRouteOutcome.InternalError)
     span.setAttribute(TraceAttr.HttpStatusCode, 500)
     span.setAttribute(TraceAttr.BillingDurationMs, duration)

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/agent-group.tsx

@@ -60,7 +60,10 @@ export function AgentGroup({
     (item): item is Extract<AgentGroupItem, { type: 'tool' }> => item.type === 'tool'
   )
   const allDone = toolItems.length > 0 && toolItems.every((t) => isToolDone(t.data.status))
-  const showDelegatingSpinner = isDelegating && !allDone
+  // Only a live turn can be delegating. Once the turn is terminal (complete,
+  // errored, or stopped) no subagent should spin — even one aborted before its
+  // first tool call, where `allDone` is false because there are no tools yet.
+  const showDelegatingSpinner = isStreaming && isDelegating && !allDone
 
   // Expand only while the turn is live and the group is still open or working.
   // Once the turn ends (isStreaming false) — or a subagent closes mid-turn — the

apps/sim/app/workspace/[workspaceId]/home/components/message-content/message-content.test.ts

@@ -3,7 +3,7 @@
  */
 import { describe, expect, it } from 'vitest'
 import type { ContentBlock } from '../../types'
-import { parseBlocks } from './message-content'
+import { parseBlocks, shouldSmoothTextSegment } from './message-content'
 
 function subagentStart(name: string, spanId: string, parentSpanId: string): ContentBlock {
   return { type: 'subagent', content: name, spanId, parentSpanId, timestamp: 1 }
@@ -182,3 +182,20 @@ describe('parseBlocks span-identity tree', () => {
     expect(groups[0].agentName).toBe('workflow')
   })
 })
+
+describe('shouldSmoothTextSegment', () => {
+  it('only smooths the trailing text segment of a live stream', () => {
+    expect(shouldSmoothTextSegment({ isStreaming: true, segmentIndex: 0, segmentCount: 2 })).toBe(
+      false
+    )
+    expect(shouldSmoothTextSegment({ isStreaming: true, segmentIndex: 1, segmentCount: 2 })).toBe(
+      true
+    )
+  })
+
+  it('never smooths completed messages', () => {
+    expect(shouldSmoothTextSegment({ isStreaming: false, segmentIndex: 0, segmentCount: 1 })).toBe(
+      false
+    )
+  })
+})

apps/sim/app/workspace/[workspaceId]/home/components/message-content/message-content.tsx

@@ -659,6 +659,18 @@ export function assistantMessageHasRenderableContent(
   return segments.length > 0
 }
 
+export function shouldSmoothTextSegment({
+  isStreaming,
+  segmentIndex,
+  segmentCount,
+}: {
+  isStreaming: boolean
+  segmentIndex: number
+  segmentCount: number
+}): boolean {
+  return isStreaming && segmentIndex === segmentCount - 1
+}
+
 interface MessageContentProps {
   blocks: ContentBlock[]
   fallbackContent: string
@@ -719,7 +731,11 @@ function MessageContentInner({
               <ChatContent
                 key={`text-${i}`}
                 content={segment.content}
-                isStreaming={isStreaming}
+                isStreaming={shouldSmoothTextSegment({
+                  isStreaming,
+                  segmentIndex: i,
+                  segmentCount: segments.length,
+                })}
                 onOptionSelect={onOptionSelect}
                 onWorkspaceResourceSelect={onWorkspaceResourceSelect}
               />

apps/sim/hooks/use-smooth-text.ts

@@ -80,6 +80,11 @@ export function useSmoothText(
   // trips React's max-update-depth guard. The running tick reads the latest content
   // from `contentRef`, so new chunks are absorbed without per-chunk teardown;
   // `hasBacklog` only flips when the reveal falls behind or catches up.
+  if (!isStreaming && effectiveRevealed !== content.length) {
+    effectiveRevealed = content.length
+    revealedRef.current = content.length
+  }
+
   const hasBacklog = effectiveRevealed < content.length
 
   useEffect(() => {