-
Notifications
You must be signed in to change notification settings - Fork 654
feat: tool truncation + auto-compaction to bash harness #2454
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 9 commits
96aa25e
25f12a6
339a785
514106c
5ea11fc
7b15d60
978e650
71d628e
f07f60a
84eae30
d9bd515
59bfc65
7c05e5c
cad252f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -12,14 +12,99 @@ | |
| from pathlib import Path | ||
|
|
||
| import httpx | ||
| from openai import AsyncOpenAI | ||
| from openai import APIError, APIStatusError, AsyncOpenAI | ||
| from tenacity import AsyncRetrying, stop_after_attempt, wait_exponential_jitter | ||
|
|
||
| SERPER_URL = "https://google.serper.dev/search" | ||
|
|
||
| MCP_CALL_ATTEMPTS = 6 | ||
| MCP_TIMEOUT = 600.0 | ||
|
|
||
| RESERVE_TOKENS = 16_384 | ||
| """Compact when this many tokens remain below the model context window.""" | ||
|
|
||
| COMPACTION_ATTEMPTS = 3 | ||
| """Checkpoint attempts before compaction fails: a rejected request falls back to the | ||
| last good snapshot; an empty or tool-calling reply is resampled.""" | ||
|
|
||
| TOOL_OUTPUT_MAX_BYTES = 20_000 | ||
| """Middle-out truncation budget for one tool result before it enters the conversation.""" | ||
|
|
||
| CHECKPOINT_COMPACTION_PROMPT = """You are performing a CONTEXT CHECKPOINT COMPACTION. Create a handoff summary for another LLM that will resume the task. | ||
|
|
||
| Include: | ||
| - Current progress and key decisions made | ||
| - Important context, constraints, or user preferences | ||
| - What remains to be done (clear next steps) | ||
| - Any critical data, examples, or references needed to continue | ||
|
|
||
| Be concise, structured, and focused on helping the next LLM seamlessly continue the work. | ||
|
|
||
| Reply with the summary as plain text. Do not call any tools - summarize from the conversation as it stands.""" | ||
|
|
||
| POST_COMPACTION_FRAMING = """Another language model started to solve this problem and produced \ | ||
| a summary of its thinking process. You also have access to the state of the tools that \ | ||
| were used by that language model. Use this to build on the work \ | ||
| that has already been done and avoid duplicating work. Here is \ | ||
| the summary produced by the other language model, use the \ | ||
| information in this summary to assist with your own analysis:""" | ||
|
|
||
| CONTEXT_OVERFLOW_MARKERS = ( | ||
| # OpenAI error code "context_length_exceeded"; OpenRouter relays the raw body. | ||
| "context_length_exceeded", | ||
| # OpenAI Responses/Completions: "Your input exceeds the context window of this model". | ||
| "exceeds the context window", | ||
| # OpenAI chat: "Input tokens exceed the configured limit of N tokens. Please reduce | ||
| # the length of the messages."; Groq words it the same way. | ||
| "reduce the length of the messages", | ||
| # vLLM: "This model's maximum context length is N tokens"; the renderers pre-flight: | ||
| # "Prompt length (N) exceeds maximum context length (M)"; Mistral uses the same words. | ||
| "maximum context length", | ||
| # Anthropic: "prompt is too long: N tokens > M maximum". | ||
| "prompt is too long", | ||
| # Anthropic byte-size overflow: HTTP 413 {"type": "request_too_large"}. | ||
| "request_too_large", | ||
| # HTTP proxies reject an oversized body with 413 "Request Entity Too Large". | ||
| "request entity too large", | ||
| # Google: "The input token count (N) exceeds the maximum number of tokens allowed (M)". | ||
| "exceeds the maximum number of tokens", | ||
| # xAI: "This model's maximum prompt length is N but the request contains M tokens". | ||
| "maximum prompt length is", | ||
| ) | ||
|
|
||
| CONTEXT_WINDOW_FIELDS = ( | ||
| "max_model_len", | ||
| "context_length", | ||
| "context_window", | ||
| "max_context_length", | ||
| ) | ||
|
|
||
|
|
||
| def default_threshold(context_window: int) -> int: | ||
| """Leave a fixed reserve below the window; small windows keep at least half.""" | ||
| return max(context_window - RESERVE_TOKENS, context_window // 2) | ||
|
|
||
|
|
||
| async def discover_threshold(client: AsyncOpenAI, model: str) -> int | None: | ||
| """The compaction threshold, when the provider's model card advertises a context window. | ||
|
|
||
| `models.list()` keeps provider extensions in each card's `model_extra`; a raw | ||
| `cast_to` parse breaks on one Python version or another.""" | ||
| try: | ||
| page = await client.models.list() | ||
| except APIError: | ||
| return None | ||
| for card in page.data: | ||
| if card.id != model: | ||
| continue | ||
| extra = card.model_extra or {} | ||
| for field in CONTEXT_WINDOW_FIELDS: | ||
| value = extra.get(field) | ||
| if isinstance(value, int) and not isinstance(value, bool) and value > 0: | ||
| return default_threshold(value) | ||
| break | ||
| return None | ||
|
|
||
|
|
||
| BASH_TOOL = { | ||
| "type": "function", | ||
|
|
@@ -178,12 +263,156 @@ def run_edit(path: str, old_str: str, new_str: str) -> str: | |
|
|
||
|
|
||
| async def chat( | ||
| client: AsyncOpenAI, model: str, messages: list[dict], tools: list[dict] | ||
| client: AsyncOpenAI, | ||
| model: str, | ||
| messages: list[dict], | ||
| tools: list[dict], | ||
| *, | ||
| tool_choice: str | None = None, | ||
| ): | ||
| completion = await client.chat.completions.create( | ||
| model=model, messages=messages, tools=tools or None | ||
| kwargs = {"model": model, "messages": messages, "tools": tools or None} | ||
| if tools and tool_choice is not None: | ||
| kwargs["tool_choice"] = tool_choice | ||
| return await client.chat.completions.create(**kwargs) | ||
|
|
||
|
|
||
| class CompactionFailed(Exception): | ||
| """Every checkpoint attempt failed - the caller ends the run cleanly instead.""" | ||
|
|
||
|
|
||
| def is_context_overflow(error: APIStatusError) -> bool: | ||
| details = f"{error} {error.body or ''}" | ||
| # An overflow is deterministic: a 400, or a 413 for a byte-size cap. | ||
| return error.status_code in (400, 413) and any( | ||
| marker in details.casefold() for marker in CONTEXT_OVERFLOW_MARKERS | ||
| ) | ||
|
|
||
|
|
||
| def compactable(messages: list[dict]) -> bool: | ||
| """Whether compaction can reclaim anything - some history beyond the task exists.""" | ||
| first_user = next( | ||
| (i for i, m in enumerate(messages) if m.get("role") == "user"), None | ||
| ) | ||
| return completion.choices[0].message | ||
| return any( | ||
| m.get("role") != "system" and i != first_user for i, m in enumerate(messages) | ||
| ) | ||
|
|
||
|
|
||
| def truncate_tool_output(text: str) -> str: | ||
| """Keep the head and tail of an oversized tool result and say what was cut.""" | ||
| data = text.encode("utf-8") | ||
| if len(data) <= TOOL_OUTPUT_MAX_BYTES: | ||
| return text | ||
| keep = TOOL_OUTPUT_MAX_BYTES // 2 | ||
| head = data[:keep].decode("utf-8", errors="ignore") | ||
| tail = data[-keep:].decode("utf-8", errors="ignore") | ||
| return ( | ||
| f"Warning: truncated output (original token count: {estimated_tokens(text)})\n" | ||
| f"Total output lines: {text.count(chr(10)) + 1}\n\n" | ||
| f"{head}\n[... {len(data) - 2 * keep} bytes truncated ...]\n{tail}" | ||
| ) | ||
|
|
||
|
|
||
| def estimated_tokens(chars: str) -> int: | ||
| """Rough token count at four characters per token.""" | ||
| return (len(chars) + 3) // 4 | ||
|
|
||
|
|
||
| def context_tokens(completion) -> int: | ||
| usage = completion.usage | ||
| if usage is None: | ||
| return 0 | ||
| return (usage.prompt_tokens or 0) + (usage.completion_tokens or 0) | ||
|
|
||
|
|
||
| class Compactor: | ||
| """Compact once and retry once when a model turn exhausts its context.""" | ||
|
|
||
| def __init__(self, client, model, tools, enabled, threshold): | ||
| self.client = client | ||
| self.model = model | ||
| self.tools = tools | ||
| self.enabled = enabled | ||
| self.threshold = threshold | ||
| self.last_good = 0 | ||
| """Message count of the newest state that passed a threshold check - by | ||
| definition a state with a full reserve of room, so a checkpoint over it fits.""" | ||
|
|
||
| def reached(self, completion, extra_tokens: int = 0) -> bool: | ||
| return ( | ||
| self.enabled | ||
| and self.threshold is not None | ||
| and context_tokens(completion) + extra_tokens >= self.threshold | ||
| ) | ||
|
|
||
| def note_good(self, messages: list[dict]) -> None: | ||
| self.last_good = len(messages) | ||
|
|
||
| async def complete(self, messages: list[dict]): | ||
| try: | ||
| completion = await chat(self.client, self.model, messages, self.tools) | ||
| except APIStatusError as error: | ||
| if ( | ||
| not self.enabled | ||
| or self.threshold is None | ||
| or not is_context_overflow(error) | ||
| or not compactable(messages) | ||
|
macroscopeapp[bot] marked this conversation as resolved.
|
||
| ): | ||
| raise | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Overflow recovery requires a thresholdMedium Severity Reactive compaction on a 400/413 overflow is skipped whenever Reviewed by Cursor Bugbot for commit f07f60a. Configure here.
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. By design: compaction requires a known context window (discovered or explicit); without one the overflow propagates. Documented scoping - training always knows the window. |
||
| else: | ||
| choice = completion.choices[0] | ||
| if choice.finish_reason != "length" or not self.reached(completion): | ||
| self.note_good(messages) | ||
| return completion, messages | ||
| if not compactable(messages): | ||
| return completion, messages | ||
|
|
||
| messages = await self.compact(messages) | ||
|
macroscopeapp[bot] marked this conversation as resolved.
|
||
| completion = await chat(self.client, self.model, messages, self.tools) | ||
| return completion, messages | ||
|
|
||
| async def compact(self, messages: list[dict]) -> list[dict]: | ||
| # A rejected checkpoint falls back to the last good snapshot (which has a | ||
| # full reserve of room, so it fits); an empty or tool-calling reply is | ||
| # resampled. Reasoning is never part of the summary. | ||
| system = [message for message in messages if message.get("role") == "system"] | ||
| base = messages | ||
| for _ in range(COMPACTION_ATTEMPTS): | ||
| checkpoint = [ | ||
| *base, | ||
| {"role": "user", "content": CHECKPOINT_COMPACTION_PROMPT}, | ||
| ] | ||
| try: | ||
| completion = await chat( | ||
| self.client, | ||
| self.model, | ||
| checkpoint, | ||
| self.tools, | ||
| tool_choice="none", | ||
| ) | ||
| except APIStatusError as error: | ||
| if not is_context_overflow(error): | ||
| raise | ||
| base = messages[: self.last_good] | ||
|
macroscopeapp[bot] marked this conversation as resolved.
macroscopeapp[bot] marked this conversation as resolved.
|
||
| continue | ||
|
cursor[bot] marked this conversation as resolved.
|
||
| message = completion.choices[0].message | ||
| # A reasoning-parsed model can put the whole reply in the reasoning | ||
| # channel ("reasoning_content" or vLLM's "reasoning"); the checkpoint | ||
| # asked for a summary, so accept it from there when content is empty. | ||
| extra = getattr(message, "model_extra", None) or {} | ||
| text = ( | ||
| (message.content or "").strip() | ||
| or str(extra.get("reasoning_content") or "").strip() | ||
| or str(extra.get("reasoning") or "").strip() | ||
| ) | ||
| if not message.tool_calls and text: | ||
| framed = POST_COMPACTION_FRAMING + "\n\n" + text | ||
| rebuilt = [*system, {"role": "user", "content": framed}] | ||
| self.note_good(rebuilt) | ||
| return rebuilt | ||
| raise CompactionFailed( | ||
| f"no usable summary after {COMPACTION_ATTEMPTS} attempts" | ||
| ) | ||
|
|
||
|
|
||
| @asynccontextmanager | ||
|
|
@@ -328,6 +557,8 @@ def parse_args() -> argparse.Namespace: | |
| parser.add_argument("--initial-messages-file", default="") | ||
| parser.add_argument("--mcp-config", default="") | ||
| parser.add_argument("--tool-interception-url", default="") | ||
| parser.add_argument("--compaction", action="store_true") | ||
| parser.add_argument("--summarize-at-tokens", type=int) | ||
| parser.add_argument("--edit", action="store_true") | ||
| parser.add_argument("--search", action="store_true") | ||
| parser.add_argument("--serper-key", default="") | ||
|
|
@@ -372,11 +603,28 @@ async def main() -> None: | |
| messages.extend(initial) | ||
| elif args.prompt: | ||
| messages.append({"role": "user", "content": args.prompt}) | ||
| compactor = Compactor( | ||
| client, | ||
| args.model, | ||
| tools, | ||
| args.compaction, | ||
| args.summarize_at_tokens, | ||
| ) | ||
| if compactor.enabled and compactor.threshold is None: | ||
| compactor.threshold = await discover_threshold(client, args.model) | ||
| while True: | ||
| message = await chat(client, args.model, messages, tools) | ||
| try: | ||
| completion, messages = await compactor.complete(messages) | ||
|
macroscopeapp[bot] marked this conversation as resolved.
|
||
| except CompactionFailed: | ||
| # The context is exhausted and could not be summarized: end the run | ||
| # cleanly with what the conversation holds - still a trainable sample. | ||
| break | ||
| choice = completion.choices[0] | ||
| message = choice.message | ||
| messages.append(message.model_dump(exclude_none=True)) | ||
| if not message.tool_calls: | ||
| break | ||
| tool_result_tokens = 0 | ||
| for call in message.tool_calls: | ||
| name = call.function.name | ||
| tool_message = { | ||
|
|
@@ -395,7 +643,11 @@ async def main() -> None: | |
| tool_message, | ||
| ) | ||
| if decision["action"] == "rewrite": | ||
| messages.append(decision["message"]) | ||
| rewritten = decision["message"] | ||
| messages.append(rewritten) | ||
| tool_result_tokens += estimated_tokens( | ||
| str(rewritten.get("content", "")) | ||
| ) | ||
| continue | ||
| try: | ||
| tool_args = json.loads(call.function.arguments or "{}") | ||
|
|
@@ -427,7 +679,7 @@ async def main() -> None: | |
| ) | ||
| else: | ||
| content = f"error: unknown tool {name!r}" | ||
| tool_message["content"] = content | ||
| tool_message["content"] = truncate_tool_output(content) | ||
|
macroscopeapp[bot] marked this conversation as resolved.
Outdated
cursor[bot] marked this conversation as resolved.
Outdated
|
||
| if args.tool_interception_url: | ||
| assert tool_client is not None | ||
| decision = await run_tool_hook( | ||
|
|
@@ -440,6 +692,14 @@ async def main() -> None: | |
| if decision["action"] == "rewrite": | ||
| tool_message = decision["message"] | ||
| messages.append(tool_message) | ||
| tool_result_tokens += estimated_tokens(str(tool_message["content"])) | ||
| if compactor.reached(completion, tool_result_tokens) and compactable(messages): | ||
| try: | ||
| messages = await compactor.compact(messages) | ||
|
macroscopeapp[bot] marked this conversation as resolved.
|
||
| except CompactionFailed: | ||
| break | ||
| else: | ||
| compactor.note_good(messages) | ||
| if tool_client is not None: | ||
| await tool_client.aclose() | ||
|
|
||
|
|
||


Uh oh!
There was an error while loading. Please reload this page.