From 9dd329d4848a1a1e7545880de823258b05896f2f Mon Sep 17 00:00:00 2001 From: Omar Baradei Date: Tue, 9 Jun 2026 16:07:30 -0700 Subject: [PATCH 1/2] fix(context): live mid-turn context estimates + compression status in desktop statusbar The context bar froze during long tool batches: tool.start/complete events carried context_compressor.last_prompt_tokens, a snapshot from the previous API call that never moves while tool results pile up mid-turn. - agent/context_compressor.py: track last_prompt_messages_len alongside last_prompt_tokens; new live_context_tokens(messages) prices the tail appended since the snapshot on top of the provider-exact base, falls back to max(base, rough) and prefers the rough estimate right after compression - agent/conversation_loop.py: record the snapshot at every last_prompt_tokens write (real usage + both preflight seeds); emit a 'Compression complete' lifecycle status after preflight compaction - agent/tool_executor.py: extract the four copy-pasted usage blocks into _context_usage_for_tool_events() which uses the live estimate - apps/desktop: handle status.update gateway events (previously dropped on the floor) in a new $sessionActivityStatus atom, cleared by the next stream activity; statusbar shows the transient status (auto-compression progress finally visible mid-session); GatewayEventPayload gains 'kind' - cli.py: stale standalone-CLI compression threshold default 0.50 -> 0.85, matching hermes_cli/config.py - tests: 9 new Python tests (live estimate math, snapshot bookkeeping, tool-event payloads), 3 new desktop tests (status.update lifecycle) Co-Authored-By: Claude Fable 5 --- agent/context_compressor.py | 49 ++++++- agent/conversation_loop.py | 27 +++- agent/tool_executor.py | 93 +++++-------- .../session/hooks/use-message-stream.test.tsx | 65 ++++++++- .../app/session/hooks/use-message-stream.ts | 24 ++++ .../app/shell/hooks/use-statusbar-items.tsx | 12 ++ apps/desktop/src/lib/chat-messages.ts | 3 + apps/desktop/src/store/session.ts | 11 ++ cli.py | 2 +- tests/agent/test_live_context_estimate.py | 130 ++++++++++++++++++ 10 files changed, 348 insertions(+), 68 deletions(-) create mode 100644 tests/agent/test_live_context_estimate.py diff --git a/agent/context_compressor.py b/agent/context_compressor.py index 8f07768982b3..8508555b0827 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -650,6 +650,10 @@ def __init__( self.last_prompt_tokens = 0 self.last_completion_tokens = 0 self.last_real_prompt_tokens = 0 + # How many conversation messages ``last_prompt_tokens`` accounts for. + # Lets ``live_context_tokens`` price only the tail appended since the + # last API call instead of freezing or re-estimating the whole list. + self.last_prompt_messages_len: Optional[int] = None self.last_compression_rough_tokens = 0 self.last_rough_tokens_when_real_prompt_fit = 0 self.awaiting_real_usage_after_compression = False @@ -681,9 +685,16 @@ def __init__( self._last_aux_model_failure_error: Optional[str] = None self._last_aux_model_failure_model: Optional[str] = None - def update_from_response(self, usage: Dict[str, Any]): - """Update tracked token usage from API response.""" + def update_from_response(self, usage: Dict[str, Any], messages_len: Optional[int] = None): + """Update tracked token usage from API response. + + ``messages_len`` is the length of the conversation list the prompt was + built from, recorded so mid-turn estimates know where the + provider-counted prefix ends. + """ self.last_prompt_tokens = usage.get("prompt_tokens", 0) + if messages_len is not None: + self.last_prompt_messages_len = messages_len self.last_completion_tokens = usage.get("completion_tokens", 0) self.last_total_tokens = usage.get("total_tokens", self.last_prompt_tokens + self.last_completion_tokens) if self.last_prompt_tokens > 0: @@ -695,6 +706,36 @@ def update_from_response(self, usage: Dict[str, Any]): self.last_rough_tokens_when_real_prompt_fit = 0 self.awaiting_real_usage_after_compression = False + def live_context_tokens(self, messages: List[Dict[str, Any]]) -> int: + """Best-effort size of the context *right now*, mid-turn. + + ``last_prompt_tokens`` is exact but frozen at the last API call (or + preflight estimate); assistant text and tool results appended since + aren't in it, so consumers that poll it mid-turn — the desktop context + bar during a long tool batch — appear stuck. Price the tail appended + since the snapshot on top of the known base instead. + """ + base = self.last_prompt_tokens or 0 + snap = self.last_prompt_messages_len + if base > 0 and snap is not None and 0 <= snap <= len(messages): + tail = 0 + if snap < len(messages): + try: + tail = estimate_messages_tokens_rough(messages[snap:]) + except Exception: + tail = 0 + return base + tail + try: + rough = estimate_messages_tokens_rough(messages) + except Exception: + rough = 0 + if self.awaiting_real_usage_after_compression: + # Right after compression the stale pre-compression base would + # overstate a freshly shrunk conversation; trust the rough + # estimate of the compressed list until real usage arrives. + return rough or base + return max(base, rough) + def should_defer_preflight_to_real_usage(self, rough_tokens: int) -> bool: """Return True when a high rough preflight estimate is known-noisy. @@ -2075,4 +2116,8 @@ def compress(self, messages: List[Dict[str, Any]], current_tokens: int = None, f ) logger.info("Compression #%d complete", self.compression_count) + # The conversation list is being replaced wholesale; the old + # messages-length snapshot no longer maps onto the new list. + self.last_prompt_messages_len = None + return compressed diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index d0bf1702ae18..4feabf0ad83a 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -117,7 +117,9 @@ def _ra(): return run_agent -def _emit_preflight_token_usage(agent: Any, request_tokens: int) -> None: +def _emit_preflight_token_usage( + agent: Any, request_tokens: int, messages_len: int | None = None +) -> None: """Send the current request-size estimate through the live usage channel.""" if request_tokens <= 0: return @@ -129,6 +131,10 @@ def _emit_preflight_token_usage(agent: Any, request_tokens: int) -> None: previous = getattr(compressor, "last_prompt_tokens", 0) or 0 if request_tokens > previous: compressor.last_prompt_tokens = request_tokens + if messages_len is not None: + # Record which conversation prefix the seeded estimate covers + # so live_context_tokens prices only the tail appended later. + compressor.last_prompt_messages_len = messages_len except Exception: logger.debug("could not update preflight context estimate", exc_info=True) @@ -678,6 +684,7 @@ def run_conversation( # would re-introduce the very desync we're avoiding. if _preflight_tokens > (_compressor.last_prompt_tokens or 0): _compressor.last_prompt_tokens = _preflight_tokens + _compressor.last_prompt_messages_len = len(messages) if _preflight_deferred: logger.info( @@ -700,6 +707,7 @@ def run_conversation( f">= {_compressor.threshold_tokens:,} threshold. " "This may take a moment." ) + _pre_compress_preflight_tokens = _preflight_tokens # May need multiple passes for very large sessions with small # context windows (each pass summarises the middle N turns). for _pass in range(3): @@ -739,6 +747,11 @@ def run_conversation( _preflight_tokens = _post_preflight_tokens if not _compressor.should_compress(_preflight_tokens): break # Under threshold or anti-thrash guard stopped it + if _preflight_tokens < _pre_compress_preflight_tokens: + agent._emit_status( + f"📦 Compression complete: ~{_pre_compress_preflight_tokens:,} " + f"→ ~{_preflight_tokens:,} tokens." + ) # Plugin hook: pre_llm_call # Fired once per turn before the tool-calling loop. Plugins can @@ -1143,7 +1156,9 @@ def run_conversation( approx_request_tokens = estimate_request_tokens_rough( api_messages, tools=agent.tools or None ) - _emit_preflight_token_usage(agent, approx_request_tokens) + _emit_preflight_token_usage( + agent, approx_request_tokens, messages_len=len(messages) + ) _runtime_context_error = _ollama_context_limit_error( agent, approx_request_tokens @@ -1995,7 +2010,9 @@ def _perform_api_call(next_api_kwargs): "cache_write_tokens": canonical_usage.cache_write_tokens, "reasoning_tokens": canonical_usage.reasoning_tokens, } - agent.context_compressor.update_from_response(usage_dict) + agent.context_compressor.update_from_response( + usage_dict, messages_len=len(messages) + ) # Cache discovered context length after successful call. # Only persist limits confirmed by the provider (parsed @@ -3278,7 +3295,9 @@ def _perform_api_call(next_api_kwargs): f"🗜️ Compressed request estimate " f"~{pre_compress_request_tokens:,} → ~{post_compress_request_tokens:,} tokens, retrying..." ) - _emit_preflight_token_usage(agent, post_compress_request_tokens) + _emit_preflight_token_usage( + agent, post_compress_request_tokens, messages_len=len(messages) + ) time.sleep(2) # Brief pause between compression retries restart_with_compressed_messages = True break diff --git a/agent/tool_executor.py b/agent/tool_executor.py index ecee62c564fe..d6914cfbe7b1 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -240,6 +240,35 @@ def _execute(next_args: dict) -> Any: return result, observed_args +def _context_usage_for_tool_events(agent, messages: list) -> dict: + """Build the usage payload attached to tool.start/complete events. + + Context numbers come from ``ContextCompressor.live_context_tokens`` so the + UI context bar keeps moving while tool results pile up mid-turn instead of + freezing at the last API call's prompt size. + """ + usage = { + "model": getattr(agent, "model", ""), + "input": getattr(agent, "session_input_tokens", 0) or 0, + "output": getattr(agent, "session_output_tokens", 0) or 0, + "total": getattr(agent, "session_total_tokens", 0) or 0, + "calls": getattr(agent, "session_api_calls", 0) or 0, + } + comp = getattr(agent, "context_compressor", None) + if comp: + try: + ctx_used = comp.live_context_tokens(messages) + except Exception: + ctx_used = getattr(comp, "last_prompt_tokens", 0) or 0 + ctx_used = ctx_used or usage["total"] or 0 + ctx_max = getattr(comp, "context_length", 0) or 0 + if ctx_max: + usage["context_used"] = ctx_used + usage["context_max"] = ctx_max + usage["context_percent"] = max(0, min(100, round(ctx_used / ctx_max * 100))) + return usage + + def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effective_task_id: str, api_call_count: int = 0) -> None: """Execute multiple tool calls concurrently using a thread pool. @@ -440,21 +469,7 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe # Inject context usage into tool events so frontend can update # the context bar in real-time during tool execution. - usage = { - "model": getattr(agent, "model", ""), - "input": getattr(agent, "session_input_tokens", 0) or 0, - "output": getattr(agent, "session_output_tokens", 0) or 0, - "total": getattr(agent, "session_total_tokens", 0) or 0, - "calls": getattr(agent, "session_api_calls", 0) or 0, - } - comp = getattr(agent, "context_compressor", None) - if comp: - ctx_used = getattr(comp, "last_prompt_tokens", 0) or usage["total"] or 0 - ctx_max = getattr(comp, "context_length", 0) or 0 - if ctx_max: - usage["context_used"] = ctx_used - usage["context_max"] = ctx_max - usage["context_percent"] = max(0, min(100, round(ctx_used / ctx_max * 100))) + usage = _context_usage_for_tool_events(agent, messages) for tc, name, args, middleware_trace, block_result, blocked_by_guardrail in parsed_calls: if block_result is not None: continue @@ -732,21 +747,7 @@ def _run_tool(index, tool_call, function_name, function_args, middleware_trace): agent._touch_activity(f"tool completed: {name} ({tool_duration:.1f}s)") # Inject context usage into tool.complete events for real-time updates. - _comp = getattr(agent, "context_compressor", None) - _usage = { - "model": getattr(agent, "model", ""), - "input": getattr(agent, "session_input_tokens", 0) or 0, - "output": getattr(agent, "session_output_tokens", 0) or 0, - "total": getattr(agent, "session_total_tokens", 0) or 0, - "calls": getattr(agent, "session_api_calls", 0) or 0, - } - if _comp: - _ctx_used = getattr(_comp, "last_prompt_tokens", 0) or _usage["total"] or 0 - _ctx_max = getattr(_comp, "context_length", 0) or 0 - if _ctx_max: - _usage["context_used"] = _ctx_used - _usage["context_max"] = _ctx_max - _usage["context_percent"] = max(0, min(100, round(_ctx_used / _ctx_max * 100))) + _usage = _context_usage_for_tool_events(agent, messages) if not blocked and agent.tool_complete_callback: try: @@ -932,21 +933,7 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe if not _execution_blocked and agent.tool_start_callback: # Inject context usage for real-time context bar updates. - _t_usage = { - "model": getattr(agent, "model", ""), - "input": getattr(agent, "session_input_tokens", 0) or 0, - "output": getattr(agent, "session_output_tokens", 0) or 0, - "total": getattr(agent, "session_total_tokens", 0) or 0, - "calls": getattr(agent, "session_api_calls", 0) or 0, - } - _t_comp = getattr(agent, "context_compressor", None) - if _t_comp: - _t_ctx_used = getattr(_t_comp, "last_prompt_tokens", 0) or _t_usage["total"] or 0 - _t_ctx_max = getattr(_t_comp, "context_length", 0) or 0 - if _t_ctx_max: - _t_usage["context_used"] = _t_ctx_used - _t_usage["context_max"] = _t_ctx_max - _t_usage["context_percent"] = max(0, min(100, round(_t_ctx_used / _t_ctx_max * 100))) + _t_usage = _context_usage_for_tool_events(agent, messages) try: agent.tool_start_callback(tool_call.id, function_name, function_args, usage=_t_usage) except Exception as cb_err: @@ -1385,21 +1372,7 @@ def _execute(next_args: dict) -> Any: if not _execution_blocked and agent.tool_complete_callback: # Inject context usage for real-time context bar updates. - _sc_usage = { - "model": getattr(agent, "model", ""), - "input": getattr(agent, "session_input_tokens", 0) or 0, - "output": getattr(agent, "session_output_tokens", 0) or 0, - "total": getattr(agent, "session_total_tokens", 0) or 0, - "calls": getattr(agent, "session_api_calls", 0) or 0, - } - _sc_comp = getattr(agent, "context_compressor", None) - if _sc_comp: - _sc_ctx_used = getattr(_sc_comp, "last_prompt_tokens", 0) or _sc_usage["total"] or 0 - _sc_ctx_max = getattr(_sc_comp, "context_length", 0) or 0 - if _sc_ctx_max: - _sc_usage["context_used"] = _sc_ctx_used - _sc_usage["context_max"] = _sc_ctx_max - _sc_usage["context_percent"] = max(0, min(100, round(_sc_ctx_used / _sc_ctx_max * 100))) + _sc_usage = _context_usage_for_tool_events(agent, messages) try: agent.tool_complete_callback(tool_call.id, function_name, function_args, function_result, usage=_sc_usage) except Exception as cb_err: diff --git a/apps/desktop/src/app/session/hooks/use-message-stream.test.tsx b/apps/desktop/src/app/session/hooks/use-message-stream.test.tsx index c1c51c2f885b..241c9ada1509 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream.test.tsx +++ b/apps/desktop/src/app/session/hooks/use-message-stream.test.tsx @@ -5,7 +5,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import type { ClientSessionState } from '@/app/types' import { createClientSessionState } from '@/lib/chat-runtime' -import { $currentUsage } from '@/store/session' +import { $currentUsage, $sessionActivityStatus } from '@/store/session' import type { RpcEvent } from '@/types/hermes' import { useMessageStream } from './use-message-stream' @@ -107,3 +107,66 @@ describe('useMessageStream token usage events', () => { expect($currentUsage.get()).toEqual({ calls: 1, input: 10, output: 5, total: 15 }) }) }) + +describe('useMessageStream status.update events', () => { + beforeEach(() => { + handleEvent = () => undefined + $sessionActivityStatus.set(null) + }) + + afterEach(() => { + cleanup() + $sessionActivityStatus.set(null) + vi.restoreAllMocks() + }) + + const statusEvent = (text: string, kind = 'lifecycle', sessionId = 'session-1') => + ({ + payload: { kind, text }, + session_id: sessionId, + type: 'status.update' + }) as RpcEvent + + it('surfaces lifecycle statuses for the active session', () => { + render() + + act(() => handleEvent(statusEvent('📦 Preflight compression: ~90,000 tokens. This may take a moment.'))) + + expect($sessionActivityStatus.get()).toEqual({ + kind: 'lifecycle', + text: '📦 Preflight compression: ~90,000 tokens. This may take a moment.' + }) + }) + + it('ignores statuses from inactive sessions and unknown kinds', () => { + render() + + act(() => handleEvent(statusEvent('background noise', 'lifecycle', 'session-2'))) + expect($sessionActivityStatus.get()).toBeNull() + + act(() => handleEvent(statusEvent('voice things', 'voice'))) + expect($sessionActivityStatus.get()).toBeNull() + }) + + it('clears on ready kind and on stream activity', () => { + render() + + act(() => handleEvent(statusEvent('⠋ compressing 120 messages', 'compressing'))) + expect($sessionActivityStatus.get()).not.toBeNull() + + act(() => handleEvent(statusEvent('ready', 'ready'))) + expect($sessionActivityStatus.get()).toBeNull() + + act(() => handleEvent(statusEvent('📦 Compression complete: ~90,000 → ~30,000 tokens.'))) + expect($sessionActivityStatus.get()).not.toBeNull() + + act(() => + handleEvent({ + payload: { text: 'hello' }, + session_id: 'session-1', + type: 'message.delta' + } as RpcEvent) + ) + expect($sessionActivityStatus.get()).toBeNull() + }) +}) diff --git a/apps/desktop/src/app/session/hooks/use-message-stream.ts b/apps/desktop/src/app/session/hooks/use-message-stream.ts index b25217d3d1b4..e3ac5b01169b 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream.ts @@ -32,6 +32,7 @@ import { setCurrentReasoningEffort, setCurrentServiceTier, setCurrentUsage, + setSessionActivityStatus, setTurnStartedAt, setYoloActive } from '@/store/session' @@ -726,6 +727,20 @@ export function useMessageStream({ if (!explicitSid || isActiveEvent) { setCurrentUsage(current => mergeTokenUsagePayload(current, event.payload as TokenUsagePayload | undefined)) } + } else if (event.type === 'status.update') { + // Gateway lifecycle statuses (auto-compression progress, background + // process notices). Statusbar is an active-session surface; the next + // stream activity (deltas / tool events) clears it. + if (!explicitSid || isActiveEvent) { + const kind = typeof payload?.kind === 'string' ? payload.kind : '' + const text = typeof payload?.text === 'string' ? payload.text.trim() : '' + + if (kind === 'ready') { + setSessionActivityStatus(null) + } else if (text && ['lifecycle', 'compressing', 'process', 'status'].includes(kind)) { + setSessionActivityStatus({ kind, text }) + } + } } else if (event.type === 'message.start') { if (!sessionId) { return @@ -754,6 +769,10 @@ export function useMessageStream({ if (sessionId) { appendAssistantDelta(sessionId, coerceGatewayText(payload?.text)) } + + if (isActiveEvent) { + setSessionActivityStatus(null) + } } else if (event.type === 'thinking.delta') { // thinking.delta carries the kawaii spinner status (face + verb from // KawaiiSpinner), not real reasoning. The bottom-of-thread loading @@ -789,6 +808,7 @@ export function useMessageStream({ if (isActiveEvent) { setTurnStartedAt(null) + setSessionActivityStatus(null) } if (payload?.usage) { @@ -805,6 +825,10 @@ export function useMessageStream({ if (payload?.usage && (!explicitSid || isActiveEvent)) { setCurrentUsage(current => ({ ...current, ...payload.usage })) } + + if (isActiveEvent) { + setSessionActivityStatus(null) + } } else if (event.type === 'tool.complete') { if (sessionId) { flushQueuedDeltas(sessionId) 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 e77ca3b99147..0fad1ea65af8 100644 --- a/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx +++ b/apps/desktop/src/app/shell/hooks/use-statusbar-items.tsx @@ -32,6 +32,7 @@ import { $currentProvider, $currentReasoningEffort, $currentUsage, + $sessionActivityStatus, $sessionStartedAt, $turnStartedAt, $workingSessionIds, @@ -91,6 +92,7 @@ export function useStatusbarItems({ const currentUsage = useStore($currentUsage) const desktopActionTasks = useStore($desktopActionTasks) const previewServerRestartStatus = useStore($previewServerRestartStatus) + const sessionActivityStatus = useStore($sessionActivityStatus) const sessionStartedAt = useStore($sessionStartedAt) const turnStartedAt = useStore($turnStartedAt) const workingSessionIds = useStore($workingSessionIds) @@ -302,6 +304,15 @@ export function useStatusbarItems({ const coreRightStatusbarItems = useMemo( () => [ + { + className: 'max-w-72', + hidden: !sessionActivityStatus, + icon: , + id: 'activity-status', + label: {sessionActivityStatus?.text ?? ''}, + title: sessionActivityStatus?.text, + variant: 'text' + }, { detail: , hidden: !busy || !turnStartedAt, @@ -383,6 +394,7 @@ export function useStatusbarItems({ currentProvider, currentReasoningEffort, modelMenuContent, + sessionActivityStatus, sessionStartedAt, showYoloToggle, toggleYolo, diff --git a/apps/desktop/src/lib/chat-messages.ts b/apps/desktop/src/lib/chat-messages.ts index c6c9cee48d8d..cde21e5fcd15 100644 --- a/apps/desktop/src/lib/chat-messages.ts +++ b/apps/desktop/src/lib/chat-messages.ts @@ -61,6 +61,9 @@ export type GatewayEventPayload = { // secret.request (skill credential capture) env_var?: string prompt?: string + // status.update (gateway lifecycle statuses: compression progress, + // background-process notices) + kind?: string } export function textPart(text: string): ChatMessagePart { diff --git a/apps/desktop/src/store/session.ts b/apps/desktop/src/store/session.ts index 7fb616be7113..30ee8ce6ca8b 100644 --- a/apps/desktop/src/store/session.ts +++ b/apps/desktop/src/store/session.ts @@ -137,6 +137,15 @@ export const $availablePersonalities = atom([]) export const $introSeed = atom(0) export const $contextSuggestions = atom([]) export const $modelPickerOpen = atom(false) +// Transient gateway lifecycle status for the ACTIVE session (auto-compression +// progress, background-process notices). Mirrors `status.update` events and is +// cleared by the next stream activity, so it only shows while nothing else +// (deltas, tool events) is moving. +export interface SessionActivityStatus { + kind: string + text: string +} +export const $sessionActivityStatus = atom(null) export const setConnection = (next: Updater) => updateAtom($connection, next) export const setGatewayState = (next: Updater) => updateAtom($gatewayState, next) @@ -181,6 +190,8 @@ export const setAvailablePersonalities = (next: Updater) => updateAtom export const setIntroSeed = (next: Updater) => updateAtom($introSeed, next) export const setContextSuggestions = (next: Updater) => updateAtom($contextSuggestions, next) export const setModelPickerOpen = (next: Updater) => updateAtom($modelPickerOpen, next) +export const setSessionActivityStatus = (next: Updater) => + updateAtom($sessionActivityStatus, next) // Watchdog tracking — when does a "working" session count as stuck? // Long-running tool calls (LLM inference, long shell commands, web fetches) diff --git a/cli.py b/cli.py index b62c27402b5e..a01c70dd44a7 100644 --- a/cli.py +++ b/cli.py @@ -415,7 +415,7 @@ def load_cli_config() -> Dict[str, Any]: }, "compression": { "enabled": True, # Auto-compress when approaching context limit - "threshold": 0.50, # Compress at 50% of model's context limit + "threshold": 0.85, # Compress at 85% of model's context limit }, "agent": { "max_turns": 90, # Default max tool-calling iterations (shared with subagents) diff --git a/tests/agent/test_live_context_estimate.py b/tests/agent/test_live_context_estimate.py new file mode 100644 index 000000000000..4895e36ea7bb --- /dev/null +++ b/tests/agent/test_live_context_estimate.py @@ -0,0 +1,130 @@ +"""Tests for mid-turn live context estimates. + +The desktop/TUI context bar reads ``context_used``/``context_percent`` off +tool.start/tool.complete event payloads. Before ``live_context_tokens`` the +value was a frozen snapshot of the previous API call's ``prompt_tokens``, so +the bar appeared stuck through long tool batches. These tests pin the new +behavior: provider-exact base plus a rough estimate of only the messages +appended since that snapshot. +""" + +import pytest +from unittest.mock import patch + +from agent.context_compressor import ContextCompressor +from agent.model_metadata import estimate_messages_tokens_rough +from agent.tool_executor import _context_usage_for_tool_events + + +@pytest.fixture() +def compressor(): + with patch("agent.context_compressor.get_model_context_length", return_value=100000): + return ContextCompressor( + model="test/model", + threshold_percent=0.85, + protect_first_n=2, + protect_last_n=2, + quiet_mode=True, + ) + + +BASE_MESSAGES = [ + {"role": "user", "content": "hello " * 50}, + {"role": "assistant", "content": "hi there " * 30}, +] + + +class TestLiveContextTokens: + def test_no_tail_returns_exact_base(self, compressor): + compressor.update_from_response( + {"prompt_tokens": 5000}, messages_len=len(BASE_MESSAGES) + ) + assert compressor.live_context_tokens(list(BASE_MESSAGES)) == 5000 + + def test_prices_tail_appended_since_snapshot(self, compressor): + compressor.update_from_response( + {"prompt_tokens": 5000}, messages_len=len(BASE_MESSAGES) + ) + tool_result = {"role": "tool", "content": "x" * 4000} + msgs = list(BASE_MESSAGES) + [tool_result] + expected_tail = estimate_messages_tokens_rough([tool_result]) + assert expected_tail > 0 + assert compressor.live_context_tokens(msgs) == 5000 + expected_tail + + def test_grows_monotonically_as_results_append(self, compressor): + compressor.update_from_response( + {"prompt_tokens": 7000}, messages_len=len(BASE_MESSAGES) + ) + msgs = list(BASE_MESSAGES) + seen = [compressor.live_context_tokens(msgs)] + for _ in range(3): + msgs.append({"role": "tool", "content": "result " * 500}) + seen.append(compressor.live_context_tokens(msgs)) + assert seen == sorted(seen) + assert seen[-1] > seen[0] + + def test_fallback_without_snapshot_takes_max(self, compressor): + msgs = [{"role": "tool", "content": "y" * 8000}] + rough = estimate_messages_tokens_rough(msgs) + + compressor.last_prompt_tokens = 50 # tiny base, no snapshot recorded + assert compressor.live_context_tokens(msgs) == rough + + compressor.last_prompt_tokens = 10_000_000 + assert compressor.live_context_tokens(msgs) == 10_000_000 + + def test_invalid_snapshot_after_compression_prefers_rough(self, compressor): + compressor.update_from_response({"prompt_tokens": 90_000}, messages_len=10) + compressor.awaiting_real_usage_after_compression = True + small = [{"role": "user", "content": "compressed summary"}] + # Snapshot (10) no longer maps onto the shrunk list; the stale 90k + # base must not win over the fresh rough estimate. + assert compressor.live_context_tokens(small) == estimate_messages_tokens_rough(small) + + def test_update_from_response_records_snapshot(self, compressor): + compressor.update_from_response({"prompt_tokens": 123}, messages_len=7) + assert compressor.last_prompt_messages_len == 7 + # Callers that don't know the list length leave the snapshot alone. + compressor.update_from_response({"prompt_tokens": 456}) + assert compressor.last_prompt_messages_len == 7 + + +class _FakeAgent: + model = "test/model" + session_input_tokens = 10 + session_output_tokens = 20 + session_total_tokens = 30 + session_api_calls = 2 + context_compressor = None + + +class TestToolEventUsage: + def test_context_fields_use_live_estimate(self, compressor): + agent = _FakeAgent() + compressor.update_from_response({"prompt_tokens": 4000}, messages_len=1) + agent.context_compressor = compressor + + msgs = [{"role": "user", "content": "q"}] + first = _context_usage_for_tool_events(agent, msgs) + assert first["context_used"] == 4000 + assert first["context_max"] == 100000 + assert first["context_percent"] == 4 + + msgs.append({"role": "tool", "content": "z" * 40_000}) + second = _context_usage_for_tool_events(agent, msgs) + assert second["context_used"] > first["context_used"] + assert second["context_percent"] == round(second["context_used"] / 100000 * 100) + + def test_no_compressor_omits_context_fields(self): + usage = _context_usage_for_tool_events(_FakeAgent(), []) + assert "context_used" not in usage + assert usage["model"] == "test/model" + assert usage["total"] == 30 + + def test_session_totals_passthrough(self, compressor): + agent = _FakeAgent() + agent.context_compressor = compressor + usage = _context_usage_for_tool_events(agent, list(BASE_MESSAGES)) + assert usage["input"] == 10 + assert usage["output"] == 20 + assert usage["calls"] == 2 From 647b1fbc209094656265a461c4332f54f0af102e Mon Sep 17 00:00:00 2001 From: Omar Baradei Date: Tue, 9 Jun 2026 16:26:37 -0700 Subject: [PATCH 2/2] =?UTF-8?q?test:=20finish=20the=200.50=E2=86=920.85=20?= =?UTF-8?q?threshold-default=20sweep;=20map=20omar@kostudios.io=20in=20AUT?= =?UTF-8?q?HOR=5FMAP?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Commit 4351acb2e raised the default compression threshold to 0.85 but left the two constructor-default tests asserting 0.50, red on main since June 8. Update them to the live default and add a dedicated small-context case so the 64K floor stays exercised. Also map omar@kostudios.io -> OmarB97 so check-attribution passes on direct (non-merge) commits. Co-Authored-By: Claude Fable 5 --- scripts/release.py | 1 + tests/agent/test_context_compressor.py | 21 ++++++++++++++------- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/scripts/release.py b/scripts/release.py index 6d2762c1435c..92d654ea46f7 100755 --- a/scripts/release.py +++ b/scripts/release.py @@ -45,6 +45,7 @@ # Auto-extracted from noreply emails + manual overrides AUTHOR_MAP = { + "omar@kostudios.io": "OmarB97", "philipadsouza@gmail.com": "PhilipAD", "zhuhaoyu0909@icloud.com": "underthestars-zhy", "raysun12142006@gmail.com": "yanxue06", diff --git a/tests/agent/test_context_compressor.py b/tests/agent/test_context_compressor.py index 5ce753864c90..4201d5f69b60 100644 --- a/tests/agent/test_context_compressor.py +++ b/tests/agent/test_context_compressor.py @@ -1576,20 +1576,27 @@ def test_ratio_clamped(self): c = ContextCompressor(model="test", quiet_mode=True, summary_target_ratio=0.95) assert c.summary_target_ratio == 0.80 - def test_default_threshold_is_50_percent(self): - """Default compression threshold should be 50%, with a 64K floor.""" + def test_default_threshold_is_85_percent(self): + """Default compression threshold should be 85%.""" with patch("agent.context_compressor.get_model_context_length", return_value=100_000): c = ContextCompressor(model="test", quiet_mode=True) - assert c.threshold_percent == 0.50 - # 50% of 100K = 50K, but the floor is 64K + assert c.threshold_percent == 0.85 + # 85% of 100K = 85K, comfortably above the 64K floor + assert c.threshold_tokens == 85_000 + + def test_threshold_floor_applies_on_small_contexts(self): + """The 64K floor wins when the percentage lands below it.""" + with patch("agent.context_compressor.get_model_context_length", return_value=70_000): + c = ContextCompressor(model="test", quiet_mode=True) + # 85% of 70K = 59.5K, raised to the 64K floor assert c.threshold_tokens == 64_000 def test_threshold_floor_does_not_apply_above_128k(self): - """On large-context models the 50% percentage is used directly.""" + """On large-context models the 85% percentage is used directly.""" with patch("agent.context_compressor.get_model_context_length", return_value=200_000): c = ContextCompressor(model="test", quiet_mode=True) - # 50% of 200K = 100K, which is above the 64K floor - assert c.threshold_tokens == 100_000 + # 85% of 200K = 170K, far above the 64K floor + assert c.threshold_tokens == 170_000 def test_default_protect_last_n_is_20(self): """Default protect_last_n should be 20."""