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, ]