perf: trim history by token budget, not message count

gcode/agent.py

@@ -90,19 +90,76 @@ def build_model(model_id: str, api_key: str):
     ).bind_tools(ALL_TOOLS)
 
 
+MAX_HISTORY_TOKENS = 12000  # approx budget for ~30 messages at ~400 tokens each
+
+
+def _estimate_tokens(msg) -> int:
+    """Heuristic token estimate for a message (len//4), with tiktoken if available."""
+    try:
+        content = getattr(msg, "content", "")
+        if isinstance(content, list):
+            # Content may be a list of parts (e.g., for tool calls)
+            text = "".join(
+                part.get("text", "") if isinstance(part, dict) else str(part) for part in content
+            )
+        else:
+            text = str(content) if content else ""
+        # Include tool_calls in estimate
+        tool_calls = getattr(msg, "tool_calls", None)
+        if tool_calls:
+            text += str(tool_calls)
+        # Try tiktoken if installed for more accurate count
+        try:
+            import tiktoken
+
+            enc = tiktoken.get_encoding("cl100k_base")
+            return len(enc.encode(text))
+        except Exception:
+            return max(1, len(text) // 4)  # tiktoken unavailable: heuristic
+    except Exception:
+        return 100  # fallback small budget
+
+
 def trim_history(messages: list) -> None:
-    """Keep the system message plus the most recent MAX_HISTORY messages.
+    """Keep history within message-count and token-budget limits.
 
-    Trims only at a settled boundary (between turns) and drops any leading
-    ToolMessages whose owning assistant message was trimmed, so the API never
-    sees an orphaned tool result.
+    Preserves the system message plus the most recent messages that fit within
+    ``MAX_HISTORY`` and ``MAX_HISTORY_TOKENS``. Trims only at a settled
+    boundary and drops any leading ToolMessages whose owning assistant was
+    trimmed, so the API never sees an orphaned tool result. A single huge
+    tool output is capped at the source (see :func:`gcode.tools.grep`).
     """
-    if len(messages) <= MAX_HISTORY + 1:
+    if len(messages) <= 1:
         return
-    tail = messages[-MAX_HISTORY:]
-    while tail and isinstance(tail[0], ToolMessage):
-        tail.pop(0)
-    messages[:] = [messages[0]] + tail
+    # Fast path: within both limits
+    if len(messages) <= MAX_HISTORY + 1:
+        total = sum(_estimate_tokens(m) for m in messages)
+        if total <= MAX_HISTORY_TOKENS:
+            return
+    # Need to trim: keep system message + most recent that fit
+    # Start from most recent and build backwards within budget
+    system = messages[0]
+    rest = messages[1:]
+    # Enforce count limit first, then token budget
+    if len(rest) > MAX_HISTORY:
+        rest = rest[-MAX_HISTORY:]
+        # Drop leading ToolMessages that would be orphaned
+        while rest and isinstance(rest[0], ToolMessage):
+            rest.pop(0)
+    # Enforce token budget by dropping oldest while over budget.
+    # Keep at least one recent turn (2 messages) if possible.
+    while len(rest) > 2:
+        if sum(_estimate_tokens(m) for m in [system] + rest) <= MAX_HISTORY_TOKENS:
+            break
+        dropped = rest.pop(0)
+        # Dropping an AIMessage that issued tool calls orphans its ToolMessages.
+        if getattr(dropped, "tool_calls", None):
+            while rest and isinstance(rest[0], ToolMessage):
+                rest.pop(0)
+    # Final orphan check: ensure rest doesn't start with ToolMessage
+    while rest and isinstance(rest[0], ToolMessage):
+        rest.pop(0)
+    messages[:] = [system] + rest
 
 
 def _stream(messages: list, model, ui) -> AIMessage:

gcode/tools.py

@@ -243,9 +243,22 @@ def grep(
             return f"No matches for {pattern!r} in {path}."
         if result.returncode != 0:
             return f"grep error: {result.stderr.strip()}"
-        return result.stdout.strip()
-
-    return _grep_python(pattern, path, glob, ignore_case)
+        out = result.stdout.strip()
+        # Cap oversized grep output so a single tool call can't blow the context window
+        if len(out) > 8000:
+            out = out[:8000] + "\n... [truncated at 8000 chars, showing first 8000]"
+        elif out.count("\n") > 200:
+            lines = out.splitlines()
+            out = "\n".join(lines[:200]) + f"\n... [truncated at 200 lines, {len(lines)} total]"
+        return out
+
+    out = _grep_python(pattern, path, glob, ignore_case)
+    if len(out) > 8000:
+        return out[:8000] + "\n... [truncated at 8000 chars, showing first 8000]"
+    if out.count("\n") > 200:
+        lines = out.splitlines()
+        return "\n".join(lines[:200]) + f"\n... [truncated at 200 lines, {len(lines)} total]"
+    return out
 
 
 def _is_binary(filepath: str) -> bool:

tests/test_agent.py

@@ -155,3 +155,77 @@ def test_print_usage_silent_without_usage():
     _print_usage(AIMessage(content="no usage here"), ui)
 
     assert not any(call[0] == "info" for call in ui.calls if isinstance(call, tuple))
+
+
+# -- trim_history -------------------------------------------------------------
+
+
+def test_trim_history_keeps_too_many_small_messages():
+    from gcode.agent import MAX_HISTORY, trim_history
+    from langchain_core.messages import HumanMessage
+
+    msgs = [HumanMessage(content=f"m{i}") for i in range(MAX_HISTORY + 5)]
+    system = HumanMessage(content="system")
+    history = [system] + msgs
+
+    trim_history(history)
+
+    assert history[0] is system
+    assert len(history) == MAX_HISTORY + 1
+
+
+def test_trim_history_trims_by_token_budget_not_count():
+    from gcode.agent import MAX_HISTORY_TOKENS, trim_history
+    from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+
+    system = HumanMessage(content="system")
+    # Varied text so the estimate is over budget whether tiktoken or len//4 is used
+    # (pure "-xxxxxxxx" compresses to almost nothing under BPE).
+    big = " ".join(f"term{i}" for i in range(MAX_HISTORY_TOKENS * 4))
+    history = [
+        system,
+        HumanMessage(content="turn one"),
+        AIMessage(
+            content="",
+            tool_calls=[{"name": "grep", "args": {}, "id": "call_1", "type": "tool_call"}],
+        ),
+        ToolMessage(content=big, tool_call_id="call_1"),
+        HumanMessage(content="turn two"),
+        AIMessage(content="reply"),
+    ]
+
+    trim_history(history)
+
+    # The oversized turn is dropped entirely (no orphaned ToolMessage).
+    assert history[0] is system
+    assert history[-1].content == "reply"
+    assert not any(isinstance(m, ToolMessage) for m in history)
+    assert len(history) <= 4
+
+
+def test_trim_history_orphaned_tool_messages_dropped():
+    from gcode.agent import MAX_HISTORY, MAX_HISTORY_TOKENS, trim_history
+    from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+
+    system = HumanMessage(content="system")
+    # Keep enough short messages that the count limit (not the token budget) trims,
+    # and make the kept tail begin with ToolMessages whose owner was trimmed out.
+    rest = [
+        AIMessage(
+            content="",
+            tool_calls=[{"name": "grep", "args": {}, "id": "call_1", "type": "tool_call"}],
+        ),
+        ToolMessage(content="res1", tool_call_id="call_1"),
+        ToolMessage(content="res2", tool_call_id="call_1"),
+    ] + [HumanMessage(content=f"m{i}") for i in range(MAX_HISTORY)]
+    history = [system] + rest
+    small = sum(len(m.content) for m in history)
+    assert small <= MAX_HISTORY_TOKENS * 4  # token budget not the deciding factor
+
+    trim_history(history)
+
+    assert history[0] is system
+    assert len(history) <= MAX_HISTORY + 1
+    assert not any(isinstance(m, ToolMessage) for m in history)
+    assert not getattr(history[1], "tool_calls", None)
+    assert history[1].content == "m0"