From 2cabeba563cfbab827583c0115ac752e52eb911a Mon Sep 17 00:00:00 2001 From: Michael Huang <154242484+michaelHMK@users.noreply.github.com> Date: Sat, 25 Jul 2026 01:57:19 +0800 Subject: [PATCH] fix(tui,desktop): refresh context usage live during active turns --- apps/desktop/e2e/sidebar-states.spec.ts | 12 +- apps/desktop/e2e/tile-unread-bug.spec.ts | 12 +- .../hooks/use-message-stream/gateway-event.ts | 16 + .../hooks/use-message-stream/usage.test.tsx | 134 ++++++++ apps/shared/src/json-rpc-gateway.ts | 1 + .../gateway/test_multiplex_busy_input_mode.py | 9 +- tests/test_tui_gateway_server.py | 306 ++++++++++++++++++ tui_gateway/server.py | 80 ++++- .../createGatewayEventHandler.test.ts | 42 +++ ui-tui/src/app/createGatewayEventHandler.ts | 15 + ui-tui/src/gatewayTypes.ts | 1 + 11 files changed, 624 insertions(+), 4 deletions(-) create mode 100644 apps/desktop/src/app/session/hooks/use-message-stream/usage.test.tsx diff --git a/apps/desktop/e2e/sidebar-states.spec.ts b/apps/desktop/e2e/sidebar-states.spec.ts index 6d8c0c2c9c..29e12debd2 100644 --- a/apps/desktop/e2e/sidebar-states.spec.ts +++ b/apps/desktop/e2e/sidebar-states.spec.ts @@ -25,6 +25,8 @@ import { /** Background-running dot aria-label (from i18n en.ts). */ const BG_DOT_LABEL = 'Background task running' +/** Foreground turn-running dot aria-label. */ +const SESSION_RUNNING_DOT_LABEL = 'Session running' /** Finished-unread dot aria-label. */ const UNREAD_DOT_LABEL = 'Finished — unread' @@ -219,12 +221,20 @@ test.describe('sidebar states — cross-session dot transition', () => { ) .toBeGreaterThan(0) - // Wait for the final answer (turn completes, but bg process still running). + // The final answer text streams before message.complete, so text visibility + // alone is not a completion barrier. Wait for the foreground-running state + // to clear before asserting the background-process state. await page.waitForFunction( (text) => (document.body.textContent ?? '').includes(text), SIDEBAR_CROSS_TEXTS.finalText, { timeout: 90_000 }, ) + await expect + .poll( + () => page.locator(`[aria-label="${SESSION_RUNNING_DOT_LABEL}"]`).count(), + { timeout: 30_000, message: 'session running dot should disappear after turn completes' }, + ) + .toBe(0) // The background dot must still be visible: the turn is done but the // process is held open by the sentinel, so this is a stable state rather diff --git a/apps/desktop/e2e/tile-unread-bug.spec.ts b/apps/desktop/e2e/tile-unread-bug.spec.ts index fa614a21d2..7a676798f5 100644 --- a/apps/desktop/e2e/tile-unread-bug.spec.ts +++ b/apps/desktop/e2e/tile-unread-bug.spec.ts @@ -33,6 +33,8 @@ import { const UNREAD_DOT_LABEL = 'Finished — unread' /** Background-running dot aria-label. */ const BG_DOT_LABEL = 'Background task running' +/** Foreground turn-running dot aria-label. */ +const SESSION_RUNNING_DOT_LABEL = 'Session running' /** Locate a session's sidebar row by its preview text. */ function sessionRow(page: import('@playwright/test').Page, text: string) { @@ -65,12 +67,20 @@ async function startTurnAndSwitchAway(page: import('@playwright/test').Page) { ) .toBeGreaterThan(0) - // Wait for the turn to complete (final answer visible). + // The final answer text streams before message.complete, so text visibility + // alone is not a completion barrier. Wait for the foreground-running state + // to clear before asserting the background-process state. await page.waitForFunction( (text) => (document.body.textContent ?? '').includes(text), SIDEBAR_CROSS_TEXTS.finalText, { timeout: 90_000 }, ) + await expect + .poll( + () => page.locator(`[aria-label="${SESSION_RUNNING_DOT_LABEL}"]`).count(), + { timeout: 30_000, message: 'session running dot should disappear after turn completes' }, + ) + .toBe(0) // The background dot must still be visible: the turn is done but the // process is held open by the sentinel, so this is a stable state rather diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts index 6cee7a2469..4592653758 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts @@ -602,6 +602,22 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { explicitSid && sessionId ? modelOptionsQueryKey(activeGatewayProfile, sessionId) : ['model-options'] }) } + } else if (event.type === 'session.usage') { + // Live usage tick emitted while a turn is mid-flight (see tui_gateway + // _start_usage_ticker) so the status-bar context window tracks growth + // during the turn instead of only jumping at message.complete. + if (payload?.usage && sessionId) { + // Per-session twin first: a focused secondary tile reads this cache, + // while the primary-only global mirrors the active session. + updateSessionState(sessionId, state => ({ + ...state, + usage: { calls: 0, input: 0, output: 0, total: 0, ...state.usage, ...payload.usage } + })) + + if (isActiveEvent) { + setCurrentUsage(current => ({ ...current, ...payload.usage })) + } + } } else if (event.type === 'message.start') { if (!sessionId) { return diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/usage.test.tsx b/apps/desktop/src/app/session/hooks/use-message-stream/usage.test.tsx new file mode 100644 index 0000000000..924e845dd0 --- /dev/null +++ b/apps/desktop/src/app/session/hooks/use-message-stream/usage.test.tsx @@ -0,0 +1,134 @@ +import { QueryClient } from '@tanstack/react-query' +import { act, cleanup, render, waitFor } from '@testing-library/react' +import { useEffect, useRef } from 'react' +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 type { RpcEvent } from '@/types/hermes' + +import { useMessageStream } from './index' + +const SID = 'session-1' +// $currentUsage mirrors the primary session; ClientSessionState.usage drives +// the same status bar when a secondary tile is focused. +const BASELINE = { calls: 2, input: 500, output: 40, total: 540 } + +let handleEvent: ((event: RpcEvent) => void) | null = null +let sessionStates = new Map() + +function Harness() { + const activeSessionIdRef = useRef(SID) + const sessionStateByRuntimeIdRef = useRef(sessionStates) + const queryClientRef = useRef(new QueryClient()) + + const stream = useMessageStream({ + activeSessionIdRef, + hydrateFromStoredSession: vi.fn(async () => undefined), + queryClient: queryClientRef.current, + refreshHermesConfig: vi.fn(async () => undefined), + refreshSessions: vi.fn(async () => undefined), + sessionStateByRuntimeIdRef, + updateSessionState: (sessionId, updater) => { + const current = sessionStateByRuntimeIdRef.current.get(sessionId) ?? createClientSessionState() + const next = updater(current) + sessionStateByRuntimeIdRef.current.set(sessionId, next) + + return next + } + }) + + useEffect(() => { + handleEvent = stream.handleGatewayEvent + }, [stream.handleGatewayEvent]) + + return null +} + +async function mountStream() { + render() + await waitFor(() => expect(handleEvent).not.toBeNull()) +} + +describe('useMessageStream status-bar usage scoping', () => { + beforeEach(() => { + handleEvent = null + sessionStates = new Map([[SID, { ...createClientSessionState(), usage: { ...BASELINE } }]]) + $currentUsage.set({ ...BASELINE }) + }) + + afterEach(() => { + cleanup() + vi.restoreAllMocks() + }) + + it('merges a live session.usage tick from the focused session', async () => { + await mountStream() + + act(() => + handleEvent!({ + payload: { usage: { context_percent: 42, input: 1200, total: 1280 } }, + session_id: SID, + type: 'session.usage' + }) + ) + + // Merge, not replace: fields absent from the tick keep their prior values. + expect($currentUsage.get()).toEqual({ ...BASELINE, context_percent: 42, input: 1200, total: 1280 }) + expect(sessionStates.get(SID)?.usage).toEqual({ + ...BASELINE, + context_percent: 42, + input: 1200, + total: 1280 + }) + }) + + it('caches a background session.usage tick without overwriting the primary status bar', async () => { + await mountStream() + + act(() => + handleEvent!({ + payload: { usage: { input: 9999, total: 9999 } }, + session_id: 'background-session', + type: 'session.usage' + }) + ) + + expect($currentUsage.get()).toEqual(BASELINE) + expect(sessionStates.get('background-session')?.usage).toEqual({ + calls: 0, + input: 9999, + output: 0, + total: 9999 + }) + }) + + it('applies message.complete usage from the focused session', async () => { + await mountStream() + + act(() => + handleEvent!({ + payload: { text: 'done', usage: { calls: 3, input: 1500, output: 90, total: 1590 } }, + session_id: SID, + type: 'message.complete' + }) + ) + + expect($currentUsage.get()).toEqual({ calls: 3, input: 1500, output: 90, total: 1590 }) + }) + + it('ignores message.complete usage from a background session', async () => { + await mountStream() + + act(() => + handleEvent!({ + payload: { text: 'done', usage: { calls: 9, input: 9999, output: 999, total: 9999 } }, + session_id: 'background-session', + type: 'message.complete' + }) + ) + + expect($currentUsage.get()).toEqual(BASELINE) + }) +}) diff --git a/apps/shared/src/json-rpc-gateway.ts b/apps/shared/src/json-rpc-gateway.ts index 4f24737ed7..acde83e96c 100644 --- a/apps/shared/src/json-rpc-gateway.ts +++ b/apps/shared/src/json-rpc-gateway.ts @@ -1,6 +1,7 @@ export type GatewayEventName = | 'gateway.ready' | 'session.info' + | 'session.usage' | 'message.start' | 'message.delta' | 'message.interim' diff --git a/tests/gateway/test_multiplex_busy_input_mode.py b/tests/gateway/test_multiplex_busy_input_mode.py index c7cfbfe1f7..3a2f901df7 100644 --- a/tests/gateway/test_multiplex_busy_input_mode.py +++ b/tests/gateway/test_multiplex_busy_input_mode.py @@ -321,8 +321,15 @@ async def test_missing_or_invalid_secondary_mode_falls_back_to_gateway_default( assert runner._busy_text_mode == "queue" -def test_profile_route_and_nonmultiplexed_resolution_preserve_boundaries(): +def test_profile_route_and_nonmultiplexed_resolution_preserve_boundaries( + tmp_path, + monkeypatch, +): runner = _runner(default_mode="interrupt") + monkeypatch.setattr( + "hermes_cli.profiles.profiles_to_serve", + lambda **_: [("research", tmp_path / "research")], + ) runner._snapshot_profile_busy_modes( "research", {"display": {"busy_input_mode": "steer"}}, diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 25a29a4d06..12e9e49e87 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -1046,6 +1046,312 @@ def test_write_json_drops_detached_ws_frames(monkeypatch): server._sessions.pop("detached-sid", None) +def test_usage_ticker_emits_wrapped_usage_payload(monkeypatch): + # The live ticker must nest the snapshot under a "usage" key, matching the + # message.complete / session.info payloads the desktop & TUI handlers read + # as payload.usage. Emitting the bare _get_usage() dict (payload.input/total + # …) silently drops every live tick on the client side. + events: list[tuple[str, str, dict]] = [] + monkeypatch.setattr( + server, "_emit", lambda event_type, sid, payload: events.append((event_type, sid, payload)) + ) + snapshot = {"input": 1200, "total": 1280} + monkeypatch.setattr(server, "_get_usage", lambda agent: dict(snapshot)) + + stop, thread = server._start_usage_ticker("sess-1", object(), interval=0.01) + # The dedup baseline is sampled synchronously inside _start_usage_ticker, + # so this mutation is guaranteed to read as the first counter movement. + snapshot["total"] = 2400 + try: + deadline = time.time() + 1.0 + while not events and time.time() < deadline: + time.sleep(0.01) + finally: + stop.set() + thread.join(timeout=2.0) + + assert events, "ticker never emitted" + event_type, sid, payload = events[0] + assert event_type == "session.usage" + assert sid == "sess-1" + assert payload == {"usage": {"input": 1200, "total": 2400}} + + +def test_usage_ticker_skips_unchanged_snapshots(monkeypatch): + # A single long API call leaves the token counters frozen for many + # intervals; the ticker must emit nothing at all (the client already has + # the turn-start values from the previous message.complete / session.info). + # Only a changed snapshot emits. + events: list[dict] = [] + monkeypatch.setattr( + server, "_emit", lambda event_type, sid, payload: events.append(payload) + ) + snapshot = {"input": 1200, "total": 1280} + monkeypatch.setattr(server, "_get_usage", lambda agent: dict(snapshot)) + + stop, thread = server._start_usage_ticker("sess-1", object(), interval=0.01) + try: + # ~15 ticks with counters frozen at the turn-start baseline: zero frames. + time.sleep(0.15) + assert events == [] + + # Counters move → the next tick emits the new snapshot. + snapshot["total"] = 2400 + deadline = time.time() + 1.0 + while not events and time.time() < deadline: + time.sleep(0.01) + finally: + stop.set() + thread.join(timeout=2.0) + + assert events == [{"usage": {"input": 1200, "total": 2400}}] + + +def test_usage_ticker_baseline_sampled_before_thread_start(monkeypatch): + """The dedup baseline must be sampled synchronously in _start_usage_ticker, + not inside the ticker thread: a late-scheduled thread would otherwise seed + itself with counters the turn's first API call already bumped, absorbing + that first growth so it never emits.""" + events: list[dict] = [] + monkeypatch.setattr( + server, "_emit", lambda event_type, sid, payload: events.append(payload) + ) + snapshot = {"input": 1200, "total": 1280} + monkeypatch.setattr(server, "_get_usage", lambda agent: dict(snapshot)) + + class _SlowStartThread(threading.Thread): + def start(self): + # Deterministic stand-in for a scheduler delay: the turn's first + # API call bumps the counters before the ticker thread ever runs. + snapshot["total"] = 2400 + super().start() + + monkeypatch.setattr(server, "_RealThread", _SlowStartThread) + + stop, thread = server._start_usage_ticker("sess-1", object(), interval=0.01) + try: + deadline = time.time() + 1.0 + while not events and time.time() < deadline: + time.sleep(0.01) + finally: + stop.set() + thread.join(timeout=2.0) + + # An in-thread seed would have read 2400 as the baseline and stayed + # silent; the synchronous seed (1280) sees it as the first growth. + assert events == [{"usage": {"input": 1200, "total": 2400}}] + + +def test_usage_ticker_stop_join_prevents_late_ticks(monkeypatch): + """The stop sequence (set + join) must guarantee no session.usage after it + returns: a tick captured mid-turn but emitted after message.complete would + roll the client's final usage back to a stale snapshot (clients merge + payload.usage unconditionally).""" + events: list[str] = [] + monkeypatch.setattr( + server, "_emit", lambda event_type, sid, payload: events.append(event_type) + ) + + in_snapshot = threading.Event() + release = threading.Event() + calls = {"n": 0} + + def _blocking_get_usage(agent): + calls["n"] += 1 + if calls["n"] == 1: + return {"total": 0} # dedup seed + # First real tick: hold it mid-snapshot so the stop lands while the + # iteration is already past the stop.wait() gate. + in_snapshot.set() + release.wait(2.0) + return {"total": 999} + + monkeypatch.setattr(server, "_get_usage", _blocking_get_usage) + + stop, thread = server._start_usage_ticker("sess-1", object(), interval=0.01) + assert in_snapshot.wait(2.0), "ticker never reached a snapshot" + + # Turn ends while the tick is mid-snapshot: run the exact stop sequence + # _run_prompt_submit uses, then emit message.complete. + stop.set() + release.set() + thread.join(timeout=2.0) + assert not thread.is_alive(), "ticker thread survived the stop sequence" + server._emit("message.complete", "sess-1", {}) + + # The in-flight tick was dropped (stop re-checked before emit), so nothing + # can land after — let alone overwrite — the final usage. + assert "session.usage" not in events + assert events[-1] == "message.complete" + + +def test_run_prompt_submit_never_ticks_after_message_complete(monkeypatch): + """End-to-end ordering through _run_prompt_submit: live session.usage ticks + happen strictly before message.complete, never after it.""" + events: list[str] = [] + tick_seen = threading.Event() + + def _record_emit(event_type, sid, payload=None): + events.append(event_type) + if event_type == "session.usage": + tick_seen.set() + + monkeypatch.setattr(server, "_emit", _record_emit) + + counter = {"n": 0} + + def _moving_usage(agent): + counter["n"] += 1 + return {"total": counter["n"]} # moves every sample → every tick emits + + monkeypatch.setattr(server, "_get_usage", _moving_usage) + real_ticker = server._start_usage_ticker + monkeypatch.setattr( + server, + "_start_usage_ticker", + lambda sid, agent, interval=1.0: real_ticker(sid, agent, interval=0.01), + ) + monkeypatch.setattr(server.threading, "Thread", _ImmediateThread) + monkeypatch.setattr(server, "make_stream_renderer", lambda cols: None) + monkeypatch.setattr(server, "render_message", lambda raw, cols: None) + monkeypatch.setattr(server, "_get_db", lambda: None) + + class _Agent: + def run_conversation( + self, + prompt, + conversation_history=None, + stream_callback=None, + persist_user_message=None, + ): + # Hold the turn open until at least one live tick has fired. + assert tick_seen.wait(5.0), "no live tick during the turn" + return {"final_response": "done", "messages": [], "completed": True} + + server._sessions["sid"] = _session(agent=_Agent()) + try: + server.handle_request( + { + "id": "1", + "method": "prompt.submit", + "params": {"session_id": "sid", "text": "hello"}, + } + ) + finally: + server._sessions.pop("sid", None) + + assert "session.usage" in events + assert "message.complete" in events + last_tick = max(i for i, e in enumerate(events) if e == "session.usage") + assert last_tick < events.index("message.complete") + + +def test_usage_ticker_unbounded_join_waits_out_blocked_emit(monkeypatch): + """A tick stalled inside _emit (a transport write can block up to + _WS_WRITE_TIMEOUT_S = 10s on a stalled event loop) must be waited out by + the stop sequence, not abandoned: stop.set() + an unbounded join may only + return after the in-flight emit has fully flushed, so nothing can land + after message.complete.""" + order: list[str] = [] + in_emit = threading.Event() + release = threading.Event() + + def _stalled_emit(event_type, sid, payload): + in_emit.set() + release.wait(10.0) # the stalled transport write + order.append(event_type) + + monkeypatch.setattr(server, "_emit", _stalled_emit) + + counter = {"n": 0} + + def _moving_usage(agent): + counter["n"] += 1 + return {"total": counter["n"]} # moves every sample → a tick emits + + monkeypatch.setattr(server, "_get_usage", _moving_usage) + + stop, thread = server._start_usage_ticker("sess-1", object(), interval=0.01) + assert in_emit.wait(2.0), "no tick got in flight" + + # Run the exact stop sequence _run_prompt_submit uses, on a side thread so + # the test can observe whether it returns while the emit is still stuck. + stopped = threading.Event() + + def _stop_sequence(): + stop.set() + thread.join() + stopped.set() + + stopper = threading.Thread(target=_stop_sequence, daemon=True) + stopper.start() + + # While the tick is stalled in the transport write, the stop sequence must + # NOT complete — a timed join returning here is exactly the bug: the + # caller would proceed to message.complete with the tick still pending. + assert not stopped.wait(0.2), "stop sequence returned with the tick still in flight" + + release.set() + assert stopped.wait(2.0), "stop sequence never completed after the emit flushed" + stopper.join(timeout=2.0) + + # The flushed tick strictly precedes anything the caller emits afterwards. + order.append("message.complete") + assert order == ["session.usage", "message.complete"] + + +def test_run_prompt_submit_joins_ticker_without_timeout(monkeypatch): + """_run_prompt_submit must join the ticker with NO timeout. A timed join + can expire while a tick sits in a stalled transport write (up to + _WS_WRITE_TIMEOUT_S = 10s) and abandon it to land after message.complete; + the wait is bounded by that same write anyway — message.complete's own + emit would stall on the same transport.""" + joins: list = [] + real_ticker = server._start_usage_ticker + + class _JoinSpy: + def __init__(self, thread): + self._thread = thread + + def join(self, timeout=None): + joins.append(timeout) + return self._thread.join(timeout) + + def __getattr__(self, name): + return getattr(self._thread, name) + + def _spying_ticker(sid, agent, interval=1.0): + stop, thread = real_ticker(sid, agent, interval=interval) + return stop, _JoinSpy(thread) + + monkeypatch.setattr(server, "_start_usage_ticker", _spying_ticker) + monkeypatch.setattr(server.threading, "Thread", _ImmediateThread) + monkeypatch.setattr(server, "_emit", lambda event_type, sid, payload=None: None) + monkeypatch.setattr(server, "make_stream_renderer", lambda cols: None) + monkeypatch.setattr(server, "render_message", lambda raw, cols: None) + monkeypatch.setattr(server, "_get_db", lambda: None) + + class _Agent: + def run_conversation( + self, prompt, conversation_history=None, stream_callback=None + ): + return {"final_response": "done", "messages": [], "completed": True} + + server._sessions["sid"] = _session(agent=_Agent()) + try: + server.handle_request( + { + "id": "1", + "method": "prompt.submit", + "params": {"session_id": "sid", "text": "hello"}, + } + ) + finally: + server._sessions.pop("sid", None) + + assert joins == [None], f"expected one unbounded join, got {joins}" + + def test_tui_verbose_tool_details_fail_closed_when_redaction_fails(monkeypatch): redact_module = types.ModuleType("agent.redact") diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 62f5725a62..00896954bc 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -10138,6 +10138,70 @@ def _plan_goal_compression_recovery( ) +# Captured at import time. Several _run_prompt_submit tests monkeypatch +# threading.Thread with a stub that runs the target synchronously to keep the +# turn deterministic. This ticker's loop only exits once the caller sets `stop` +# *after* run_conversation returns, so running it inline would spin forever. +# It's a non-critical, fire-and-forget background poller, so it always uses a +# real daemon thread regardless of any such patch. +_RealThread = threading.Thread + + +def _start_usage_ticker( + sid: str, agent, interval: float = 1.0 +) -> tuple[threading.Event, threading.Thread]: + """Push live usage snapshots while a turn runs. + + The desktop/TUI status-bar context-window figure is otherwise refreshed + only at ``message.complete``, so it stays frozen for the whole (often + multi-minute, multi-tool) turn. On the standard chat-completions path the + agent's token counters grow after every internal API call, so this daemon + emits a lightweight ``session.usage`` event every ``interval`` seconds and + the bar tracks context growth live. (The codex app-server runtime folds + usage into the counters only at turn end — codex_runtime. + _record_codex_app_server_usage — so it gets no mid-turn ticks; its final + value still lands via ``message.complete``.) + + The caller must set the returned Event AND join the returned thread + before emitting ``message.complete``: a tick that survived past it would + roll the client's final usage back to a stale mid-turn snapshot. + """ + stop = threading.Event() + + # Sample the dedup baseline BEFORE the thread starts: the client already + # has the turn-start values from the previous message.complete / + # session.info. Seeding here (not in the thread) guarantees the baseline + # predates the turn's first API call — a late-scheduled thread would + # otherwise absorb that first counter growth and never emit it. + try: + baseline: dict | None = _get_usage(agent) + except Exception: + baseline = None + + def _loop() -> None: + last = baseline + while not stop.wait(interval): + try: + usage = _get_usage(agent) + if usage == last: + # Counters frozen (e.g. one long API call in flight) — + # skip the redundant frame so idle ticks don't re-render + # the client status bar every second. + continue + last = usage + if stop.is_set(): + # Turn ended while snapshotting — drop the tick; + # message.complete carries the authoritative usage. + break + _emit("session.usage", sid, {"usage": usage}) + except Exception: + pass + + thread = _RealThread(target=_loop, daemon=True) + thread.start() + return stop, thread + + def _run_prompt_submit( rid, sid: str, @@ -10455,7 +10519,21 @@ def _run_prompt_submit( agent._on_session_title = lambda t, _src, _k=_title_key: _emit( "session.title", sid, {"session_id": _k, "title": t} ) - result = agent.run_conversation(run_message, **run_kwargs) + _usage_stop, _usage_thread = _start_usage_ticker(sid, agent) + try: + result = agent.run_conversation(run_message, **run_kwargs) + finally: + # Stop AND join before anything below emits: an in-flight tick + # surviving past message.complete would roll the client's final + # usage back to a stale mid-turn snapshot. The join is + # deliberately unbounded — once stop is set it only ever waits + # out one in-flight _get_usage/_emit, and the worst case there + # (a stalled transport write, up to _WS_WRITE_TIMEOUT_S) would + # stall the message.complete emit below just the same. A + # timed-out join would abandon the tick to land after + # message.complete. + _usage_stop.set() + _usage_thread.join() if display_kind and isinstance(text, str): db = getattr(agent, "_session_db", None) current_session_id = getattr(agent, "session_id", None) or session.get("session_key") diff --git a/ui-tui/src/__tests__/createGatewayEventHandler.test.ts b/ui-tui/src/__tests__/createGatewayEventHandler.test.ts index a49a7cec49..2b8fc76184 100644 --- a/ui-tui/src/__tests__/createGatewayEventHandler.test.ts +++ b/ui-tui/src/__tests__/createGatewayEventHandler.test.ts @@ -5,6 +5,7 @@ import { getOverlayState, patchOverlayState, resetOverlayState } from '../app/ov import { turnController } from '../app/turnController.js' import { getTurnState, resetTurnState } from '../app/turnStore.js' import { getUiState, patchUiState, resetUiState } from '../app/uiStore.js' +import { ZERO } from '../domain/usage.js' import { estimateTokensRough } from '../lib/text.js' import type { Msg } from '../types.js' @@ -1938,6 +1939,47 @@ describe('createGatewayEventHandler', () => { }) }) + describe('session.usage', () => { + it('merges a live usage tick into uiState (payload.usage shape, see tui_gateway _start_usage_ticker)', () => { + patchUiState({ sid: 'sess-1' }) + const onEvent = createGatewayEventHandler(buildCtx([])) + + onEvent({ + payload: { usage: { calls: 3, context_percent: 42, input: 1200, output: 80, total: 1280 } }, + session_id: 'sess-1', + type: 'session.usage' + } as any) + + expect(getUiState().usage).toMatchObject({ context_percent: 42, input: 1200, total: 1280 }) + }) + + it('keeps existing usage fields when the tick only carries a subset', () => { + patchUiState({ sid: 'sess-1', usage: { calls: 2, input: 500, output: 40, total: 540 } }) + const onEvent = createGatewayEventHandler(buildCtx([])) + + onEvent({ + payload: { usage: { context_percent: 55 } }, + session_id: 'sess-1', + type: 'session.usage' + } as any) + + expect(getUiState().usage).toMatchObject({ context_percent: 55, input: 500, total: 540 }) + }) + + it('drops a tick for a non-focused session', () => { + patchUiState({ sid: 'focused', usage: ZERO }) + const onEvent = createGatewayEventHandler(buildCtx([])) + + onEvent({ + payload: { usage: { input: 9999, total: 9999 } }, + session_id: 'background', + type: 'session.usage' + } as any) + + expect(getUiState().usage).toEqual(ZERO) + }) + }) + describe('message.interim', () => { it('finalizes an interim segment without settling the turn', () => { const appended: Msg[] = [] diff --git a/ui-tui/src/app/createGatewayEventHandler.ts b/ui-tui/src/app/createGatewayEventHandler.ts index 9507cb0854..798c6e4726 100644 --- a/ui-tui/src/app/createGatewayEventHandler.ts +++ b/ui-tui/src/app/createGatewayEventHandler.ts @@ -782,6 +782,21 @@ export function createGatewayEventHandler(ctx: GatewayEventHandlerContext): (ev: return } + case 'session.usage': { + // Live usage tick while a turn runs (see tui_gateway + // _start_usage_ticker) — keeps the status-bar context window current + // mid-turn instead of only at message.complete. The session filter at + // the top of this handler already dropped ticks for non-focused + // sessions. + const usage = ev.payload?.usage + + if (usage) { + patchUiState(state => ({ ...state, usage: { ...state.usage, ...usage } })) + } + + return + } + case 'thinking.delta': { if (!getUiState().busy) { return diff --git a/ui-tui/src/gatewayTypes.ts b/ui-tui/src/gatewayTypes.ts index ba8e85c42a..d5ab544ab5 100644 --- a/ui-tui/src/gatewayTypes.ts +++ b/ui-tui/src/gatewayTypes.ts @@ -739,4 +739,5 @@ export type GatewayEvent = session_id?: string type: 'message.complete' } + | { payload?: { usage?: Usage }; session_id?: string; type: 'session.usage' } | { payload?: { message?: string }; session_id?: string; type: 'error' }