From b17658525fcc27fafa6598279cbe465cd35fc755 Mon Sep 17 00:00:00 2001
From: ytahdn <1294726970@qq.com>
Date: Thu, 17 Sep 2026 07:35:31 +0000
Subject: [PATCH] fix(web-shell): restore context usage from persisted replay
 (#12070)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: 钉萁 <dingqi.jww@alibaba-inc.com>
---
 .../client/daemon/session/mappers.test.ts     | 33 +++++++++++++++++++
 .../client/daemon/session/mappers.ts          |  7 ++--
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/packages/web-shell/client/daemon/session/mappers.test.ts b/packages/web-shell/client/daemon/session/mappers.test.ts
index 41d03d2f976..2a2dc11b60f 100644
--- a/packages/web-shell/client/daemon/session/mappers.test.ts
+++ b/packages/web-shell/client/daemon/session/mappers.test.ts
@@ -399,6 +399,39 @@ describe('getReplayTokenCount', () => {
     ).toBe(23_000);
   });
 
+  it('restores flat persisted usage and ignores newer subagent usage', () => {
+    const usage = { inputTokens: 23_000, totalTokens: 25_000 };
+    const events: DaemonEvent[] = [
+      usageEvent(1, { inputTokens: 11_000 }),
+      {
+        id: 2,
+        v: 1,
+        type: 'session_update',
+        data: {
+          sessionUpdate: 'agent_message_chunk',
+          content: { type: 'text', text: '' },
+          _meta: { usage },
+        },
+      },
+      {
+        id: 3,
+        v: 1,
+        type: 'session_update',
+        data: {
+          sessionUpdate: 'agent_message_chunk',
+          content: { type: 'text', text: '' },
+          _meta: {
+            parentToolCallId: 'agent-1',
+            usage: { inputTokens: 1_000 },
+          },
+        },
+      },
+    ];
+
+    expect(getReplayTokenUsage(events)).toEqual(usage);
+    expect(getReplayTokenCount(events)).toBe(23_000);
+  });
+
   it('returns the latest structured usage fields', () => {
     expect(
       getReplayTokenUsage([
diff --git a/packages/web-shell/client/daemon/session/mappers.ts b/packages/web-shell/client/daemon/session/mappers.ts
index f0cf932d7e8..4ff72d64a3f 100644
--- a/packages/web-shell/client/daemon/session/mappers.ts
+++ b/packages/web-shell/client/daemon/session/mappers.ts
@@ -5,7 +5,10 @@
  */
 
 import type { Dispatch, SetStateAction } from 'react';
-import { parseDaemonBackgroundTurn } from '@qwen-code/sdk/daemon';
+import {
+  getSessionUpdatePayload,
+  parseDaemonBackgroundTurn,
+} from '@qwen-code/sdk/daemon';
 import type {
   DaemonAvailableCommand,
   DaemonEvent,
@@ -826,7 +829,7 @@ export function getReplayTokenUsage(
     try {
       const event = events[i];
       if (event.type !== 'session_update') continue;
-      const update = getRecord(getRecord(event.data)?.['update']);
+      const update = getSessionUpdatePayload(event.data);
       const tokenUsage = getUsageTokenUsage(update);
       if (tokenUsage) return tokenUsage;
     } catch {