fix: fallback ws usage totals (#54940) (thanks @lyfuci)

src/agents/openai-ws-message-conversion.ts

@@ -540,6 +540,7 @@ export function buildAssistantMessageFromResponse(
   const hasToolCalls = content.some((part) => part.type === "toolCall");
   const stopReason: StopReason = hasToolCalls ? "toolUse" : "stop";
   const normalizedUsage = normalizeUsage(response.usage);
+  const rawTotalTokens = normalizedUsage?.total;
 
   const message = buildAssistantMessage({
     model: modelInfo,
@@ -548,7 +549,7 @@ export function buildAssistantMessageFromResponse(
     usage: buildUsageWithNoCost({
       input: normalizedUsage?.input ?? 0,
       output: normalizedUsage?.output ?? 0,
-      totalTokens: normalizedUsage?.total ?? 0,
+      totalTokens: rawTotalTokens && rawTotalTokens > 0 ? rawTotalTokens : undefined,
     }),
   });
 

src/agents/openai-ws-stream.test.ts

@@ -874,6 +874,33 @@ describe("buildAssistantMessageFromResponse", () => {
     expect(msg.usage.totalTokens).toBe(55);
   });
 
+  it("falls back to normalized input and output when total_tokens is missing", () => {
+    const response = makeResponseObject("resp_5c", "Hello");
+    response.usage = {
+      prompt_tokens: 10,
+      completion_tokens: 5,
+    };
+
+    const msg = buildAssistantMessageFromResponse(response, modelInfo);
+    expect(msg.usage.input).toBe(10);
+    expect(msg.usage.output).toBe(5);
+    expect(msg.usage.totalTokens).toBe(15);
+  });
+
+  it("falls back to normalized input and output when total_tokens is zero", () => {
+    const response = makeResponseObject("resp_5d", "Hello");
+    response.usage = {
+      input_tokens: 10,
+      output_tokens: 5,
+      total_tokens: 0,
+    };
+
+    const msg = buildAssistantMessageFromResponse(response, modelInfo);
+    expect(msg.usage.input).toBe(10);
+    expect(msg.usage.output).toBe(5);
+    expect(msg.usage.totalTokens).toBe(15);
+  });
+
   it("sets model/provider/api from modelInfo", () => {
     const response = makeResponseObject("resp_6", "Hi");
     const msg = buildAssistantMessageFromResponse(response, modelInfo);