From 81599b2af9512b2b8078216af12c3423da7175d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A8=8B=E5=BA=8F=E5=91=98=E9=98=BF=E6=B1=9F=28Relakkes?= =?UTF-8?q?=29?= Date: Fri, 12 Jun 2026 17:57:47 +0800 Subject: [PATCH] fix: recover context meter right after compaction (#743) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The context usage indicator stayed at the pre-compact percentage (e.g. 100%) until the next API response arrived. Two stacked causes: - The CLI anchors the displayed total to max(local estimate, last API usage) so the meter never drops mid-turn — but preserved messages (SM-compact / partial compact) still carried the pre-compact usage, pinning the meter after compaction. buildPostCompactMessages now zeroes token usage on preserved assistant copies, the established stale placeholder that getCurrentUsage() skips, covering every compaction path in one place. Originals are not mutated, so transcript and cost accounting are unaffected. - Right after compaction the CLI is often still busy, so the indicator's refresh timed out ("Request timed out after 30s") and the catch kept the stale context on screen with no later retry (auto refresh is throttled to 10s and stops once the session goes idle). compact_boundary now bumps a per-session compactCount; the indicator force-refreshes on that nonce — bypassing the throttle and any in-flight request that may still hold pre-compact data — and retries once after 5s if the forced refresh fails. Tested: bun test src/services/compact/ Tested: cd desktop && bun run test -- src/components/chat/ContextUsageIndicator.test.tsx src/stores/chatStore.test.ts Tested: cd desktop && bun run lint Confidence: high Scope-risk: narrow --- desktop/src/components/chat/ChatInput.tsx | 1 + .../chat/ContextUsageIndicator.test.tsx | 92 +++++++++++++ .../components/chat/ContextUsageIndicator.tsx | 55 ++++++-- desktop/src/stores/chatStore.test.ts | 13 ++ desktop/src/stores/chatStore.ts | 9 ++ src/services/compact/compact.test.ts | 126 ++++++++++++++++++ src/services/compact/compact.ts | 34 ++++- 7 files changed, 320 insertions(+), 10 deletions(-) create mode 100644 src/services/compact/compact.test.ts diff --git a/desktop/src/components/chat/ChatInput.tsx b/desktop/src/components/chat/ChatInput.tsx index f200db2a..54d0ce13 100644 --- a/desktop/src/components/chat/ChatInput.tsx +++ b/desktop/src/components/chat/ChatInput.tsx @@ -1276,6 +1276,7 @@ export function ChatInput({ variant = 'default', compact = false }: ChatInputPro runtimeSelectionKey={runtimeSelectionKey} fallbackModelLabel={runtimeModelLabel} compact={useCompactControls} + refreshNonce={sessionState?.compactCount ?? 0} /> )} {!isMemberSession && activeTabId && ( diff --git a/desktop/src/components/chat/ContextUsageIndicator.test.tsx b/desktop/src/components/chat/ContextUsageIndicator.test.tsx index cdc67efc..00295a28 100644 --- a/desktop/src/components/chat/ContextUsageIndicator.test.tsx +++ b/desktop/src/components/chat/ContextUsageIndicator.test.tsx @@ -265,4 +265,96 @@ describe('ContextUsageIndicator request behavior', () => { expect(sessionsApiMock.getInspection).toHaveBeenCalledTimes(1) expect(screen.queryByText('90%')).not.toBeInTheDocument() }) + + it('forces a fresh inspection when refreshNonce bumps after a compaction (#743)', async () => { + // First request hangs — simulates an auto refresh that started just + // before the compact boundary and would resolve with pre-compact data. + const first = deferred() + sessionsApiMock.getInspection + .mockImplementationOnce(() => first.promise) + .mockResolvedValueOnce({ + ...baseInspection, + context: { ...baseInspection.context, totalTokens: 9_000, percentage: 5 }, + }) + + const { rerender } = render( + , + ) + + await act(async () => { + await Promise.resolve() + }) + expect(sessionsApiMock.getInspection).toHaveBeenCalledTimes(1) + + rerender( + , + ) + + // The forced refresh bypasses both the auto-refresh throttle and the + // in-flight request reuse. + await waitFor(() => { + expect(sessionsApiMock.getInspection).toHaveBeenCalledTimes(2) + }) + await waitFor(() => { + expect(screen.getByTestId('context-usage-indicator')).toHaveTextContent('5%') + }) + }) + + it('retries the forced refresh once when it fails right after compaction (#743)', async () => { + vi.useFakeTimers() + try { + sessionsApiMock.getInspection + .mockResolvedValueOnce(baseInspection) + .mockRejectedValueOnce(new Error('Request timed out after 30s')) + .mockResolvedValueOnce({ + ...baseInspection, + context: { ...baseInspection.context, totalTokens: 9_000, percentage: 5 }, + }) + + const { rerender } = render( + , + ) + await act(async () => { + await vi.advanceTimersByTimeAsync(0) + }) + expect(sessionsApiMock.getInspection).toHaveBeenCalledTimes(1) + + rerender( + , + ) + await act(async () => { + await vi.advanceTimersByTimeAsync(0) + }) + expect(sessionsApiMock.getInspection).toHaveBeenCalledTimes(2) + + // The CLI was still busy and the forced refresh failed — one delayed + // retry recovers the meter instead of leaving the stale percentage. + await act(async () => { + await vi.advanceTimersByTimeAsync(5_000) + }) + expect(sessionsApiMock.getInspection).toHaveBeenCalledTimes(3) + } finally { + vi.useRealTimers() + } + }) }) diff --git a/desktop/src/components/chat/ContextUsageIndicator.tsx b/desktop/src/components/chat/ContextUsageIndicator.tsx index 489b71c4..f907094e 100644 --- a/desktop/src/components/chat/ContextUsageIndicator.tsx +++ b/desktop/src/components/chat/ContextUsageIndicator.tsx @@ -12,11 +12,21 @@ type Props = { fallbackModelLabel?: string draft?: boolean compact?: boolean + /** + * Bump to force an immediate refresh that bypasses the auto-refresh + * throttle and any in-flight (possibly pre-compact) request. Used after + * context compaction so the meter recovers right away (#743). + */ + refreshNonce?: number } const ACTIVE_REFRESH_MS = 30_000 const CONTEXT_REQUEST_TIMEOUT_MS = 20_000 const AUTO_REFRESH_MIN_INTERVAL_MS = 10_000 +// Right after a compaction the CLI may still be busy finishing the turn, so +// the forced refresh can time out — retry once instead of keeping the stale +// pre-compact percentage on screen. +const FORCED_REFRESH_RETRY_MS = 5_000 function formatNumber(value: number | undefined) { return new Intl.NumberFormat().format(value ?? 0) @@ -67,6 +77,7 @@ export function ContextUsageIndicator({ fallbackModelLabel, draft = false, compact = false, + refreshNonce = 0, }: Props) { const t = useTranslation() const [context, setContext] = useState(null) @@ -78,29 +89,31 @@ export function ContextUsageIndicator({ const [mobileDetailsOpen, setMobileDetailsOpen] = useState(false) const requestSeq = useRef(0) const contextIdentityRef = useRef('') - const inFlightRequestRef = useRef | null>(null) + const inFlightRequestRef = useRef | null>(null) const inFlightIdentityRef = useRef(null) const lastAutoRefreshAtRef = useRef(0) - const refresh = useCallback(async (mode: 'auto' | 'manual' = 'manual') => { + const refresh = useCallback(async (mode: 'auto' | 'manual' | 'force' = 'manual'): Promise => { if (!sessionId || draft) { setLoading(false) - return + return false } if (mode === 'auto' && !isDocumentVisible()) { setLoading(false) - return + return false } if (mode === 'auto' && Date.now() - lastAutoRefreshAtRef.current < AUTO_REFRESH_MIN_INTERVAL_MS) { - return inFlightRequestRef.current ?? undefined + return inFlightRequestRef.current ?? false } if (typeof sessionsApi.getInspection !== 'function') { setLoading(false) - return + return false } const activeSessionId = sessionId const activeContextIdentity = `${activeSessionId}:${runtimeSelectionKey}` - if (inFlightRequestRef.current && inFlightIdentityRef.current === activeContextIdentity) { + // 'force' must not reuse an in-flight request: one started just before a + // compact boundary would resolve with the pre-compact context. + if (mode !== 'force' && inFlightRequestRef.current && inFlightIdentityRef.current === activeContextIdentity) { return inFlightRequestRef.current } const seq = requestSeq.current + 1 @@ -114,7 +127,7 @@ export function ContextUsageIndicator({ timeout: CONTEXT_REQUEST_TIMEOUT_MS, }) .then((inspection) => { - if (seq !== requestSeq.current || activeContextIdentity !== contextIdentityRef.current) return + if (seq !== requestSeq.current || activeContextIdentity !== contextIdentityRef.current) return false const nextContext = inspection.context ?? inspection.contextEstimate ?? null const nextSource = inspection.context ? 'live' : inspection.contextEstimate ? 'estimate' : null const usageModel = inspection.usage?.models.find((model) => firstNonEmpty(model.displayName, model.model)) ?? null @@ -129,10 +142,12 @@ export function ContextUsageIndicator({ ) ?? null) setError(nextContext ? null : inspection.errors?.context ?? null) setUpdatedAt(Date.now()) + return nextContext !== null }) .catch((err) => { - if (seq !== requestSeq.current || activeContextIdentity !== contextIdentityRef.current) return + if (seq !== requestSeq.current || activeContextIdentity !== contextIdentityRef.current) return false setError(err instanceof Error ? err.message : String(err)) + return false }) .finally(() => { if (inFlightRequestRef.current === request) { @@ -146,6 +161,28 @@ export function ContextUsageIndicator({ return request }, [draft, runtimeSelectionKey, sessionId]) + // After a compaction the context shrinks server-side but nothing else + // re-reads it promptly (auto refreshes are throttled and stop once the + // session goes idle), leaving the pre-compact percentage on screen (#743). + // Force a fresh request, and retry once if the CLI was still busy. + const lastRefreshNonceRef = useRef(refreshNonce) + useEffect(() => { + if (refreshNonce === lastRefreshNonceRef.current) return + lastRefreshNonceRef.current = refreshNonce + let cancelled = false + let retryTimer: ReturnType | null = null + void refresh('force').then((ok) => { + if (ok || cancelled) return + retryTimer = setTimeout(() => { + void refresh('force') + }, FORCED_REFRESH_RETRY_MS) + }) + return () => { + cancelled = true + if (retryTimer) clearTimeout(retryTimer) + } + }, [refresh, refreshNonce]) + useEffect(() => { const contextIdentity = `${sessionId}:${runtimeSelectionKey}` const identityChanged = contextIdentityRef.current !== contextIdentity diff --git a/desktop/src/stores/chatStore.test.ts b/desktop/src/stores/chatStore.test.ts index 3b99de6f..60b31d25 100644 --- a/desktop/src/stores/chatStore.test.ts +++ b/desktop/src/stores/chatStore.test.ts @@ -2599,6 +2599,19 @@ describe('chatStore history mapping', () => { preTokens: 120000, }, ]) + // The context usage indicator watches this counter to force an + // immediate post-compact refresh (#743). The seeded session state above + // intentionally lacks compactCount (legacy persisted shape) — the bump + // must tolerate that. + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.compactCount).toBe(1) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'system_notification', + subtype: 'compact_boundary', + message: 'Context compacted', + data: { trigger: 'manual' }, + }) + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.compactCount).toBe(2) }) it('attaches compact summary content to the latest compact card', () => { diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index 9abf540a..77b728a4 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -88,6 +88,13 @@ export type PerSessionState = { request: ComputerUsePermissionRequest } | null tokenUsage: TokenUsage + /** + * Bumped each time a compact boundary arrives. The context usage indicator + * watches this to force an immediate re-read of the (now much smaller) + * context instead of waiting for the next API response (#743). + * Optional: legacy persisted sessions predate the field. + */ + compactCount?: number /** * Characters streamed by the assistant during the current turn (text, * thinking, tool input). ÷4 approximates output tokens for the streaming @@ -127,6 +134,7 @@ const DEFAULT_SESSION_STATE: PerSessionState = { pendingPermission: null, pendingComputerUsePermission: null, tokenUsage: { input_tokens: 0, output_tokens: 0 }, + compactCount: 0, streamingResponseChars: 0, elapsedSeconds: 0, statusVerb: '', @@ -1923,6 +1931,7 @@ export const useChatStore = create((set, get) => ({ update((session) => ({ chatState: session.chatState === 'compacting' ? 'thinking' : session.chatState, statusVerb: session.chatState === 'compacting' ? '' : session.statusVerb, + compactCount: (session.compactCount ?? 0) + 1, messages: appendOrUpdateTailCompactSummary( session.messages, { diff --git a/src/services/compact/compact.test.ts b/src/services/compact/compact.test.ts new file mode 100644 index 00000000..4576be77 --- /dev/null +++ b/src/services/compact/compact.test.ts @@ -0,0 +1,126 @@ +import { describe, expect, test } from 'bun:test' + +import { buildPostCompactMessages, type CompactionResult } from './compact.js' +import { getCurrentUsage } from '../../utils/tokens.js' +import type { Message } from '../../types/message.js' + +const PRE_COMPACT_USAGE = { + input_tokens: 150_000, + output_tokens: 900, + cache_creation_input_tokens: 2_000, + cache_read_input_tokens: 120_000, + service_tier: 'standard', +} + +function makeBoundary(): CompactionResult['boundaryMarker'] { + return { + type: 'system', + subtype: 'compact_boundary', + content: 'Conversation compacted', + isMeta: false, + timestamp: new Date().toISOString(), + uuid: '00000000-0000-0000-0000-000000000001', + level: 'info', + compactMetadata: { trigger: 'manual', preTokens: 150_000 }, + } as unknown as CompactionResult['boundaryMarker'] +} + +function makeSummaryMessage(): CompactionResult['summaryMessages'][number] { + return { + type: 'user', + uuid: '00000000-0000-0000-0000-000000000002', + timestamp: new Date().toISOString(), + isCompactSummary: true, + message: { role: 'user', content: 'This session is being continued…' }, + } as unknown as CompactionResult['summaryMessages'][number] +} + +function makePreservedAssistant(): Message { + return { + type: 'assistant', + uuid: '00000000-0000-0000-0000-000000000003', + timestamp: new Date().toISOString(), + message: { + id: 'msg_old', + role: 'assistant', + model: 'mock-model', + content: [{ type: 'text', text: 'old reply kept after compact' }], + stop_reason: 'end_turn', + usage: { ...PRE_COMPACT_USAGE }, + }, + } as unknown as Message +} + +function makePreservedUser(): Message { + return { + type: 'user', + uuid: '00000000-0000-0000-0000-000000000004', + timestamp: new Date().toISOString(), + message: { role: 'user', content: 'kept user message' }, + } as unknown as Message +} + +function makeResult(messagesToKeep?: Message[]): CompactionResult { + return { + boundaryMarker: makeBoundary(), + summaryMessages: [makeSummaryMessage()], + attachments: [], + hookResults: [], + ...(messagesToKeep ? { messagesToKeep } : {}), + } +} + +describe('buildPostCompactMessages stale-usage stripping (#743)', () => { + test('zeroes provider usage on preserved assistant messages', () => { + const kept = makePreservedAssistant() + const result = buildPostCompactMessages(makeResult([kept, makePreservedUser()])) + + const assistant = result.find(m => m.type === 'assistant') as { + message: { usage: Record } + } + expect(assistant.message.usage).toMatchObject({ + input_tokens: 0, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + }) + // Non-token usage metadata survives the strip. + expect(assistant.message.usage.service_tier).toBe('standard') + }) + + test('does not mutate the original preserved message', () => { + const kept = makePreservedAssistant() + buildPostCompactMessages(makeResult([kept])) + + expect( + (kept as unknown as { message: { usage: { input_tokens: number } } }) + .message.usage.input_tokens, + ).toBe(150_000) + }) + + test('keeps ordering and passes non-assistant messages through untouched', () => { + const keptUser = makePreservedUser() + const result = buildPostCompactMessages(makeResult([keptUser])) + + expect(result[0]?.type).toBe('system') + expect(result[1]?.type).toBe('user') + expect(result[2]).toBe(keptUser) + }) + + test('post-compact view no longer anchors getCurrentUsage to the pre-compact request size', () => { + const result = buildPostCompactMessages( + makeResult([makePreservedUser(), makePreservedAssistant()]), + ) + + // getCurrentUsage skips zeroed usage (its stale placeholder convention), + // so the context meter falls back to the local estimate and recovers + // immediately instead of staying pinned at the pre-compact 100%. + expect(getCurrentUsage(result)).toBeNull() + }) + + test('handles results without messagesToKeep', () => { + const result = buildPostCompactMessages(makeResult()) + expect(result).toHaveLength(2) + expect(getCurrentUsage(result)).toBeNull() + }) +}) diff --git a/src/services/compact/compact.ts b/src/services/compact/compact.ts index f8f86eaf..fad233e4 100644 --- a/src/services/compact/compact.ts +++ b/src/services/compact/compact.ts @@ -322,6 +322,38 @@ export type RecompactionInfo = { querySource?: QuerySource } +/** + * Strip provider-reported usage off preserved assistant messages. After a + * compaction the last pre-compact API usage no longer describes the active + * context, but getCurrentUsage() walks backwards for the newest non-zero + * usage and calculateCurrentContextTokenTotal() takes + * max(local estimate, API usage) — so a stale anchor pins the context meter + * at the pre-compact level until the next API response arrives (#743). + * Zeroed usage is the established "stale/unavailable" placeholder that + * getCurrentUsage() skips. Copies, never mutates: the originals stay intact + * for transcript/cost accounting. + */ +function stripStaleUsageFromPreservedMessages(messages: Message[]): Message[] { + return messages.map(message => { + if (message.type !== 'assistant') return message + const usage = message.message?.usage + if (!usage) return message + return { + ...message, + message: { + ...message.message, + usage: { + ...usage, + input_tokens: 0, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + }, + }, + } + }) +} + /** * Build the base post-compact messages array from a CompactionResult. * This ensures consistent ordering across all compaction paths. @@ -331,7 +363,7 @@ export function buildPostCompactMessages(result: CompactionResult): Message[] { return [ result.boundaryMarker, ...result.summaryMessages, - ...(result.messagesToKeep ?? []), + ...stripStaleUsageFromPreservedMessages(result.messagesToKeep ?? []), ...result.attachments, ...result.hookResults, ]