From 39720ff74054748bd006ff3e234b469aa6e1e35c Mon Sep 17 00:00:00 2001 From: Johannes du Plessis Date: Mon, 10 Aug 2026 15:18:17 -0700 Subject: [PATCH 1/4] feat(code): add `/context` usage report --- libs/code/COMMANDS.md | 3 +- libs/code/deepagents_code/app.py | 29 +++++- libs/code/deepagents_code/command_registry.py | 6 ++ .../tui/widgets/context_usage.py | 89 +++++++++++++++++++ .../tui/widgets/startup_tip.py | 1 + .../tui/widgets/test_context_usage.py | 19 ++++ 6 files changed, 142 insertions(+), 5 deletions(-) create mode 100644 libs/code/deepagents_code/tui/widgets/context_usage.py create mode 100644 libs/code/tests/unit_tests/tui/widgets/test_context_usage.py diff --git a/libs/code/COMMANDS.md b/libs/code/COMMANDS.md index e08e68c694e..7465902b689 100644 --- a/libs/code/COMMANDS.md +++ b/libs/code/COMMANDS.md @@ -8,7 +8,7 @@ Regenerate this file with `make commands-catalog` after changing command names, aliases, descriptions, visibility, or hidden-command metadata. -## Public (38) +## Public (39) | Command | Aliases | Description | | --- | --- | --- | @@ -18,6 +18,7 @@ aliases, descriptions, visibility, or hidden-command metadata. | `/auto-update` | | Turn automatic updates on or off | | `/changelog` | | Open the changelog in a browser | | `/clear` | | Clear the chat and start a new thread | +| `/context` | | Show current context window usage | | `/copy` | | Copy the latest assistant message to clipboard | | `/cost` | | Show estimated thread cost | | `/docs` | | Open the docs | diff --git a/libs/code/deepagents_code/app.py b/libs/code/deepagents_code/app.py index 51c971ef9b3..de2b146854d 100644 --- a/libs/code/deepagents_code/app.py +++ b/libs/code/deepagents_code/app.py @@ -109,6 +109,7 @@ ) from deepagents_code.tui.widgets._links import open_url_async from deepagents_code.tui.widgets.chat_input import ChatInput +from deepagents_code.tui.widgets.context_usage import build_context_usage_markdown from deepagents_code.tui.widgets.goal_status import GoalStatusPanel from deepagents_code.tui.widgets.loading import LoadingWidget from deepagents_code.tui.widgets.message_store import ( @@ -14109,6 +14110,22 @@ async def _handle_command(self, command: str) -> None: timeout=5, markup=False, ) + elif cmd == "/context": + await self._mount_message(UserMessage(command)) + conversation_tokens = await self._get_conversation_token_count() + model_spec = self._effective_model_spec() or settings.model_name + await self._mount_message( + AppMessage( + build_context_usage_markdown( + context_tokens=self._context_tokens, + conversation_tokens=conversation_tokens, + context_limit=settings.model_context_limit, + model_spec=model_spec, + approximate=self._tokens_approximate, + ), + markdown=True, + ) + ) elif cmd == "/tokens": await self._mount_message(UserMessage(command)) if self._context_tokens > 0: @@ -14824,10 +14841,14 @@ async def _get_conversation_token_count(self) -> int | None: if not state or not state.values: return None messages = state.values.get("messages", []) - if not messages: + if not isinstance(messages, list) or not messages: return None - return count_tokens_approximately(messages) - except Exception: # best-effort for /tokens display + effective = _effective_conversation( + messages, + state.values.get("_summarization_event"), + ) + return count_tokens_approximately(effective) + except Exception: # best-effort for context-usage displays logger.debug("Failed to retrieve conversation token count", exc_info=True) return None @@ -15051,7 +15072,7 @@ async def _handle_offload(self) -> None: ) ) - self._on_tokens_update(tokens_after) + self._on_tokens_update(tokens_after, approximate=True) except Exception as exc: # surface offload errors to user logger.exception("Offload failed") diff --git a/libs/code/deepagents_code/command_registry.py b/libs/code/deepagents_code/command_registry.py index 5fc61fb7dfd..e2e462c6991 100644 --- a/libs/code/deepagents_code/command_registry.py +++ b/libs/code/deepagents_code/command_registry.py @@ -127,6 +127,12 @@ def to_entry(self) -> CommandEntry: description="Copy the latest assistant message to clipboard", bypass_tier=BypassTier.SIDE_EFFECT_FREE, ), + SlashCommand( + name="/context", + description="Show current context window usage", + bypass_tier=BypassTier.QUEUED, + hidden_keywords="tokens window usage remaining offload compact", + ), SlashCommand( name="/cost", description="Show estimated thread cost", diff --git a/libs/code/deepagents_code/tui/widgets/context_usage.py b/libs/code/deepagents_code/tui/widgets/context_usage.py new file mode 100644 index 00000000000..6490bcc8899 --- /dev/null +++ b/libs/code/deepagents_code/tui/widgets/context_usage.py @@ -0,0 +1,89 @@ +"""Presentation for the `/context` slash command.""" + +from __future__ import annotations + +from deepagents_code._markdown import escape_markdown +from deepagents_code._session_stats import format_token_count + + +def build_context_usage_markdown( + *, + context_tokens: int, + conversation_tokens: int | None, + context_limit: int | None, + model_spec: str | None, + approximate: bool, +) -> str: + """Build a context report that distinguishes measured and estimated usage. + + Args: + context_tokens: Latest total for the active model context. + conversation_tokens: Approximate tokens in the effective message history. + context_limit: Active model's input-token limit, when known. + model_spec: Active model identifier. + approximate: Whether `context_tokens` is stale or locally estimated. + + Returns: + Markdown for the context-window report. + """ + if context_limit is not None and context_limit <= 0: + context_limit = None + lines = ["## Context usage"] + if model_spec: + lines.append(escape_markdown(model_spec)) + + if context_tokens <= 0: + lines.extend(("", "No usage reported yet.")) + if context_limit is not None: + lines.append( + f"**Context window:** {format_token_count(context_limit)} tokens" + ) + else: + lines.append("Context window limit unavailable for this model.") + else: + count_prefix = "~" if approximate else "" + current_tokens = format_token_count(context_tokens) + lines.append("") + if context_limit is not None: + percent = context_tokens / context_limit * 100 + remaining = max(0, context_limit - context_tokens) + limit = format_token_count(context_limit) + current = f"**Current:** {count_prefix}{current_tokens} / {limit} tokens" + lines.extend( + ( + f"{current} ({percent:.1f}%)", + f"**Remaining:** {format_token_count(remaining)} tokens", + ) + ) + else: + lines.extend( + ( + f"**Current:** {count_prefix}{current_tokens} tokens", + "Context window limit unavailable for this model.", + ) + ) + + if conversation_tokens is not None: + conversation = min(max(0, conversation_tokens), context_tokens) + fixed = context_tokens - conversation + lines.extend(("", "### Estimated composition")) + for label, tokens in ( + ("System prompt + tools", fixed), + ("Conversation", conversation), + ): + percent_suffix = ( + f" ({tokens / context_limit * 100:.1f}%)" if context_limit else "" + ) + count = format_token_count(tokens) + lines.append(f"- **{label}:** ~{count} tokens{percent_suffix}") + + lines.extend( + ("", "**Automatic offload:** enabled; use `/offload` to compact sooner.") + ) + if context_tokens > 0: + lines.append( + "*Current total is approximate; composition is estimated.*" + if approximate + else "*Current total is provider-reported; composition is estimated.*" + ) + return "\n".join(lines) diff --git a/libs/code/deepagents_code/tui/widgets/startup_tip.py b/libs/code/deepagents_code/tui/widgets/startup_tip.py index 0d0aebe7977..68ff19fba82 100644 --- a/libs/code/deepagents_code/tui/widgets/startup_tip.py +++ b/libs/code/deepagents_code/tui/widgets/startup_tip.py @@ -20,6 +20,7 @@ "Use @ to reference files and / for commands": 3, "Try /threads to resume a previous conversation": 2, "Use /offload to summarize older messages and free up the context window": 2, + "Use /context to see context window usage and remaining space": 1, "Use /copy to copy the latest message": 3, "Use /cost to see a breakdown of estimated spend": 1, "Use /tools to list the tools available to the agent": 1, diff --git a/libs/code/tests/unit_tests/tui/widgets/test_context_usage.py b/libs/code/tests/unit_tests/tui/widgets/test_context_usage.py new file mode 100644 index 00000000000..6cf7b7f541d --- /dev/null +++ b/libs/code/tests/unit_tests/tui/widgets/test_context_usage.py @@ -0,0 +1,19 @@ +"""Tests for the context-usage report.""" + +from deepagents_code.tui.widgets.context_usage import build_context_usage_markdown + + +def test_reports_capacity_and_estimated_breakdown() -> None: + content = build_context_usage_markdown( + context_tokens=20_000, + conversation_tokens=15_000, + context_limit=100_000, + model_spec="anthropic:claude-sonnet", + approximate=False, + ) + + assert "20.0K / 100.0K tokens (20.0%)" in content + assert "**Remaining:** 80.0K tokens" in content + assert "**System prompt + tools:** ~5.0K tokens (5.0%)" in content + assert "**Conversation:** ~15.0K tokens (15.0%)" in content + assert "**Automatic offload:** enabled" in content From 65aa051f890afde3e257313e53303bcfffe79e3c Mon Sep 17 00:00:00 2001 From: Johannes du Plessis Date: Mon, 10 Aug 2026 15:48:02 -0700 Subject: [PATCH 2/4] fix(code): prefer reported context usage --- libs/code/deepagents_code/app.py | 60 +++++++++++++------ .../tui/widgets/context_usage.py | 14 +++-- libs/code/tests/unit_tests/test_app.py | 26 ++++++++ 3 files changed, 79 insertions(+), 21 deletions(-) diff --git a/libs/code/deepagents_code/app.py b/libs/code/deepagents_code/app.py index de2b146854d..6b1a69067b1 100644 --- a/libs/code/deepagents_code/app.py +++ b/libs/code/deepagents_code/app.py @@ -14112,16 +14112,32 @@ async def _handle_command(self, command: str) -> None: ) elif cmd == "/context": await self._mount_message(UserMessage(command)) - conversation_tokens = await self._get_conversation_token_count() + ( + reported_tokens, + conversation_tokens, + ) = await self._get_context_usage_counts() + if reported_tokens is not None: + context_tokens = reported_tokens + approximate = False + elif self._context_tokens > 0 and not self._tokens_approximate: + context_tokens = self._context_tokens + approximate = False + elif conversation_tokens is not None or self._context_tokens > 0: + context_tokens = None + conversation_tokens = conversation_tokens or self._context_tokens + approximate = True + else: + context_tokens = 0 + approximate = False model_spec = self._effective_model_spec() or settings.model_name await self._mount_message( AppMessage( build_context_usage_markdown( - context_tokens=self._context_tokens, + context_tokens=context_tokens, conversation_tokens=conversation_tokens, context_limit=settings.model_context_limit, model_spec=model_spec, - approximate=self._tokens_approximate, + approximate=approximate, ), markdown=True, ) @@ -14821,36 +14837,46 @@ async def _has_conversation_messages(self) -> bool: ) return True - async def _get_conversation_token_count(self) -> int | None: - """Return the approximate conversation-only token count. + async def _get_context_usage_counts(self) -> tuple[int | None, int | None]: + """Read provider-reported total and estimated conversation usage together. Returns: - Token count as an integer, or `None` if state is unavailable. + Pair of provider-reported context tokens and approximate effective + conversation tokens. Either value is `None` when unavailable. """ if not self._agent: - return None + return None, None try: - from langchain_core.messages.utils import ( - count_tokens_approximately, - ) + from langchain_core.messages.utils import count_tokens_approximately config: RunnableConfig = { "configurable": {"thread_id": self._lc_thread_id}, } state = await self._agent.aget_state(config) if not state or not state.values: - return None - messages = state.values.get("messages", []) + return None, None + values = dict(state.values) + reported = _persisted_context_tokens(values) or None + messages = values.get("messages", []) if not isinstance(messages, list) or not messages: - return None + return reported, None effective = _effective_conversation( messages, - state.values.get("_summarization_event"), + values.get("_summarization_event"), ) - return count_tokens_approximately(effective) + return reported, count_tokens_approximately(effective) except Exception: # best-effort for context-usage displays - logger.debug("Failed to retrieve conversation token count", exc_info=True) - return None + logger.debug("Failed to retrieve context usage", exc_info=True) + return None, None + + async def _get_conversation_token_count(self) -> int | None: + """Return the approximate conversation-only token count. + + Returns: + Token count as an integer, or `None` if state is unavailable. + """ + _, conversation = await self._get_context_usage_counts() + return conversation async def _handle_offload(self) -> None: """Offload older messages to free context window space. diff --git a/libs/code/deepagents_code/tui/widgets/context_usage.py b/libs/code/deepagents_code/tui/widgets/context_usage.py index 6490bcc8899..ba01f964459 100644 --- a/libs/code/deepagents_code/tui/widgets/context_usage.py +++ b/libs/code/deepagents_code/tui/widgets/context_usage.py @@ -8,7 +8,7 @@ def build_context_usage_markdown( *, - context_tokens: int, + context_tokens: int | None, conversation_tokens: int | None, context_limit: int | None, model_spec: str | None, @@ -17,7 +17,8 @@ def build_context_usage_markdown( """Build a context report that distinguishes measured and estimated usage. Args: - context_tokens: Latest total for the active model context. + context_tokens: Latest total for the active model context, or `None` when + no reliable total is available. conversation_tokens: Approximate tokens in the effective message history. context_limit: Active model's input-token limit, when known. model_spec: Active model identifier. @@ -32,7 +33,12 @@ def build_context_usage_markdown( if model_spec: lines.append(escape_markdown(model_spec)) - if context_tokens <= 0: + if context_tokens is None: + lines.extend(("", "Current total unavailable.")) + if conversation_tokens is not None: + estimate = format_token_count(max(0, conversation_tokens)) + lines.append(f"**Conversation estimate:** ~{estimate} tokens") + elif context_tokens <= 0: lines.extend(("", "No usage reported yet.")) if context_limit is not None: lines.append( @@ -80,7 +86,7 @@ def build_context_usage_markdown( lines.extend( ("", "**Automatic offload:** enabled; use `/offload` to compact sooner.") ) - if context_tokens > 0: + if context_tokens is not None and context_tokens > 0: lines.append( "*Current total is approximate; composition is estimated.*" if approximate diff --git a/libs/code/tests/unit_tests/test_app.py b/libs/code/tests/unit_tests/test_app.py index a92864bb5af..1a87dab0928 100644 --- a/libs/code/tests/unit_tests/test_app.py +++ b/libs/code/tests/unit_tests/test_app.py @@ -179,6 +179,32 @@ def test_strips_provider_prefix( assert _display_model_label(spec) == expected +class TestContextCommand: + """Tests for the `/context` command's data source.""" + + async def test_prefers_checkpoint_total_after_offload( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + app = DeepAgentsApp() + app._context_tokens = 200 + app._tokens_approximate = True + get_counts = AsyncMock(return_value=(1_000, 200)) + mount_message = AsyncMock() + monkeypatch.setattr(app, "_get_context_usage_counts", get_counts) + monkeypatch.setattr(app, "_mount_message", mount_message) + + with patch("deepagents_code.config.settings") as mock_settings: + mock_settings.model_provider = "anthropic" + mock_settings.model_name = "claude-sonnet" + mock_settings.model_context_limit = 2_000 + await app._handle_command("/context") + + report = mount_message.await_args_list[-1].args[0] + assert isinstance(report, AppMessage) + assert "**Current:** 1.0K / 2.0K tokens (50.0%)" in str(report._content) + assert "**System prompt + tools:** ~800 tokens (40.0%)" in str(report._content) + + class TestWhatsNewMessage: """Tests for the post-upgrade banner content.""" From f19d325f8870bb8999a954177090b6fbd97f52e8 Mon Sep 17 00:00:00 2001 From: Johannes du Plessis Date: Mon, 10 Aug 2026 16:23:32 -0700 Subject: [PATCH 3/4] feat(code): render `/context` as color-coded modal --- libs/code/deepagents_code/app.py | 27 ++- .../tui/widgets/context_usage.py | 95 ---------- .../tui/widgets/context_usage/__init__.py | 5 + .../tui/widgets/context_usage/_models.py | 93 ++++++++++ .../tui/widgets/context_usage/_screen.py | 103 +++++++++++ .../tui/widgets/context_usage/_widgets.py | 167 ++++++++++++++++++ libs/code/tests/unit_tests/test_app.py | 21 ++- .../tui/widgets/test_context_usage.py | 19 -- 8 files changed, 395 insertions(+), 135 deletions(-) delete mode 100644 libs/code/deepagents_code/tui/widgets/context_usage.py create mode 100644 libs/code/deepagents_code/tui/widgets/context_usage/__init__.py create mode 100644 libs/code/deepagents_code/tui/widgets/context_usage/_models.py create mode 100644 libs/code/deepagents_code/tui/widgets/context_usage/_screen.py create mode 100644 libs/code/deepagents_code/tui/widgets/context_usage/_widgets.py delete mode 100644 libs/code/tests/unit_tests/tui/widgets/test_context_usage.py diff --git a/libs/code/deepagents_code/app.py b/libs/code/deepagents_code/app.py index 6b1a69067b1..dc15a16a1fa 100644 --- a/libs/code/deepagents_code/app.py +++ b/libs/code/deepagents_code/app.py @@ -109,7 +109,7 @@ ) from deepagents_code.tui.widgets._links import open_url_async from deepagents_code.tui.widgets.chat_input import ChatInput -from deepagents_code.tui.widgets.context_usage import build_context_usage_markdown +from deepagents_code.tui.widgets.context_usage import ContextUsageScreen from deepagents_code.tui.widgets.goal_status import GoalStatusPanel from deepagents_code.tui.widgets.loading import LoadingWidget from deepagents_code.tui.widgets.message_store import ( @@ -14111,7 +14111,6 @@ async def _handle_command(self, command: str) -> None: markup=False, ) elif cmd == "/context": - await self._mount_message(UserMessage(command)) ( reported_tokens, conversation_tokens, @@ -14129,19 +14128,19 @@ async def _handle_command(self, command: str) -> None: else: context_tokens = 0 approximate = False - model_spec = self._effective_model_spec() or settings.model_name - await self._mount_message( - AppMessage( - build_context_usage_markdown( - context_tokens=context_tokens, - conversation_tokens=conversation_tokens, - context_limit=settings.model_context_limit, - model_spec=model_spec, - approximate=approximate, - ), - markdown=True, - ) + screen = ContextUsageScreen( + context_tokens=context_tokens, + conversation_tokens=conversation_tokens, + context_limit=settings.model_context_limit, + model_spec=self._effective_model_spec() or settings.model_name, + approximate=approximate, ) + + def handle_result(_result: None) -> None: + if self._chat_input: + self._chat_input.focus_input() + + self.push_screen(screen, handle_result) elif cmd == "/tokens": await self._mount_message(UserMessage(command)) if self._context_tokens > 0: diff --git a/libs/code/deepagents_code/tui/widgets/context_usage.py b/libs/code/deepagents_code/tui/widgets/context_usage.py deleted file mode 100644 index ba01f964459..00000000000 --- a/libs/code/deepagents_code/tui/widgets/context_usage.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Presentation for the `/context` slash command.""" - -from __future__ import annotations - -from deepagents_code._markdown import escape_markdown -from deepagents_code._session_stats import format_token_count - - -def build_context_usage_markdown( - *, - context_tokens: int | None, - conversation_tokens: int | None, - context_limit: int | None, - model_spec: str | None, - approximate: bool, -) -> str: - """Build a context report that distinguishes measured and estimated usage. - - Args: - context_tokens: Latest total for the active model context, or `None` when - no reliable total is available. - conversation_tokens: Approximate tokens in the effective message history. - context_limit: Active model's input-token limit, when known. - model_spec: Active model identifier. - approximate: Whether `context_tokens` is stale or locally estimated. - - Returns: - Markdown for the context-window report. - """ - if context_limit is not None and context_limit <= 0: - context_limit = None - lines = ["## Context usage"] - if model_spec: - lines.append(escape_markdown(model_spec)) - - if context_tokens is None: - lines.extend(("", "Current total unavailable.")) - if conversation_tokens is not None: - estimate = format_token_count(max(0, conversation_tokens)) - lines.append(f"**Conversation estimate:** ~{estimate} tokens") - elif context_tokens <= 0: - lines.extend(("", "No usage reported yet.")) - if context_limit is not None: - lines.append( - f"**Context window:** {format_token_count(context_limit)} tokens" - ) - else: - lines.append("Context window limit unavailable for this model.") - else: - count_prefix = "~" if approximate else "" - current_tokens = format_token_count(context_tokens) - lines.append("") - if context_limit is not None: - percent = context_tokens / context_limit * 100 - remaining = max(0, context_limit - context_tokens) - limit = format_token_count(context_limit) - current = f"**Current:** {count_prefix}{current_tokens} / {limit} tokens" - lines.extend( - ( - f"{current} ({percent:.1f}%)", - f"**Remaining:** {format_token_count(remaining)} tokens", - ) - ) - else: - lines.extend( - ( - f"**Current:** {count_prefix}{current_tokens} tokens", - "Context window limit unavailable for this model.", - ) - ) - - if conversation_tokens is not None: - conversation = min(max(0, conversation_tokens), context_tokens) - fixed = context_tokens - conversation - lines.extend(("", "### Estimated composition")) - for label, tokens in ( - ("System prompt + tools", fixed), - ("Conversation", conversation), - ): - percent_suffix = ( - f" ({tokens / context_limit * 100:.1f}%)" if context_limit else "" - ) - count = format_token_count(tokens) - lines.append(f"- **{label}:** ~{count} tokens{percent_suffix}") - - lines.extend( - ("", "**Automatic offload:** enabled; use `/offload` to compact sooner.") - ) - if context_tokens is not None and context_tokens > 0: - lines.append( - "*Current total is approximate; composition is estimated.*" - if approximate - else "*Current total is provider-reported; composition is estimated.*" - ) - return "\n".join(lines) diff --git a/libs/code/deepagents_code/tui/widgets/context_usage/__init__.py b/libs/code/deepagents_code/tui/widgets/context_usage/__init__.py new file mode 100644 index 00000000000..75319793177 --- /dev/null +++ b/libs/code/deepagents_code/tui/widgets/context_usage/__init__.py @@ -0,0 +1,5 @@ +"""Color-coded context-window visualization for `/context`.""" + +from deepagents_code.tui.widgets.context_usage._screen import ContextUsageScreen + +__all__ = ["ContextUsageScreen"] diff --git a/libs/code/deepagents_code/tui/widgets/context_usage/_models.py b/libs/code/deepagents_code/tui/widgets/context_usage/_models.py new file mode 100644 index 00000000000..12e83e4dcd0 --- /dev/null +++ b/libs/code/deepagents_code/tui/widgets/context_usage/_models.py @@ -0,0 +1,93 @@ +"""Data model for the context-usage visualization.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +_ColorRole = Literal["warning", "primary", "secondary", "accent", "muted"] + + +@dataclass(frozen=True, slots=True) +class _Category: + label: str + tokens: int + color: _ColorRole + + +@dataclass(frozen=True, slots=True) +class _Snapshot: + context_tokens: int | None + context_limit: int | None + conversation_tokens: int | None + model_spec: str | None + approximate: bool + categories: tuple[_Category, ...] + + @classmethod + def from_usage( + cls, + *, + context_tokens: int | None, + conversation_tokens: int | None, + context_limit: int | None, + model_spec: str | None, + approximate: bool, + ) -> _Snapshot: + total = None if context_tokens is None else max(0, context_tokens) + conversation = ( + None if conversation_tokens is None else max(0, conversation_tokens) + ) + limit = ( + context_limit if context_limit is not None and context_limit > 0 else None + ) + categories: list[_Category] = [] + + if total is None: + if conversation: + categories.append( + _Category("Conversation estimate", conversation, "primary") + ) + if limit is not None: + categories.append( + _Category( + "Unreported capacity", + max(0, limit - (conversation or 0)), + "accent", + ) + ) + elif total > 0: + if conversation is None: + categories.append(_Category("Used context", total, "secondary")) + else: + conversation = min(conversation, total) + fixed = total - conversation + if fixed: + categories.append( + _Category("System prompt + tools", fixed, "warning") + ) + if conversation: + categories.append( + _Category("Conversation", conversation, "primary") + ) + + if total is not None and limit is not None: + categories.append(_Category("Free space", max(0, limit - total), "muted")) + + return cls( + context_tokens=total, + context_limit=limit, + conversation_tokens=conversation, + model_spec=model_spec, + approximate=approximate, + categories=tuple(categories), + ) + + @property + def scale_tokens(self) -> int: + categorized = sum(category.tokens for category in self.categories) + return max(self.context_limit or 0, categorized, 1) + + @property + def displayed_usage(self) -> int: + return self.context_tokens or self.conversation_tokens or 0 diff --git a/libs/code/deepagents_code/tui/widgets/context_usage/_screen.py b/libs/code/deepagents_code/tui/widgets/context_usage/_screen.py new file mode 100644 index 00000000000..d289e1555b4 --- /dev/null +++ b/libs/code/deepagents_code/tui/widgets/context_usage/_screen.py @@ -0,0 +1,103 @@ +"""Modal shell for the context-usage visualization.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from textual.binding import Binding, BindingType +from textual.containers import Vertical +from textual.screen import ModalScreen +from textual.widgets import Static + +from deepagents_code import theme +from deepagents_code.config import is_ascii_mode +from deepagents_code.tui.widgets.context_usage._models import _Snapshot +from deepagents_code.tui.widgets.context_usage._widgets import ( + _ContextBar, + _ContextHeader, + _ContextLegend, +) + +if TYPE_CHECKING: + from textual.app import ComposeResult + + +class ContextUsageScreen(ModalScreen[None]): + """Modal visualization of the current model context window.""" + + BINDINGS: ClassVar[list[BindingType]] = [ + Binding("escape", "close", "Close", show=False) + ] + + CSS = """ + ContextUsageScreen { align: center middle; } + + ContextUsageScreen > Vertical { + width: 94%; + max-width: 120; + height: auto; + max-height: 90%; + background: $surface; + border: solid $primary; + padding: 1 2; + } + + ContextUsageScreen _ContextHeader { height: auto; min-height: 2; } + ContextUsageScreen _ContextBar { height: 2; margin: 1 0; } + ContextUsageScreen _ContextLegend { height: auto; margin-top: 1; } + + ContextUsageScreen .context-usage-help { + height: 1; + color: $text-muted; + margin-top: 2; + } + """ + + def __init__( + self, + *, + context_tokens: int | None, + conversation_tokens: int | None, + context_limit: int | None, + model_spec: str | None, + approximate: bool, + ) -> None: + """Initialize the modal from the latest usage measurements. + + Args: + context_tokens: Reliable total usage, or `None` when unavailable. + conversation_tokens: Estimated effective conversation usage. + context_limit: Configured model context limit. + model_spec: Active model identifier. + approximate: Whether the displayed usage is approximate. + """ + super().__init__() + self._snapshot = _Snapshot.from_usage( + context_tokens=context_tokens, + conversation_tokens=conversation_tokens, + context_limit=context_limit, + model_spec=model_spec, + approximate=approximate, + ) + + def compose(self) -> ComposeResult: + """Compose the context header, bar, legend, and close hint. + + Yields: + Widgets that make up the visualization. + """ + with Vertical(): + yield _ContextHeader(self._snapshot) + yield _ContextBar(self._snapshot) + yield _ContextLegend(self._snapshot) + yield Static("Esc to close", classes="context-usage-help") + + def on_mount(self) -> None: + """Use an ASCII border when the terminal cannot render Unicode.""" + if is_ascii_mode(): + panel = self.query_one(Vertical) + panel.styles.border = ("ascii", theme.get_theme_colors(self).primary) + + def action_close(self) -> None: + """Dismiss the context visualization.""" + self.dismiss(None) diff --git a/libs/code/deepagents_code/tui/widgets/context_usage/_widgets.py b/libs/code/deepagents_code/tui/widgets/context_usage/_widgets.py new file mode 100644 index 00000000000..90180d9836c --- /dev/null +++ b/libs/code/deepagents_code/tui/widgets/context_usage/_widgets.py @@ -0,0 +1,167 @@ +"""Responsive render widgets for the context-usage modal.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from textual.content import Content +from textual.widgets import Static + +from deepagents_code import theme +from deepagents_code._session_stats import format_token_count +from deepagents_code.config import get_glyphs + +if TYPE_CHECKING: + from textual.widget import Widget + + from deepagents_code.tui.widgets.context_usage._models import ( + _ColorRole, + _Snapshot, + ) + + +def _compact_tokens(tokens: int) -> str: + return format_token_count(tokens).replace(".0K", "K").replace(".0M", "M") + + +def _category_color(widget: Widget, role: _ColorRole) -> str: + colors = theme.get_theme_colors(widget) + return { + "warning": colors.warning, + "primary": colors.primary, + "secondary": colors.secondary, + "accent": colors.accent, + "muted": colors.muted, + }[role] + + +def _allocate_widths(snapshot: _Snapshot, width: int) -> list[int]: + categories = snapshot.categories + nonempty = [index for index, category in enumerate(categories) if category.tokens] + widths = [0] * len(categories) + if not nonempty or width <= 0: + return widths + if width <= len(nonempty): + for index in nonempty[:width]: + widths[index] = 1 + return widths + + remaining = width - len(nonempty) + raw = [ + category.tokens / snapshot.scale_tokens * remaining for category in categories + ] + for index in nonempty: + widths[index] = 1 + int(raw[index]) + remainder = width - sum(widths) + order = sorted(nonempty, key=lambda index: raw[index] % 1, reverse=True) + for offset in range(remainder): + widths[order[offset % len(order)]] += 1 + return widths + + +def _scale_line(width: int, maximum: int) -> str: + line = [" "] * width + for fraction in (0.0, 0.25, 0.5, 0.75, 1.0): + label = _compact_tokens(round(maximum * fraction)) + start = min(round((width - 1) * fraction), max(0, width - len(label))) + visible = label[: width - start] + line[start : start + len(visible)] = visible + return "".join(line) + + +class _ContextHeader(Static): + def __init__(self, snapshot: _Snapshot) -> None: + super().__init__() + self._snapshot = snapshot + + def render(self) -> Content: + colors = theme.get_theme_colors(self) + glyphs = get_glyphs() + model = self._snapshot.model_spec or "Unknown model" + maximum = ( + f"{_compact_tokens(self._snapshot.context_limit)} Max" + if self._snapshot.context_limit is not None + else "Max unavailable" + ) + left = Content.assemble( + ("Context", f"bold {colors.primary}"), + (f" {glyphs.bullet} ", colors.muted), + model, + " ", + maximum, + ) + used = _compact_tokens(self._snapshot.displayed_usage) + prefix = ( + "~" + if self._snapshot.approximate or self._snapshot.context_tokens is None + else "" + ) + right_text = f"{prefix}{used}" + if self._snapshot.context_limit is not None: + right_text += f" / {_compact_tokens(self._snapshot.context_limit)}" + right = Content(right_text) + if self._snapshot.context_tokens is not None and self._snapshot.context_limit: + percent = self._snapshot.context_tokens / self._snapshot.context_limit * 100 + right = Content.assemble(right, (f" {percent:.1f}%", colors.success)) + + gap = self.content_size.width - left.cell_length - right.cell_length + title = ( + Content.assemble(left, " " * gap, right) + if gap > 0 + else Content.assemble(left, "\n", right) + ) + subtitle = Content.styled("Current context usage by category.", colors.muted) + return Content("\n").join((title, subtitle)) + + +class _ContextBar(Static): + def __init__(self, snapshot: _Snapshot) -> None: + super().__init__() + self._snapshot = snapshot + + def render(self) -> Content: + width = max(self.content_size.width, 1) + glyph = get_glyphs().box_horizontal + segments = [ + Content.styled(glyph * segment_width, _category_color(self, category.color)) + for category, segment_width in zip( + self._snapshot.categories, + _allocate_widths(self._snapshot, width), + strict=True, + ) + if segment_width + ] + bar = Content.assemble(*segments) + scale = Content.styled( + _scale_line(width, self._snapshot.scale_tokens), + theme.get_theme_colors(self).muted, + ) + return Content("\n").join((bar, scale)) + + +class _ContextLegend(Static): + def __init__(self, snapshot: _Snapshot) -> None: + super().__init__() + self._snapshot = snapshot + + def render(self) -> Content: + colors = theme.get_theme_colors(self) + glyphs = get_glyphs() + marker = glyphs.box_horizontal * 2 + width = max(self.content_size.width, 1) + rows: list[Content] = [] + for category in self._snapshot.categories: + color = _category_color(self, category.color) + percent = category.tokens / self._snapshot.scale_tokens * 100 + value = ( + f"{_compact_tokens(category.tokens)} {glyphs.bullet} {percent:.1f}%" + ) + label_style = ( + colors.muted if category.color == "muted" else colors.foreground + ) + left = Content.assemble((marker, color), " ", (category.label, label_style)) + gap = max(1, width - left.cell_length - len(value)) + rows.append(Content.assemble(left, " " * gap, (value, colors.muted))) + if not rows: + rows.append(Content.styled("No context usage reported yet.", colors.muted)) + return Content("\n").join(rows) diff --git a/libs/code/tests/unit_tests/test_app.py b/libs/code/tests/unit_tests/test_app.py index 1a87dab0928..4146de1c314 100644 --- a/libs/code/tests/unit_tests/test_app.py +++ b/libs/code/tests/unit_tests/test_app.py @@ -189,20 +189,27 @@ async def test_prefers_checkpoint_total_after_offload( app._context_tokens = 200 app._tokens_approximate = True get_counts = AsyncMock(return_value=(1_000, 200)) - mount_message = AsyncMock() + push_screen = MagicMock() monkeypatch.setattr(app, "_get_context_usage_counts", get_counts) - monkeypatch.setattr(app, "_mount_message", mount_message) + monkeypatch.setattr(app, "push_screen", push_screen) - with patch("deepagents_code.config.settings") as mock_settings: + with ( + patch("deepagents_code.app.ContextUsageScreen") as screen_type, + patch("deepagents_code.config.settings") as mock_settings, + ): mock_settings.model_provider = "anthropic" mock_settings.model_name = "claude-sonnet" mock_settings.model_context_limit = 2_000 await app._handle_command("/context") - report = mount_message.await_args_list[-1].args[0] - assert isinstance(report, AppMessage) - assert "**Current:** 1.0K / 2.0K tokens (50.0%)" in str(report._content) - assert "**System prompt + tools:** ~800 tokens (40.0%)" in str(report._content) + assert screen_type.call_args.kwargs == { + "context_tokens": 1_000, + "conversation_tokens": 200, + "context_limit": 2_000, + "model_spec": "anthropic:claude-sonnet", + "approximate": False, + } + assert push_screen.call_args.args[0] is screen_type.return_value class TestWhatsNewMessage: diff --git a/libs/code/tests/unit_tests/tui/widgets/test_context_usage.py b/libs/code/tests/unit_tests/tui/widgets/test_context_usage.py deleted file mode 100644 index 6cf7b7f541d..00000000000 --- a/libs/code/tests/unit_tests/tui/widgets/test_context_usage.py +++ /dev/null @@ -1,19 +0,0 @@ -"""Tests for the context-usage report.""" - -from deepagents_code.tui.widgets.context_usage import build_context_usage_markdown - - -def test_reports_capacity_and_estimated_breakdown() -> None: - content = build_context_usage_markdown( - context_tokens=20_000, - conversation_tokens=15_000, - context_limit=100_000, - model_spec="anthropic:claude-sonnet", - approximate=False, - ) - - assert "20.0K / 100.0K tokens (20.0%)" in content - assert "**Remaining:** 80.0K tokens" in content - assert "**System prompt + tools:** ~5.0K tokens (5.0%)" in content - assert "**Conversation:** ~15.0K tokens (15.0%)" in content - assert "**Automatic offload:** enabled" in content From 944cc19dd5830cfef25d9653bc9f1665a8431c6e Mon Sep 17 00:00:00 2001 From: Johannes du Plessis <51395795+johannes117@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:59:35 +0000 Subject: [PATCH 4/4] fix(code): defer `/context` focus and trim modal Co-authored-by: open-swe[bot] --- libs/code/deepagents_code/app.py | 54 ++---- .../tui/widgets/context_usage.py | 169 ++++++++++++++++++ .../tui/widgets/context_usage/__init__.py | 5 - .../tui/widgets/context_usage/_models.py | 93 ---------- .../tui/widgets/context_usage/_screen.py | 103 ----------- .../tui/widgets/context_usage/_widgets.py | 167 ----------------- libs/code/tests/unit_tests/test_app.py | 57 +++--- 7 files changed, 214 insertions(+), 434 deletions(-) create mode 100644 libs/code/deepagents_code/tui/widgets/context_usage.py delete mode 100644 libs/code/deepagents_code/tui/widgets/context_usage/__init__.py delete mode 100644 libs/code/deepagents_code/tui/widgets/context_usage/_models.py delete mode 100644 libs/code/deepagents_code/tui/widgets/context_usage/_screen.py delete mode 100644 libs/code/deepagents_code/tui/widgets/context_usage/_widgets.py diff --git a/libs/code/deepagents_code/app.py b/libs/code/deepagents_code/app.py index dc15a16a1fa..48dfaeaac20 100644 --- a/libs/code/deepagents_code/app.py +++ b/libs/code/deepagents_code/app.py @@ -14111,36 +14111,26 @@ async def _handle_command(self, command: str) -> None: markup=False, ) elif cmd == "/context": - ( - reported_tokens, - conversation_tokens, - ) = await self._get_context_usage_counts() - if reported_tokens is not None: - context_tokens = reported_tokens - approximate = False - elif self._context_tokens > 0 and not self._tokens_approximate: - context_tokens = self._context_tokens - approximate = False - elif conversation_tokens is not None or self._context_tokens > 0: - context_tokens = None + context_tokens, conversation_tokens = await self._get_context_usage_counts() + if context_tokens is None and not self._tokens_approximate: + context_tokens = self._context_tokens or None + approximate = context_tokens is None and bool( + conversation_tokens or self._context_tokens + ) + if approximate: conversation_tokens = conversation_tokens or self._context_tokens - approximate = True - else: + elif context_tokens is None: context_tokens = 0 - approximate = False - screen = ContextUsageScreen( - context_tokens=context_tokens, - conversation_tokens=conversation_tokens, - context_limit=settings.model_context_limit, - model_spec=self._effective_model_spec() or settings.model_name, - approximate=approximate, + self.push_screen( + ContextUsageScreen( + context_tokens=context_tokens, + conversation_tokens=conversation_tokens, + context_limit=settings.model_context_limit, + model_spec=self._effective_model_spec() or settings.model_name, + approximate=approximate, + ), + lambda _result: self._focus_chat_input_after_refresh(), ) - - def handle_result(_result: None) -> None: - if self._chat_input: - self._chat_input.focus_input() - - self.push_screen(screen, handle_result) elif cmd == "/tokens": await self._mount_message(UserMessage(command)) if self._context_tokens > 0: @@ -14843,18 +14833,12 @@ async def _get_context_usage_counts(self) -> tuple[int | None, int | None]: Pair of provider-reported context tokens and approximate effective conversation tokens. Either value is `None` when unavailable. """ - if not self._agent: + if not self._agent or not self._lc_thread_id: return None, None try: from langchain_core.messages.utils import count_tokens_approximately - config: RunnableConfig = { - "configurable": {"thread_id": self._lc_thread_id}, - } - state = await self._agent.aget_state(config) - if not state or not state.values: - return None, None - values = dict(state.values) + values = await self._get_thread_state_values(self._lc_thread_id) reported = _persisted_context_tokens(values) or None messages = values.get("messages", []) if not isinstance(messages, list) or not messages: diff --git a/libs/code/deepagents_code/tui/widgets/context_usage.py b/libs/code/deepagents_code/tui/widgets/context_usage.py new file mode 100644 index 00000000000..3be48f46aef --- /dev/null +++ b/libs/code/deepagents_code/tui/widgets/context_usage.py @@ -0,0 +1,169 @@ +"""Color-coded context-window visualization for `/context`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from textual.binding import Binding, BindingType +from textual.containers import VerticalScroll +from textual.content import Content +from textual.screen import ModalScreen +from textual.widgets import Static + +from deepagents_code import theme +from deepagents_code._session_stats import format_token_count +from deepagents_code.config import get_glyphs, is_ascii_mode + +if TYPE_CHECKING: + from textual.app import ComposeResult + + +class _ContextUsage(Static): + def __init__( + self, + *, + context_tokens: int | None, + conversation_tokens: int | None, + context_limit: int | None, + model_spec: str | None, + approximate: bool, + ) -> None: + super().__init__() + self._total = None if context_tokens is None else max(0, context_tokens) + self._conversation = ( + None if conversation_tokens is None else max(0, conversation_tokens) + ) + self._limit = context_limit if context_limit and context_limit > 0 else None + self._model = model_spec or "Unknown model" + self._approximate = approximate + + def render(self) -> Content: + colors = theme.get_theme_colors(self) + glyphs = get_glyphs() + usage = self._total if self._total is not None else self._conversation or 0 + maximum = format_token_count(self._limit) if self._limit else "unavailable" + prefix = "~" if self._approximate or self._total is None else "" + right = f"{prefix}{format_token_count(usage)} / {maximum}" + if self._total is not None and self._limit: + right += f" {self._total / self._limit * 100:.1f}%" + left = Content.assemble( + ("Context", f"bold {colors.primary}"), + (f" {glyphs.bullet} ", colors.muted), + self._model, + ) + gap = self.content_size.width - left.cell_length - len(right) + header = Content.assemble(left, " " * max(1, gap), (right, colors.muted)) + + categories: list[tuple[str, int, str]] = [] + if self._total is None: + if self._conversation: + categories.append( + ("Conversation estimate", self._conversation, colors.primary) + ) + elif self._total: + if self._conversation is None: + categories.append(("Used context", self._total, colors.secondary)) + else: + conversation = min(self._conversation, self._total) + if overhead := self._total - conversation: + categories.append( + ("System prompt + tools", overhead, colors.warning) + ) + if conversation: + categories.append(("Conversation", conversation, colors.primary)) + if self._total is not None and self._limit: + categories.append( + ("Free space", max(0, self._limit - self._total), colors.muted) + ) + + scale = max(self._limit or 0, sum(tokens for _, tokens, _ in categories), 1) + width = max(self.content_size.width, 1) + used = 0 + segments: list[Content] = [] + for _label, tokens, color in categories: + end = round(min(scale, used + tokens) / scale * width) + start = round(min(scale, used) / scale * width) + segments.append( + Content.styled(glyphs.box_horizontal * (end - start), color) + ) + used += tokens + bar = Content.assemble(*segments) + + rows: list[Content] = [] + marker = glyphs.box_horizontal * 2 + for label, tokens, color in categories: + percent = tokens / scale * 100 + value = f"{format_token_count(tokens)} {glyphs.bullet} {percent:.1f}%" + item = Content.assemble((marker, color), " ", label) + rows.append( + Content.assemble( + item, + " " * max(1, width - item.cell_length - len(value)), + (value, colors.muted), + ) + ) + if not rows: + rows.append(Content.styled("No context usage reported yet.", colors.muted)) + elif self._total is None: + rows.append(Content.styled("Total usage unavailable.", colors.muted)) + return Content("\n").join((header, bar, *rows)) + + +class ContextUsageScreen(ModalScreen[None]): + """Modal visualization of the current model context window.""" + + BINDINGS: ClassVar[list[BindingType]] = [ + Binding("escape", "close", "Close", show=False) + ] + CSS = """ + ContextUsageScreen { align: center middle; } + ContextUsageScreen > VerticalScroll { + width: 94%; max-width: 120; height: auto; max-height: 90%; + background: $surface; border: solid $primary; padding: 1 2; + } + ContextUsageScreen _ContextUsage { height: auto; } + ContextUsageScreen .context-usage-help { + height: 1; color: $text-muted; margin-top: 2; + } + """ + + def __init__( + self, + *, + context_tokens: int | None, + conversation_tokens: int | None, + context_limit: int | None, + model_spec: str | None, + approximate: bool, + ) -> None: + """Initialize the modal from the latest usage measurements.""" + super().__init__() + self._usage = _ContextUsage( + context_tokens=context_tokens, + conversation_tokens=conversation_tokens, + context_limit=context_limit, + model_spec=model_spec, + approximate=approximate, + ) + + def compose(self) -> ComposeResult: + """Compose the context visualization and close hint. + + Yields: + Widgets that make up the modal. + """ + with VerticalScroll(): + yield self._usage + yield Static("Esc to close", classes="context-usage-help") + + def on_mount(self) -> None: + """Use an ASCII border when the terminal cannot render Unicode.""" + if is_ascii_mode(): + self.query_one(VerticalScroll).styles.border = ( + "ascii", + theme.get_theme_colors(self).primary, + ) + + def action_close(self) -> None: + """Dismiss the context visualization.""" + self.dismiss(None) diff --git a/libs/code/deepagents_code/tui/widgets/context_usage/__init__.py b/libs/code/deepagents_code/tui/widgets/context_usage/__init__.py deleted file mode 100644 index 75319793177..00000000000 --- a/libs/code/deepagents_code/tui/widgets/context_usage/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Color-coded context-window visualization for `/context`.""" - -from deepagents_code.tui.widgets.context_usage._screen import ContextUsageScreen - -__all__ = ["ContextUsageScreen"] diff --git a/libs/code/deepagents_code/tui/widgets/context_usage/_models.py b/libs/code/deepagents_code/tui/widgets/context_usage/_models.py deleted file mode 100644 index 12e83e4dcd0..00000000000 --- a/libs/code/deepagents_code/tui/widgets/context_usage/_models.py +++ /dev/null @@ -1,93 +0,0 @@ -"""Data model for the context-usage visualization.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Literal - -_ColorRole = Literal["warning", "primary", "secondary", "accent", "muted"] - - -@dataclass(frozen=True, slots=True) -class _Category: - label: str - tokens: int - color: _ColorRole - - -@dataclass(frozen=True, slots=True) -class _Snapshot: - context_tokens: int | None - context_limit: int | None - conversation_tokens: int | None - model_spec: str | None - approximate: bool - categories: tuple[_Category, ...] - - @classmethod - def from_usage( - cls, - *, - context_tokens: int | None, - conversation_tokens: int | None, - context_limit: int | None, - model_spec: str | None, - approximate: bool, - ) -> _Snapshot: - total = None if context_tokens is None else max(0, context_tokens) - conversation = ( - None if conversation_tokens is None else max(0, conversation_tokens) - ) - limit = ( - context_limit if context_limit is not None and context_limit > 0 else None - ) - categories: list[_Category] = [] - - if total is None: - if conversation: - categories.append( - _Category("Conversation estimate", conversation, "primary") - ) - if limit is not None: - categories.append( - _Category( - "Unreported capacity", - max(0, limit - (conversation or 0)), - "accent", - ) - ) - elif total > 0: - if conversation is None: - categories.append(_Category("Used context", total, "secondary")) - else: - conversation = min(conversation, total) - fixed = total - conversation - if fixed: - categories.append( - _Category("System prompt + tools", fixed, "warning") - ) - if conversation: - categories.append( - _Category("Conversation", conversation, "primary") - ) - - if total is not None and limit is not None: - categories.append(_Category("Free space", max(0, limit - total), "muted")) - - return cls( - context_tokens=total, - context_limit=limit, - conversation_tokens=conversation, - model_spec=model_spec, - approximate=approximate, - categories=tuple(categories), - ) - - @property - def scale_tokens(self) -> int: - categorized = sum(category.tokens for category in self.categories) - return max(self.context_limit or 0, categorized, 1) - - @property - def displayed_usage(self) -> int: - return self.context_tokens or self.conversation_tokens or 0 diff --git a/libs/code/deepagents_code/tui/widgets/context_usage/_screen.py b/libs/code/deepagents_code/tui/widgets/context_usage/_screen.py deleted file mode 100644 index d289e1555b4..00000000000 --- a/libs/code/deepagents_code/tui/widgets/context_usage/_screen.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Modal shell for the context-usage visualization.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, ClassVar - -from textual.binding import Binding, BindingType -from textual.containers import Vertical -from textual.screen import ModalScreen -from textual.widgets import Static - -from deepagents_code import theme -from deepagents_code.config import is_ascii_mode -from deepagents_code.tui.widgets.context_usage._models import _Snapshot -from deepagents_code.tui.widgets.context_usage._widgets import ( - _ContextBar, - _ContextHeader, - _ContextLegend, -) - -if TYPE_CHECKING: - from textual.app import ComposeResult - - -class ContextUsageScreen(ModalScreen[None]): - """Modal visualization of the current model context window.""" - - BINDINGS: ClassVar[list[BindingType]] = [ - Binding("escape", "close", "Close", show=False) - ] - - CSS = """ - ContextUsageScreen { align: center middle; } - - ContextUsageScreen > Vertical { - width: 94%; - max-width: 120; - height: auto; - max-height: 90%; - background: $surface; - border: solid $primary; - padding: 1 2; - } - - ContextUsageScreen _ContextHeader { height: auto; min-height: 2; } - ContextUsageScreen _ContextBar { height: 2; margin: 1 0; } - ContextUsageScreen _ContextLegend { height: auto; margin-top: 1; } - - ContextUsageScreen .context-usage-help { - height: 1; - color: $text-muted; - margin-top: 2; - } - """ - - def __init__( - self, - *, - context_tokens: int | None, - conversation_tokens: int | None, - context_limit: int | None, - model_spec: str | None, - approximate: bool, - ) -> None: - """Initialize the modal from the latest usage measurements. - - Args: - context_tokens: Reliable total usage, or `None` when unavailable. - conversation_tokens: Estimated effective conversation usage. - context_limit: Configured model context limit. - model_spec: Active model identifier. - approximate: Whether the displayed usage is approximate. - """ - super().__init__() - self._snapshot = _Snapshot.from_usage( - context_tokens=context_tokens, - conversation_tokens=conversation_tokens, - context_limit=context_limit, - model_spec=model_spec, - approximate=approximate, - ) - - def compose(self) -> ComposeResult: - """Compose the context header, bar, legend, and close hint. - - Yields: - Widgets that make up the visualization. - """ - with Vertical(): - yield _ContextHeader(self._snapshot) - yield _ContextBar(self._snapshot) - yield _ContextLegend(self._snapshot) - yield Static("Esc to close", classes="context-usage-help") - - def on_mount(self) -> None: - """Use an ASCII border when the terminal cannot render Unicode.""" - if is_ascii_mode(): - panel = self.query_one(Vertical) - panel.styles.border = ("ascii", theme.get_theme_colors(self).primary) - - def action_close(self) -> None: - """Dismiss the context visualization.""" - self.dismiss(None) diff --git a/libs/code/deepagents_code/tui/widgets/context_usage/_widgets.py b/libs/code/deepagents_code/tui/widgets/context_usage/_widgets.py deleted file mode 100644 index 90180d9836c..00000000000 --- a/libs/code/deepagents_code/tui/widgets/context_usage/_widgets.py +++ /dev/null @@ -1,167 +0,0 @@ -"""Responsive render widgets for the context-usage modal.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -from textual.content import Content -from textual.widgets import Static - -from deepagents_code import theme -from deepagents_code._session_stats import format_token_count -from deepagents_code.config import get_glyphs - -if TYPE_CHECKING: - from textual.widget import Widget - - from deepagents_code.tui.widgets.context_usage._models import ( - _ColorRole, - _Snapshot, - ) - - -def _compact_tokens(tokens: int) -> str: - return format_token_count(tokens).replace(".0K", "K").replace(".0M", "M") - - -def _category_color(widget: Widget, role: _ColorRole) -> str: - colors = theme.get_theme_colors(widget) - return { - "warning": colors.warning, - "primary": colors.primary, - "secondary": colors.secondary, - "accent": colors.accent, - "muted": colors.muted, - }[role] - - -def _allocate_widths(snapshot: _Snapshot, width: int) -> list[int]: - categories = snapshot.categories - nonempty = [index for index, category in enumerate(categories) if category.tokens] - widths = [0] * len(categories) - if not nonempty or width <= 0: - return widths - if width <= len(nonempty): - for index in nonempty[:width]: - widths[index] = 1 - return widths - - remaining = width - len(nonempty) - raw = [ - category.tokens / snapshot.scale_tokens * remaining for category in categories - ] - for index in nonempty: - widths[index] = 1 + int(raw[index]) - remainder = width - sum(widths) - order = sorted(nonempty, key=lambda index: raw[index] % 1, reverse=True) - for offset in range(remainder): - widths[order[offset % len(order)]] += 1 - return widths - - -def _scale_line(width: int, maximum: int) -> str: - line = [" "] * width - for fraction in (0.0, 0.25, 0.5, 0.75, 1.0): - label = _compact_tokens(round(maximum * fraction)) - start = min(round((width - 1) * fraction), max(0, width - len(label))) - visible = label[: width - start] - line[start : start + len(visible)] = visible - return "".join(line) - - -class _ContextHeader(Static): - def __init__(self, snapshot: _Snapshot) -> None: - super().__init__() - self._snapshot = snapshot - - def render(self) -> Content: - colors = theme.get_theme_colors(self) - glyphs = get_glyphs() - model = self._snapshot.model_spec or "Unknown model" - maximum = ( - f"{_compact_tokens(self._snapshot.context_limit)} Max" - if self._snapshot.context_limit is not None - else "Max unavailable" - ) - left = Content.assemble( - ("Context", f"bold {colors.primary}"), - (f" {glyphs.bullet} ", colors.muted), - model, - " ", - maximum, - ) - used = _compact_tokens(self._snapshot.displayed_usage) - prefix = ( - "~" - if self._snapshot.approximate or self._snapshot.context_tokens is None - else "" - ) - right_text = f"{prefix}{used}" - if self._snapshot.context_limit is not None: - right_text += f" / {_compact_tokens(self._snapshot.context_limit)}" - right = Content(right_text) - if self._snapshot.context_tokens is not None and self._snapshot.context_limit: - percent = self._snapshot.context_tokens / self._snapshot.context_limit * 100 - right = Content.assemble(right, (f" {percent:.1f}%", colors.success)) - - gap = self.content_size.width - left.cell_length - right.cell_length - title = ( - Content.assemble(left, " " * gap, right) - if gap > 0 - else Content.assemble(left, "\n", right) - ) - subtitle = Content.styled("Current context usage by category.", colors.muted) - return Content("\n").join((title, subtitle)) - - -class _ContextBar(Static): - def __init__(self, snapshot: _Snapshot) -> None: - super().__init__() - self._snapshot = snapshot - - def render(self) -> Content: - width = max(self.content_size.width, 1) - glyph = get_glyphs().box_horizontal - segments = [ - Content.styled(glyph * segment_width, _category_color(self, category.color)) - for category, segment_width in zip( - self._snapshot.categories, - _allocate_widths(self._snapshot, width), - strict=True, - ) - if segment_width - ] - bar = Content.assemble(*segments) - scale = Content.styled( - _scale_line(width, self._snapshot.scale_tokens), - theme.get_theme_colors(self).muted, - ) - return Content("\n").join((bar, scale)) - - -class _ContextLegend(Static): - def __init__(self, snapshot: _Snapshot) -> None: - super().__init__() - self._snapshot = snapshot - - def render(self) -> Content: - colors = theme.get_theme_colors(self) - glyphs = get_glyphs() - marker = glyphs.box_horizontal * 2 - width = max(self.content_size.width, 1) - rows: list[Content] = [] - for category in self._snapshot.categories: - color = _category_color(self, category.color) - percent = category.tokens / self._snapshot.scale_tokens * 100 - value = ( - f"{_compact_tokens(category.tokens)} {glyphs.bullet} {percent:.1f}%" - ) - label_style = ( - colors.muted if category.color == "muted" else colors.foreground - ) - left = Content.assemble((marker, color), " ", (category.label, label_style)) - gap = max(1, width - left.cell_length - len(value)) - rows.append(Content.assemble(left, " " * gap, (value, colors.muted))) - if not rows: - rows.append(Content.styled("No context usage reported yet.", colors.muted)) - return Content("\n").join(rows) diff --git a/libs/code/tests/unit_tests/test_app.py b/libs/code/tests/unit_tests/test_app.py index 4146de1c314..28757ae5cae 100644 --- a/libs/code/tests/unit_tests/test_app.py +++ b/libs/code/tests/unit_tests/test_app.py @@ -179,37 +179,32 @@ def test_strips_provider_prefix( assert _display_model_label(spec) == expected -class TestContextCommand: - """Tests for the `/context` command's data source.""" - - async def test_prefers_checkpoint_total_after_offload( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - app = DeepAgentsApp() - app._context_tokens = 200 - app._tokens_approximate = True - get_counts = AsyncMock(return_value=(1_000, 200)) - push_screen = MagicMock() - monkeypatch.setattr(app, "_get_context_usage_counts", get_counts) - monkeypatch.setattr(app, "push_screen", push_screen) - - with ( - patch("deepagents_code.app.ContextUsageScreen") as screen_type, - patch("deepagents_code.config.settings") as mock_settings, - ): - mock_settings.model_provider = "anthropic" - mock_settings.model_name = "claude-sonnet" - mock_settings.model_context_limit = 2_000 - await app._handle_command("/context") - - assert screen_type.call_args.kwargs == { - "context_tokens": 1_000, - "conversation_tokens": 200, - "context_limit": 2_000, - "model_spec": "anthropic:claude-sonnet", - "approximate": False, - } - assert push_screen.call_args.args[0] is screen_type.return_value +async def test_context_prefers_checkpoint_total_after_offload( + monkeypatch: pytest.MonkeyPatch, +) -> None: + app = DeepAgentsApp() + app._context_tokens = 200 + app._tokens_approximate = True + monkeypatch.setattr( + app, "_get_context_usage_counts", AsyncMock(return_value=(1_000, 200)) + ) + push_screen = MagicMock() + monkeypatch.setattr(app, "push_screen", push_screen) + focus = MagicMock() + monkeypatch.setattr(app, "_focus_chat_input_after_refresh", focus) + + with ( + patch("deepagents_code.app.ContextUsageScreen") as screen_type, + patch("deepagents_code.config.settings") as settings, + ): + settings.model_provider = "anthropic" + settings.model_name = "claude-sonnet" + settings.model_context_limit = 2_000 + await app._handle_command("/context") + + assert screen_type.call_args.kwargs["context_tokens"] == 1_000 + push_screen.call_args.args[1](None) + focus.assert_called_once_with() class TestWhatsNewMessage: