diff --git a/README.md b/README.md index 0e741c9..dfdaab2 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ The model gets a single built-in tool, `ipython`: a persistent IPython kernel fo For convenience, rlm ships built-in *skills* that can be enabled per session via the runtime contract's `skills` list (off by default): `edit` (single-occurrence string replacement), `search` (web search via Serper, needs `SERPER_API_KEY`), and `fetch` (retrieve a URL as cleaned text). Enabled skills are pre-imported into the IPython kernel like any other skill (see [Skills](#skills)), so the agent calls `await edit(path=..., old_str=..., new_str=...)`, `await search(query=...)`, or `await fetch(url=...)`. `fetch` also exists as a native builtin tool with the same semantics, for tool-calling runs (opt-in via `RLM_BUILTIN_TOOLS`). -Context is reclaimed automatically: when a turn's prompt token count crosses the policy's `summarize_at_tokens`, the engine compacts the conversation into a summary and continues on a fresh branch. The IPython kernel keeps running across the compaction, so REPL state survives (see [Compaction](#compaction)). +Context compaction is optional. When enabled, the engine compacts when 16k tokens remain below an advertised model context window. The policy can set an explicit `summarize_at_tokens` threshold. The IPython kernel keeps running across compaction, so REPL state survives (see [Compaction](#compaction)). Inside the IPython session, a callable `rlm` is pre-injected into the namespace. When recursion is allowed, the model can call `await rlm(...)` to spawn sub-agents. Skills supplied by the host environment (see [Skills](#skills)) are importable directly by name, e.g. `import websearch`. @@ -114,9 +114,11 @@ Recursive calls are created by a session-local supervisor rather than by the IPy ## Compaction -There is no model-driven compaction tool. Compaction is automatic: set the policy's `summarize_at_tokens` and, once a turn's prompt token count reaches that threshold, the engine asks the model for a handoff summary and resumes the task on a fresh branch seeded with that summary. The original task prompt is dropped — the summary carries the goal forward. +There is no model-driven compaction tool. Set the policy's `compaction` field to enable compaction. The engine reads the model context window from the provider's `/models` response and compacts when 16k tokens remain below it; small windows keep at least half. Set `summarize_at_tokens` to pin the threshold explicitly. Without a known window or explicit threshold, compaction stays off and an overflow propagates. A tool result larger than 20KB is truncated to its head and tail before it enters the conversation, with a warning naming the original size. -The IPython kernel keeps running across the compaction, so all variables, imports, and in-memory data are preserved; the model is told to mention important variable names in its summary so the resumed branch knows what's available. With `summarize_at_tokens` unset, no auto-compaction occurs. +The engine asks the model for a plain-text handoff summary and resumes the task on a fresh branch seeded with that summary; reasoning is never part of it. A provider overflow (a 400 or 413 naming a context limit) triggers the same compaction reactively from the current state. A rejected checkpoint request falls back to the last state that passed a threshold check - by definition a state with a full reserve of room - and an empty or tool-calling reply is resampled; after three failed attempts the run ends cleanly with what the conversation holds. An overflow with no history beyond the task propagates: the task alone approaches the window and there is nothing to reclaim. + +The IPython kernel keeps running across the compaction, so all variables, imports, and in-memory data are preserved. The model is told to mention important variable names in its summary so the resumed branch knows what is available. The same policy applies to the main agent and all recursive agents. ## Session Directory @@ -255,4 +257,3 @@ Install dev dependencies and run the suite: uv sync --group dev uv run pytest tests/ ``` - diff --git a/src/rlm/acp.py b/src/rlm/acp.py index 3d264df..eae034b 100644 --- a/src/rlm/acp.py +++ b/src/rlm/acp.py @@ -98,6 +98,7 @@ class _LimitsSnapshot(_ContractModel): max_concurrent_subagents: int = Field(gt=0) max_subagent_calls: int = Field(gt=0) max_tokens: int | None = Field(default=None, gt=0) + compaction: bool summarize_at_tokens: int | None = Field(default=None, gt=0) max_compactions: int | None = Field(default=None, gt=0) allow_git: bool diff --git a/src/rlm/compaction.py b/src/rlm/compaction.py new file mode 100644 index 0000000..cf8f064 --- /dev/null +++ b/src/rlm/compaction.py @@ -0,0 +1,161 @@ +"""Context checkpoint helpers.""" + +from collections.abc import Mapping +from typing import Any + +from openai import APIError, APIStatusError, AsyncOpenAI + +CHECKPOINT_PROMPT = """You are performing a CONTEXT CHECKPOINT COMPACTION. Create a handoff summary another LLM can ACT on immediately to resume the task. + +It MUST contain, as fenced code blocks (not prose): +- The exact shell/test command(s) to reproduce and verify — copy-pasteable, with the real path and test filter +- Any edit still to apply, as the concrete `await edit(path=..., old_str=..., new_str=...)` call + +Then: +- A NUMBERED list of remaining next steps +- Current progress, key decisions, and constraints + +Be concise and concrete: prefer runnable commands over descriptions. + +Reply with the summary as plain text. Do not call any tools - summarize from the conversation as it stands.""" + +REPL_NOTE = ( + "\n\n" + "Note: the IPython kernel stays running across this compaction. " + "All variables, imports, and in-memory data are preserved. " + "Mention important variable names and what they contain so the " + "next LLM knows what's available." +) + +SUMMARY_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:""" + +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.""" + +_CONTEXT_FIELDS = ( + "max_model_len", + "context_length", + "context_window", + "max_context_length", +) +_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", +) +_window_cache: dict[tuple[str, str], int | None] = {} + + +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 _OVERFLOW_MARKERS + ) + + +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) + + +def _model_context_window(payload: Mapping[str, Any], model: str) -> int | None: + card = next( + ( + item + for item in payload.get("data") or [] + if isinstance(item, Mapping) and item.get("id") == model + ), + None, + ) + if card is None: + return None + for field in _CONTEXT_FIELDS: + value = card.get(field) + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + return value + return None + + +async def discover_threshold(client: AsyncOpenAI, model: str) -> int | None: + key = (str(client.base_url), model) + if key not in _window_cache: + try: + # `models.list()` keeps provider extensions in each card's `model_extra`; + # a raw `cast_to` parse breaks on one Python version or another. + page = await client.models.list() + payload = { + "data": [ + {"id": card.id, **(card.model_extra or {})} for card in page.data + ] + } + except (APIError, AttributeError): + # A transient listing failure must not disable compaction for the + # rest of the process - leave the cache empty so the next engine retries. + return None + _window_cache[key] = _model_context_window(payload, model) + window = _window_cache[key] + return default_threshold(window) if window is not None else None + + +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 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) + ) diff --git a/src/rlm/config.py b/src/rlm/config.py index abd0bf6..2d1a1bd 100644 --- a/src/rlm/config.py +++ b/src/rlm/config.py @@ -54,7 +54,8 @@ class ExecutionPolicy(_ConfigModel): max_depth: int = Field(default=1, ge=0) exec_timeout: int = Field(default=300, gt=0) max_tokens: int | None = Field(default=None, gt=0) - summarize_at_tokens: int | None = Field(default=256_000, gt=0) + compaction: bool = False + summarize_at_tokens: int | None = Field(default=None, gt=0) max_compactions: int | None = Field(default=None, gt=0) max_concurrent_subagents: int = Field(default=4, gt=0) max_subagent_calls: int = Field(default=64, gt=0) diff --git a/src/rlm/engine.py b/src/rlm/engine.py index dd33445..2034fd7 100644 --- a/src/rlm/engine.py +++ b/src/rlm/engine.py @@ -11,8 +11,9 @@ import uuid from copy import deepcopy from pathlib import Path +from typing import Any -from openai import AsyncOpenAI, BadRequestError +from openai import APIStatusError, AsyncOpenAI from rlm.client import ( call_with_retries, @@ -20,6 +21,18 @@ make_client, model_call_headers, ) +from rlm.compaction import ( + COMPACTION_ATTEMPTS, + CHECKPOINT_PROMPT, + CompactionFailed, + REPL_NOTE, + SUMMARY_FRAMING, + compactable, + discover_threshold, + estimated_tokens, + is_context_overflow, + truncate_tool_output, +) from rlm.config import RuntimeConfig from rlm.semantic import SemanticEdgeTracker from rlm.mcp import MCPServer, validate_mcp_servers @@ -48,85 +61,6 @@ logger = logging.getLogger(__name__) -# Hard cap for a single tool result entering the conversation (the session log -# keeps the full output). Head+tail with a warning naming the original size — -# Codex's output policy, matching verifiers' bash-harness truncation. -TOOL_OUTPUT_MAX_BYTES = 10_000 - - -def estimated_tokens(chars: str) -> int: - """Rough token count at four characters per token.""" - return (len(chars) + 3) // 4 - - -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}" - ) - - -# Injected as a user message when the branch's context size reaches the -# compaction threshold. The model's next reply is expected to be a -# plain-text handoff summary; any tool calls it emits are ignored and -# the message is compacted in place of them. -CHECKPOINT_COMPACTION_PROMPT = ( - "You are performing a CONTEXT CHECKPOINT COMPACTION. " - "Create a handoff summary another LLM can ACT on immediately to resume the task.\n" - "\n" - "It MUST contain, as fenced code blocks (not prose):\n" - "- The exact shell/test command(s) to reproduce and verify — copy-pasteable, " - "with the real path and test filter\n" - "- Any edit still to apply, as the concrete " - "`await edit(path=..., old_str=..., new_str=...)` call\n" - "\n" - "Then:\n" - "- A NUMBERED list of remaining next steps\n" - "- Current progress, key decisions, and constraints\n" - "\n" - "Be concise and concrete: prefer runnable commands over descriptions." -) - -# Appended to the checkpoint prompt when the IPython REPL is active. -REPL_RESTART_NOTE = ( - "\n\n" - "Note: the IPython kernel stays running across this compaction. " - "All variables, imports, and in-memory data are preserved. " - "Mention important variable names and what they contain so the " - "next LLM knows what's available." -) - -# Wrapper text that frames the summary as the sole user-facing context -# for the post-compaction branch. The original task prompt is dropped; -# the summary is responsible for carrying the goal. -POST_COMPACTION_FRAMING = ( - "Another language model started to solve this problem and produced " - "a summary of its thinking process. 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:" -) - - -def _is_request_too_large(e: BadRequestError) -> bool: - """True if a 400 says the request outgrew the serving limit, so the session can end - with its graceful hard-ceiling stop instead of erroring. Matches the proxy's - "Request Entity Too Large" body and vLLM's context-length phrasing - ("This model's maximum context length is ..." / "Prompt length ... exceeds the - maximum context length").""" - haystack = f"{e} {getattr(e, 'body', '') or ''}".lower() - return ( - "request entity too large" in haystack or "maximum context length" in haystack - ) - def _parse_tool_call_args(raw: str) -> tuple[dict | None, dict | None]: """Parse a tool-call arguments blob. Returns (args, error_info). @@ -185,6 +119,7 @@ def __init__( self.model = config.model self.cwd = cwd or os.getcwd() self.exec_timeout = config.policy.exec_timeout + self.compaction = config.policy.compaction self.summarize_at_tokens = config.policy.summarize_at_tokens self.max_compactions = config.policy.max_compactions self.system_prompt_path = config.system_prompt_path @@ -222,6 +157,11 @@ def __init__( self._owns_supervisor = False self._total_usage = TokenUsage() self._last_prompt_tokens = 0 + 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.""" + self._compacted = False + self._last_call_id: str | None = None # Metrics self._metrics = RLMMetrics() @@ -292,10 +232,16 @@ async def prompt(self, prompt: str) -> RLMResult: ) raise messages_before = self._messages[:1] + last_good_before = len(messages_before) else: messages_before = list(self._messages) + last_good_before = self._last_good self._messages.append({"role": "user", "content": prompt}) + # This turn's opening state is the floor for checkpoint fallbacks: + # a fallback must never drop the newest user instruction. + self._last_good = len(self._messages) branch_start_before = self._branch_start_turn + compacted_before = self._compacted semantic_edges_before = self._semantic_edges.checkpoint(self._invocation_id) turn_before = self._turn usage_before = TokenUsage( @@ -326,6 +272,8 @@ async def prompt(self, prompt: str) -> RLMResult: # kernel/tool side effects, and the append-only audit log describe work # that really ran and remain part of session accounting. self._messages[:] = messages_before + self._last_good = last_good_before + self._compacted = compacted_before self._branch_start_turn = branch_start_before self._semantic_edges.restore(self._invocation_id, semantic_edges_before) self._turn = turn_before @@ -348,6 +296,9 @@ async def prompt(self, prompt: str) -> RLMResult: async def _start(self, prompt: str) -> None: """Initialize the session, tools, conversation, and persistent kernel.""" + if self.compaction and self.summarize_at_tokens is None: + self.summarize_at_tokens = await discover_threshold(self.client, self.model) + self._ensure_session() self.session.write_meta( @@ -415,6 +366,9 @@ async def _start(self, prompt: str) -> None: {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}, ] + # The initial conversation is the floor for checkpoint fallbacks: a + # first-turn checkpoint must never retry from an empty base. + self._last_good = len(self._messages) self._started = True except BaseException: self._repl.shutdown() @@ -434,34 +388,15 @@ async def _run_loop(self) -> RLMResult: for turn in itertools.count(self._turn): self._turn = turn + 1 - # Call LLM - request_id = self._semantic_edges.start_request(self._invocation_id) - call_id = request_id - request_kwargs = { - "model": self.model, - "messages": messages, - "extra_headers": model_call_headers(request_id), - } - if self._active_tool_schemas: - request_kwargs["tools"] = self._active_tool_schemas - request_kwargs["parallel_tool_calls"] = False try: - response = await call_with_retries( - self.client.chat.completions.create, - **request_kwargs, - ) - except BaseException as exc: - self._semantic_edges.fail_request(request_id) - if isinstance(exc, BadRequestError) and _is_request_too_large(exc): - self._metrics.stop_reason = "request_too_large" - final_text = "[request body too large]" - break - raise - self._semantic_edges.finish_request(request_id) - usage = extract_usage(response) - self._total_usage.prompt_tokens += usage.prompt_tokens - self._total_usage.completion_tokens += usage.completion_tokens - self._last_prompt_tokens = usage.prompt_tokens + response, usage = await self._complete(messages, turn) + except CompactionFailed: + # The context is exhausted and could not be summarized: end the run + # cleanly with what the conversation holds - still a trainable sample. + self._metrics.stop_reason = "compaction_failed" + final_text = "[context exhausted: compaction failed]" + break + call_id = self._last_call_id self._metrics.turns_since_last_compaction = ( turn + 1 - self._branch_start_turn @@ -600,37 +535,21 @@ async def _run_loop(self) -> RLMResult: result = tool_result.content self.session.log_tool_result(turn, tool_name, result, duration) + content = truncate_tool_output(result) messages.append( { "role": "tool", "tool_call_id": tc.id, - "content": truncate_tool_output(result), + "content": content, } ) - # Auto-compaction: if this turn's prompt_tokens reached the - # configured threshold, ask the model for a handoff summary and - # rebuild the branch around it. Fires at most once per loop - # iteration; the compaction op takes its own LLM call. A - # max_compactions cap, once hit, disables further compaction so - # the context grows to the model's natural limit. - if ( - self.summarize_at_tokens is not None - and usage.prompt_tokens >= self.summarize_at_tokens - and ( - self.max_compactions is None - or self._metrics.num_compactions < self.max_compactions - ) - ): + if self._should_compact(messages, usage, content): try: - await self._compact_branch( - messages, turn, self._active_tool_schemas - ) - except BadRequestError as e: - if not _is_request_too_large(e): - raise - self._metrics.stop_reason = "request_too_large" - final_text = "[request body too large]" + await self._compact_branch(messages, turn) + except CompactionFailed: + self._metrics.stop_reason = "compaction_failed" + final_text = "[context exhausted: compaction failed]" break result = RLMResult( @@ -737,21 +656,119 @@ def _programmatic_tool_call_stats( child = child.merge(trusted_child) return direct, child, child_aggregate.num_sessions + def _can_compact(self) -> bool: + return self.compaction and ( + self.max_compactions is None + or self._metrics.num_compactions < self.max_compactions + ) + + def _should_compact( + self, messages: list[dict], usage: TokenUsage, extra_text: str = "" + ) -> bool: + if self.summarize_at_tokens is None or not self._can_compact(): + return False + if not compactable(messages): + return False + tokens = usage.total + estimated_tokens(extra_text) + return tokens >= self.summarize_at_tokens + + async def _call_model( + self, + messages: list[dict], + *, + checkpoint: bool = False, + compaction_id: str | None = None, + ) -> tuple[Any, TokenUsage]: + request_id = self._semantic_edges.start_request( + self._invocation_id, compaction_id=compaction_id + ) + request: dict = { + "model": self.model, + "messages": messages, + "extra_headers": model_call_headers(request_id), + } + if self._active_tool_schemas: + request["tools"] = self._active_tool_schemas + if checkpoint: + request["tool_choice"] = "none" + else: + request["parallel_tool_calls"] = False + + try: + response = await call_with_retries( + self.client.chat.completions.create, **request + ) + except BaseException: + self._semantic_edges.fail_request(request_id) + raise + self._semantic_edges.finish_request(request_id) + usage = extract_usage(response) + self._total_usage.prompt_tokens += usage.prompt_tokens + self._total_usage.completion_tokens += usage.completion_tokens + if not checkpoint: + self._last_prompt_tokens = usage.prompt_tokens + self._last_call_id = request_id + return response, usage + + async def _complete( + self, messages: list[dict], turn: int + ) -> tuple[Any, TokenUsage]: + """Complete one turn, with at most one compact-and-retry cycle.""" + try: + response, usage = await self._call_model(messages) + except APIStatusError as error: + if ( + self.summarize_at_tokens is None + or not self._can_compact() + or not is_context_overflow(error) + ): + raise + if not compactable(messages): + if self._compacted: + # The conversation is already a compaction floor and still + # overflows - out of moves, end cleanly. + raise CompactionFailed( + "the compacted conversation still overflows" + ) from error + raise + else: + choice = response.choices[0] + if ( + self.summarize_at_tokens is not None + and usage.total < self.summarize_at_tokens + ): + # Usage-verified: this exact prompt was accepted with a full + # reserve of room, so it is a safe checkpoint fallback. + self._last_good = len(messages) + if choice.finish_reason != "length" or not self._should_compact( + messages, usage + ): + return response, usage + + await self._compact_branch(messages, turn) + try: + return await self._call_model(messages) + except APIStatusError as error: + # The rebuilt conversation is sized to fit, so this is out of moves. + if is_context_overflow(error): + raise CompactionFailed( + "the rebuilt conversation still overflows" + ) from error + raise + async def _compact_branch( self, messages: list[dict], turn: int, - active_tools: list[dict], ) -> None: """Ask the model for a handoff summary and rebuild ``messages``. Called in-place: mutates ``messages`` to ``[system, user(framing + - summary)]`` and restarts the ipython kernel. The LLM call for the + summary)]`` while preserving the IPython kernel. The LLM call for the summary is housekeeping, not a work turn, but its tokens land in ``_total_usage`` for cost accounting. - ``active_tools`` is forwarded as ``tools=`` with - ``tool_choice="none"`` so the rendered system prompt matches + Active tools are forwarded with ``tool_choice="none"`` so the system prompt matches regular turns (vLLM's chat-completions layer injects the tools block into the system message only when ``tools=`` is set). With a matching system prompt, prime-rl's RL trajectory walker keeps @@ -760,60 +777,65 @@ async def _compact_branch( keeps the original "text-only summary" behaviour by forbidding tool calls on this turn. """ - # Measure what's about to be dropped BEFORE appending the - # checkpoint prompt — otherwise the prompt's own chars get - # counted as "dropped conversation content", inflating the - # metric and the session log's dropped_chars field. dropped_chars = _count_messages_chars(messages[1:]) turns_since_last = turn + 1 - self._branch_start_turn - # Append the checkpoint prompt and ask the model for a text-only - # summary turn. Tools are advertised to the server (so the system - # prompt renders identically to regular turns) but - # ``tool_choice="none"`` forbids the model from calling any. - # Warn about the REPL restart only when a kernel is actually running. - checkpoint_prompt = CHECKPOINT_COMPACTION_PROMPT + checkpoint_prompt = CHECKPOINT_PROMPT if self._repl is not None: - checkpoint_prompt += REPL_RESTART_NOTE - messages.append({"role": "user", "content": checkpoint_prompt}) + checkpoint_prompt += REPL_NOTE compaction = self._semantic_edges.begin_compaction(self._invocation_id) - request_id = self._semantic_edges.start_request( - self._invocation_id, - compaction_id=compaction.compaction_id, - ) - request_kwargs: dict = { - "model": self.model, - "messages": messages, - "extra_headers": model_call_headers(request_id), - } - if active_tools: - request_kwargs["tools"] = active_tools - request_kwargs["tool_choice"] = "none" try: - response = await call_with_retries( - self.client.chat.completions.create, - **request_kwargs, - ) - usage = extract_usage(response) - summary_text = response.choices[0].message.content or "" + # 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. + base = messages + summary_text = "" + for _ in range(COMPACTION_ATTEMPTS): + checkpoint = [ + *base, + {"role": "user", "content": checkpoint_prompt}, + ] + try: + response, usage = await self._call_model( + checkpoint, + checkpoint=True, + compaction_id=compaction.compaction_id, + ) + except APIStatusError as e: + if not is_context_overflow(e): + raise + base = messages[: self._last_good] + continue + message = response.choices[0].message + # Reasoning never enters the summary: only the reply's final text + # counts, so a reply that lives entirely in the reasoning channel + # is resampled like an empty one. + text = (message.content or "").strip() + if not message.tool_calls and text: + summary_text = text + break + # An unusable reply finished its request without failing it, so the + # compaction still holds the claim - release it for the resample. + self._semantic_edges.release_summary_request(compaction.compaction_id) + if not summary_text: + raise CompactionFailed( + f"no usable summary after {COMPACTION_ATTEMPTS} attempts" + ) except BaseException as exc: - self._semantic_edges.fail_request(request_id) self._semantic_edges.finish_compaction( compaction.compaction_id, "cancelled" if isinstance(exc, asyncio.CancelledError) else "failed", ) - messages.pop() raise - self._semantic_edges.finish_request(request_id) - self._total_usage.prompt_tokens += usage.prompt_tokens - self._total_usage.completion_tokens += usage.completion_tokens system_msg = messages[0] - compacted_user_content = POST_COMPACTION_FRAMING + "\n\n" + summary_text + compacted_user_content = SUMMARY_FRAMING + "\n\n" + summary_text messages[:] = [ system_msg, {"role": "user", "content": compacted_user_content}, ] + self._last_good = len(messages) + self._compacted = True self._semantic_edges.finish_compaction(compaction.compaction_id, "completed") # Log the compaction for traceability. @@ -884,7 +906,8 @@ def execution_snapshot(self) -> dict: "max_concurrent_subagents": self.runtime_config.policy.max_concurrent_subagents, "max_subagent_calls": self.runtime_config.policy.max_subagent_calls, "max_tokens": self.runtime_config.policy.max_tokens, - "summarize_at_tokens": self.runtime_config.policy.summarize_at_tokens, + "compaction": self.compaction, + "summarize_at_tokens": self.summarize_at_tokens, "max_compactions": self.runtime_config.policy.max_compactions, "allow_git": self.runtime_config.policy.allow_git, }, diff --git a/src/rlm/semantic.py b/src/rlm/semantic.py index 6df6b1f..5731aef 100644 --- a/src/rlm/semantic.py +++ b/src/rlm/semantic.py @@ -46,6 +46,7 @@ class _Session: class _Request: session_id: str inbound_edges: list[_PendingEdge] + compaction_id: str | None = None @dataclass @@ -113,7 +114,7 @@ def start_request( inbound.append(_PendingEdge(session.last_request_id, "continuation")) request_id = uuid.uuid4().hex - self._requests[request_id] = _Request(session_id, inbound) + self._requests[request_id] = _Request(session_id, inbound, compaction_id) if compaction_id is not None: compaction = self._compactions[compaction_id] if compaction.session_id != session_id: @@ -139,6 +140,17 @@ def fail_request(self, request_id: str) -> None: request = self._requests.pop(request_id) session = self._sessions[request.session_id] session.pending_edges = request.inbound_edges + session.pending_edges + # A failed summary attempt releases its compaction, so a retry can claim it. + if request.compaction_id is not None: + compaction = self._compactions.get(request.compaction_id) + if compaction is not None and compaction.summary_request_id == request_id: + compaction.summary_request_id = None + + def release_summary_request(self, compaction_id: str) -> None: + """Unbind a compaction's summary request so a resampled attempt can claim it.""" + compaction = self._compactions.get(compaction_id) + if compaction is not None: + compaction.summary_request_id = None def begin_compaction(self, session_id: str) -> Compaction: compaction_id = uuid.uuid4().hex diff --git a/tests/conftest.py b/tests/conftest.py index 4f35734..0a40ebf 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -90,6 +90,7 @@ def model_dump(self, exclude_none: bool = True) -> dict[str, Any]: @dataclass class DummyChoice: message: DummyMessage + finish_reason: str = "stop" @dataclass diff --git a/tests/test_acp.py b/tests/test_acp.py index 4f284c3..dcca78c 100644 --- a/tests/test_acp.py +++ b/tests/test_acp.py @@ -23,7 +23,12 @@ RLMACPAgent, ) from rlm.engine import RLMEngine -from rlm.config import ExecutionPolicy, InvocationContext, ProviderConfig, RuntimeConfig +from rlm.config import ( + ExecutionPolicy, + InvocationContext, + ProviderConfig, + RuntimeConfig, +) from rlm.mcp import MCPHTTPServer, MCPStdioServer from rlm.session import Session from rlm.types import RLMResult, TokenUsage @@ -43,6 +48,7 @@ def _runtime_metadata(**overrides: Any) -> dict[str, Any]: "max_depth": 0, "exec_timeout": 300, "max_tokens": None, + "compaction": False, "summarize_at_tokens": None, "max_compactions": None, "max_concurrent_subagents": 4, @@ -140,6 +146,7 @@ def execution_snapshot(self) -> dict[str, Any]: "max_concurrent_subagents": 4, "max_subagent_calls": 64, "max_tokens": None, + "compaction": False, "summarize_at_tokens": None, "max_compactions": None, "allow_git": False, @@ -318,7 +325,7 @@ async def flaky_first_call(**kwargs): model="test-model", provider=ProviderConfig(base_url=None, api_key="test-key"), invocation=InvocationContext(), - policy=ExecutionPolicy(summarize_at_tokens=1, max_depth=0), + policy=ExecutionPolicy(compaction=True, summarize_at_tokens=1, max_depth=0), ) engine = RLMEngine( client=client, # type: ignore[arg-type] @@ -417,7 +424,7 @@ async def test_compaction_counts_seed_prompt(session): ] try: - await engine._compact_branch(messages, turn=0, active_tools=[]) + await engine._compact_branch(messages, turn=0) finally: await engine.aclose() @@ -479,7 +486,7 @@ async def test_failed_prompt_restores_pre_compaction_context(session): model="test-model", provider=ProviderConfig(base_url=None, api_key="test-key"), invocation=InvocationContext(), - policy=ExecutionPolicy(summarize_at_tokens=1), + policy=ExecutionPolicy(compaction=True, summarize_at_tokens=1), ) engine = RLMEngine( client=client, # type: ignore[arg-type] @@ -822,6 +829,7 @@ def make_engine(**kwargs): agent = RLMACPAgent() agent.on_connect(_Client()) # type: ignore[arg-type] runtime_metadata = _runtime_metadata() + runtime_metadata[RUNTIME_METADATA_KEY]["policy"]["compaction"] = True runtime_metadata[RUNTIME_METADATA_KEY]["policy"]["summarize_at_tokens"] = 1 created = await agent.new_session(str(tmp_path), **runtime_metadata) diff --git a/tests/test_compaction.py b/tests/test_compaction.py new file mode 100644 index 0000000..6d2d38e --- /dev/null +++ b/tests/test_compaction.py @@ -0,0 +1,269 @@ +from __future__ import annotations + +from copy import deepcopy +from types import SimpleNamespace +from typing import Any + +import httpx +from openai import BadRequestError +import pytest + +from conftest import ( + DummyChoice, + DummyClient, + DummyMessage, + DummyResponse, + DummyToolCall, + DummyUsage, +) +from rlm.compaction import is_context_overflow +from rlm.config import ( + ExecutionPolicy, + InvocationContext, + ProviderConfig, + RuntimeConfig, +) +from rlm.engine import RLMEngine +from rlm.session import Session +from rlm.supervisor import SessionTreeSupervisor + + +def _response( + message: DummyMessage, + *, + prompt_tokens: int = 1, + completion_tokens: int = 1, + finish_reason: str = "stop", +) -> DummyResponse: + return DummyResponse( + choices=[DummyChoice(message=message, finish_reason=finish_reason)], + usage=DummyUsage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ), + ) + + +def _overflow() -> BadRequestError: + response = httpx.Response( + 400, + request=httpx.Request("POST", "http://interceptor/v1/chat/completions"), + ) + return BadRequestError( + "This model's maximum context length is 4096 tokens.", + response=response, + body={ + "error": {"message": "This model's maximum context length is 4096 tokens."} + }, + ) + + +def test_overflow_detection_is_status_gated(): + response = httpx.Response( + 400, + request=httpx.Request("POST", "http://interceptor/v1/chat/completions"), + ) + error = BadRequestError( + "maximum context length is 32,768 tokens", + response=response, + body={"error": {"message": "maximum context length is 32,768 tokens"}}, + ) + + assert is_context_overflow(error) + + +class _ScriptedClient(DummyClient): + def __init__( + self, + actions: list[DummyResponse | BaseException], + *, + max_model_len: int | None = None, + ): + super().__init__([]) + self.actions = list(actions) + self.max_model_len = max_model_len + self.base_url = f"http://scripted-{id(self)}" + + @property + def models(self): + outer = self + + class _Models: + async def list(self): + extra = ( + {"max_model_len": outer.max_model_len} + if outer.max_model_len is not None + else {} + ) + card = SimpleNamespace(id="test-model", model_extra=extra) + return SimpleNamespace(data=[card]) + + return _Models() + + async def create(self, **kwargs: Any) -> DummyResponse: + self.calls.append(deepcopy(kwargs)) + if not self.actions: + raise AssertionError("script exhausted") + action = self.actions.pop(0) + if isinstance(action, BaseException): + raise action + return action + + +def _config( + *, + max_depth: int = 0, + summarize_at_tokens: int | None = None, + compaction: bool = True, +): + return RuntimeConfig( + model="test-model", + provider=ProviderConfig(base_url=None, api_key="test-key"), + invocation=InvocationContext(), + policy=ExecutionPolicy( + max_depth=max_depth, + max_concurrent_subagents=max(4, max_depth), + compaction=compaction, + summarize_at_tokens=summarize_at_tokens, + ), + ) + + +async def test_tool_result_overflow_compacts_and_retries(session): + client = _ScriptedClient( + [ + _response( + DummyMessage( + tool_calls=[DummyToolCall("ipython", {"code": "print('x' * 4000)"})] + ) + ), + _overflow(), + _overflow(), + _response(DummyMessage(content="summary")), + _response(DummyMessage(content="done")), + ], + max_model_len=32_768, + ) + engine = RLMEngine( + client=client, # type: ignore[arg-type] + session=session, + runtime_config=_config(), + ) + + try: + result = await engine.run("produce a large tool result") + finally: + engine.close() + + assert result.answer == "done" + assert engine._metrics.num_compactions == 1 + assert client.calls[3]["tool_choice"] == "none" + # The retried work call runs on the rebuilt branch: system + framed summary. + retry_messages = client.calls[4]["messages"] + assert len(retry_messages) == 2 + assert "summary" in retry_messages[1]["content"] + + +async def test_context_overflow_propagates_when_compaction_is_disabled(session): + client = _ScriptedClient([_overflow()]) + engine = RLMEngine( + client=client, # type: ignore[arg-type] + session=session, + runtime_config=_config(compaction=False), + ) + + try: + with pytest.raises(BadRequestError): + await engine.run("overflow without compaction") + finally: + engine.close() + + assert engine._metrics.num_compactions == 0 + + +async def test_decode_context_limit_uses_discovered_threshold(session): + client = _ScriptedClient( + [ + _response( + DummyMessage( + tool_calls=[DummyToolCall("ipython", {"code": "print('ready')"})] + ) + ), + _response( + DummyMessage(content="partial decode"), + prompt_tokens=95, + completion_tokens=5, + finish_reason="length", + ), + _response(DummyMessage(content="summary")), + _response(DummyMessage(content="done")), + ], + max_model_len=112, + ) + engine = RLMEngine( + client=client, # type: ignore[arg-type] + session=session, + runtime_config=_config(), + ) + + try: + result = await engine.run("fill the remaining context") + finally: + engine.close() + + assert result.answer == "done" + assert engine._metrics.num_compactions == 1 + checkpoint_messages = client.calls[2]["messages"] + assert all( + message.get("content") != "partial decode" for message in checkpoint_messages + ) + + +async def test_subagent_recovers_from_context_overflow(tmp_path): + clients: list[_ScriptedClient] = [] + engines: list[RLMEngine] = [] + + def engine_factory(**kwargs: Any) -> RLMEngine: + client = _ScriptedClient( + [ + _response( + DummyMessage( + tool_calls=[DummyToolCall("ipython", {"code": "print('hi')"})] + ) + ), + _overflow(), + _response(DummyMessage(content="summary")), + _response(DummyMessage(content="child done")), + ], + max_model_len=32_768, + ) + engine = RLMEngine(client=client, **kwargs) # type: ignore[arg-type] + clients.append(client) + engines.append(engine) + return engine + + config = _config(max_depth=1) + root = Session(tmp_path / "root") + supervisor = SessionTreeSupervisor( + root_session=root, + runtime_config=config, + cwd=str(tmp_path), + engine_factory=engine_factory, + ) + await supervisor.start() + scope = await supervisor.open_scope(supervisor.root_id) + endpoint = supervisor.endpoint_for(supervisor.root_id) + try: + task = await supervisor._start_child( + endpoint.capability, scope, "recover in the child" + ) + result = await task + finally: + await supervisor.close_scope(scope) + await supervisor.aclose() + root.close() + + assert result.answer == "child done" + assert engines[0].depth == 1 + assert engines[0]._metrics.num_compactions == 1 + assert len(clients[0].calls) == 4 diff --git a/tests/test_config.py b/tests/test_config.py index beddc2f..17bb91d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -18,6 +18,7 @@ def test_runtime_config_redacts_secrets(): max_depth=4, exec_timeout=30, max_tokens=500, + compaction=True, summarize_at_tokens=2048, max_compactions=3, allow_git=True, @@ -27,7 +28,23 @@ def test_runtime_config_redacts_secrets(): search_api_key="search-secret", ) + assert config.model == "override" + assert config.invocation == InvocationContext(depth=2) assert config.invocation.child() == InvocationContext(depth=3) + assert config.policy == ExecutionPolicy( + max_depth=4, + exec_timeout=30, + max_tokens=500, + compaction=True, + summarize_at_tokens=2048, + max_compactions=3, + max_concurrent_subagents=4, + max_subagent_calls=64, + allow_git=True, + ) + assert config.skills == ("search", "edit") + assert config.kernel_env == (("TASK_TOKEN", "task-secret"),) + assert config.search_api_key == "search-secret" assert "task-secret" not in repr(config) assert "search-secret" not in repr(config) assert "test-key" not in repr(config.provider) @@ -45,7 +62,8 @@ def test_policy_rejects_unsafe_recursive_values_and_reserved_headers(): make_client(provider) -def test_default_policy_enables_compaction_and_recursion(): +def test_default_policy_enables_recursion_but_not_compaction(): policy = ExecutionPolicy() - assert policy.summarize_at_tokens == 256_000 + assert policy.compaction is False + assert policy.summarize_at_tokens is None assert policy.max_depth == 1 diff --git a/tests/test_tools.py b/tests/test_tools.py index a0b9457..8b11e7a 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -226,7 +226,7 @@ def test_kernel_receives_policy_env(session): def test_truncate_tool_output_caps_and_reports(): - from rlm.engine import TOOL_OUTPUT_MAX_BYTES, truncate_tool_output + from rlm.compaction import TOOL_OUTPUT_MAX_BYTES, truncate_tool_output small = "x" * 100 assert truncate_tool_output(small) == small