diff --git a/gcode/agent.py b/gcode/agent.py index 05e260a..f1de590 100644 --- a/gcode/agent.py +++ b/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: diff --git a/gcode/tools.py b/gcode/tools.py index e919f37..dba66b9 100644 --- a/gcode/tools.py +++ b/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: diff --git a/tests/test_agent.py b/tests/test_agent.py index a14f71a..cb4174a 100644 --- a/tests/test_agent.py +++ b/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"