Merge branch 'dev' into zen-accounting-latency

Move stream-end and non-stream tracking/billing round trips into waitUntil so the client stream closes at the last token, and add duration.preflight, duration.sticky_set, and duration.accounting metrics for the previously unmeasured blocks.

packages/console/app/src/routes/zen/util/handler.ts

@@ -43,7 +43,7 @@ import { createRateLimiter as createKeyRateLimiter } from "./keyRateLimiter"
 import { createTrialLimiter } from "./trialLimiter"
 import { createStickyTracker } from "./stickyProviderTracker"
 import { LiteData } from "@opencode-ai/console-core/lite.js"
-import { Resource } from "@opencode-ai/console-resource"
+import { Resource, waitUntil } from "@opencode-ai/console-resource"
 import { i18n, type Key } from "~/i18n"
 import { localeFromRequest } from "~/lib/language"
 import { createModelTpmLimiter } from "./modelTpmLimiter"
@@ -96,6 +96,7 @@ export async function handler(
   ]
 
   try {
+    const timestampHandlerStart = Date.now()
     const url = input.request.url
     const body = await input.request.json()
     const model = opts.parseModel(url, body)
@@ -288,10 +289,21 @@ export async function handler(
       return { providerInfo, reqBody, res, startTimestamp }
     }
 
+    // Time spent on pre-flight checks (body parse, limiter checks, auth,
+    // sticky lookup) before the first upstream dispatch. Not covered by
+    // time_to_first_byte, which starts at dispatch.
+    logger.metric({ "duration.preflight": Date.now() - timestampHandlerStart })
     const { providerInfo, reqBody, res, startTimestamp } = await retriableRequest()
 
-    // Store sticky provider
+    // Store sticky provider. Awaited on purpose: subsequent requests of the
-    if (res.status === 200) await stickyTracker?.set(providerInfo.id)
+    // same session must observe it (strict sticky providers depend on it),
+    // but it sits between upstream response headers and the client's first
+    // byte, so measure it.
+    if (res.status === 200 && stickyTracker) {
+      const timestampStickySet = Date.now()
+      await stickyTracker.set(providerInfo.id)
+      logger.metric({ "duration.sticky_set": Date.now() - timestampStickySet })
+    }
 
     // Temporarily change 404 to 400 status code b/c solid start automatically override 404 response
     const resStatus = res.status === 404 ? 400 : res.status
@@ -309,18 +321,28 @@ export async function handler(
     // Handle non-streaming response
     if (!isStream || [400, 404, 429].includes(res.status)) {
       const json = await res.json()
-      await rateLimiter?.track()
       const usage = providerInfo.extractUsage(json)
-      if (usage) {
+      const usageInfo = usage ? providerInfo.normalizeUsage(usage) : undefined
-        const usageInfo = providerInfo.normalizeUsage(usage)
+      const costInfo = usageInfo ? calculateCost(modelInfo, usageInfo) : undefined
-        const costInfo = calculateCost(modelInfo, usageInfo)
+      if (costInfo) json.cost = calculateOccurredCost(billingSource, costInfo)
+      // Cost math above is pure; the accounting round trips below must not
+      // delay the response, so run them after the response is sent.
+      const timestampAccountingStart = Date.now()
+      waitUntil(
+        (async () => {
+          await rateLimiter?.track()
+          if (usageInfo && costInfo) {
             await trialLimiter?.track(usageInfo)
             await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)
             await providerBudgetTracker?.track(providerInfo.id, providerInfo.budgetPriority, costInfo.totalCostInCent)
             await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
             await reload(billingSource, authInfo, costInfo)
-        json.cost = calculateOccurredCost(billingSource, costInfo)
           }
+          logger.metric({ "duration.accounting": Date.now() - timestampAccountingStart })
+        })().catch((error) => {
+          logger.metric({ "error.accounting": error instanceof Error ? error.message : String(error) })
+        }),
+      )
       if (res.status === 400) {
         logger.metric({ "error.response": JSON.stringify(json) })
       }
@@ -363,11 +385,20 @@ export async function handler(
                   response_length: responseLength,
                   "timestamp.last_byte": timestampLastByte,
                 })
-                await rateLimiter?.track()
                 const usage = usageParser.retrieve()
-                if (usage) {
+                const usageInfo = usage ? providerInfo.normalizeUsage(usage) : undefined
-                  const usageInfo = providerInfo.normalizeUsage(usage)
+                const costInfo = usageInfo ? calculateCost(modelInfo, usageInfo) : undefined
-                  const costInfo = calculateCost(modelInfo, usageInfo)
+                if (costInfo) {
+                  c.enqueue(encoder.encode(buildCostChunk(opts.format, calculateOccurredCost(billingSource, costInfo))))
+                }
+                // Close the client stream before accounting: the tracking and
+                // billing round trips below used to hold the stream open after
+                // the last token, which the client observes as a hang.
+                c.close()
+                waitUntil(
+                  (async () => {
+                    await rateLimiter?.track()
+                    if (usageInfo && costInfo) {
                       await trialLimiter?.track(usageInfo)
                       await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)
                       await modelTpsLimiter?.track(
@@ -385,10 +416,12 @@ export async function handler(
                       )
                       await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
                       await reload(billingSource, authInfo, costInfo)
-                  const cost = calculateOccurredCost(billingSource, costInfo)
-                  c.enqueue(encoder.encode(buildCostChunk(opts.format, cost)))
                     }
-                c.close()
+                    logger.metric({ "duration.accounting": Date.now() - timestampLastByte })
+                  })().catch((error) => {
+                    logger.metric({ "error.accounting": error instanceof Error ? error.message : String(error) })
+                  }),
+                )
                 return
               }