fix: code review

Author: omeraplak

packages/core/src/utils/usage-normalizer.spec.ts

@@ -124,6 +124,32 @@ describe("normalizeFinishUsageStream", () => {
     expect(normalized[2].totalUsage).toEqual(lastStepUsage);
   });
 
+  it("uses finish metadata to override totalUsage when finish-step metadata is missing", async () => {
+    const lastStepUsage: LanguageModelUsage = {
+      inputTokens: 9,
+      outputTokens: 4,
+      totalTokens: 13,
+    };
+    const totalUsage: LanguageModelUsage = {
+      inputTokens: 20,
+      outputTokens: 10,
+      totalTokens: 30,
+    };
+    const parts = [
+      { type: "finish-step", usage: lastStepUsage },
+      {
+        type: "finish",
+        finishReason: "stop",
+        totalUsage,
+        providerMetadata: { anthropic: { model: "claude" } },
+      },
+    ];
+
+    const normalized = await collectStream(normalizeFinishUsageStream(toAsync(parts)));
+
+    expect(normalized[1].totalUsage).toEqual(lastStepUsage);
+  });
+
   it("overrides totalUsage when cache fields indicate anthropic usage", async () => {
     const lastStepUsage = {
       inputTokens: 8,

packages/core/src/utils/usage-normalizer.ts

@@ -60,6 +60,15 @@ export async function* normalizeFinishUsageStream<T extends StreamPartWithUsage>
       }
     }
 
+    if (part.type === "finish" && !useLastStepUsage) {
+      if (shouldUseLastStepUsage(part.providerMetadata, lastStepUsage)) {
+        useLastStepUsage = true;
+        if (part.usage) {
+          lastStepUsage = part.usage;
+        }
+      }
+    }
+
     if (part.type === "finish" && useLastStepUsage && lastStepUsage) {
       if (part.totalUsage !== undefined) {
         yield { ...part, totalUsage: lastStepUsage };