perf: trim history by token budget, not message count

gcode/agent.py

@@ -41,19 +41,87 @@ 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  # type: ignore
+
+            enc = tiktoken.get_encoding("cl100k_base")
+            return len(enc.encode(text))
+        except Exception:
+            pass
+        return max(1, len(text) // 4)
+    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 rest and sum(_estimate_tokens(m) for m in [system] + rest) > MAX_HISTORY_TOKENS:
+        # Drop oldest message in rest, but avoid orphaning ToolMessages
+        # If oldest is AIMessage with tool_calls, also drop its ToolMessages
+        if len(rest) <= 2:
+            break
+        dropped = rest.pop(0)
+        # If we dropped an AIMessage that had tool_calls, also drop its ToolMessages
+        # that immediately follow (they are now orphaned)
+        while rest and isinstance(rest[0], ToolMessage):
+            # Check if this ToolMessage belonged to the dropped AIMessage
+            # Heuristic: if dropped was AIMessage with tool_calls, drop all leading ToolMessages
+            if getattr(dropped, "tool_calls", None):
+                rest.pop(0)
+            else:
+                break
+            # Avoid dropping too many and breaking the budget loop
+            if not getattr(dropped, "tool_calls", None):
+                break
+    # 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: