feat: show token usage and rate-limit footer per turn

gcode/agent.py

@@ -12,6 +12,55 @@
 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 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 = 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 {}
+            cand: dict = {}
+            if isinstance(meta, dict):
+                cand = meta.get("token_usage") or meta.get("usage") or {}
+            usage = cand if isinstance(cand, dict) else None
+        if not usage or all(v is None for v in usage.values()):
+            return
+        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 = []
+        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 (TypeError, ValueError):
+                parts.append(f"cost {cost}")
+        if parts:
+            ui.info(f"[dim]Usage: {'  ·  '.join(parts)}[/dim]")
+    except Exception:
+        return
+
+
 def build_model(model_id: str, api_key: str):
     """Build a ChatOpenAI model bound to all GCode tools.
 
@@ -82,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,
     )
 
@@ -128,6 +181,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 +203,7 @@ def run_turn(user_input: str, messages: list, model, ui) -> None:
             break
 
         messages.append(response)
+        _print_usage(response, ui)
 
     if errored:
         return

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))