fix(claude): retain measured usage after a steered query exits

apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts

@@ -4506,11 +4506,16 @@ describe("ClaudeAdapterV2 background wake turns", () => {
 
   it.effect.each(
     (["aborted_tools", "aborted_streaming"] as const).flatMap((terminalReason) =>
-      [true, false].map((steered) => ({ terminalReason, steered })),
+      [
+        { steered: false, queryExits: false },
+        { steered: true, queryExits: false },
+        { steered: true, queryExits: true },
+      ].map((scenario) => ({ terminalReason, ...scenario })),
     ),
-  )("handles $terminalReason with active steering=$steered", ({ terminalReason, steered }) =>
+  )("handles $terminalReason (steered=$steered, exits=$queryExits)", (scenario) =>
     Effect.scoped(
       Effect.gen(function* () {
+        const { terminalReason, steered, queryExits } = scenario;
         const harness = yield* makeWakeHarness;
         const idAllocator = yield* IdAllocator.IdAllocatorV2;
         const attemptId = RunAttemptId.make("attempt-steering-abort");
@@ -4542,8 +4547,7 @@ describe("ClaudeAdapterV2 background wake turns", () => {
           });
           assert.equal(harness.offeredMessages[1]?.priority, "now");
         }
-        yield* Queue.offer(
-          harness.sdkMessages,
+        yield* harness.offerAndWait(
           makeResultFrame({
             uuid: "00000000-0000-4000-8000-000000000901",
             result: "",
@@ -4556,7 +4560,9 @@ describe("ClaudeAdapterV2 background wake turns", () => {
             },
           }),
         );
-        if (steered) {
+        if (queryExits) {
+          yield* Queue.shutdown(harness.sdkMessages);
+        } else if (steered) {
           yield* Queue.offer(harness.sdkMessages, wakeAssistant);
           yield* Queue.offer(
             harness.sdkMessages,
@@ -4573,16 +4579,20 @@ describe("ClaudeAdapterV2 background wake turns", () => {
           );
         }
         const terminal = yield* Queue.take(harness.terminalReceipts);
-        assert.equal(terminal.status, steered ? "completed" : "interrupted");
+        assert.equal(
+          terminal.status,
+          queryExits ? "failed" : steered ? "completed" : "interrupted",
+        );
         if (steered) {
-          assert.isTrue(
-            harness.events.some(
-              (event) =>
-                event.type === "turn_item.updated" &&
-                event.turnItem.type === "assistant_message" &&
-                event.turnItem.text === WAKE_ASSISTANT_TEXT,
-            ),
-          );
+          if (!queryExits)
+            assert.isTrue(
+              harness.events.some(
+                (event) =>
+                  event.type === "turn_item.updated" &&
+                  event.turnItem.type === "assistant_message" &&
+                  event.turnItem.text === WAKE_ASSISTANT_TEXT,
+              ),
+            );
           const completed = harness.events.findLast(
             (event) => event.type === "provider_turn.updated",
           );
@@ -4592,12 +4602,12 @@ describe("ClaudeAdapterV2 background wake turns", () => {
               : undefined,
             {
               usageScope: "main_agent",
-              usageStatus: "complete",
+              usageStatus: queryExits ? "partial" : "complete",
               hasSubagents: false,
-              inputTokens: 603 + 2847 + 11455,
-              outputTokens: 46,
-              cachedInputTokens: 11455,
-              cacheCreationTokens: 2847,
+              inputTokens: queryExits ? 3100 : 603 + 2847 + 11455,
+              outputTokens: queryExits ? 40 : 46,
+              cachedInputTokens: queryExits ? 2000 : 11455,
+              cacheCreationTokens: queryExits ? 500 : 2847,
             },
           );
         }

apps/server/src/provider/ClaudeTurnTokenUsage.ts

@@ -15,12 +15,10 @@ const BILLED_USAGE_FIELDS = [
 ] as const;
 
 function addKnownCount(finalCount: unknown, earlierCounts: ReadonlyArray<unknown>): unknown {
-  const known = finiteNonNegativeInteger(finalCount);
-  if (known === undefined) return finalCount;
-  return earlierCounts.reduce<number>(
-    (sum, count) => sum + (finiteNonNegativeInteger(count) ?? 0),
-    known,
-  );
+  const counts = [finalCount, ...earlierCounts]
+    .map(finiteNonNegativeInteger)
+    .filter((count) => count !== undefined);
+  return counts.length === 0 ? finalCount : counts.reduce((sum, count) => sum + count, 0);
 }
 
 // Steering aborts the current segment; each SDK result reports only that segment's usage.
@@ -38,18 +36,16 @@ function withEarlierSegmentUsage(
     );
   }
   const outputDetails = usage.output_tokens_details as Record<string, unknown> | undefined;
-  if (outputDetails != null) {
-    total.output_tokens_details = {
-      ...outputDetails,
-      thinking_tokens: addKnownCount(
-        outputDetails.thinking_tokens,
-        segments.map(
-          (segment) =>
-            (segment.output_tokens_details as Record<string, unknown> | undefined)?.thinking_tokens,
-        ),
+  total.output_tokens_details = {
+    ...outputDetails,
+    thinking_tokens: addKnownCount(
+      outputDetails?.thinking_tokens,
+      segments.map(
+        (segment) =>
+          (segment.output_tokens_details as Record<string, unknown> | undefined)?.thinking_tokens,
       ),
-    };
-  }
+    ),
+  };
   return total;
 }
 
@@ -60,15 +56,7 @@ export function normalizeClaudeTurnTokenUsage(
   earlierSegmentUsage: ReadonlyArray<unknown> = [],
 ): TurnTokenUsage {
   const finalUsage = result?.usage as Record<string, unknown> | undefined;
-  const usage = finalUsage && withEarlierSegmentUsage(finalUsage, earlierSegmentUsage);
-  if (!usage) {
-    return {
-      usageStatus: "unavailable",
-      usageScope: "main_agent",
-      hasSubagents,
-    };
-  }
-
+  const usage = withEarlierSegmentUsage(finalUsage ?? {}, earlierSegmentUsage);
   const uncachedInputTokens = finiteNonNegativeInteger(usage.input_tokens);
   const cachedInputTokens = finiteNonNegativeInteger(usage.cache_read_input_tokens);
   const cacheCreationTokens = finiteNonNegativeInteger(usage.cache_creation_input_tokens);
@@ -116,8 +104,14 @@ export function normalizeClaudeTurnTokenUsage(
   if (
     terminalStatus === "completed" &&
     result?.subtype === "success" &&
+    finiteNonNegativeInteger(finalUsage?.input_tokens) !== undefined &&
+    finiteNonNegativeInteger(finalUsage?.output_tokens) !== undefined &&
     inputTokens !== undefined &&
-    rawOutputTokens !== undefined
+    rawOutputTokens !== undefined &&
+    BILLED_USAGE_FIELDS.every(
+      (field) =>
+        finalUsage?.[field] == null || finiteNonNegativeInteger(finalUsage[field]) !== undefined,
+    )
   ) {
     return {
       ...commonUsage,

apps/server/src/provider/TurnTokenUsage.test.ts

@@ -130,4 +130,93 @@ describe("billable native turn usage", () => {
       ),
     ).toMatchObject({ usageStatus: "partial", inputTokens: 4, outputTokens: 2 });
   });
+
+  it.each(["failed", "interrupted", "completed"] as const)(
+    "retains measured Claude steering usage as partial after %s without final usage",
+    (status) => {
+      expect(
+        normalizeClaudeTurnTokenUsage(
+          status === "completed" ? { subtype: "success" } : undefined,
+          true,
+          status,
+          [
+            { input_tokens: 100, cache_read_input_tokens: 40, output_tokens: 20 },
+            { input_tokens: 5, cache_read_input_tokens: 2, output_tokens: 3 },
+          ],
+        ),
+      ).toEqual({
+        usageStatus: "partial",
+        usageScope: "main_agent",
+        hasSubagents: true,
+        inputTokens: 147,
+        cachedInputTokens: 42,
+        outputTokens: 23,
+      });
+    },
+  );
+
+  it.each([undefined, null, {}])(
+    "keeps unavailable Claude usage unavailable without measured segments (%s)",
+    (usage) => {
+      expect(
+        normalizeClaudeTurnTokenUsage({ subtype: "success", usage }, false, "completed")
+          .usageStatus,
+      ).toBe("unavailable");
+    },
+  );
+
+  it("retains measured thinking tokens without a final result", () => {
+    expect(
+      normalizeClaudeTurnTokenUsage(undefined, false, "failed", [
+        { input_tokens: 4, output_tokens: 3, output_tokens_details: { thinking_tokens: 2 } },
+      ]),
+    ).toMatchObject({
+      usageStatus: "partial",
+      inputTokens: 4,
+      outputTokens: 3,
+      reasoningTokens: 2,
+    });
+  });
+
+  it.each([undefined, null])(
+    "retains earlier cache counts when a steering segment reports %s",
+    (cache) => {
+      expect(
+        normalizeClaudeTurnTokenUsage(undefined, false, "failed", [
+          { input_tokens: 100, cache_read_input_tokens: 40, output_tokens: 20 },
+          { input_tokens: 5, cache_read_input_tokens: cache, output_tokens: 3 },
+        ]),
+      ).toMatchObject({
+        usageStatus: "partial",
+        inputTokens: 145,
+        cachedInputTokens: 40,
+        outputTokens: 23,
+      });
+    },
+  );
+
+  it.each(["cache_read_input_tokens", "cache_creation_input_tokens"])(
+    "keeps Claude usage partial when the final %s is invalid",
+    (field) => {
+      expect(
+        normalizeClaudeTurnTokenUsage(
+          { subtype: "success", usage: { input_tokens: 5, output_tokens: 3, [field]: -1 } },
+          false,
+          "completed",
+          [{ input_tokens: 100, output_tokens: 20, [field]: 40 }],
+        ),
+      ).toMatchObject({ usageStatus: "partial", inputTokens: 145, outputTokens: 23 });
+    },
+  );
+
+  it("adds Claude steering segments to the final measured result once", () => {
+    expect(
+      normalizeClaudeTurnTokenUsage(
+        { subtype: "success", usage: { input_tokens: 5, output_tokens: 3 } },
+        false,
+        "completed",
+        [{ input_tokens: 4, output_tokens: 2 }],
+      ),
+    ).toMatchObject({ usageStatus: "complete", inputTokens: 9, outputTokens: 5 });
+  });
 });