feat: tool truncation + auto-compaction to bash harness

verifiers/v1/harnesses/bash/harness.py

@@ -2,6 +2,9 @@
 import os
 from pathlib import Path
 
+from pydantic import PositiveInt
+from pydantic_config import BaseConfig
+
 from verifiers.v1.clients import ModelContext
 from verifiers.v1.configs.harness import HarnessConfig
 from verifiers.v1.dialects.chat import message_to_wire
@@ -27,6 +30,14 @@
 )
 
 
+class CompactionConfig(BaseConfig):
+    """Context compaction policy for the bash agent loop."""
+
+    summarize_at_tokens: PositiveInt | None = None
+    """Compact at this token count. When unset, compact when 16k tokens remain below the
+    model context window when the provider advertises it."""
+
+
 class BashHarnessConfig(HarnessConfig):
     edit: bool = True
     """Offer the local `edit` tool (single-occurrence string replacement in a file) alongside
@@ -37,6 +48,9 @@ class BashHarnessConfig(HarnessConfig):
     eval environment; the key is handed to the program over argv (like the interception secret) so
     the agent's `bash` subprocesses don't inherit it."""
 
+    compaction: CompactionConfig | None = None
+    """Context compaction policy. Set an empty config to use automatic thresholds."""
+
 
 class BashHarness(Harness[BashHarnessConfig]):
     APPENDS_SYSTEM_PROMPT = True
@@ -77,6 +91,11 @@ async def launch(
         ]
         if tool_interception_url:
             args.append(f"--tool-interception-url={tool_interception_url}")
+        if self.config.compaction is not None:
+            args.append("--compaction")
+            threshold = self.config.compaction.summarize_at_tokens
+            if threshold is not None:
+                args.append(f"--summarize-at-tokens={threshold}")
         if self.config.edit:
             args.append("--edit")
         if self.config.search:

verifiers/v1/harnesses/bash/program.py

@@ -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,153 @@ 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 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}"
     )
-    return completion.choices[0].message
+
+
+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)
+            ):
+                raise
+        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)
+        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]
+                continue
+            message = completion.choices[0].message
+            # A reasoning-parsed model can put the whole reply in the reasoning
+            # channel; the checkpoint asked for a summary, so accept it from
+            # there when content is empty.
+            text = (message.content or "").strip() or (
+                getattr(message, "reasoning_content", None) 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 +554,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 +600,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)
+        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 +640,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 +676,7 @@ async def main() -> None:
                     )
                 else:
                     content = f"error: unknown tool {name!r}"
-            tool_message["content"] = content
+            tool_message["content"] = truncate_tool_output(content)
             if args.tool_interception_url:
                 assert tool_client is not None
                 decision = await run_tool_hook(
@@ -440,6 +689,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)
+            except CompactionFailed:
+                break
+        else:
+            compactor.note_good(messages)
     if tool_client is not None:
         await tool_client.aclose()