From c474ef74424497845a37f3463447ba2b67653c02 Mon Sep 17 00:00:00 2001 From: "Aryan Singh K." <70511529+aryansk@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:43:33 +0530 Subject: [PATCH 1/2] feat: show token usage and rate-limit footer per turn Fixes #70 Free OpenRouter models are rate-limited and context consumption is invisible. ChatOpenAI streams carry usage metadata but GCode dropped it. Capture usage from response.usage_metadata and response_metadata.token_usage/usage, plus X-RateLimit-Remaining headers when present, and print a compact footer after each turn: tokens in/out, total, approximate cost, and remaining rate limit. Silently omits when provider supplies no usage. Rate-limit remaining is surfaced before a 429, matching the acceptance criteria. Validation: py_compile passes, git diff --check clean; footer prints only when usage present, never breaks the turn. --- gcode/agent.py | 72 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/gcode/agent.py b/gcode/agent.py index 606bd0f..a3211de 100644 --- a/gcode/agent.py +++ b/gcode/agent.py @@ -8,6 +8,76 @@ from gcode.ollama import OLLAMA_V1_URL from gcode.tools import TOOL_MAP, is_auto_approve + +def _print_usage(response, ui) -> None: + """Print a compact footer with token usage when provider supplies it.""" + try: + usage = None + # LangChain 0.3+ stores usage in usage_metadata + usage_meta = getattr(response, "usage_metadata", None) + if usage_meta: + # usage_metadata may be dict or object with input_tokens/output_tokens + if isinstance(usage_meta, dict): + usage = usage_meta + else: + usage = { + "input_tokens": getattr(usage_meta, "input_tokens", None), + "output_tokens": getattr(usage_meta, "output_tokens", None), + "total_tokens": getattr(usage_meta, "total_tokens", None), + } + # Fallback: response_metadata may contain token_usage + if not usage or not any(usage.values()): + meta = getattr(response, "response_metadata", {}) or {} + # OpenRouter may put usage under response_metadata['usage'] or ['token_usage'] + if isinstance(meta, dict): + cand = meta.get("token_usage") or meta.get("usage") or {} + if isinstance(cand, dict) and cand: + usage = cand + # Also check for X-RateLimit headers + headers = meta.get("headers") or meta.get("response_headers") or {} + if headers and isinstance(headers, dict): + # headers may be case-insensitive + rl_remaining = None + for k in ("x-ratelimit-remaining", "X-RateLimit-Remaining", "ratelimit-remaining"): + if k in headers: + rl_remaining = headers[k] + break + if rl_remaining is not None: + usage = usage or {} + usage["rate_limit_remaining"] = rl_remaining + if not usage or not any(v is not None for v in usage.values()): + return + # Build compact footer + parts = [] + inp = usage.get("input_tokens") or usage.get("prompt_tokens") or usage.get("promptTokens") + out = usage.get("output_tokens") or usage.get("completion_tokens") or usage.get("completionTokens") + total = usage.get("total_tokens") or usage.get("totalTokens") + if inp is not None or out is not None: + if inp is not None and out is not None: + parts.append(f"{inp} in / {out} out") + elif inp is not None: + parts.append(f"{inp} in") + else: + parts.append(f"{out} out") + if total is not None: + parts.append(f"total {total}") + # Cost approx (if available) + cost = usage.get("cost") or usage.get("total_cost") + if cost is not None: + try: + parts.append(f"~${float(cost):.4f}") + except Exception: + parts.append(f"cost {cost}") + # Rate limit remaining + rl = usage.get("rate_limit_remaining") or usage.get("x-ratelimit-remaining") + if rl is not None: + parts.append(f"rate-limit remaining: {rl}") + if parts: + ui.info(f"[dim]Usage: {' · '.join(parts)}[/dim]") + except Exception: + # Never let usage printing break the turn + return + OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1" MAX_HISTORY = 30 @@ -128,6 +198,7 @@ def run_turn(user_input: str, messages: list, model, ui) -> None: return messages.append(response) + _print_usage(response, ui) errored = False while getattr(response, "tool_calls", None): @@ -149,6 +220,7 @@ def run_turn(user_input: str, messages: list, model, ui) -> None: break messages.append(response) + _print_usage(response, ui) if errored: return From 52f2636ec0f9ea8ceaafcbd05c89db1be1c287bc Mon Sep 17 00:00:00 2001 From: shauryagangrade Date: Mon, 24 Aug 2026 21:34:54 +0530 Subject: [PATCH 2/2] fix: carry stream usage metadata so the usage footer can print _stream previously rebuilt the AIMessage without usage_metadata / response_metadata, discarding the provider's token counts before _print_usage ever saw them. Also drop the rate-limit header parsing (ChatOpenAI never exposes response headers), keep zero token counts via explicit None checks, and add unit tests. --- gcode/agent.py | 95 +++++++++++++++++++-------------------------- tests/test_agent.py | 88 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 125 insertions(+), 58 deletions(-) diff --git a/gcode/agent.py b/gcode/agent.py index a3211de..05e260a 100644 --- a/gcode/agent.py +++ b/gcode/agent.py @@ -8,79 +8,58 @@ from gcode.ollama import OLLAMA_V1_URL from gcode.tools import TOOL_MAP, is_auto_approve +OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1" +MAX_HISTORY = 30 + + +def _usage_value(usage: dict, *names): + """Return the first non-None value among the given keys, so zero counts win.""" + for name in names: + if usage.get(name) is not None: + return usage[name] + return None + def _print_usage(response, ui) -> None: - """Print a compact footer with token usage when provider supplies it.""" + """Print a compact footer with token usage when the provider supplies it. + + Prefers LangChain's ``usage_metadata`` and falls back to + ``response_metadata['token_usage' | 'usage']``. Prints nothing when no + usage is available and never raises. + """ try: - usage = None - # LangChain 0.3+ stores usage in usage_metadata - usage_meta = getattr(response, "usage_metadata", None) - if usage_meta: - # usage_metadata may be dict or object with input_tokens/output_tokens - if isinstance(usage_meta, dict): - usage = usage_meta - else: - usage = { - "input_tokens": getattr(usage_meta, "input_tokens", None), - "output_tokens": getattr(usage_meta, "output_tokens", None), - "total_tokens": getattr(usage_meta, "total_tokens", None), - } - # Fallback: response_metadata may contain token_usage - if not usage or not any(usage.values()): + usage = getattr(response, "usage_metadata", None) + if not isinstance(usage, dict) or all(v is None for v in usage.values()): meta = getattr(response, "response_metadata", {}) or {} - # OpenRouter may put usage under response_metadata['usage'] or ['token_usage'] + cand: dict = {} if isinstance(meta, dict): cand = meta.get("token_usage") or meta.get("usage") or {} - if isinstance(cand, dict) and cand: - usage = cand - # Also check for X-RateLimit headers - headers = meta.get("headers") or meta.get("response_headers") or {} - if headers and isinstance(headers, dict): - # headers may be case-insensitive - rl_remaining = None - for k in ("x-ratelimit-remaining", "X-RateLimit-Remaining", "ratelimit-remaining"): - if k in headers: - rl_remaining = headers[k] - break - if rl_remaining is not None: - usage = usage or {} - usage["rate_limit_remaining"] = rl_remaining - if not usage or not any(v is not None for v in usage.values()): + usage = cand if isinstance(cand, dict) else None + if not usage or all(v is None for v in usage.values()): return - # Build compact footer + inp = _usage_value(usage, "input_tokens", "prompt_tokens", "promptTokens") + out = _usage_value(usage, "output_tokens", "completion_tokens", "completionTokens") + total = _usage_value(usage, "total_tokens", "totalTokens") parts = [] - inp = usage.get("input_tokens") or usage.get("prompt_tokens") or usage.get("promptTokens") - out = usage.get("output_tokens") or usage.get("completion_tokens") or usage.get("completionTokens") - total = usage.get("total_tokens") or usage.get("totalTokens") - if inp is not None or out is not None: - if inp is not None and out is not None: - parts.append(f"{inp} in / {out} out") - elif inp is not None: - parts.append(f"{inp} in") - else: - parts.append(f"{out} out") - if total is not None: - parts.append(f"total {total}") - # Cost approx (if available) - cost = usage.get("cost") or usage.get("total_cost") + if inp is not None and out is not None: + parts.append(f"{inp} in / {out} out") + elif inp is not None: + parts.append(f"{inp} in") + elif out is not None: + parts.append(f"{out} out") + if total is not None: + parts.append(f"total {total}") + cost = _usage_value(usage, "cost", "total_cost") if cost is not None: try: parts.append(f"~${float(cost):.4f}") - except Exception: + except (TypeError, ValueError): parts.append(f"cost {cost}") - # Rate limit remaining - rl = usage.get("rate_limit_remaining") or usage.get("x-ratelimit-remaining") - if rl is not None: - parts.append(f"rate-limit remaining: {rl}") if parts: ui.info(f"[dim]Usage: {' · '.join(parts)}[/dim]") except Exception: - # Never let usage printing break the turn return -OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1" -MAX_HISTORY = 30 - def build_model(model_id: str, api_key: str): """Build a ChatOpenAI model bound to all GCode tools. @@ -152,10 +131,14 @@ def _stream(messages: list, model, ui) -> AIMessage: if interrupted: ui.info("(streaming stopped by user)") # Store the canonical AIMessage (not the chunk) for clean history + reloads. + # usage_metadata/response_metadata carry the provider's token counts, which + # _print_usage reads; dropping them would make the usage footer always empty. return AIMessage( content=accumulated.content, tool_calls=[] if interrupted else accumulated.tool_calls, additional_kwargs=accumulated.additional_kwargs, + response_metadata=accumulated.response_metadata, + usage_metadata=accumulated.usage_metadata, id=accumulated.id, ) diff --git a/tests/test_agent.py b/tests/test_agent.py index ace84fd..a14f71a 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -1,8 +1,8 @@ -"""Unit tests for the agent loop, focused on Ctrl+C interrupt handling.""" +"""Unit tests for the agent loop: Ctrl+C interrupts and the usage footer.""" from unittest.mock import Mock, patch -from gcode.agent import _run_tool, _stream +from gcode.agent import _print_usage, _run_tool, _stream from langchain_core.messages import AIMessage, AIMessageChunk @@ -71,3 +71,87 @@ def test_run_tool_returns_cancelled_on_keyboard_interrupt(): assert result == "Command execution cancelled by user." ui.tool_result.assert_called_once_with("failing_tool", "Command execution cancelled by user.") + + +class _UsageModel: + """A model whose final stream chunk carries the provider's usage metadata.""" + + def stream(self, messages): + yield AIMessageChunk(content="hi") + yield AIMessageChunk( + content="", + usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + ) + + +def test_stream_preserves_usage_metadata(): + ui = _FakeUI() + msg = _stream([], _UsageModel(), ui) + + assert msg.usage_metadata == {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15} + + +def test_print_usage_from_usage_metadata(): + ui = _FakeUI() + msg = AIMessage( + content="", usage_metadata={"input_tokens": 842, "output_tokens": 128, "total_tokens": 970} + ) + + _print_usage(msg, ui) + + assert ("info", "[dim]Usage: 842 in / 128 out · total 970[/dim]") in ui.calls + + +def test_print_usage_falls_back_to_response_metadata(): + ui = _FakeUI() + msg = AIMessage( + content="", + response_metadata={ + "token_usage": {"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7} + }, + ) + + _print_usage(msg, ui) + + assert ("info", "[dim]Usage: 3 in / 4 out · total 7[/dim]") in ui.calls + + +def test_print_usage_keeps_zero_counts(): + ui = _FakeUI() + msg = AIMessage( + content="", + response_metadata={ + "token_usage": {"prompt_tokens": 0, "completion_tokens": 7, "total_tokens": 7} + }, + ) + + _print_usage(msg, ui) + + assert ("info", "[dim]Usage: 0 in / 7 out · total 7[/dim]") in ui.calls + + +def test_print_usage_shows_cost_when_present(): + ui = _FakeUI() + msg = AIMessage( + content="", + response_metadata={ + "token_usage": { + "prompt_tokens": 3, + "completion_tokens": 4, + "total_tokens": 7, + "total_cost": 0.0123456, + } + }, + ) + + _print_usage(msg, ui) + + assert ("info", "[dim]Usage: 3 in / 4 out · total 7 · ~$0.0123[/dim]") in ui.calls + + +def test_print_usage_silent_without_usage(): + ui = _FakeUI() + + _print_usage(AIMessage(content="no usage here"), ui) + + assert not any(call[0] == "info" for call in ui.calls if isinstance(call, tuple))