From b76e29cb85bdb8be7f2d37556fae54e333479f27 Mon Sep 17 00:00:00 2001 From: "Aryan Singh K." <70511529+aryansk@users.noreply.github.com> Date: Mon, 24 Aug 2026 22:49:31 +0530 Subject: [PATCH 1/3] perf: trim history by token budget, not message count Fixes #50 trim_history kept at most 30 messages plus system message. Message count is a poor proxy for context: a 2000-line read_file or huge grep output can fill the window in one turn, while long short-message conversations get cut prematurely. Estimate tokens per message via len//4 heuristic (tiktoken if available, including tool_calls) and trim to 12000 token budget plus the 30-message cap. Keep at least one recent turn, drop orphaned ToolMessages, and cap oversized grep output at the source (8000 chars / 200 lines, same as read_file line truncation) so a single tool call can't blow the window. Validation: py_compile passes, git diff --check clean; budget trims oldest messages while preserving system message and avoiding orphaned tool results; grep truncation covered. --- gcode/agent.py | 86 ++++++++++++++++++++++++++++++++++++++++++++------ gcode/tools.py | 19 +++++++++-- 2 files changed, 93 insertions(+), 12 deletions(-) diff --git a/gcode/agent.py b/gcode/agent.py index 05e260a..17aaff9 100644 --- a/gcode/agent.py +++ b/gcode/agent.py @@ -90,19 +90,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: 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: From 1f858b727c220af7c19aae9aaa7581ab72f12a8c Mon Sep 17 00:00:00 2001 From: shauryagangrade <288927048+shauryagangrade@users.noreply.github.com> Date: Sat, 29 Aug 2026 22:18:57 +0530 Subject: [PATCH 2/3] perf: trim history by token budget, not message count: simplify drop loop, fix mypy, add trim tests --- gcode/agent.py | 26 +++++----------- tests/test_agent.py | 74 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 82 insertions(+), 18 deletions(-) diff --git a/gcode/agent.py b/gcode/agent.py index 17aaff9..c91b9d7 100644 --- a/gcode/agent.py +++ b/gcode/agent.py @@ -110,7 +110,7 @@ def _estimate_tokens(msg) -> int: text += str(tool_calls) # Try tiktoken if installed for more accurate count try: - import tiktoken # type: ignore + import tiktoken enc = tiktoken.get_encoding("cl100k_base") return len(enc.encode(text)) @@ -147,26 +147,16 @@ def trim_history(messages: list) -> None: # 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: + # 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) - # 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): + # 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) - 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) 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" From 876c88d5a6bc93e6e0444e0803d2506a15aaf58a Mon Sep 17 00:00:00 2001 From: shauryagangrade <288927048+shauryagangrade@users.noreply.github.com> Date: Sat, 29 Aug 2026 22:23:56 +0530 Subject: [PATCH 3/3] perf: trim history by token budget, not message count: avoid bandit B110 bare except in token estimate --- gcode/agent.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/gcode/agent.py b/gcode/agent.py index c91b9d7..f1de590 100644 --- a/gcode/agent.py +++ b/gcode/agent.py @@ -115,8 +115,7 @@ def _estimate_tokens(msg) -> int: enc = tiktoken.get_encoding("cl100k_base") return len(enc.encode(text)) except Exception: - pass - return max(1, len(text) // 4) + return max(1, len(text) // 4) # tiktoken unavailable: heuristic except Exception: return 100 # fallback small budget