fix: isolate usage tests and preserve stream usage context

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>

Author: cuipengfei

src/lib/token-usage/index.ts

@@ -138,10 +138,16 @@ export function recordTokenUsageEvent(input: TokenUsageEventInput): void {
 export function createTokenUsageRecorder(
   options: TokenUsageRecorderOptions,
 ): (usage: UsageTokens) => void {
+  const store = requestContext.getStore()
+  const traceId = options.traceId ?? store?.traceId
+  const sessionId = options.sessionId ?? store?.sessionAffinity
+
   return (usage) => {
     recordTokenUsageEvent({
       ...usage,
       ...options,
+      sessionId,
+      traceId,
     })
   }
 }

src/routes/messages/api-flows.ts

@@ -85,6 +85,8 @@ export interface FlowBaseOptions {
   logger: ConsolaInstance
   subagentMarker?: SubagentMarker | null
   requestId: string
+  requestSessionAffinity?: string
+  requestTraceId?: string
   sessionId?: string
   compactType?: CompactType
 }
@@ -103,14 +105,24 @@ export const handleWithChatCompletions = async (
   anthropicPayload: AnthropicMessagesPayload,
   options: FlowBaseOptions,
 ) => {
-  const { logger, subagentMarker, requestId, sessionId, compactType } = options
+  const {
+    logger,
+    requestSessionAffinity,
+    requestTraceId,
+    subagentMarker,
+    requestId,
+    sessionId,
+    compactType,
+  } = options
   const openAIPayload = translateToOpenAI(anthropicPayload)
   prepareCopilotChatCompletionsPayload(openAIPayload)
-  const recordUsage = createCopilotUsageRecorder({
+  const recordUsage = createMessagesFlowUsageRecorder({
+    anthropicPayload,
     endpoint: "chat_completions",
     fallbackSessionId: sessionId,
     model: openAIPayload.model,
-    payload: anthropicPayload,
+    requestSessionAffinity,
+    requestTraceId,
   })
   debugJson(logger, "Translated OpenAI request payload:", openAIPayload)
 
@@ -130,6 +142,23 @@ export const handleWithChatCompletions = async (
     })
   }
 
+  return handleChatCompletionsStream({
+    c,
+    logger,
+    model: openAIPayload.model,
+    recordUsage,
+    response,
+  })
+}
+
+const handleChatCompletionsStream = (options: {
+  c: Context
+  logger: ConsolaInstance
+  model: string
+  recordUsage: (usage: UsageTokens) => void
+  response: AsyncIterable<{ data?: string }>
+}) => {
+  const { c, logger, model, recordUsage, response } = options
   logger.debug("Streaming response from Copilot")
   applyForwardableResponseHeaders(c, getAttachedResponseHeaders(response), {
     "content-type": null,
@@ -180,7 +209,7 @@ export const handleWithChatCompletions = async (
       const premium = await resolvePremiumInfo(response, "messages/chat-stream")
       writeStreamLog(
         {
-          model: openAIPayload.model,
+          model,
           chunks: chunkCount,
           done: true,
           premium,
@@ -226,7 +255,14 @@ export const handleWithResponsesApi = async (
   anthropicPayload: AnthropicMessagesPayload,
   options: ResponsesFlowOptions,
 ) => {
-  const { logger, selectedModel, compactType, ...requestOptions } = options
+  const {
+    logger,
+    requestSessionAffinity,
+    requestTraceId,
+    selectedModel,
+    compactType,
+    ...requestOptions
+  } = options
 
   const responsesPayload =
     translateAnthropicMessagesToResponsesPayload(anthropicPayload)
@@ -235,6 +271,8 @@ export const handleWithResponsesApi = async (
     fallbackSessionId: requestOptions.sessionId,
     model: responsesPayload.model,
     payload: anthropicPayload,
+    requestSessionAffinity,
+    requestTraceId,
   })
 
   applyResponsesApiContextManagement(
@@ -304,6 +342,8 @@ export const handleWithMessagesApi = async (
   const {
     logger,
     anthropicBetaHeader,
+    requestSessionAffinity,
+    requestTraceId,
     subagentMarker,
     selectedModel,
     requestId,
@@ -317,6 +357,8 @@ export const handleWithMessagesApi = async (
     fallbackSessionId: sessionId,
     model: anthropicPayload.model,
     payload: anthropicPayload,
+    requestSessionAffinity,
+    requestTraceId,
   })
 
   debugJson(logger, "Translated Messages payload:", anthropicPayload)
@@ -475,6 +517,7 @@ const createNativeStreamBody = (options: {
   const { body, model, premium, logger, recordUsage } = options
   let chunkCount = 0
   let usage: UsageTokens = {}
+  let usageRecorded = false
   let buffer = ""
   const decoder = new TextDecoder()
   const reader = body.getReader()
@@ -491,7 +534,9 @@ const createNativeStreamBody = (options: {
         }
 
         writeStreamLog({ model, chunks: chunkCount, done: true, premium }, true)
-        recordUsage(usage)
+        if (!usageRecorded) {
+          recordUsage(usage)
+        }
 
         controller.close()
         return
@@ -526,6 +571,10 @@ const createNativeStreamBody = (options: {
             normalizeAnthropicUsage(parsedEvent.usage),
           )
         }
+        if (!usageRecorded && parsedEvent?.type === "message_delta") {
+          recordUsage(usage)
+          usageRecorded = true
+        }
         return data !== "" && data !== "[DONE]"
       }).length
       if (newEvents > 0) {
@@ -625,12 +674,33 @@ const createCopilotUsageRecorder = (options: {
   fallbackSessionId?: string
   model: string
   payload: AnthropicMessagesPayload
+  requestSessionAffinity?: string
+  requestTraceId?: string
 }): ((usage: UsageTokens) => void) =>
   createCopilotTokenUsageRecorder({
     endpoint: options.endpoint,
     fallbackSessionId: options.fallbackSessionId,
     model: options.model,
-    sessionId: getMetadataSessionId(options.payload),
+    sessionId:
+      options.requestSessionAffinity ?? getMetadataSessionId(options.payload),
+    traceId: options.requestTraceId,
+  })
+
+const createMessagesFlowUsageRecorder = (options: {
+  anthropicPayload: AnthropicMessagesPayload
+  endpoint: TokenUsageEndpoint
+  fallbackSessionId?: string
+  model: string
+  requestSessionAffinity?: string
+  requestTraceId?: string
+}) =>
+  createCopilotUsageRecorder({
+    endpoint: options.endpoint,
+    fallbackSessionId: options.fallbackSessionId,
+    model: options.model,
+    payload: options.anthropicPayload,
+    requestSessionAffinity: options.requestSessionAffinity,
+    requestTraceId: options.requestTraceId,
   })
 
 const getMetadataSessionId = (

src/routes/messages/handler.ts

@@ -109,6 +109,8 @@ export async function handleCompletion(c: Context) {
         subagentMarker,
         selectedModel,
         requestId,
+        requestSessionAffinity: c.req.header("x-session-affinity"),
+        requestTraceId: c.req.header("x-trace-id"),
         sessionId,
         compactType,
         logger,
@@ -124,6 +126,8 @@ export async function handleCompletion(c: Context) {
         subagentMarker,
         selectedModel,
         requestId,
+        requestSessionAffinity: c.req.header("x-session-affinity"),
+        requestTraceId: c.req.header("x-trace-id"),
         sessionId,
         compactType,
         logger,
@@ -137,6 +141,8 @@ export async function handleCompletion(c: Context) {
     {
       subagentMarker,
       requestId,
+      requestSessionAffinity: c.req.header("x-session-affinity"),
+      requestTraceId: c.req.header("x-trace-id"),
       sessionId,
       compactType,
       logger,

tests/native-messages-handler.test.ts

@@ -11,13 +11,18 @@ import * as modelsModule from "~/lib/models"
 import * as rateLimitModule from "~/lib/rate-limit"
 import { attachResponseHeaders } from "~/lib/response-headers"
 import { state } from "~/lib/state"
+import { closeUsageStore } from "~/lib/token-usage"
+import { traceIdMiddleware } from "~/lib/trace"
 import {
   getInitiatorFromPayload,
   isClaudeModel,
 } from "~/routes/messages/api-flows"
 import { handleCompletion } from "~/routes/messages/handler"
+import { tokenUsageRoute } from "~/routes/token-usage/route"
 import * as createMessagesModule from "~/services/copilot/create-messages"
 
+const DB_PATH_ENV = "COPILOT_API_SQLITE_DB_PATH"
+
 const createSseResponse = (events: Array<string>): Response => {
   const encoder = new TextEncoder()
 
@@ -296,6 +301,7 @@ describe("native handler", () => {
 
   test("logs anthropic effort mapping for config fallback", async () => {
     const app = new Hono()
+    app.use("*", traceIdMiddleware)
     app.post("/v1/messages", (c) => handleCompletion(c))
 
     const res = await app.request("/v1/messages", {
@@ -449,4 +455,63 @@ describe("native handler", () => {
       "rem=74.9&rst=2026-04-27T00%3A00%3A00Z",
     )
   })
+
+  test("records usage from native stream message events", async () => {
+    process.env[DB_PATH_ENV] = ":memory:"
+    await closeUsageStore()
+    createMessagesSpy.mockResolvedValueOnce(
+      createSseResponse([
+        'event: message_start\ndata: {"type":"message_start","message":{"id":"msg-usage","type":"message","role":"assistant","model":"claude-opus-4","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":37,"output_tokens":1}}}\n\n',
+        'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"max_tokens","stop_sequence":null},"usage":{"input_tokens":37,"output_tokens":96}}\n\n',
+        'event: message_stop\ndata: {"type":"message_stop"}\n\n',
+      ]),
+    )
+
+    const app = new Hono()
+    app.post("/v1/messages", (c) => handleCompletion(c))
+    app.route("/token-usage", tokenUsageRoute)
+
+    const res = await app.request("/v1/messages", {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "x-session-affinity": "native-stream-test",
+        "x-trace-id": "native-stream-usage-test",
+      },
+      body: JSON.stringify({
+        model: "claude-opus-4",
+        stream: true,
+        max_tokens: 1024,
+        messages: [{ role: "user", content: "hi" }],
+      } satisfies AnthropicMessagesPayload),
+    })
+
+    expect(res.status).toBe(200)
+    await res.text()
+
+    const eventsResponse = await app.request(
+      "/token-usage/events?period=day&page=1&page_size=10",
+    )
+    const page = (await eventsResponse.json()) as {
+      items: Array<{
+        input_tokens: number
+        output_tokens: number
+        session_id: string
+        total_tokens: number
+        trace_id: string
+      }>
+    }
+
+    expect(page.items).toHaveLength(1)
+    expect(page.items[0]).toMatchObject({
+      input_tokens: 37,
+      output_tokens: 96,
+      session_id: "native-stream-test",
+      total_tokens: 133,
+      trace_id: "native-stream-usage-test",
+    })
+
+    await closeUsageStore()
+    Reflect.deleteProperty(process.env, DB_PATH_ENV)
+  })
 })

tests/responses-handler.test.ts

@@ -32,6 +32,13 @@ import * as createChatCompletionsModule from "~/services/copilot/create-chat-com
 import * as createMessagesModule from "~/services/copilot/create-messages"
 import * as createResponsesModule from "~/services/copilot/create-responses"
 
+const DB_PATH_ENV = "COPILOT_API_SQLITE_DB_PATH"
+
+const isolateTokenUsageStore = async () => {
+  process.env[DB_PATH_ENV] = ":memory:"
+  await closeUsageStore()
+}
+
 const responseResult: ResponsesResult = {
   id: "resp-1",
   object: "response",
@@ -338,6 +345,15 @@ const restoreRoutingTestSpies = (
   spies.createMessagesSpy?.mockRestore()
 }
 
+beforeEach(async () => {
+  await isolateTokenUsageStore()
+})
+
+afterEach(async () => {
+  await closeUsageStore()
+  Reflect.deleteProperty(process.env, DB_PATH_ENV)
+})
+
 test("normalizes ANSI-colored info logs before assertions", () => {
   expect(
     normalizeInfoCall(
@@ -1187,7 +1203,6 @@ describe("responses handler token usage", () => {
     rateLimitWait: state.rateLimitWait,
     verbose: state.verbose,
   }
-  const dbPathEnv = "COPILOT_API_SQLITE_DB_PATH"
   const createResponsesMock = mock((() =>
     Promise.resolve(
       createStreamResponse([]),
@@ -1198,8 +1213,7 @@ describe("responses handler token usage", () => {
   let getConfigSpy: ReturnType<typeof spyOn<typeof configModule, "getConfig">>
 
   beforeEach(async () => {
-    process.env[dbPathEnv] = ":memory:"
-    await closeUsageStore()
+    await isolateTokenUsageStore()
 
     state.copilotToken = "test-token"
     state.manualApprove = false
@@ -1226,7 +1240,7 @@ describe("responses handler token usage", () => {
     createResponsesSpy.mockRestore()
     getConfigSpy.mockRestore()
     await closeUsageStore()
-    Reflect.deleteProperty(process.env, dbPathEnv)
+    Reflect.deleteProperty(process.env, DB_PATH_ENV)
 
     state.copilotToken = originalState.copilotToken
     state.manualApprove = originalState.manualApprove