diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 4c105d4906e6..33aa074ecbd9 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -57,6 +57,19 @@ COMPACTION_STATUS = ( f"🗜️ {COMPACTION_STATUS_MARKER} — summarizing earlier conversation so I can continue..." ) +COMPACTION_DONE_STATUS = "✓ Context compaction complete — continuing turn..." + + +def _emit_compaction_done(agent: Any) -> None: + """Emit the structured terminal edge for a started compaction.""" + status_callback = getattr(agent, "status_callback", None) + if not status_callback: + return + try: + status_callback("compacted", COMPACTION_DONE_STATUS) + except Exception: + logger.debug("status_callback error in compaction completion", exc_info=True) + def _builtin_memory_prompt_snapshot(agent: Any) -> Optional[Tuple[str, str]]: """Return the built-in memory text that can affect a system prompt. @@ -1025,6 +1038,14 @@ def compress_context( focus_topic, ) agent._emit_status(COMPACTION_STATUS) + _compaction_done_emitted = False + + def _complete_compaction_lifecycle() -> None: + nonlocal _compaction_done_emitted + if _compaction_done_emitted: + return + _compaction_done_emitted = True + _emit_compaction_done(agent) # ── Compression lock ──────────────────────────────────────────────── # Atomic, state.db-backed lock per session_id. Without this, two @@ -1174,12 +1195,14 @@ def compress_context( split_status="aborted", failure_class="lock_contended", ) + _complete_compaction_lifecycle() return messages, _existing_sp _lock_released = False def _release_lock() -> None: """Release the lock keyed on the OLD session_id (before rotation).""" nonlocal _lock_released + _complete_compaction_lifecycle() if _lock_released: return _lock_released = True @@ -1876,6 +1899,15 @@ def _compress_context_via_codex_app_server( except Exception: pass + _compaction_done_emitted = False + + def _complete_compaction_lifecycle() -> None: + nonlocal _compaction_done_emitted + if _compaction_done_emitted: + return + _compaction_done_emitted = True + _emit_compaction_done(agent) + _activity_heartbeat: Optional[_CompressionActivityHeartbeat] = None try: _activity_heartbeat = _CompressionActivityHeartbeat(agent).start() @@ -1883,6 +1915,7 @@ def _compress_context_via_codex_app_server( except BaseException: if _activity_heartbeat is not None: _activity_heartbeat.stop("context compression failed") + _complete_compaction_lifecycle() raise if getattr(result, "interrupted", False) or getattr(result, "error", None): @@ -1907,6 +1940,7 @@ def _compress_context_via_codex_app_server( existing_prompt = getattr(agent, "_cached_system_prompt", None) if not existing_prompt: existing_prompt = agent._build_system_prompt(system_message) + _complete_compaction_lifecycle() return messages, existing_prompt try: @@ -1946,6 +1980,7 @@ def _compress_context_via_codex_app_server( existing_prompt = getattr(agent, "_cached_system_prompt", None) if not existing_prompt: existing_prompt = agent._build_system_prompt(system_message) + _complete_compaction_lifecycle() return messages, existing_prompt @@ -2206,6 +2241,7 @@ def try_shrink_image_parts_in_messages( __all__ = [ "COMPACTION_STATUS", + "COMPACTION_DONE_STATUS", "COMPACTION_STATUS_MARKER", "check_compression_model_feasibility", "replay_compression_warning", diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/compaction-event.test.tsx b/apps/desktop/src/app/session/hooks/use-message-stream/compaction-event.test.tsx index a403bde23800..9fcabadcd1ee 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream/compaction-event.test.tsx +++ b/apps/desktop/src/app/session/hooks/use-message-stream/compaction-event.test.tsx @@ -79,4 +79,14 @@ describe('useMessageStream compaction lifecycle', () => { expect($compactingSessions.get()).toEqual({ [OTHER_SID]: true }) }) + + it('clears the compaction phase on the structured completion edge', async () => { + await mountStream() + setSessionCompacting(OTHER_SID, true) + + emit('status.update', { kind: 'compacting' }) + emit('status.update', { kind: 'compacted' }) + + expect($compactingSessions.get()).toEqual({ [OTHER_SID]: true }) + }) }) diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts index 5b8fccd8bb00..173e356958de 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts @@ -763,6 +763,9 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { if (sessionId && payload?.kind === 'compacting') { setSessionCompacting(sessionId, true) compactedTurnRef.current.add(sessionId) + } else if (sessionId && payload?.kind === 'compacted') { + setSessionCompacting(sessionId, false) + compactedTurnRef.current.delete(sessionId) } else if (sessionId && payload?.kind === 'process') { // The gateway's notification poller announces background process // completions / watch matches here — re-sync the status stack. diff --git a/apps/desktop/src/app/shell/context-usage-panel.test.tsx b/apps/desktop/src/app/shell/context-usage-panel.test.tsx new file mode 100644 index 000000000000..28d72e861a9e --- /dev/null +++ b/apps/desktop/src/app/shell/context-usage-panel.test.tsx @@ -0,0 +1,93 @@ +import { act, cleanup, render, waitFor } from '@testing-library/react' +import { useState } from 'react' +import { afterEach, describe, expect, it, vi } from 'vitest' + +import type { ContextBreakdown, UsageStats } from '@/types/hermes' + +import { ContextUsagePanel } from './context-usage-panel' + +const initialUsage: UsageStats = { + calls: 1, + context_max: 272_000, + context_percent: 47, + context_used: 128_200, + input: 0, + output: 0, + total: 0 +} + +const breakdown: ContextBreakdown = { + categories: [{ color: 'teal', id: 'conversation', label: 'Conversation', tokens: 241_400 }], + context_max: 272_000, + context_percent: 89, + context_used: 241_400, + estimated_total: 286_600, + model: 'test-model' +} + +afterEach(() => { + cleanup() + vi.restoreAllMocks() +}) + +describe('ContextUsagePanel', () => { + it('publishes once without refetching when publication recreates the callback', async () => { + const requestGateway = vi.fn().mockResolvedValue(breakdown) + const published = vi.fn() + const renderedUsage: UsageStats[] = [] + + function Harness() { + const [currentUsage, setCurrentUsage] = useState(initialUsage) + renderedUsage.push(currentUsage) + + return ( + { + published(snapshot) + setCurrentUsage(current => ({ ...current, ...snapshot })) + }} + requestGateway={requestGateway} + sessionId="runtime-1" + /> + ) + } + + render() + + await waitFor(() => { + expect(published).toHaveBeenCalledWith({ + context_max: 272_000, + context_percent: 89, + context_used: 241_400 + }) + expect(renderedUsage.at(-1)?.context_used).toBe(241_400) + }) + await act(async () => {}) + + expect(requestGateway).toHaveBeenCalledTimes(1) + expect(requestGateway).toHaveBeenCalledWith('session.context_breakdown', { session_id: 'runtime-1' }) + }) + + it('refetches when the session or gateway requester changes', async () => { + const firstGateway = vi.fn().mockResolvedValue(breakdown) + const secondGateway = vi.fn().mockResolvedValue(breakdown) + + const { rerender } = render( + + ) + + await waitFor(() => expect(firstGateway).toHaveBeenCalledTimes(1)) + + rerender() + + await waitFor(() => { + expect(firstGateway).toHaveBeenCalledTimes(2) + expect(firstGateway).toHaveBeenLastCalledWith('session.context_breakdown', { session_id: 'runtime-2' }) + }) + + rerender() + + await waitFor(() => expect(secondGateway).toHaveBeenCalledTimes(1)) + }) +}) diff --git a/apps/desktop/src/app/shell/context-usage-panel.tsx b/apps/desktop/src/app/shell/context-usage-panel.tsx index 5a243c0913f7..03ddd21a50cd 100644 --- a/apps/desktop/src/app/shell/context-usage-panel.tsx +++ b/apps/desktop/src/app/shell/context-usage-panel.tsx @@ -1,4 +1,4 @@ -import { useEffect, useMemo, useState } from 'react' +import { useEffect, useMemo, useRef, useState } from 'react' import { useI18n } from '@/i18n' import { compactNumber } from '@/lib/format' @@ -7,15 +7,18 @@ import type { ContextBreakdown, ContextUsageCategory, UsageStats } from '@/types interface ContextUsagePanelProps { currentUsage: UsageStats + onUsageSnapshot?: (usage: Pick) => void requestGateway: (method: string, params?: Record) => Promise sessionId: string | null } -export function ContextUsagePanel({ currentUsage, requestGateway, sessionId }: ContextUsagePanelProps) { +export function ContextUsagePanel({ currentUsage, onUsageSnapshot, requestGateway, sessionId }: ContextUsagePanelProps) { const { t } = useI18n() const copy = t.shell.statusbar.contextUsagePanel const [breakdown, setBreakdown] = useState(null) const [loading, setLoading] = useState(false) + const onUsageSnapshotRef = useRef(onUsageSnapshot) + onUsageSnapshotRef.current = onUsageSnapshot useEffect(() => { if (!sessionId) { @@ -32,6 +35,11 @@ export function ContextUsagePanel({ currentUsage, requestGateway, sessionId }: C .then(data => { if (!cancelled) { setBreakdown(data) + onUsageSnapshotRef.current?.({ + context_max: data.context_max, + context_percent: data.context_percent, + context_used: data.context_used + }) } }) .catch(() => { diff --git a/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx b/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx index f66e40c4aafe..62e469e8a082 100644 --- a/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx +++ b/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx @@ -1,5 +1,5 @@ import { useStore } from '@nanostores/react' -import { useMemo } from 'react' +import { useCallback, useMemo } from 'react' import type { CommandCenterSection } from '@/app/command-center' import { $terminalTakeover, setTerminalTakeover } from '@/app/right-sidebar/store' @@ -27,7 +27,8 @@ import { $sessions, $sessionStartedAt, $turnStartedAt, - sessionMatchesStoredId + sessionMatchesStoredId, + setCurrentUsage } from '@/store/session' import { $focusedRuntimeId, $focusedSessionState, $focusedStoredSessionId } from '@/store/session-states' import { $subagentsBySession, activeSubagentCount, failedSubagentCount } from '@/store/subagents' @@ -40,7 +41,7 @@ import { $updateStatus, openUpdateOverlayFor } from '@/store/updates' -import type { StatusResponse } from '@/types/hermes' +import type { StatusResponse, UsageStats } from '@/types/hermes' import { CRON_ROUTE, SETTINGS_ROUTE } from '../../routes' import type { StatusbarItem } from '../statusbar-controls' @@ -144,6 +145,14 @@ export function useStatusbarItems({ const contextUsage = useMemo(() => usageContextLabel(currentUsage), [currentUsage]) const contextBar = useMemo(() => contextBarLabel(currentUsage), [currentUsage]) + + const publishContextUsage = useCallback( + (snapshot: Pick) => { + setCurrentUsage(current => ({ ...current, ...snapshot })) + }, + [] + ) + const approvalModeItem = useApprovalModeStatusbarItem(activeGatewayProfile, requestGateway) const gatewayMenuContent = useMemo( @@ -455,7 +464,12 @@ export function useStatusbarItems({ menuAlign: 'end', menuClassName: 'w-auto border-(--ui-stroke-secondary) p-0', menuContent: ( - + ), title: copy.openContextUsage, variant: 'menu' @@ -496,6 +510,7 @@ export function useStatusbarItems({ contextUsage, copy, currentUsage, + publishContextUsage, requestGateway, sessionStartedAt, gatewayState, diff --git a/tests/run_agent/test_413_compression.py b/tests/run_agent/test_413_compression.py index d7c2b72e35b1..11c8cecc19fa 100644 --- a/tests/run_agent/test_413_compression.py +++ b/tests/run_agent/test_413_compression.py @@ -16,6 +16,7 @@ from unittest.mock import MagicMock, patch from agent.context_compressor import SUMMARY_PREFIX +from agent.conversation_compression import COMPACTION_DONE_STATUS, COMPACTION_STATUS from run_agent import AIAgent import run_agent @@ -599,9 +600,50 @@ class TestPreflightCompression: ] assert new_system_prompt == "You are helpful." build_prompt.assert_not_called() - assert events[0][0] == "lifecycle" - assert "Compacting context" in events[0][1] - assert events[1] == ("compress", "started") + assert events == [ + ("lifecycle", COMPACTION_STATUS), + ("compress", "started"), + ("compacted", COMPACTION_DONE_STATUS), + ] + + def test_compress_context_emits_one_terminal_status_when_lock_is_unavailable(self, agent): + """A rejected lock must retire the started desktop compaction phase.""" + agent.compression_enabled = False + agent.session_id = "session-with-contended-lock" + agent._session_db = SimpleNamespace( + get_compression_lock_holder=lambda _session_id: "other-agent", + try_acquire_compression_lock=lambda *_args, **_kwargs: False, + ) + events = [] + agent.status_callback = lambda event, message: events.append((event, message)) + messages = [{"role": "user", "content": "hello"}] + + compressed, prompt = agent._compress_context(messages, "system prompt", force=True) + + assert compressed is messages + assert prompt == "You are helpful." + assert [event for event, _ in events] == ["lifecycle", "warn", "compacted"] + assert events[-1] == ("compacted", COMPACTION_DONE_STATUS) + + def test_compress_context_emits_one_terminal_status_after_an_abort(self, agent): + """An aborted summary must retire the started desktop compaction phase.""" + agent.compression_enabled = False + events = [] + agent.status_callback = lambda event, message: events.append((event, message)) + messages = [{"role": "user", "content": "hello"}] + + def _abort_compression(current_messages, **_kwargs): + agent.context_compressor._last_compress_aborted = True + agent.context_compressor._last_summary_error = "auxiliary model unavailable" + return current_messages + + with patch.object(agent.context_compressor, "compress", side_effect=_abort_compression): + compressed, prompt = agent._compress_context(messages, "system prompt", force=True) + + assert compressed is messages + assert prompt == "You are helpful." + assert [event for event, _ in events] == ["lifecycle", "warn", "compacted"] + assert events[-1] == ("compacted", COMPACTION_DONE_STATUS) def test_compression_reuses_cached_prompt_when_memory_snapshot_is_unchanged(self, agent): """A memory reload without new injected text must keep the cache prefix.""" diff --git a/tests/run_agent/test_codex_app_server_compaction.py b/tests/run_agent/test_codex_app_server_compaction.py index 98f8cb8059a4..bf14a88e78fe 100644 --- a/tests/run_agent/test_codex_app_server_compaction.py +++ b/tests/run_agent/test_codex_app_server_compaction.py @@ -4,7 +4,7 @@ from types import SimpleNamespace import pytest from agent.codex_runtime import _record_codex_app_server_compaction -from agent.conversation_compression import COMPACTION_STATUS, compress_context +from agent.conversation_compression import COMPACTION_DONE_STATUS, COMPACTION_STATUS, compress_context from agent.transports.codex_app_server_session import TurnResult @@ -63,6 +63,8 @@ class DummyAgent: awaiting_real_usage_after_compression=False, ) self.statuses = [] + self.status_events = [] + self.status_callback = lambda kind, text: self.status_events.append((kind, text)) self.warnings = [] self.events = [] self.built_prompts = [] @@ -74,9 +76,11 @@ class DummyAgent: def _emit_status(self, message): self.statuses.append(message) + self.status_callback("lifecycle", message) def _emit_warning(self, message): self.warnings.append(message) + self.status_callback("warn", message) def _build_system_prompt(self, system_message): self.built_prompts.append(system_message) @@ -156,6 +160,10 @@ def test_codex_app_server_manual_compression_routes_to_codex_thread(): assert prompt == "cached prompt" assert agent._codex_session.calls == 1 assert agent.context_compressor.compression_count == 1 + assert agent.status_events == [ + ("lifecycle", COMPACTION_STATUS), + ("compacted", COMPACTION_DONE_STATUS), + ] assert agent.context_compressor.last_compression_rough_tokens == 100000 # This minimal fake compressor does not implement update_from_response(), # so the runtime preserves its existing pending-usage bookkeeping here. @@ -219,6 +227,11 @@ def test_codex_app_server_compression_failure_preserves_bookkeeping(): assert agent.warnings assert agent.touch_calls[0] == "context compression started" assert agent.touch_calls[-1] == "context compression failed" + assert agent.status_events == [ + ("lifecycle", COMPACTION_STATUS), + ("warn", "⚠ Codex app-server compaction failed: compact failed"), + ("compacted", COMPACTION_DONE_STATUS), + ] def test_codex_app_server_native_compaction_notice_emits_status_and_event(): diff --git a/tests/tui_gateway/test_compaction_status.py b/tests/tui_gateway/test_compaction_status.py index 0e98bde18547..4f3468a2261e 100644 --- a/tests/tui_gateway/test_compaction_status.py +++ b/tests/tui_gateway/test_compaction_status.py @@ -48,6 +48,14 @@ def test_compaction_lifecycle_is_retagged(server, monkeypatch): assert events == [{"kind": "compacting", "text": COMPACTION_STATUS}] +def test_compaction_completion_status_is_preserved(server, monkeypatch): + from agent.conversation_compression import COMPACTION_DONE_STATUS + + events = _capture(server, monkeypatch) + server._status_update("sid", "compacted", COMPACTION_DONE_STATUS) + + assert events == [{"kind": "compacted", "text": COMPACTION_DONE_STATUS}] + def test_other_lifecycle_status_stays_lifecycle(server, monkeypatch): events = _capture(server, monkeypatch) server._status_update("sid", "lifecycle", "❌ Rate limited after 5 retries")