Skip to content

[codex] Stream live token usage into desktop context meter - #39201

Closed
OmarB97 wants to merge 4 commits into
NousResearch:mainfrom
OmarB97:fix/context-usage-realtime-upstream
Closed

[codex] Stream live token usage into desktop context meter#39201
OmarB97 wants to merge 4 commits into
NousResearch:mainfrom
OmarB97:fix/context-usage-realtime-upstream

Conversation

@OmarB97

@OmarB97 OmarB97 commented Jun 4, 2026

Copy link
Copy Markdown
Contributor

Summary

  • emit token usage snapshots from the agent loop after provider usage is recorded and bridge them through the gateway as token.usage events
  • emit a preflight token.usage snapshot before each API call using the already-computed request-size estimate, so providers that only report usage at completion still move the desktop context meter during the active turn
  • consume those events in the desktop session hook so the context meter can move before message.complete
  • add focused backend, desktop, and TUI coverage, plus a small spawn stdio typing fix so ui-tui type-checks cleanly

Why

The first live-usage patch handled final/canonical provider usage, but the official macOS app could still sit at 0/131.1k while a turn was running because many providers expose usage only at completion. The new preflight snapshot sends the current request/context estimate at call start, then the final provider usage corrects totals afterward.

Fork / install follow-up

Verification

  • python3 -m pytest tests/tui_gateway/test_token_usage_events.py
  • python3 -m pytest tests/run_agent/test_run_agent.py::TestRunConversation::test_preflight_token_usage_emits_before_api_response tests/run_agent/test_run_agent.py::TestRunConversation::test_stop_finish_reason_returns_response
  • python3 -m py_compile agent/conversation_loop.py
  • npx vitest run --environment jsdom src/lib/token-usage.test.ts src/app/session/hooks/use-message-stream.test.tsx from apps/desktop
  • npm run type-check --workspace ui-tui
  • Installed-code gateway canary on macOS: token.usage arrived at ~2.1s and message.complete arrived at ~11.9s for the same turn, proving the desktop can receive usage before end-of-turn.

Omar Baradei added 3 commits June 4, 2026 09:24
…ime context bar

After every API call that updates session token counters, emit a structured
token.usage event through the existing status_callback → _status_update →
WebSocket event pipe. This gives the desktop app's context bar component
(@nous-research/ui) the raw material to animate smoothly during a turn
instead of jumping to the final value at end-of-turn.

Changes:
- run_agent.py: add _emit_token_usage() method that JSON-encodes
  token counts and context dimensions, then fires via status_callback
- conversation_loop.py: call _emit_token_usage() after session_*_tokens +=
- tui_gateway/server.py: handle kind=token_usage in _status_update by
  parsing the JSON payload and re-emitting as a typed 'token.usage' event

Token.usage payload:
  { input_tokens, output_tokens, total_tokens, context_tokens,
    context_length, context_pct }

The frontend side (@nous-research/ui npm package) requires a separate
upstream change to subscribe to 'token.usage' events for the context bar.
@OmarB97
OmarB97 marked this pull request as ready for review June 4, 2026 16:44
@alt-glitch alt-glitch added type/feature New feature or request P3 Low — cosmetic, nice to have comp/agent Core agent runtime: loop, agent_init, prompt builder, context-compression, responses endpoint comp/tui Terminal UI (ui-tui/ + tui_gateway/) labels Jun 4, 2026
@OmarB97

OmarB97 commented Jun 5, 2026

Copy link
Copy Markdown
Contributor Author

Review: approve with suggestions

Reviewed the diff for live token usage streaming into the desktop/TUI context meter. The data flow is complete and correct: backend _emit_token_usagestatus_callback("token_usage", ...)tui_gateway/server.py intercepts → token.usage event → desktop use-message-stream.ts / TUI createGatewayEventHandler.ts → usage store update.

Verified

  • last_prompt_tokens is initialized to 0 in ContextCompressor.__init__ (line 650) — direct access in the post-API emit is safe.
  • Preflight emit fires before the first API call with estimate_request_tokens_rough, solving the "0/131k" problem.
  • Post-API emit fires after every successful _perform_api_call with actual usage from canonical_usage.
  • Both desktop and TUI side have defensive parsing (finiteNumber, nonNegativeNumber, positiveNumber) that rejects NaN/null/malformed values.
  • Session-scoped filtering: the desktop handler only updates the usage bar for the active session.

Non-blocking suggestions

  1. tui_gateway/server.py: Consider adding isinstance(payload, dict) guard before _emit("token.usage", ...) — prevents forwarding non-dict JSON values (already in the working tree but not in this diff).

  2. tests: Consider a test for the post-API-call _emit_token_usage path in test_run_agent.py, to mirror test_preflight_token_usage_emits_before_api_response.

  3. Style nit: import json inside the method body in tui_gateway/server.py is unconventional but functional.

Decision

Approve — no blocker-severity findings. The PR achieves the stated goal of streaming live context usage into the desktop/TUI meter.

@OmarB97

OmarB97 commented Jun 9, 2026

Copy link
Copy Markdown
Contributor Author

Superseded by #39370, which carries this slice (gateway emission, preflight context estimate, desktop + TUI consumption) rebased on current main. Two micro-slices not carried there (scripts/release.py author-email map; execFileNoThrow.ts SpawnOptions stdio typing fix) — will rehome in a future small PR if still wanted. Closing to reduce reviewer surface.

@OmarB97 OmarB97 closed this Jun 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

comp/agent Core agent runtime: loop, agent_init, prompt builder, context-compression, responses endpoint comp/tui Terminal UI (ui-tui/ + tui_gateway/) P3 Low — cosmetic, nice to have type/feature New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants