From a59b3b18a9f01297a5b8cb6c9ae49afe16d5fd08 Mon Sep 17 00:00:00 2001 From: Zijian Xu Date: Sun, 4 Oct 2026 12:49:11 -0400 Subject: [PATCH 1/2] fix(chat): refresh context usage between model calls Change-Id: I423982a12d024afebc8da1e67836722a6e2aae3b --- components/AppShell.workspace-memory.test.mjs | 1 + docs/agents/sessions.md | 2 +- hooks/context-usage.test.mjs | 94 +++++++++++++++++++ hooks/useAgentSession.ts | 53 ++++++++--- 4 files changed, 138 insertions(+), 12 deletions(-) create mode 100644 hooks/context-usage.test.mjs diff --git a/components/AppShell.workspace-memory.test.mjs b/components/AppShell.workspace-memory.test.mjs index fe3f86f792..b173717c15 100644 --- a/components/AppShell.workspace-memory.test.mjs +++ b/components/AppShell.workspace-memory.test.mjs @@ -115,6 +115,7 @@ test("New restores the draft after session navigation and workspace auto-restore // Run the actual hook cleanup with the outgoing mount's captured draft key. const makeCleanup = vm.runInContext(stripTypeScriptTypes(`((isNew, newSessionDraftKey) => { const sessionHookMountedRef = { current: true }; + const contextUsageRequestIdRef = { current: 0 }; const newSessionPromotedRef = { current: false }; const sessionIdRef = { current: null }; const dataRef = { current: null }; diff --git a/docs/agents/sessions.md b/docs/agents/sessions.md index ddcf2744f8..6d4ce384fa 100644 --- a/docs/agents/sessions.md +++ b/docs/agents/sessions.md @@ -47,7 +47,7 @@ On mount `useAgentSession` loads the history, then `GET /api/sessions/[id]/state - The sidebar polls `/api/agent/running` every 2.5 s while the tab is visible; the session-list response is the initial fallback. - `invalidateSessionListCache()` bumps the generation but **keeps** the previous scan, fresh only while its generation matches. Callers needing only metadata (search hits to sidebar rows) pass `listAllSessions({ allowStale: true })` to read it while it rebuilds in the background, accepting that a seconds-old session is missing. - `useAgentSession` treats per-session SSE as primary and opens it before each prompt. `prompt_done` completes the UI stage and notification at once, but the stream stays open for the next prompt: the selected session's while selected (`scheduleEventStreamClose()` skips it), any other's for a 30-second grace window. `agent_start` cancels the close timer; `agent_settled` finishes extension-injected runs that have no wrapper-level `prompt_done` and starts a fresh grace window. Never close on the first `agent_end`: retries, compaction and extension-queued messages continue the same logical prompt. -- While a run is active, `useAgentSession` polls `GET /api/agent/[id]` and reconciles on `visibilitychange` / `online`, for terminal events missed by background tabs or half-open connections. +- While a run is active, `useAgentSession` polls `GET /api/agent/[id]` and reconciles on `visibilitychange` / `online`, for terminal events missed by background tabs or half-open connections. Busy replies also refresh context usage; each completed assistant message requests an immediate usage read between model calls. Usage reads from mount, message completion, reconciliation and `agent_end` share a monotonic request id and check the current session/run and mounted hook, so delayed reads cannot overwrite newer usage. Cleanup invalidates pending usage reads. - Prompt runs carry a monotonic run id; late SSE or reconciliation answers from an old run must be ignored, or they resurrect stale streaming bubbles. - Every SSE (re)connection is gated on `sessionHookMountedRef`. Under React Strict Mode (`next dev`) the mount-only effect's cleanup clears it and restores it only after the warm-session effect re-runs, so that effect must re-assert it before `maintainEventsConnected()`, or a dev tab never opens its stream. diff --git a/hooks/context-usage.test.mjs b/hooks/context-usage.test.mjs new file mode 100644 index 0000000000..1e8c2e39e6 --- /dev/null +++ b/hooks/context-usage.test.mjs @@ -0,0 +1,94 @@ +import assert from "node:assert/strict"; +import { readFile } from "node:fs/promises"; +import test from "node:test"; +import { Script, createContext } from "node:vm"; +import ts from "typescript"; + +const source = ts.createSourceFile("useAgentSession.ts", await readFile(new URL("./useAgentSession.ts", import.meta.url), "utf8"), ts.ScriptTarget.Latest, true); +const nodes = []; +function visit(node) { nodes.push(node); ts.forEachChild(node, visit); } +visit(source); +function callback(name) { + const node = nodes.find((node) => ts.isVariableDeclaration(node) && node.name.getText(source) === name); + assert.ok(node, `missing ${name}`); + return new Script(ts.transpileModule(`(${node.initializer.arguments[0].getText(source)})`, { compilerOptions: { target: ts.ScriptTarget.ESNext } }).outputText); +} +function setup() { + const writes = []; + const requests = []; + const context = createContext({ + sessionHookMountedRef: { current: true }, sessionIdRef: { current: "a" }, + promptRunIdRef: { current: 1 }, contextUsageRequestIdRef: { current: 0 }, + agentRunningRef: { current: true }, sdkAgentActiveRef: { current: true }, rpcPromptPendingRef: { current: true }, + setContextUsage: (value) => writes.push(value), + fetch: (url) => { const request = Promise.withResolvers(); requests.push({ ...request, url }); return request.promise; }, + syncLiveModel() {}, setIsCompacting() {}, setAutoCompactionEnabled() {}, setQueuedMessages() {}, + normalizeQueuedMessages: (value) => value, finishPromptWithoutStream: () => { throw new Error("busy run settled"); }, + }); + context.applyContextUsage = callback("applyContextUsage").runInContext(context); + context.refreshContextUsage = callback("refreshContextUsage").runInContext(context); + context.reconcileAgentState = callback("reconcileAgentState").runInContext(context); + const reply = (index, usage, busy = true) => requests[index].resolve(Response.json({ running: busy, state: { contextUsage: usage, isStreaming: busy, isPromptRunning: busy } })); + return { context, writes, requests, reply }; +} +const usage = (tokens) => ({ percent: tokens / 100, contextWindow: 10_000, tokens }); + +test("a busy reconciliation refreshes context usage without settling the run", async () => { + const state = setup(); + const pending = state.context.reconcileAgentState("a"); + state.reply(0, usage(100)); + await pending; + assert.deepEqual(state.writes, [usage(100)]); + assert.equal(state.context.agentRunningRef.current, true); +}); + +test("a newer assistant usage read wins over a delayed poll", async () => { + const state = setup(); + const poll = state.context.reconcileAgentState("a"); + const refresh = state.context.refreshContextUsage("a"); + state.reply(1, usage(200)); await refresh; + state.reply(0, usage(100)); await poll; + assert.deepEqual(state.writes, [usage(200)]); +}); + +test("usage replies from an old run, another session or an unmounted hook are ignored", async () => { + for (const invalidate of [ + (ctx) => { ctx.promptRunIdRef.current++; }, + (ctx) => { ctx.sessionIdRef.current = "b"; }, + (ctx) => { ctx.sessionHookMountedRef.current = false; }, + (ctx) => { ctx.contextUsageRequestIdRef.current++; }, + ]) { + const state = setup(); + const pending = state.context.refreshContextUsage("a"); + invalidate(state.context); state.reply(0, usage(100)); await pending; + assert.deepEqual(state.writes, []); + } +}); + +test("failed reads preserve usage and a subsequent read can recover", async () => { + const state = setup(); + const failed = state.context.refreshContextUsage("a"); + state.requests[0].resolve(new Response("Unavailable", { status: 503 })); await failed; + const broken = state.context.refreshContextUsage("a"); + state.requests[1].reject(new TypeError("offline")); await broken; + assert.deepEqual(state.writes, []); + const recovered = state.context.refreshContextUsage("a"); + state.reply(2, null); await recovered; + assert.deepEqual(state.writes, [null]); +}); + +test("only completed assistant messages trigger an immediate usage read", () => { + const messageEnd = nodes.find((node) => ts.isCaseClause(node) && node.expression.getText(source) === '"message_end"'); + const script = new Script(ts.transpileModule(`(() => { switch(event.type) { ${messageEnd.getText(source)} } })()`, { compilerOptions: { target: ts.ScriptTarget.ESNext } }).outputText); + for (const role of ["assistant", "toolResult", "user", "system"]) { + const reads = []; + script.runInNewContext({ + event: { type: "message_end", message: { role, content: [] } }, + agentRunningRef: { current: true }, sessionIdRef: { current: "a" }, optimisticUserMessageKeyRef: { current: null }, + isSystemMessageEvent: (event) => event.message.role === "system", + normalizeToolCalls: (message) => message, userMessageKey: () => "user", setMessages() {}, dispatch() {}, setAgentPhase() {}, + refreshContextUsage: (sid) => reads.push(sid), + }); + assert.deepEqual(reads, role === "assistant" ? ["a"] : []); + } +}); diff --git a/hooks/useAgentSession.ts b/hooks/useAgentSession.ts index 0fb352e355..70ecb12301 100644 --- a/hooks/useAgentSession.ts +++ b/hooks/useAgentSession.ts @@ -350,6 +350,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { const [liveThinkingLevel, setLiveThinkingLevel] = useState(null); const [retryInfo, setRetryInfo] = useState<{ attempt: number; maxAttempts: number; errorMessage?: string } | null>(null); const [contextUsage, setContextUsage] = useState<{ percent: number | null; contextWindow: number; tokens: number | null } | null>(null); + const contextUsageRequestIdRef = useRef(0); const [systemPrompt, setSystemPrompt] = useState(null); const [forkingEntryId, setForkingEntryId] = useState(null); const [currentModelOverride, setCurrentModelOverride] = useState<{ provider: string; modelId: string } | null>(null); @@ -575,6 +576,25 @@ export function useAgentSession(opts: UseAgentSessionOptions) { } satisfies SessionStatsInfo; }, [messages, sessionStatsOverride, contextUsage, data?.context.messages, data?.filePath, data?.totalActiveMs, data?.stats, session?.id, session?.name]); + const applyContextUsage = useCallback((state: AgentStateResponse | undefined, sid: string, runId: number, requestId: number) => { + if (!sessionHookMountedRef.current || sessionIdRef.current !== sid + || promptRunIdRef.current !== runId || contextUsageRequestIdRef.current !== requestId) return; + if (state?.contextUsage !== undefined) setContextUsage(state.contextUsage ?? null); + }, []); + + const refreshContextUsage = useCallback(async (sid: string) => { + const runId = promptRunIdRef.current; + const requestId = ++contextUsageRequestIdRef.current; + try { + const res = await fetch(`/api/agent/${encodeURIComponent(sid)}`); + if (!res.ok) return; + const data = await res.json() as { state?: AgentStateResponse }; + applyContextUsage(data.state, sid, runId, requestId); + } catch { + // A later message or the running-state poll retries the usage read. + } + }, [applyContextUsage]); + const loadSession = useCallback(async (sid: string, showLoading = false, includeState = false, options?: { force?: boolean }) => { // Single-flight: concurrent reads for the same session (mount + SSE settle + // reconcile) share one request unless the caller forces a fresh read. @@ -689,6 +709,8 @@ export function useAgentSession(opts: UseAgentSessionOptions) { if (!includeState) return null; try { + const runId = promptRunIdRef.current; + const usageRequestId = ++contextUsageRequestIdRef.current; const stateRes = await fetch(`/api/sessions/${encodeURIComponent(sid)}/state`); if (!stateRes.ok) throw new Error(`HTTP ${stateRes.status}`); const agentState = await stateRes.json() as { running: boolean; state?: AgentStateResponse }; @@ -697,7 +719,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { const liveState = agentState.state; syncLiveModel(liveState); if (liveState) { - if (liveState.contextUsage !== undefined) setContextUsage(liveState.contextUsage ?? null); + applyContextUsage(liveState, sid, runId, usageRequestId); if (liveState.systemPrompt !== undefined) setSystemPrompt(liveState.systemPrompt ?? null); if (liveState.extensionStatuses !== undefined) setExtensionStatuses(liveState.extensionStatuses ?? []); if (liveState.extensionWidgets !== undefined) setExtensionWidgets(liveState.extensionWidgets ?? []); @@ -723,7 +745,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { if (loadFlightsRef.current.get(flightKey) === flight) loadFlightsRef.current.delete(flightKey); }); return await flight; - }, [setToolPresetState, syncLiveModel]); + }, [applyContextUsage, setToolPresetState, syncLiveModel]); const loadContext = useCallback(async (sid: string, leafId: string | null, before?: string | null, options?: { tail?: number; signal?: AbortSignal }) => { try { @@ -1280,6 +1302,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { const reconcileAgentState = useCallback(async (sid: string) => { if (!agentRunningRef.current || sessionIdRef.current !== sid) return; const runId = promptRunIdRef.current; + const usageRequestId = ++contextUsageRequestIdRef.current; try { const res = await fetch(`/api/agent/${encodeURIComponent(sid)}`); if (!res.ok) return; @@ -1289,6 +1312,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { // flight) — everything in it is stale, drop it. if (sessionIdRef.current !== sid || promptRunIdRef.current !== runId) return; const state = data.state; + applyContextUsage(state, sid, runId, usageRequestId); syncLiveModel(state); // Mirror compaction state unconditionally: a missed compaction_end // would otherwise leave the "Stop compaction" UI stuck. No state @@ -1305,7 +1329,6 @@ export function useAgentSession(opts: UseAgentSessionOptions) { } if (!agentRunningRef.current) return; if (state) { - if (state.contextUsage !== undefined) setContextUsage(state.contextUsage ?? null); if (state.systemPrompt !== undefined) setSystemPrompt(state.systemPrompt ?? null); if (state.extensionStatuses !== undefined) setExtensionStatuses(state.extensionStatuses ?? []); if (state.extensionWidgets !== undefined) setExtensionWidgets(state.extensionWidgets ?? []); @@ -1314,7 +1337,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { } catch { // Network still down — the next poll / visibility / online tick retries. } - }, [finishPromptWithoutStream, syncLiveModel]); + }, [applyContextUsage, finishPromptWithoutStream, syncLiveModel]); // Recovery net for missed SSE events: while the agent is running, verify // against the server periodically and whenever the tab returns to the @@ -1380,12 +1403,16 @@ export function useAgentSession(opts: UseAgentSessionOptions) { setRetryInfo(null); dispatch({ type: "end" }); if (sessionIdRef.current) { - loadSession(sessionIdRef.current); - fetch(`/api/agent/${encodeURIComponent(sessionIdRef.current)}`) - .then((r) => r.json()) - .then((d: { state?: AgentStateResponse }) => { + const sid = sessionIdRef.current; + const runId = promptRunIdRef.current; + const usageRequestId = ++contextUsageRequestIdRef.current; + loadSession(sid); + fetch(`/api/agent/${encodeURIComponent(sid)}`) + .then((r) => r.ok ? r.json() : null) + .then((d: { state?: AgentStateResponse } | null) => { + if (!d || !sessionHookMountedRef.current || sessionIdRef.current !== sid || promptRunIdRef.current !== runId) return; syncLiveModel(d.state); - if (d.state?.contextUsage !== undefined) setContextUsage(d.state.contextUsage ?? null); + applyContextUsage(d.state, sid, runId, usageRequestId); if (d.state?.systemPrompt !== undefined) setSystemPrompt(d.state.systemPrompt ?? null); if (d.state?.extensionStatuses !== undefined) setExtensionStatuses(d.state.extensionStatuses ?? []); if (d.state?.extensionWidgets !== undefined) setExtensionWidgets(d.state.extensionWidgets ?? []); @@ -1506,6 +1533,10 @@ export function useAgentSession(opts: UseAgentSessionOptions) { }); } else if (completed) { setMessages((prev) => [...prev, normalizeToolCalls(completed)]); + if (completed.role === "assistant") { + const sid = sessionIdRef.current; + if (sid) void refreshContextUsage(sid); + } } dispatch({ type: "end" }); setAgentPhase({ kind: "waiting_model" }); @@ -1617,7 +1648,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { setExtensionDialogs((queue) => removeExtensionUiRequest(queue, event.id as string)); break; } - }, [addNotice, cancelEventStreamGrace, handleExtensionUiRequest, loadSession, notifyPromptStage, onAgentEnd, scheduleEventStreamClose, scrollToBottom, settleUiStage, syncLiveModel]); + }, [addNotice, applyContextUsage, cancelEventStreamGrace, handleExtensionUiRequest, loadSession, notifyPromptStage, onAgentEnd, refreshContextUsage, scheduleEventStreamClose, scrollToBottom, settleUiStage, syncLiveModel]); handleAgentEventRef.current = handleAgentEvent; const handleSend = useCallback(async (message: string, images?: AttachedImage[]) => { @@ -2451,7 +2482,6 @@ export function useAgentSession(opts: UseAgentSessionOptions) { } if (agentState?.state) { if (agentState.state.isCompacting !== undefined) setIsCompacting(agentState.state.isCompacting); - if (agentState.state.contextUsage !== undefined) setContextUsage(agentState.state.contextUsage ?? null); if (agentState.state.systemPrompt !== undefined) setSystemPrompt(agentState.state.systemPrompt ?? null); if (agentState.state.extensionStatuses !== undefined) setExtensionStatuses(agentState.state.extensionStatuses ?? []); if (agentState.state.extensionWidgets !== undefined) setExtensionWidgets(agentState.state.extensionWidgets ?? []); @@ -2461,6 +2491,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { } return () => { sessionHookMountedRef.current = false; + contextUsageRequestIdRef.current += 1; const abandonedDraftKey = isNew ? newSessionDraftKey : null; if (abandonedDraftKey) { queueMicrotask(() => { From 514df1789062694da5e6d833d740839b572ecde1 Mon Sep 17 00:00:00 2001 From: Alex Yang Date: Wed, 7 Oct 2026 20:59:59 +0900 Subject: [PATCH 2/2] fix(chat): keep an older usage read when a newer one fails The shared request id let any newer usage read invalidate older ones, even when the newer read failed (non-ok reply or network error). A good reply from an in-flight poll was then dropped, and usage stayed stale until the next assistant message, poll tick or agent_end. Track the last applied id instead: a reply applies only when it is newer than the last one applied, so delayed replies still cannot overwrite newer usage. Each mount marks the reads started before it as stale, which keeps the unmount invalidation without reading a ref in effect cleanup. Co-Authored-By: Claude Opus 5.5 --- components/AppShell.workspace-memory.test.mjs | 1 - docs/agents/sessions.md | 2 +- hooks/context-usage.test.mjs | 13 +++++++++++-- hooks/useAgentSession.ts | 9 +++++++-- 4 files changed, 19 insertions(+), 6 deletions(-) diff --git a/components/AppShell.workspace-memory.test.mjs b/components/AppShell.workspace-memory.test.mjs index b173717c15..fe3f86f792 100644 --- a/components/AppShell.workspace-memory.test.mjs +++ b/components/AppShell.workspace-memory.test.mjs @@ -115,7 +115,6 @@ test("New restores the draft after session navigation and workspace auto-restore // Run the actual hook cleanup with the outgoing mount's captured draft key. const makeCleanup = vm.runInContext(stripTypeScriptTypes(`((isNew, newSessionDraftKey) => { const sessionHookMountedRef = { current: true }; - const contextUsageRequestIdRef = { current: 0 }; const newSessionPromotedRef = { current: false }; const sessionIdRef = { current: null }; const dataRef = { current: null }; diff --git a/docs/agents/sessions.md b/docs/agents/sessions.md index 6d4ce384fa..eaadeebb05 100644 --- a/docs/agents/sessions.md +++ b/docs/agents/sessions.md @@ -47,7 +47,7 @@ On mount `useAgentSession` loads the history, then `GET /api/sessions/[id]/state - The sidebar polls `/api/agent/running` every 2.5 s while the tab is visible; the session-list response is the initial fallback. - `invalidateSessionListCache()` bumps the generation but **keeps** the previous scan, fresh only while its generation matches. Callers needing only metadata (search hits to sidebar rows) pass `listAllSessions({ allowStale: true })` to read it while it rebuilds in the background, accepting that a seconds-old session is missing. - `useAgentSession` treats per-session SSE as primary and opens it before each prompt. `prompt_done` completes the UI stage and notification at once, but the stream stays open for the next prompt: the selected session's while selected (`scheduleEventStreamClose()` skips it), any other's for a 30-second grace window. `agent_start` cancels the close timer; `agent_settled` finishes extension-injected runs that have no wrapper-level `prompt_done` and starts a fresh grace window. Never close on the first `agent_end`: retries, compaction and extension-queued messages continue the same logical prompt. -- While a run is active, `useAgentSession` polls `GET /api/agent/[id]` and reconciles on `visibilitychange` / `online`, for terminal events missed by background tabs or half-open connections. Busy replies also refresh context usage; each completed assistant message requests an immediate usage read between model calls. Usage reads from mount, message completion, reconciliation and `agent_end` share a monotonic request id and check the current session/run and mounted hook, so delayed reads cannot overwrite newer usage. Cleanup invalidates pending usage reads. +- While a run is active, `useAgentSession` polls `GET /api/agent/[id]` and reconciles on `visibilitychange` / `online`, for terminal events missed by background tabs or half-open connections. Busy replies also refresh context usage; each completed assistant message requests an immediate usage read between model calls. Usage reads from mount, message completion, reconciliation and `agent_end` share a monotonic request id and check the current session/run and mounted hook. A reply applies only when its id is newer than the last applied one, so a delayed read cannot overwrite newer usage, and a newer read that fails does not discard an older one still in flight. Each mount marks the reads started before it as stale. - Prompt runs carry a monotonic run id; late SSE or reconciliation answers from an old run must be ignored, or they resurrect stale streaming bubbles. - Every SSE (re)connection is gated on `sessionHookMountedRef`. Under React Strict Mode (`next dev`) the mount-only effect's cleanup clears it and restores it only after the warm-session effect re-runs, so that effect must re-assert it before `maintainEventsConnected()`, or a dev tab never opens its stream. diff --git a/hooks/context-usage.test.mjs b/hooks/context-usage.test.mjs index 1e8c2e39e6..4d06e5ca94 100644 --- a/hooks/context-usage.test.mjs +++ b/hooks/context-usage.test.mjs @@ -18,7 +18,7 @@ function setup() { const requests = []; const context = createContext({ sessionHookMountedRef: { current: true }, sessionIdRef: { current: "a" }, - promptRunIdRef: { current: 1 }, contextUsageRequestIdRef: { current: 0 }, + promptRunIdRef: { current: 1 }, contextUsageRequestIdRef: { current: 0 }, contextUsageAppliedIdRef: { current: 0 }, agentRunningRef: { current: true }, sdkAgentActiveRef: { current: true }, rpcPromptPendingRef: { current: true }, setContextUsage: (value) => writes.push(value), fetch: (url) => { const request = Promise.withResolvers(); requests.push({ ...request, url }); return request.promise; }, @@ -51,12 +51,21 @@ test("a newer assistant usage read wins over a delayed poll", async () => { assert.deepEqual(state.writes, [usage(200)]); }); +test("a failed newer read does not discard an older poll's usage", async () => { + const state = setup(); + const poll = state.context.reconcileAgentState("a"); + const refresh = state.context.refreshContextUsage("a"); + state.requests[1].resolve(new Response("Unavailable", { status: 503 })); await refresh; + state.reply(0, usage(100)); await poll; + assert.deepEqual(state.writes, [usage(100)]); +}); + test("usage replies from an old run, another session or an unmounted hook are ignored", async () => { for (const invalidate of [ (ctx) => { ctx.promptRunIdRef.current++; }, (ctx) => { ctx.sessionIdRef.current = "b"; }, (ctx) => { ctx.sessionHookMountedRef.current = false; }, - (ctx) => { ctx.contextUsageRequestIdRef.current++; }, + (ctx) => { ctx.contextUsageAppliedIdRef.current = ctx.contextUsageRequestIdRef.current; }, ]) { const state = setup(); const pending = state.context.refreshContextUsage("a"); diff --git a/hooks/useAgentSession.ts b/hooks/useAgentSession.ts index 70ecb12301..2056912b30 100644 --- a/hooks/useAgentSession.ts +++ b/hooks/useAgentSession.ts @@ -351,6 +351,9 @@ export function useAgentSession(opts: UseAgentSessionOptions) { const [retryInfo, setRetryInfo] = useState<{ attempt: number; maxAttempts: number; errorMessage?: string } | null>(null); const [contextUsage, setContextUsage] = useState<{ percent: number | null; contextWindow: number; tokens: number | null } | null>(null); const contextUsageRequestIdRef = useRef(0); + // Highest request id whose reply was applied. A reply applies only when it + // is newer, so a failed newer read never discards an older good one. + const contextUsageAppliedIdRef = useRef(0); const [systemPrompt, setSystemPrompt] = useState(null); const [forkingEntryId, setForkingEntryId] = useState(null); const [currentModelOverride, setCurrentModelOverride] = useState<{ provider: string; modelId: string } | null>(null); @@ -578,7 +581,8 @@ export function useAgentSession(opts: UseAgentSessionOptions) { const applyContextUsage = useCallback((state: AgentStateResponse | undefined, sid: string, runId: number, requestId: number) => { if (!sessionHookMountedRef.current || sessionIdRef.current !== sid - || promptRunIdRef.current !== runId || contextUsageRequestIdRef.current !== requestId) return; + || promptRunIdRef.current !== runId || requestId <= contextUsageAppliedIdRef.current) return; + contextUsageAppliedIdRef.current = requestId; if (state?.contextUsage !== undefined) setContextUsage(state.contextUsage ?? null); }, []); @@ -2427,6 +2431,8 @@ export function useAgentSession(opts: UseAgentSessionOptions) { // Load session on mount useEffect(() => { sessionHookMountedRef.current = true; + // Usage reads started before a remount are stale. + contextUsageAppliedIdRef.current = contextUsageRequestIdRef.current; if (session) { sessionIdRef.current = session.id; // Snapshot fast path: show the cached history window immediately, then @@ -2491,7 +2497,6 @@ export function useAgentSession(opts: UseAgentSessionOptions) { } return () => { sessionHookMountedRef.current = false; - contextUsageRequestIdRef.current += 1; const abandonedDraftKey = isNew ? newSessionDraftKey : null; if (abandonedDraftKey) { queueMicrotask(() => {