diff --git a/clients/terminal/tui/tui_app.py b/clients/terminal/tui/tui_app.py index 15cf1b1..0081743 100644 --- a/clients/terminal/tui/tui_app.py +++ b/clients/terminal/tui/tui_app.py @@ -169,6 +169,7 @@ async def attach_session(self, session_id: str) -> None: self._ctx.session_id = session_id history: list[dict] = [] + session: dict = {} try: session = api.get_session(session_id) self._ctx.profile_id = session.get("profile_id") or settings.default_profile_id @@ -196,6 +197,12 @@ # Replay the session's past conversation before the connection banner # so a resumed/switched session shows its history instead of a blank chat. self._chat_panel.load_history(history) + # Seed the context gauge immediately on resume/switch so the bar isn't + # blank until the first stream_end of the new turn reports tokens. + self._status_bar.set_context( + session.get("context_token_count"), + session.get("max_context_tokens"), + ) if cwd := self._ctx.cwd: self._chat_panel.handle_ws_event({"type": "status", "content": f"cwd: {cwd}"}) @@ -353,6 +360,12 @@ self._input_box.set_placeholder("Ask anything...") self._input_box.focus_input() self._status_bar.activity_stop() + # stream_end reports the prompt-token size of the last LLM call — + # how full the context window was for this turn. Update the gauge. + if msg_type == "stream_end": + self._status_bar.set_context( + payload.get("context_tokens"), payload.get("max_context_tokens") + ) self._update_footer_bindings() elif msg_type == "model_info": # The model that actually served this turn (resolved by the backend, @@ -362,6 +375,16 @@ if model: self._status_panel.set_model(model) return + elif msg_type in ("compression_started", "context_compressed"): + # Compression carries a fresh context-token count (post-compress on + # context_compressed, pre-compress on compression_started). Update + # the gauge; not a chat item. + self._status_bar.set_context( + payload.get("context_tokens"), payload.get("max_context_tokens") + ) + if msg_type == "compression_started": + self._status_bar.activity_label("compressing") + return elif msg_type == "stream_delta" and self._streaming: # First real output: the agent is now responding, not just thinking. self._status_bar.activity_label("responding") diff --git a/clients/terminal/tui/widgets/status_bar.py b/clients/terminal/tui/widgets/status_bar.py index 82a14d5..ac95c1b 100644 --- a/clients/terminal/tui/widgets/status_bar.py +++ b/clients/terminal/tui/widgets/status_bar.py @@ -1,24 +1,70 @@ """Bottom status line for the TUI. -Hosts the activity indicator (left) and a dim key-combo summary (right). The -indicator is the live "agent is working" signal; the combos keep the line -useful while idle. Sits at the very bottom, replacing Textual's Footer so the -working indicator and the key hints share one line. +Hosts the activity indicator (left), the context-window fill gauge, and a dim +key-combo summary (right). The indicator is the live "agent is working" signal; +the context gauge shows how full the LLM context window is (absolute tokens + +percent); the combos keep the line useful while idle. Sits at the very bottom, +replacing Textual's Footer so these share one line. """ from __future__ import annotations +from rich.text import Text from textual.app import ComposeResult from textual.containers import Horizontal from textual.widgets import Static +from clients.terminal.tui.themes import get_active_theme from clients.terminal.tui.widgets.activity_indicator import ActivityIndicator _COMBOS = "Ctrl+P palette · Ctrl+X Q quit · Esc stop" +def _humanize(n: int) -> str: + """Compact token count: 12345 -> '12.3k', 32768 -> '32.8k', 500 -> '500'.""" + if n < 1000: + return str(n) + if n < 1_000_000: + return f"{n / 1000:.1f}k" + return f"{n / 1_000_000:.2f}M" + + +class ContextFill(Static): + """Context-window fill: ``ctx used/max · pct%``, colored by fill level.""" + + def __init__(self) -> None: + super().__init__("—") + self._used: int | None = None + self._max: int | None = None + + def set_tokens(self, used: int | None, total: int | None) -> None: + # None means the backend didn't report; keep the last known value. + if used is not None: + self._used = used + if total is not None: + self._max = total + if not self.is_mounted: + return + if self._used is None or not self._max: + self.update(Text("—", style=get_active_theme().text_dim.hex)) + return + pct = round(self._used / self._max * 100) + color = self._color_for(pct) + self.update( + Text(f"{_humanize(self._used)}/{_humanize(self._max)} · {pct}%", style=color) + ) + + def _color_for(self, pct: int) -> str: + theme = get_active_theme() + if pct >= 90: + return theme.error.hex + if pct >= 70: + return theme.warning.hex + return theme.text_dim.hex + + class StatusBar(Horizontal): - """One-line bottom bar: activity indicator + key combos.""" + """One-line bottom bar: activity indicator + context fill + key combos.""" DEFAULT_CSS = """ StatusBar { @@ -31,6 +77,11 @@ width: auto; height: 1; } + StatusBar ContextFill { + width: auto; + height: 1; + margin-left: 1; + } StatusBar .combos { color: $tui-text-dim; text-align: right; @@ -42,10 +93,16 @@ def __init__(self) -> None: super().__init__() self._activity = ActivityIndicator() + # NOTE: do not name this attribute ``_context`` — that shadows + # ``MessagePump._context``, the internal context manager Textual wraps + # around message processing. Shadowing it (e.g. with a widget) hangs the + # app's message pump at mount time. + self._ctx_fill = ContextFill() self._combos = Static(_COMBOS, classes="combos") def compose(self) -> ComposeResult: yield self._activity + yield self._ctx_fill yield self._combos def activity_start(self, label: str) -> None: @@ -55,4 +112,7 @@ self._activity.set_label(label) def activity_stop(self) -> None: - self._activity.stop() \ No newline at end of file + self._activity.stop() + + def set_context(self, used: int | None, total: int | None) -> None: + self._ctx_fill.set_tokens(used, total) \ No newline at end of file diff --git a/tests/clients/test_tui_app.py b/tests/clients/test_tui_app.py index dd1aae4..ca4d765 100644 --- a/tests/clients/test_tui_app.py +++ b/tests/clients/test_tui_app.py @@ -608,3 +608,155 @@ assert "earlier answer" in contents # The context-only message (is_display=False) was skipped. assert "context-only" not in contents + + +@pytest.mark.anyio +async def test_stream_end_updates_context_fill() -> None: + """A stream_end event carrying context_tokens/max_context_tokens seeds the gauge.""" + async with NaviCodeTui(new_session=True).run_test() as pilot: + await pilot.pause() + ctx = pilot.app.query_one("StatusBar")._ctx_fill + # No data yet → placeholder. + assert ctx._used is None + pilot.app.on_ws_event(WsEvent({ + "type": "stream_end", "content": "", + "context_tokens": 12000, + "max_context_tokens": 32000, + })) + await pilot.pause() + assert ctx._used == 12000 + assert ctx._max == 32000 + # 12000/32000 = 37.5%, which rounds to 38% (banker's rounding to even). + assert "38%" in str(ctx.render()) + + +@pytest.mark.anyio +async def test_context_compressed_updates_fill() -> None: + """context_compressed reports the post-compression token count.""" + async with NaviCodeTui(new_session=True).run_test() as pilot: + await pilot.pause() + ctx = pilot.app.query_one("StatusBar")._ctx_fill + pilot.app.on_ws_event(WsEvent({ + "type": "stream_end", "content": "", + "context_tokens": 30000, + "max_context_tokens": 32000, + })) + await pilot.pause() + assert ctx._used == 30000 + + pilot.app.on_ws_event(WsEvent({ + "type": "context_compressed", + "context_tokens": 8000, + "max_context_tokens": 32000, + })) + await pilot.pause() + assert ctx._used == 8000 + + +@pytest.mark.anyio +async def test_stream_stopped_does_not_touch_context() -> None: + """stream_stopped carries no context fields; the gauge keeps its last value.""" + async with NaviCodeTui(new_session=True).run_test() as pilot: + await pilot.pause() + ctx = pilot.app.query_one("StatusBar")._ctx_fill + pilot.app.on_ws_event(WsEvent({ + "type": "stream_end", "content": "", + "context_tokens": 5000, + "max_context_tokens": 32000, + })) + await pilot.pause() + assert ctx._used == 5000 + + pilot.app.on_ws_event(WsEvent({"type": "stream_stopped"})) + await pilot.pause() + # Unchanged — stream_stopped has no token fields. + assert ctx._used == 5000 + + +@pytest.mark.anyio +async def test_context_fill_keeps_last_on_none() -> None: + """A context report missing token fields keeps the last known value.""" + async with NaviCodeTui(new_session=True).run_test() as pilot: + await pilot.pause() + ctx = pilot.app.query_one("StatusBar")._ctx_fill + pilot.app.on_ws_event(WsEvent({ + "type": "stream_end", "content": "", + "context_tokens": 10000, + "max_context_tokens": 32000, + })) + await pilot.pause() + pilot.app.on_ws_event(WsEvent({ + "type": "stream_end", "content": "", + "context_tokens": None, + "max_context_tokens": None, + })) + await pilot.pause() + assert ctx._used == 10000 + assert ctx._max == 32000 + + +@pytest.mark.anyio +async def test_context_fill_color_thresholds() -> None: + """The gauge color escalates as the window fills (dim → warning → error).""" + async with NaviCodeTui(new_session=True).run_test() as pilot: + await pilot.pause() + ctx = pilot.app.query_one("StatusBar")._ctx_fill + from clients.terminal.tui.themes import get_active_theme + t = get_active_theme() + + # Under 70% → dim. + assert ctx._color_for(50) == t.text_dim.hex + # 70–89% → warning. + assert ctx._color_for(75) == t.warning.hex + # 90%+ → error. + assert ctx._color_for(95) == t.error.hex + + +@pytest.mark.anyio +async def test_context_fill_k_notation() -> None: + """Token counts render in compact k/M notation.""" + from clients.terminal.tui.widgets.status_bar import _humanize + assert _humanize(500) == "500" + assert _humanize(12000) == "12.0k" + assert _humanize(1_500_000) == "1.50M" + + +@pytest.mark.anyio +async def test_attach_session_seeds_context_fill(monkeypatch: pytest.MonkeyPatch) -> None: + """On resume the gauge is seeded from get_session so it isn't blank pre-turn.""" + import clients.terminal.api as api_module + import clients.terminal.tui.tui_app as tui_app_module + + class FakeBridge: + def __init__(self, app, session_id, cwd=None) -> None: + self.client = None + + async def start(self) -> None: + pass + + async def stop(self) -> None: + pass + + monkeypatch.setattr(tui_app_module, "WsBridge", FakeBridge) + monkeypatch.setattr(api_module, "get_session", lambda sid: { + "session_id": sid, + "profile_id": "navi_code", + "messages": [], + "context_token_count": 9000, + "max_context_tokens": 32000, + }) + + async with NaviCodeTui(new_session=True).run_test() as pilot: + await pilot.pause() + ctx = pilot.app.query_one("StatusBar")._ctx_fill + # _startup already ran attach_session (against the monkeypatched + # get_session), so the gauge is seeded — not blank — before any turn. + assert ctx._used == 9000 + assert ctx._max == 32000 + assert "28%" in str(ctx.render()) + + # An explicit re-attach re-seeds from the same session response. + await pilot.app.attach_session("sess-resume") + await pilot.pause() + assert ctx._used == 9000 + assert "28%" in str(ctx.render())