diff --git a/README.md b/README.md index aeca97f..e569daa 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,8 @@ Version `0.1.x` contains the converged foundation layer: - compaction policy and deterministic compactor; - context-budget estimation and non-blocking tokenizer access; - message trimming; +- bounded LLM summaries, projected-token triggers, and tiered compaction; +- session-isolated, content-addressed spill storage for recoverable tool output; - provider-neutral LLM response, stream, and client contracts; - loop configuration, lifecycle contexts, observer protocol, intervention merging, and observer dispatch helpers. @@ -37,9 +39,9 @@ Those revisions are provenance, not runtime dependencies. AgentCore tests and builds without either product checkout. The shared loop accepts product decisions through explicit hooks. Products own -execution-context storage, endpoint registries, metering, deadline policy, -result spilling and aggregate budgets; they no longer need private copies of -the orchestration or parsing engines. +execution-context storage, endpoint registries, metering, deadline policy, and +sandbox mount policy; they no longer need private copies of the orchestration, +parsing, context-management, or spill-storage engines. ## Repository boundary @@ -51,8 +53,9 @@ Code belongs in AgentCore when it: - accepts product behavior through typed inputs or explicit hooks; - has tests that run without either product repository installed. -Provider clients, environment/config loading, session affinity, persistence, -workflow assembly, and product-specific observers remain in their product. +Provider clients, environment/config loading, session affinity, checkpoint +persistence, workflow assembly, sandbox mounting, and product-specific +observers remain in their product. ## Development @@ -113,15 +116,17 @@ edit both products' core copies, that is evidence it belongs here. ## Migration plan -1. **Foundation** (this version): messages, token estimation, compaction, - context budget, trimming. +1. **Foundation** (complete): messages, token estimation, deterministic + compaction, context budget, trimming. 2. **Runtime contracts** (complete): LLM protocols, loop types, errors, retry classification, and explicit product hooks are shared. 3. **LLM runtime** (complete): binding, calls, streaming, response normalization, runaway recovery, and the public `llm_client` facade. 4. **Agent loop** (complete in AgentCore): model/tool parsing, hook-driven tool execution, and `agent_loop` orchestration. -5. Remove product compatibility facades after downstream imports have moved to +5. **Context management Phase 2** (complete in AgentCore): calibrated trigger, + bounded LLM summary, tier selection, observability, and spill storage. +6. Remove product compatibility facades after downstream imports have moved to `agent_core`. Each slice must leave product CI green and must not depend on a floating branch. diff --git a/agent_core/runtime/__init__.py b/agent_core/runtime/__init__.py index db41bd5..2801880 100644 --- a/agent_core/runtime/__init__.py +++ b/agent_core/runtime/__init__.py @@ -1 +1,5 @@ """Shared runtime implementations.""" + +from agent_core.runtime.spill import SpillStore + +__all__ = ["SpillStore"] diff --git a/agent_core/runtime/loop/__init__.py b/agent_core/runtime/loop/__init__.py index 42e0c1d..6653ac9 100644 --- a/agent_core/runtime/loop/__init__.py +++ b/agent_core/runtime/loop/__init__.py @@ -1,10 +1,15 @@ """Shared loop foundation primitives.""" from agent_core.runtime.loop.agent_loop import AgentLoopHooks, run_agent_loop +from agent_core.runtime.loop.budget_consistency import ( + COMPACTION_TRIGGER_RATIO, + check_context_budget, +) from agent_core.runtime.loop.compact import ( DefaultCompactionPolicy, DefaultMessageCompactor, ) +from agent_core.runtime.loop.compact_llm import LLMSummaryCompactor from agent_core.runtime.loop.llm_client import ( RUNAWAY_STATE_KEY, TRUNCATION_CONTINUATION_GUIDANCE, @@ -36,6 +41,13 @@ NativeMessageNormalizer, configure_model_registry, ) +from agent_core.runtime.loop.tiered_compact import ( + DEFAULT_TRIGGER_RATIO, + InputTokenGauge, + InputTokenThresholdPolicy, + TieredCompactor, + compaction_trigger_tokens, +) from agent_core.runtime.loop.tool_call_parser import ( DefaultToolCallParser, MultiFormatToolCallParser, @@ -47,6 +59,8 @@ ) __all__ = [ + "COMPACTION_TRIGGER_RATIO", + "DEFAULT_TRIGGER_RATIO", "RUNAWAY_STATE_KEY", "TRUNCATION_CONTINUATION_GUIDANCE", "AgentLoopHooks", @@ -56,10 +70,13 @@ "DefaultToolCallParser", "DefaultToolResultPostProcessor", "HistoryPolicy", + "InputTokenGauge", + "InputTokenThresholdPolicy", "LLMCallExhausted", "LLMDeadlineExceeded", "LLMReasoningRunaway", "LLMStreamStalled", + "LLMSummaryCompactor", "MessageTrimmer", "ModelProfile", "MultiFormatToolCallParser", @@ -67,12 +84,15 @@ "NullTrimmer", "TaskBoundaryTrimmer", "ThinkTagSplitter", + "TieredCompactor", "ToolExecutionHooks", "bind_max_tokens", "bind_session_id", "bind_temperature", "bind_tools", "call_llm", + "check_context_budget", + "compaction_trigger_tokens", "configure_model_registry", "execute_tools", "extract_final_content", diff --git a/agent_core/runtime/loop/budget_consistency.py b/agent_core/runtime/loop/budget_consistency.py new file mode 100644 index 0000000..3fcbd34 --- /dev/null +++ b/agent_core/runtime/loop/budget_consistency.py @@ -0,0 +1,128 @@ +"""Consistency of the three token budgets a loop is configured with. + +``config/sglang/README.md`` already describes the context/output relationship, +and ``docker/sglang-doctor.sh`` checks it — but only for the sglang compose path, +expressed in ``SGLANG_*`` variables, by a script an operator has to remember to +run. The values that actually reach the loop are the profile's +``max_len`` / ``max_input_tokens`` and the LLM's ``max_tokens``, and those can be +set straight through ``OPENAI_CONTEXT_WINDOW`` / ``OPENAI_MAX_INPUT_TOKENS`` / +``OPENAI_MAX_TOKENS`` against any endpoint, with nothing checking them at all. + +A violation is not visibly a misconfiguration. It surfaces as a provider +rejection mid-run, or as a reasoning watchdog that cannot reliably pre-empt the +provider's output cap. Both read as a runtime fault. + +Warnings only. A deploy that is running today with an unusual combination must +keep running; the point is to say so once, at startup, in terms of the knob to +turn. +""" + +from __future__ import annotations + +import logging + +from agent_core.runtime.loop.tiered_compact import ( + DEFAULT_TRIGGER_RATIO as COMPACTION_TRIGGER_RATIO, +) +from agent_core.runtime.loop.tiered_compact import compaction_trigger_tokens + +logger = logging.getLogger(__name__) + +__all__ = ["COMPACTION_TRIGGER_RATIO", "check_context_budget"] + +# The ratio of ``max_len`` at which tiered compaction triggers is re-exported, +# not restated: a second literal here could silently drift from the trigger, and +# a check that disagreed with the mechanism it describes would be worse than no +# check. A host that passes its own ``ratio`` to ``compaction_trigger_tokens`` +# must pass the same one here. + + +def check_context_budget( + *, + max_len: int, + max_input_tokens: int | None, + max_tokens: int | None, + reasoning_only_max_tokens: int | None = None, + ratio: float = COMPACTION_TRIGGER_RATIO, + label: str, +) -> list[str]: + """Warn about token budgets that cannot all hold at once. + + Returns the warning strings, so a caller (or a test) can assert on them + rather than scraping the log. Anything unset or non-positive is skipped: a + missing bound is a deliberate configuration, not an inconsistency. + + ``ratio`` must be the one the loop's compaction trigger was built with; the + margin reported below is meaningless against any other. + """ + problems: list[str] = [] + + # This relationship does not depend on knowing the context window. Keep it + # outside the max_len-gated checks so profiles without tiered compaction + # still get the warning. + if ( + max_tokens is not None + and max_tokens > 0 + and reasoning_only_max_tokens is not None + and reasoning_only_max_tokens > 0 + and reasoning_only_max_tokens >= max_tokens + ): + # The reasoning-only watchdog cancels a stream that is still ONLY + # reasoning. At or above max_tokens it cannot reliably pre-empt the + # provider's completion cap, so the runaway path handles the reply + # instead — a different mechanism with a different recovery. + problems.append( + f"{label}: reasoning_only_max_tokens " + f"({reasoning_only_max_tokens:,}) is not below max_tokens " + f"({max_tokens:,}); the reasoning watchdog cannot reliably fire " + f"before the provider truncates the reply" + ) + + if max_len <= 0: + for problem in problems: + logger.warning("%s", problem) + return problems + + trigger = compaction_trigger_tokens(max_len, ratio) + + if ( + max_input_tokens is not None + and max_input_tokens > 0 + and max_tokens is not None + and max_tokens > 0 + and max_input_tokens + max_tokens > max_len + ): + # A prompt sitting at the guard limit plus a full-length completion does + # not fit what the endpoint serves. Nothing catches this before the + # provider does, and by then the turn is lost. + problems.append( + f"{label}: max_input_tokens ({max_input_tokens:,}) + max_tokens " + f"({max_tokens:,}) = {max_input_tokens + max_tokens:,} exceeds " + f"max_len ({max_len:,}); a full-length reply to a prompt at the " + f"input guard cannot fit the served context" + ) + + if max_tokens is not None and max_tokens > 0 and max_len - trigger > 0: + # The margin the compaction trigger has before the hard wall, and what + # one full-length reply costs against it. INFORMATIONAL, deliberately not + # a threshold: whether a turn actually crosses this is dynamic (reply + + # its tool result + whatever the reasoning guard really allowed), and no + # static rule separates a safe configuration from an unsafe one. + # + # It is logged because the diagnosis that needed it (PR #66) had to be + # done by hand: 51 of 51 dead sub-agents had a last-successful prompt + # BELOW the trigger and a next request past the wall — a whole turn's + # growth fitting inside the margin. `should_compact` reads the request + # already sent, so it cannot see that coming; the standing fix is to + # project the next request, not to bound this ratio. + margin = max_len - trigger + logger.info( + "%s: compaction trigger %s leaves %s tokens before max_len %s; " + "one full-length reply (max_tokens %s) is %.0f%% of that margin", + label, f"{trigger:,}", f"{margin:,}", f"{max_len:,}", + f"{max_tokens:,}", 100.0 * max_tokens / margin, + ) + + for problem in problems: + logger.warning("%s", problem) + return problems diff --git a/agent_core/runtime/loop/compact.py b/agent_core/runtime/loop/compact.py index afea3b1..496e497 100644 --- a/agent_core/runtime/loop/compact.py +++ b/agent_core/runtime/loop/compact.py @@ -31,6 +31,7 @@ "compact_messages", "compress_tool_results", "estimate_tokens", + "tool_names_by_call_id", ] # --------------------------------------------------------------------------- @@ -89,7 +90,7 @@ ) -def _tool_names_by_call_id(messages: list[Message]) -> dict[str, str]: +def tool_names_by_call_id(messages: list[Message]) -> dict[str, str]: """Map ``tool_call_id`` → tool name from AIMessage ``tool_calls``. A ``ToolMessage`` carries no tool name, so any name-keyed policy has to @@ -112,6 +113,10 @@ def _tool_names_by_call_id(messages: list[Message]) -> dict[str, str]: return out +# Compatibility for product facades that predate the public spelling. +_tool_names_by_call_id = tool_names_by_call_id + + def _condense(content: str, max_chars: int) -> str: """Head + tail + URLs of *content*, never longer than the original.""" prefix = f"[Compressed tool result: {len(content):,} characters]\n" @@ -169,7 +174,7 @@ def compress_tool_results( if protect_max_chars is not None and protect_max_chars < 200: raise ValueError("protect_max_chars must be at least 200") - id_to_name = _tool_names_by_call_id(messages) if protect_tool_names else {} + id_to_name = tool_names_by_call_id(messages) if protect_tool_names else {} compacted: list[Message] = [] for message in messages: clone = message.copy() @@ -464,7 +469,7 @@ def compact( return messages id_to_name = ( - _tool_names_by_call_id(messages) if self._protect or self._spill is not None else {} + tool_names_by_call_id(messages) if self._protect or self._spill is not None else {} ) out: list[Message] = [] diff --git a/agent_core/runtime/loop/compact_llm.py b/agent_core/runtime/loop/compact_llm.py new file mode 100644 index 0000000..2f58f41 --- /dev/null +++ b/agent_core/runtime/loop/compact_llm.py @@ -0,0 +1,612 @@ +"""Compact message history with an LLM and a bounded rollback fallback. + +If summarization fails, the compactor keeps only the system message, a +placeholder, and the latest user query so the next turn fits its budget. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import inspect +import logging +import re +import time +from collections.abc import Awaitable, Callable +from typing import Any, cast + +from agent_core.messages import ( + Message, + is_tool_msg, + text_of, + user_msg, +) +from agent_core.runtime.loop.compact import ( + SPILL_MANIFEST_HEADER, + compress_tool_results, + estimate_tokens, + tool_names_by_call_id, +) +from agent_core.runtime.loop.compact import ( + compact_messages as _string_slice_compact, +) +from agent_core.runtime.loop.summary_prompt import ( + compaction_prompt, + format_conversation_for_summary, +) +from agent_core.runtime.retriable import ( + is_context_length_error, + is_transient_network, +) + +__all__ = [ + "CompactionEventEmitter", + "LLMSummaryCompactor", + "SummaryPromptBuilder", + "is_transient_summary_error", +] + +logger = logging.getLogger(__name__) + +# Strong references to in-flight fire-and-forget event emits. asyncio only holds +# a weak reference to a running task, so without this a compaction event could be +# garbage collected before it is delivered. +_PENDING_EMITS: set[asyncio.Task[Any]] = set() + + +# Substrings that identify a *deterministic* summariser failure — retrying one +# of these burns the retry budget on a guaranteed second failure. The dominant +# case is the summariser being handed more than it can read, which is exactly +# what happens on the runs that need compaction most. +_DETERMINISTIC_ERROR_MARKERS = ( + "context length", + "context_length", + "maximum context", + "too many tokens", + "too large", + "invalid_request", + "invalid request", + "bad request", + "unauthorized", + "forbidden", + "not found", + "model_not_found", +) + +# Substrings that identify a *transient* failure worth one more attempt. +_TRANSIENT_ERROR_MARKERS = ( + "timeout", + "timed out", + "connection", + "temporarily", + "unavailable", + "overloaded", + "rate limit", + "rate_limit", + "too many requests", +) + +# Status codes, kept apart from the substrings above because a bare number is +# not a substring worth matching: provider messages carry request ids, durations +# and quota hints, so "Rate limit reached — try again in 400ms" used to classify +# a textbook retriable 429 as a permanent 400 and skip its whole retry budget. +# Matched only as a standalone token, and only after the exception's own status +# field — which is authoritative when the provider exposes one. +_DETERMINISTIC_STATUS_CODES = frozenset({400, 401, 403, 404, 413, 422}) +_TRANSIENT_STATUS_CODES = frozenset({408, 425, 429, 500, 502, 503, 504, 529}) +_STATUS_ATTRS = ("status_code", "status", "http_status") +# Neither side of a code may abut a word character, a dot or a percent sign: +# ``400ms``, ``req_400abc``, ``0.400`` and ``400%`` are not status codes. +_CODE_PATTERNS = { + code: re.compile(rf"(? int | None: + """The HTTP status an exception carries, if it names one at all.""" + + candidates: list[Any] = [getattr(exc, attr, None) for attr in _STATUS_ATTRS] + response = getattr(exc, "response", None) + if response is not None: + candidates.extend(getattr(response, attr, None) for attr in _STATUS_ATTRS) + for value in candidates: + if isinstance(value, bool): + continue + if isinstance(value, str): + value = value.strip() + if not value.isdigit(): + continue + value = int(value) + if isinstance(value, int) and 100 <= value <= 599: + return value + return None + + +def _mentions_status(text: str, codes: frozenset[int]) -> bool: + return any(_CODE_PATTERNS[code].search(text) for code in codes) + + +# Backoff between transient summary retries: ``min(base * attempt, cap)``. +# Module constants for the same reason ``_runaway.py`` keeps its own — the values +# are a judgement about endpoint hiccups, not a per-call parameter, and a test +# has to be able to shorten them without waiting out a real backoff. +_RETRY_BACKOFF_BASE_S = 2.0 +_RETRY_BACKOFF_CAP_S = 5.0 + + +def is_transient_summary_error(exc: BaseException) -> bool: + """Whether a summariser failure is worth retrying. + + Retrying a deterministic error — a 4xx, and above all "context too large" — + is guaranteed to fail again while consuming the retry budget. A single + summariser call can block for the full LLM timeout (600s in the shipped + profiles), so getting this wrong costs minutes, not milliseconds. That + asymmetry is why the classification exists at all rather than a flat retry + count. + + Unrecognised errors are treated as transient: one extra attempt is cheaper + than losing the whole research history to the fallback layout. The bias is + deliberate and bounded on the other side by ``retry_total_timeout_s``. + """ + if isinstance(exc, (asyncio.TimeoutError, ConnectionError, TimeoutError)): + return True + if isinstance(exc, asyncio.CancelledError): + return False + # Keep this ordering aligned with the main LLM retry path. Context overflow + # is deterministic even when a provider's message mentions transport-like + # details. Conversely, OpenAI-compatible gateways sometimes wrap an + # upstream 5xx/timeout in an HTTP 400 response (``new_api_error`` / + # ``bad_response_status_code``); the wrapper status must not suppress the + # transient signal carried in the structured body. + if is_context_length_error(exc): + return False + if is_transient_network(exc): + return True + # A structured status beats any reading of the message text. + status = _status_code(exc) + if status is not None: + if status in _TRANSIENT_STATUS_CODES: + return True + if status in _DETERMINISTIC_STATUS_CODES: + return False + text = f"{type(exc).__name__}: {exc}".lower() + for marker in _DETERMINISTIC_ERROR_MARKERS: + if marker in text: + return False + for marker in _TRANSIENT_ERROR_MARKERS: + if marker in text: + return True + # Codes last, and transient first among them: a rate limit is a 4xx, so a + # message naming both ("429 ... retry after 400ms") must not read as + # permanent. + if _mentions_status(text, _TRANSIENT_STATUS_CODES): + return True + return not _mentions_status(text, _DETERMINISTIC_STATUS_CODES) + + +# Type alias for the optional event emitter callback. Signature matches +# what the SDK's event sink expects: a payload dict; the caller wraps +# the actual transport (event_store.append / stdout JSONL / no-op). +CompactionEventEmitter = Callable[[dict[str, Any]], Awaitable[None] | None] +SummaryPromptBuilder = Callable[[list[Message]], str] + + +def _last_human_message(messages: list[Message]) -> Message | None: + for msg in reversed(messages): + if msg.get("role") == "user": + return msg + return None + + +class LLMSummaryCompactor: + """Compactor that delegates summarization to an LLM call. + + Implements the ``MessageCompactor`` Protocol with an *async* compact + method. The agent loop awaits the result transparently — see the + awaitable-aware call site in ``agent_loop.py`` step 10. + + Failure handling + ---------------- + - **Summary LLM raises**: log, emit ``compaction`` event with + ``rolled_back=True``, return rollback layout + ``[system, "[Compaction failed — earlier turns dropped]", last_user]``. + - **Summary returns empty / whitespace**: same as raise — rollback + with placeholder text. + - **Successful summary**: emit ``compaction`` event with + ``rolled_back=False``, return ``[system, HumanMessage(summary), recent]``. + + Args: + summary_llm: any object with ``async def chat(messages)`` + returning an object with ``.content`` (see ``_generate_summary``). + ``None`` makes the + compactor delegate to the deterministic string-slice path + (no LLM call, no rollback) — useful when no summarizer is + available but you still want a Protocol-conforming compactor. + emit_event: optional callback that receives the ``compaction`` + event payload. The SDK's stdout JSONL emitter wires this in; + other callers pass ``None`` and the event is just logged. + failure_fallback: ``"panic"`` preserves the legacy tiny rollback; + ``"deterministic"`` keeps a bounded string-slice summary instead. + max_transient_retries: extra attempts for a *transient* summariser + failure. **Defaults to zero**, so every existing direct caller + (``apodex.session``, ``apodex.task_runner``, ``session_history``) + keeps its historical single-attempt behaviour. Those callers do not + pass ``retry_total_timeout_s``, which would make any retry they + inherited unbounded in time — the reason this is opt-in rather + than on by default. Deterministic failures never retry. + retry_total_timeout_s: wall-clock ceiling across ALL attempts, backoff + included. Enforced with ``asyncio.wait_for`` around each attempt, + not merely checked between them: a summariser call carries its own + (much larger) transport timeout, so a between-attempts check is not + a ceiling. ``None`` disables the bound, which is only safe when + ``max_transient_retries`` is zero. + """ + + def __init__( + self, + *, + summary_llm: Any = None, + emit_event: CompactionEventEmitter | None = None, + failure_fallback: str = "panic", + max_transient_retries: int = 0, + retry_total_timeout_s: float | None = None, + prompt_builder: SummaryPromptBuilder = compaction_prompt, + ) -> None: + self._summary_llm = summary_llm + self._emit_event = emit_event + self._failure_fallback = failure_fallback + self._max_transient_retries = max(0, int(max_transient_retries)) + self._retry_total_timeout_s = retry_total_timeout_s + self._prompt_builder = prompt_builder + if self._max_transient_retries and not retry_total_timeout_s: + # Not fatal — a caller may genuinely own the outer deadline — but it + # is the shape of the defect this parameter pair exists to prevent, + # so it must not pass silently. + logger.warning( + "LLMSummaryCompactor: max_transient_retries=%d without " + "retry_total_timeout_s — retries are unbounded in time", + self._max_transient_retries, + ) + + async def compact( + self, + messages: list[Message], + keep_recent: int, + *, + compress_all_tool_results: bool = False, + preserve_tool_names: frozenset[str] = frozenset(), + ) -> list[Message]: + # No LLM configured → behave as the string-slice fallback. This + # keeps the SDK construction simple: callers can pass + # ``LLMSummaryCompactor()`` and trust it never crashes on a + # missing summarizer. + source_messages = ( + compress_tool_results(messages) if compress_all_tool_results else messages + ) + if self._summary_llm is None: + return self._string_slice(source_messages, keep_recent) + + sys_msgs, middle, recent = self._partition(source_messages, keep_recent) + if not middle: + return source_messages # nothing to summarize + + # Compute once; reused by both the success and rollback paths so + # we don't re-walk the full message list per emit. + tokens_before = estimate_tokens(source_messages) + + summary_text, attempts, failure_reason, failure_error = ( + await self._generate_summary_with_retry( + middle, preserve_tool_names=preserve_tool_names, + ) + ) + + if failure_reason: + logger.warning( + "LLMSummaryCompactor: %s after %d attempt(s) (%s) — rolling back", + failure_reason, attempts, failure_error or "no detail", + ) + return await self._rollback( + source_messages, + sys_msgs, + len(middle), + keep_recent=keep_recent, + tokens_before=tokens_before, + reason=failure_reason, + error=failure_error, + attempts=attempts, + ) + + if not summary_text: + logger.warning( + "LLMSummaryCompactor: empty summary text — rolling back", + ) + return await self._rollback( + source_messages, + sys_msgs, + len(middle), + keep_recent=keep_recent, + tokens_before=tokens_before, + reason="empty_summary", + attempts=attempts, + ) + + summary_msg = user_msg( + "[Compacted summary of earlier turns — older raw messages " + "have been replaced by this rollup. Continue from here.]\n\n" + + summary_text + ) + new_messages: list[Message] = [*sys_msgs, summary_msg, *recent] + await self._emit( + { + "rolled_back": False, + "messages_before": len(messages), + "messages_after": len(new_messages), + "tokens_before": tokens_before, + "tokens_after": estimate_tokens(new_messages), + "compactor": "llm", + "attempts": attempts, + # The summary IS the compaction's product: every token the model + # keeps of the replaced turns passes through here. Without it an + # observer can report that a compaction happened and how much it + # freed, but nothing about whether what survived was worth + # keeping — which is the only question a reader of the record + # actually has. + "summary": summary_text, + } + ) + return new_messages + + @staticmethod + def _partition( + messages: list[Message], + keep_recent: int, + ) -> tuple[list[Message], list[Message], list[Message]]: + """Split into ``(system_prefix, to_summarize, kept_recent)``.""" + sys_msgs: list[Message] = [] + rest: list[Message] = [] + for msg in messages: + if msg.get("role") == "system" and not rest: + sys_msgs.append(msg) + else: + rest.append(msg) + + if len(rest) <= keep_recent: + return sys_msgs, [], rest + + split_idx = len(rest) - keep_recent + # Avoid an orphan ToolMessage at the head of the kept window: the + # matching AIMessage(tool_calls=[...]) would otherwise be left + # in the middle and Azure rejects an orphan tool_call_id with 400. + forward = split_idx + while forward < len(rest) and is_tool_msg(rest[forward]): + forward += 1 + if forward < len(rest): + split_idx = forward + else: + # The whole tail is tool results — a parallel tool-call turn that + # emitted at least ``keep_recent`` of them. Walking forward runs off + # the end and leaves the last result orphaned (the bug this guard + # exists to prevent), so walk BACK to the assistant message that + # owns the calls and keep that turn whole instead. The kept window + # grows past ``keep_recent``; ``keep_recent`` is a floor on how much + # recent history survives, not a cap. + while split_idx > 0 and is_tool_msg(rest[split_idx]): + split_idx -= 1 + + return sys_msgs, rest[:split_idx], rest[split_idx:] + + async def _generate_summary_with_retry( + self, + middle: list[Message], + *, + preserve_tool_names: frozenset[str] = frozenset(), + ) -> tuple[str, int, str, str]: + """Summarize with bounded retries → ``(text, attempts, reason, error)``. + + ``reason`` is empty on success. Only transient failures are retried, and + the sequence is additionally bounded by ``retry_total_timeout_s`` so a + stuck endpoint cannot stall the loop for ``retries × llm_timeout``. + + A permanent failure returns ``llm_error_permanent`` rather than + ``llm_error``: both roll back identically, but a reader of the record + needs to know the difference between "the endpoint hiccuped and we ran + out of attempts" and "this summary can never succeed at this size", + because only the second one means the relief target is unreachable. + """ + deadline = ( + time.monotonic() + self._retry_total_timeout_s + if self._retry_total_timeout_s is not None + and self._retry_total_timeout_s > 0 + else None + ) + attempts = 0 + last_error = "" + while True: + remaining = deadline - time.monotonic() if deadline is not None else None + if remaining is not None and remaining <= 0: + return ( + "", + attempts, + "llm_error", + last_error or "summary retry deadline exceeded", + ) + + attempts += 1 + try: + if remaining is None: + summary = await self._generate_summary( + middle, preserve_tool_names=preserve_tool_names, + ) + else: + # The deadline is a ceiling for the whole sequence, not a + # flag checked after an individual call has already spent + # its own transport timeout. + summary = await asyncio.wait_for( + self._generate_summary( + middle, preserve_tool_names=preserve_tool_names, + ), + timeout=remaining, + ) + return summary, attempts, "", "" + except asyncio.CancelledError: + raise + except Exception as exc: + last_error = str(exc) or type(exc).__name__ + if not is_transient_summary_error(exc): + logger.info( + "LLMSummaryCompactor: deterministic summariser error, " + "not retrying (%s)", exc, + ) + return "", attempts, "llm_error_permanent", last_error + exhausted = attempts > self._max_transient_retries + out_of_time = deadline is not None and time.monotonic() >= deadline + if exhausted or out_of_time: + return "", attempts, "llm_error", last_error + # Short fixed backoff: the failure modes worth retrying here are + # endpoint hiccups, and the call itself already took seconds. + delay = min(_RETRY_BACKOFF_BASE_S * attempts, _RETRY_BACKOFF_CAP_S) + if deadline is not None and deadline - time.monotonic() <= delay: + # Refuse a retry whose backoff alone would consume the rest + # of the budget. Returning now is faster and keeps the + # advertised ceiling strict rather than approximate. + return "", attempts, "llm_error", last_error + await asyncio.sleep(delay) + + async def _generate_summary( + self, + to_summarize: list[Message], + *, + preserve_tool_names: frozenset[str] = frozenset(), + ) -> str: + id_to_name = ( + tool_names_by_call_id(to_summarize) if preserve_tool_names else {} + ) + preserved_ids = frozenset( + call_id + for call_id, name in id_to_name.items() + if name in preserve_tool_names + ) + # Drop the spill index, the way ``compact_messages`` already does on the + # deterministic path: it is a list of paths, there is nothing in it to + # summarize, and ``TieredCompactor`` re-attaches the real one afterwards + # from the refs it collected. Nothing breaks if a summary now quotes the + # header — no code reads an index back out of prose — but spending + # summarizer budget on it is still waste. The header check is the legacy + # path for a history checkpointed before the field existed. + conversation = format_conversation_for_summary( + [ + message + for message in to_summarize + if not message.get("spill_refs") + and SPILL_MANIFEST_HEADER not in text_of(message.get("content")) + ], + preserve_tool_result_ids=preserved_ids, + ) + prompt = self._prompt_builder(to_summarize).format(conversation=conversation) + resp = await self._summary_llm.chat([user_msg(prompt)]) + text: Any = getattr(resp, "content", None) or "" + if isinstance(text, list): + # Anthropic-style content blocks: list of {"type":"text","text":...} + parts: list[str] = [] + for block in cast(list[object], text): + if isinstance(block, dict): + value = cast(dict[str, Any], block).get("text", "") + parts.append(str(value)) + else: + parts.append(str(block)) + text = "".join(parts) + if not isinstance(text, str): + text = str(text) + return text.strip() + + async def _rollback( + self, + messages: list[Message], + sys_msgs: list[Message], + dropped: int, + *, + keep_recent: int, + tokens_before: int, + reason: str, + error: str | None = None, + attempts: int = 1, + ) -> list[Message]: + """Use deterministic compaction when requested, else panic truncation. + + The legacy mode keeps only the system prefix, a failure marker, and the + latest user message. Tiered compaction opts into deterministic slicing, + whose size is verified by the caller before it is selected. + """ + if self._failure_fallback == "deterministic": + new_messages = _string_slice_compact(messages, keep_recent) + else: + last_user = _last_human_message(messages) + placeholder = user_msg( + f"[Compaction failed — {dropped} earlier turns dropped. " + f"Reason: {reason}.]" + ) + new_messages = [*sys_msgs, placeholder] + if last_user is not None and last_user not in sys_msgs: + new_messages.append(last_user) + + payload: dict[str, Any] = { + "rolled_back": True, + "rollback_reason": reason, + "messages_before": len(messages), + "messages_after": len(new_messages), + "tokens_before": tokens_before, + "tokens_after": estimate_tokens(new_messages), + "compactor": "llm", + "attempts": attempts, + } + if error: + payload["error"] = error + await self._emit(payload) + return new_messages + + def _string_slice( + self, + messages: list[Message], + keep_recent: int, + ) -> list[Message]: + new_messages = _string_slice_compact(messages, keep_recent) + # Synchronous emit: schedule async path on a dummy loop only when + # an emitter exists. Most callers pass ``emit_event=None`` (SDK + # default), so the common case is a pure-sync no-op. + if self._emit_event is not None: + try: + rv = self._emit_event( + { + "rolled_back": False, + "messages_before": len(messages), + "messages_after": len(new_messages), + "tokens_before": estimate_tokens(messages), + "tokens_after": estimate_tokens(new_messages), + "compactor": "string", + } + ) + if inspect.isawaitable(rv): + # Best-effort fire-and-forget. The string-slice path + # is sync so we cannot reliably await; the caller + # should wire async emitters at the LLM-summary path. + with contextlib.suppress(RuntimeError): + # Keep a reference: a bare ensure_future can be garbage + # collected before it runs. Discarded via the callback + # once it settles. + task = asyncio.ensure_future(rv) + _PENDING_EMITS.add(task) + task.add_done_callback(_PENDING_EMITS.discard) + except Exception: + logger.debug("compaction event emitter raised", exc_info=True) + return new_messages + + async def _emit(self, payload: dict[str, Any]) -> None: + if self._emit_event is None: + return + try: + rv = self._emit_event(payload) + if inspect.isawaitable(rv): + await rv + except Exception: + logger.debug("compaction event emitter raised", exc_info=True) diff --git a/agent_core/runtime/loop/summary_prompt.py b/agent_core/runtime/loop/summary_prompt.py new file mode 100644 index 0000000..10f8bb2 --- /dev/null +++ b/agent_core/runtime/loop/summary_prompt.py @@ -0,0 +1,300 @@ +"""Shared structured-summary prompt for message-history compaction.""" + +from __future__ import annotations + +from collections.abc import Callable + +from agent_core.messages import Message, text_of + +__all__ = [ + "COMPACTION_PROMPT", + "HANDOFF_COMPACTION_PROMPT", + "RESEARCH_COMPACTION_PROMPT", + "compaction_prompt", + "format_conversation_for_summary", +] + + +# Shaped for research / QA: the units it preserves are candidates, sources and +# queries. See ``HANDOFF_COMPACTION_PROMPT`` for the long-run coding shape, and +# ``compaction_prompt`` for how one is chosen. +RESEARCH_COMPACTION_PROMPT = """\ +Summarize the conversation that follows so the assistant agent can \ +continue from a smaller context. The ORIGINAL question is in the \ +system message and is NOT being summarized — it remains accessible. + +PRESERVE EXACTLY (do NOT compress these): +- All entity names ever mentioned as candidates (people, places, \ +organizations, works, songs, etc). +- Candidates that were RULED OUT and the specific reason \ +(so future turns don't re-explore them). +- Source URLs and titles already consulted (so future turns don't \ +re-fetch the same pages). +- The exact search queries already issued, verbatim (so future turns \ +don't re-run them). List them even when they returned nothing useful. +- Verified facts that confirm or contradict each clue against each \ +candidate. +- Any partial answer hypotheses that are still in play. +- Numeric values, dates, or quoted phrases that came up in evidence. + +OK TO COMPRESS: +- Tool call mechanics (drop the JSON args, EXCEPT a search query — that \ +one is preserved above; keep the substantive result). +- The RESULTS of a search that returned nothing useful (keep the query \ +itself under PRESERVE above; the finding compresses to "no useful hits"). +- Reasoning deliberation prose (keep conclusions, drop intermediate \ +musings). +- Repeated framing or restatements of the question. + +OUTPUT FORMAT — write a concise structured summary that the agent can \ +read at the start of its next turn: + +## Investigation so far + + +## Candidates +- Confirmed / strong: +- Ruled out: +- Still open: + +## Sources consulted +- — + +## Queries already run +- — + +## Verified facts +- + +## Next steps suggested by current state + + +Conversation to summarize: +{conversation}""" + + +# Shaped for long-run coding and file work, where what the next turn needs is +# not a list of candidates but the exact state of the work: which commands ran, +# which paths they touched, what they actually returned, and what is still +# unverified. Written first-person on the same reasoning that kimi-code's +# compaction instruction uses — a note the agent writes to itself, not a report +# about someone else's work — because a third-party summary drops the operational +# detail (the exact command, the exact error line) that a resuming agent has to +# re-derive by re-running things. +HANDOFF_COMPACTION_PROMPT = """\ +You are about to run out of context. Write a note to YOURSELF so you can pick \ +this task up after the conversation below is discarded. + +Write it first person, present tense, as your own continuing train of thought — \ +not a third-party report about someone else's work. Write it in the language the \ +conversation has been using. The next turn will see your most recent user \ +messages and this note, and nothing else: every tool call and tool result below \ +will be gone. + +PRESERVE EXACTLY — these cannot be re-derived cheaply, and a resuming agent that \ +guesses them will redo work or corrupt state: +- The exact commands that were run, verbatim, and whether each succeeded. +- The exact file paths that were read or written, and what changed in each. +- What the results actually SAID: the concrete values returned, the key output \ +lines, the exact error text, the signature or schema a lookup revealed. Not \ +"the tests failed" but which test and what the assertion said. +- Any recovery path already mentioned in the conversation (a spilled tool result \ +under a .spill or spill directory). Reproduce the path character for character; \ +it is how the discarded detail is retrieved. +- Decisions already settled, kept SEPARATE from questions still open, so the \ +next turn neither reopens a closed choice nor treats an undecided one as decided. +- Anything an earlier step CLAIMED but never verified — "tests pass", "the fix \ +works", "the file was created". Say plainly that it is unverified. + +OK TO COMPRESS OR DROP: +- Intermediate attempts that were superseded. Keep only the final working \ +version of any code or command. +- Errors already diagnosed and fixed. +- Deliberation prose: keep the conclusion, drop the reasoning that led to it. +- Repeated restatements of the task. + +Then set out the forward plan, and invest in it: right now you hold more context \ +on this task than you ever will again. Give the exact next command or tool call, \ +the remaining sequence after it, the decisions you have already made for those \ +later steps, and the obstacles you can already foresee with how you mean to \ +handle each. Anything you settle here is one less thing to rediscover. + +Name what you still do NOT know that the next step depends on — files referenced \ +but never read, an API assumed but never inspected, a question the user has not \ +answered — so the next turn goes and checks instead of assuming. + +Be concise and proportional: a long multi-step task earns detail, a nearly \ +finished one earns a few sentences. Do not pad. + +Conversation to summarize: +{conversation}""" + + +# Backwards-compatible name for the research shape, which was the only one. +COMPACTION_PROMPT = RESEARCH_COMPACTION_PROMPT + +_PROMPT_STYLES = { + "research": RESEARCH_COMPACTION_PROMPT, + "handoff": HANDOFF_COMPACTION_PROMPT, +} + +# ``ToolMeta.category`` values that mean the agent was doing research: reading +# the web. Everything else that a tool call can be — compute, file, search +# (glob/grep over the local filesystem), finance sandboxes — is work on a +# machine, whose state is what a resumed run needs to know about. +_RESEARCH_CATEGORIES = frozenset({"web"}) +# Categories that say nothing either way and must not cast a vote. +_NEUTRAL_CATEGORIES = frozenset({"meta", "orchestration", ""}) + + +def _tool_call_names(messages: list[Message]) -> list[str]: + """Every tool name called across ``messages``, in order, with repeats. + + Accepts both shapes a call can arrive in: the wire form + ``{"function": {"name": ...}}`` that ``Message.tool_calls`` carries, and the + flattened ``{"name": ...}`` that some observers and tests use. + """ + names: list[str] = [] + for message in messages: + for call in message.get("tool_calls") or []: + function = call.get("function") + name = "" + if function: + name = str(function.get("name") or "") + if not name: + name = str(call.get("name") or "") + if name: + names.append(name) + return names + + +def _is_machine_work( + messages: list[Message], + tool_category: Callable[[str], str], +) -> bool: + """Whether this conversation is work on a machine rather than web research. + + Majority vote over the categorised tool calls, ties and empties going to + research — the incumbent shape. Deliberately conservative in that direction: + a research run misclassified as coding would lose the candidate and query + preservation that is the whole point of the research prompt, while a coding + run misclassified as research merely keeps the summary it has had all along. + + The signal is the tool MIX, not any single call: a coding task legitimately + fetches documentation, and a research task legitimately runs ``bash`` to + tabulate what it found. + """ + machine = research = 0 + for name in _tool_call_names(messages): + category = tool_category(name) + if category in _NEUTRAL_CATEGORIES: + continue + if category in _RESEARCH_CATEGORIES: + research += 1 + else: + machine += 1 + return machine > research + + +def compaction_prompt( + messages: list[Message] | None = None, + *, + style: str = "auto", + tool_category: Callable[[str], str] | None = None, +) -> str: + """Return a product-neutral summary prompt. + + Hosts may select ``research`` or ``handoff`` directly. ``auto`` dispatches + only when the host supplies a tool-category callback; AgentCore deliberately + has no dependency on a product plugin registry. Missing or invalid policy + falls back to the established research prompt — as does ``auto`` itself + without a callback or without ``messages``, so the default is identical to + ``research`` for a host that wires neither. + + ``auto`` (the default) dispatches on the conversation's own tool mix, for the + reason the apex A/B surfaced: at identical task success the handoff shape + produced smaller prompts on a coding benchmark, because the research shape's + candidate / source / query sections are empty or padded there. The converse + risk is real too, which is why this routes rather than replaces — on a + web-dominated conversation ``auto`` returns exactly the prompt that was in + use before it existed. Without ``messages`` there is nothing to dispatch on, + so it also falls back to research. + """ + style = style.strip().lower() + if style == "auto": + return ( + HANDOFF_COMPACTION_PROMPT + if messages and tool_category and _is_machine_work(messages, tool_category) + else RESEARCH_COMPACTION_PROMPT + ) + return _PROMPT_STYLES.get(style, RESEARCH_COMPACTION_PROMPT) + + +# Per-call cap on rendered tool arguments. Generous enough for a search +# batch or a file path, small enough that a hundred turns of ``bash`` +# payloads cannot dominate the summarizer's window. +_TOOL_ARGS_MAX_CHARS = 300 + + +def _render_tool_calls(message: Message) -> str: + """Render an assistant message's tool calls as ``-> name(args)`` lines. + + Without this the summarizer never sees a single tool argument, because + only ``role`` and ``content`` were rendered — and a search query lives in + ``tool_calls[i]["function"]["arguments"]``. The prompt above asks for the + exact queries already issued under PRESERVE EXACTLY, and a model asked to + preserve something absent from its input will invent it. A fabricated + "already run" list is worse than no list: it steers later turns away from + searches the agent never actually tried. + + The tool result text cannot substitute. ``web_search``'s single-query path + formats results without echoing ``q``, and ``web_search_aligned`` never + echoes it at all — while the *empty*-result strings DO carry the query, so + relying on results would preserve exactly the failed searches and lose the + useful ones. + """ + calls = message.get("tool_calls") or [] + lines: list[str] = [] + for call in calls: + function = call.get("function") or {} + name = str(function.get("name") or "") + if not name: + continue + args = str(function.get("arguments") or "") + if len(args) > _TOOL_ARGS_MAX_CHARS: + args = args[:_TOOL_ARGS_MAX_CHARS] + "…" + lines.append(f"-> {name}({args})") + return "\n".join(lines) + + +def format_conversation_for_summary( + messages: list[Message], + *, + preserve_tool_result_ids: frozenset[str] = frozenset(), +) -> str: + """Render messages as a plain-text dialogue for the summarizer LLM. + + Long tool results (>4000 chars) are truncated to 3500 chars + ellipsis + so a single noisy turn doesn't dominate the summarizer's input window. + ``preserve_tool_result_ids`` exempts structured fan-in whose middle must + remain visible; callers use it only for explicitly protected tool names. + + Assistant tool calls are rendered with their (truncated) arguments — see + :func:`_render_tool_calls` for why the prompt's query-preservation rule + depends on it. + """ + parts: list[str] = [] + for m in messages: + role = m.get("role", "") or "" + content = text_of(m.get("content")) + preserve_full = ( + role == "tool" + and str(m.get("tool_call_id") or "") in preserve_tool_result_ids + ) + if len(content) > 4000 and not preserve_full: + content = content[:3500] + "\n...[truncated]..." + rendered_calls = _render_tool_calls(m) + if rendered_calls: + content = f"{content}\n{rendered_calls}" if content else rendered_calls + parts.append(f"[{role}]\n{content}") + return "\n\n".join(parts) diff --git a/agent_core/runtime/loop/tiered_compact.py b/agent_core/runtime/loop/tiered_compact.py new file mode 100644 index 0000000..fe52308 --- /dev/null +++ b/agent_core/runtime/loop/tiered_compact.py @@ -0,0 +1,625 @@ +"""Tiered, threshold-triggered context compaction for long tool-heavy loops.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from typing import Any + +from agent_core.loop_types import ( + BaseObserver, + CompactionEvent, + LoopConfig, + TurnContext, +) +from agent_core.messages import Message, is_tool_msg, text_of, user_msg +from agent_core.runtime.loop.compact import ( + INPUT_ESTIMATE_KEY, + KeepLastNToolResultsCompactor, + compress_tool_results, + estimate_tokens, + tool_names_by_call_id, +) +from agent_core.runtime.loop.compact import ( + SPILL_MANIFEST_HEADER as _SPILL_MANIFEST_HEADER, +) +from agent_core.runtime.loop.compact_llm import LLMSummaryCompactor, SummaryPromptBuilder +from agent_core.runtime.loop.summary_prompt import compaction_prompt +from agent_core.runtime.spill import SpillStore + +__all__ = [ + "InputTokenGauge", + "InputTokenThresholdPolicy", + "TieredCompactor", + "compaction_trigger_tokens", +] + +logger = logging.getLogger(__name__) + +# Fraction of ``max_len`` at which tiered compaction fires. At 262,144 this +# leaves ~52k tokens before the endpoint's hard ceiling — a margin one assistant +# turn can cross on its own once reasoning is replayed into history as +# ```` (measured max 130,965 characters ≈ 33k tokens, plus a tool result). +# +# 0.8 is KEPT, and widening it was measured and REJECTED: on the long-reasoning +# subset (160 trials each) 0.65 compacted in 72.6% of trials versus 44.4% at 0.8 +# and scored 49.7% versus 50.0% — no gain, just more history discarded and 4% +# fewer searches per trial. Making the trigger look at the request about to be +# sent (see :class:`InputTokenThresholdPolicy`) is what removed the failures; the +# margin was never the binding constraint. Hosts may pass an explicit ratio so +# the question can be revisited without forking this implementation. +DEFAULT_TRIGGER_RATIO = 0.8 + + +def compaction_trigger_tokens(max_len: int, ratio: float = DEFAULT_TRIGGER_RATIO) -> int: + """The absolute token threshold tiered compaction should fire at. + + Keeping the ratio explicit makes environment/config ownership a host concern + while every construction site shares validation and rounding semantics. + """ + if not 0.0 < ratio < 1.0: + raise ValueError(f"ratio must be between 0 and 1 (got {ratio})") + return int(max_len * ratio) + + +def _spill_can_recover(spill: Callable[[str, str], str | None] | None) -> bool: + """Whether a spill callback can actually hand the model back a path. + + ``SpillStore.spill_compacted_body`` writes with ``require_visible=True`` and + so returns ``None`` for every body when the store has no ``visible_root`` — + a configuration ``docs/context-management-boundary.md`` explicitly allows for + a backend that cannot name host paths. A candidate selected on the strength + of "a spill callback exists" would then shorten results with nothing + recoverable behind them, so the question has to be whether the callback can + produce a ref, not whether it is set. + + A bound method of an unmounted store answers for itself; any other callable + is the host's own and is taken at its word. + """ + if spill is None: + return False + owner = getattr(spill, "__self__", None) + if isinstance(owner, SpillStore): + return owner.visible_root is not None + return True + + +_FULL_TEXT_PREFIX = "[Full text] " +_SPILL_MANIFEST_MAX_PATHS = 20 +_SPILL_MANIFEST_MAX_CHARS = 3_000 + + +class InputTokenGauge(BaseObserver): + """Record the REAL input tokens of each LLM call for the compaction policy. + + The loop's ``CompactionPolicy.should_compact`` only receives the loop's own + token *estimate*; this gauge exposes the actual ``prompt_tokens`` so the + threshold can be expressed against the true context window. ``critical`` so + the update lands INLINE (before the turn's compaction step), not as a + fire-and-forget task that could race it. Read-only — never intervenes. + """ + + critical: bool = True + + def __init__(self) -> None: + self.tokens = 0 + self.estimate = 0 + + async def on_llm_response(self, ctx: TurnContext) -> None: + u = ctx.usage or {} + # Normalised usage uses ``prompt_tokens``; accept raw Anthropic + # ``input_tokens`` as a fallback. + self.tokens = int(u.get("prompt_tokens") or u.get("input_tokens") or 0) + # Both sides of the scale must describe the SAME request. Estimating + # ``ctx.messages`` here cannot: the loop appends this turn's assistant + # reply before building the context, and the per-call system addendum + # only ever exists inside ``messages_for_call``. The loop publishes the + # estimate of the list it actually sent; 0 when absent, which leaves the + # scale at 1.0 rather than silently understating it. + self.estimate = int(ctx.metadata.get(INPUT_ESTIMATE_KEY, 0) or 0) + return None + + async def on_loop_start(self, config: LoopConfig) -> None: + self.tokens = 0 + self.estimate = 0 + + def real_to_estimate_scale(self) -> float: + """``real / estimate`` for the SAME request; 1.0 before the first reply. + + Both readers of this ratio want it measured on one snapshot — that is the + whole reason :attr:`estimate` exists rather than being recomputed at the + point of use. They clamp differently, so clamping is left to them. + """ + if self.tokens > 0 and self.estimate > 0: + return self.tokens / self.estimate + return 1.0 + + +class InputTokenThresholdPolicy: + """Compact when the request ABOUT TO BE SENT would exceed ``limit``. + + The gauge alone is not enough, and reading only it killed sub-agents. The + gauge describes the request already SENT, while ``should_compact`` runs at + turn end — after this turn's assistant message and tool results were + appended — so the next request is strictly larger. Trigger on the gauge alone + and a single turn can step from "under the limit" straight past the + endpoint's hard ceiling. + + Measured over 51 sub-agents killed by ``llm_error``: 51/51 had a last + successful prompt just UNDER the 209,715 trigger (median 207,803) and then + issued a request of 264,743 tokens median, which the endpoint rejected with + HTTP 400 — losing the sub-agent and every report it had gathered. One + assistant turn can add 40k+ tokens once reasoning is replayed as ````, + so the trigger-to-ceiling margin is not something a reactive read can cover. + + So both are consulted: the gauge for what the endpoint really charged, and the + loop's own estimate for what is about to go out. The estimate is uncalibrated + (tiktoken cl100k_base standing in for the served tokenizer, measured to + UNDER-state this checkpoint by ~14%), so it is scaled by the ratio the last + call revealed, clamped to ``[1.0, 3.0]``: only under-estimation can cost a + request here, so the floor matters and the ceiling is a sanity bound. + """ + + _MAX_SCALE = 3.0 + + def __init__(self, gauge: InputTokenGauge, limit: int) -> None: + self._gauge = gauge + self._limit = limit + + def should_compact( + self, turn: int, messages: list[Message], estimated_tokens: int, + ) -> bool: + real = self._gauge.tokens + # Same-snapshot ratio, straight off the gauge. Deriving it here from + # ``messages`` instead would reintroduce the defect the gauge's + # ``estimate`` was added to fix: ``messages`` is the turn-END history, so + # dividing by it inflates the denominator with the very tool results that + # make this turn dangerous, understating the scale exactly when it is + # needed most. + scale = min(max(self._gauge.real_to_estimate_scale(), 1.0), self._MAX_SCALE) + projected = int(estimated_tokens * scale) + fire = max(real, projected) > self._limit + if fire: + logger.info( + "compaction trigger: real=%d projected=%d (est=%d scale=%.2f) " + "limit=%d turn=%d", + real, projected, estimated_tokens, scale, self._limit, turn, + ) + return fire + + +class TieredCompactor: + """Compact old results, then summarize real history only when still needed.""" + + def __init__( + self, + *, + keep_tool_result: int, + summary_llm: Any, + relief_target: int, + protect_tool_names: frozenset[str] = frozenset(), + gauge: InputTokenGauge | None = None, + spill: Callable[[str, str], str | None] | None = None, + spill_store: SpillStore | None = None, + summary_retries: int = 2, + summary_retry_timeout_s: float | None = None, + prompt_builder: SummaryPromptBuilder = compaction_prompt, + ) -> None: + if spill is not None and spill_store is not None: + raise ValueError("pass spill or spill_store, not both") + spill_callback = spill or ( + spill_store.spill_compacted_body if spill_store is not None else None + ) + self._tier1 = KeepLastNToolResultsCompactor( + keep_tool_result=keep_tool_result, + protect_tool_names=protect_tool_names, + spill=spill_callback, + ) + # Tier 1 is not the only candidate that can win, and the others rewrite + # the SAME protected results — which Tier 1 also keeps out of its spill + # path, so nothing would be recoverable afterwards. + self._protect = frozenset(protect_tool_names) + self._spill = spill_callback + self._spill_recoverable = _spill_can_recover(spill_callback) + if spill_callback is not None and not self._spill_recoverable: + logger.info( + "TieredCompactor: spill store has no visible_root — unseen tool " + "results stay verbatim, as if no spill were configured", + ) + # ``emit_event`` has existed on LLMSummaryCompactor all along with no + # caller, so the summary it produced reached nothing. Used here as the + # internal channel that carries the summary text up to ``last_event``; + # broadcasting to observers is the loop's job, not the summariser's. + # Retries are on here and off for the bare callers of + # ``LLMSummaryCompactor``: this is the only path where losing the summary + # has a cost worth two extra attempts. Pass + # ``summary_retry_timeout_s=llm_timeout`` and the whole sequence fits + # inside the budget a single call was already allowed to spend, so + # retrying adds no worst-case latency to the turn. + # + # No ceiling → no retries, whatever ``summary_retries`` says. A retry + # count without a time bound is the defect this pair exists to prevent: + # one summariser call can occupy the full LLM timeout, so "two more + # attempts" reads as a small allowance and spends up to 30 minutes. The + # two knobs are therefore not independent, and the safe reading of a + # half-specified pair is the conservative one. + retries = summary_retries if (summary_retry_timeout_s or 0) > 0 else 0 + if retries != summary_retries: + logger.info( + "TieredCompactor: summary_retries=%d ignored — no " + "summary_retry_timeout_s to bound it", + summary_retries, + ) + self._tier2 = LLMSummaryCompactor( + summary_llm=summary_llm, + failure_fallback="deterministic", + emit_event=self._capture_tier2, + max_transient_retries=retries, + retry_total_timeout_s=summary_retry_timeout_s, + prompt_builder=prompt_builder, + ) + self._tier2_payload: dict[str, Any] = {} + #: What the most recent ``compact`` selected. Read by the agent loop to + #: notify ``on_compaction`` observers; ``None`` until the first call. + self.last_event: CompactionEvent | None = None + self._relief_target = relief_target + # Optional: the same gauge that drives the trigger. When present, the + # relief check is expressed in REAL tokens instead of the raw estimate. + self._gauge = gauge + self._last_no_relief_estimate = 0 + + @staticmethod + def _spill_refs(messages: list[Message]) -> list[str]: + """Collect unique spill paths from the messages that carry them. + + Reads ``Message.spill_refs``, which both producers set: the Tier 1 tool + placeholder and the recovery index below. This replaced two heuristics + that recovered the same paths from prose — one deciding whether a bullet + looked like a spill path, one deciding whether a block of text WAS an + index rather than a summary quoting one. Both existed only because the + index was prose in a user message that the summarizer can echo; with the + paths carried as data, an echoed header is just words. + """ + refs: list[str] = [] + for message in messages: + carried = message.get("spill_refs") + if isinstance(carried, list): + for path in carried: + if path and path not in refs: + refs.append(path) + continue + # Legacy: a Tier 1 placeholder from a history checkpointed before the + # field existed. A fixed prefix on its own line, not a shape guess. + content = text_of(message.get("content")) + if message.get("role") != "tool" or _FULL_TEXT_PREFIX not in content: + continue + for line in content.splitlines(): + if line.startswith(_FULL_TEXT_PREFIX): + path = line[len(_FULL_TEXT_PREFIX):].strip() + if path and path not in refs: + refs.append(path) + return refs + + def _spill_protected_results(self, messages: list[Message]) -> list[str]: + """Persist protected fan-in bodies before Tier 2 replaces them. + + Tier 1 deliberately leaves these results inline and therefore never + creates recovery files for them. Once a Tier 2 summary wins, however, + the original tool messages disappear just like every other old result. + Spill them at that transition so summary omissions remain recoverable. + """ + if self._spill is None or not self._protect: + return [] + id_to_name = tool_names_by_call_id(messages) + refs: list[str] = [] + for message in messages: + if not is_tool_msg(message): + continue + name = id_to_name.get(message.get("tool_call_id", ""), "") + if name not in self._protect: + continue + try: + path = self._spill(name, text_of(message.get("content"))) + except Exception: + path = None + if path and path not in refs: + refs.append(path) + return refs + + @staticmethod + def _latest_tool_result_ids(messages: list[Message]) -> frozenset[str]: + """Tool-result ids produced by the latest assistant tool-call turn. + + Compaction runs at turn end, immediately after those results are + appended and before the model has seen them. A spill-less candidate + must therefore keep them verbatim; shrinking them would discard unseen + evidence rather than compacting history the model already consumed. + """ + for message in reversed(messages): + if message.get("role") != "assistant": + continue + ids: set[str] = { + str(call.get("id") or "") + for call in message.get("tool_calls") or [] + if call.get("id") + } + return frozenset(ids) + return frozenset() + + def _spill_changed_tool_results( + self, + source: list[Message], + transformed: list[Message], + ) -> list[str]: + """Persist every tool body a selected candidate shortened. + + The cheap compression candidates operate on the original history, not + Tier 1, so they can shorten results Tier 1 deliberately kept as recent. + Those results may never have been shown to the model. Spill the source + body before the transformed candidate becomes the final history. + Content-addressed spill naming makes re-spilling an older result cheap + and idempotent. + """ + if self._spill is None: + return [] + id_to_name = tool_names_by_call_id(source) + refs: list[str] = [] + for original, replacement in zip(source, transformed, strict=True): + if not is_tool_msg(original) or not is_tool_msg(replacement): + continue + original_body = text_of(original.get("content")) + if original_body == text_of(replacement.get("content")): + continue + name = id_to_name.get(original.get("tool_call_id", ""), "tool") + try: + path = self._spill(name, original_body) + except Exception: + path = None + if path and path not in refs: + refs.append(path) + return refs + + @staticmethod + def _with_spill_manifest( + messages: list[Message], refs: list[str], + ) -> list[Message]: + """Keep a bounded session-local recovery index when Tier 1 is replaced. + + The index is its own message, carrying the paths in ``spill_refs`` and + rendering them as text for the model. Keeping it separate from the + summary is what removes a whole class of bug: it used to be appended into + the summary message, so replacing it meant finding where the index + started inside prose the summarizer wrote — and the summarizer sees the + previous index and quotes its header. + """ + if not refs: + return messages + selected: list[str] = [] + used = 0 + for path in reversed(refs): + cost = len(path) + 3 + if len(selected) >= _SPILL_MANIFEST_MAX_PATHS or used + cost > _SPILL_MANIFEST_MAX_CHARS: + break + selected.append(path) + used += cost + selected.reverse() + index: Message = user_msg( + _SPILL_MANIFEST_HEADER + "\n" + "\n".join(f"- {path}" for path in selected), + ) + index["spill_refs"] = selected + + out: list[Message] = [] + replaced = False + for message in messages: + # Exactly one message can be the index: the one that says it is. + if not replaced and message.get("role") == "user" and message.get("spill_refs"): + out.append(index) + replaced = True + continue + out.append(message) + if replaced: + return out + + insert_at = 0 + while insert_at < len(out) and out[insert_at].get("role") == "system": + insert_at += 1 + out.insert(insert_at, index) + return out + + def _capture_tier2(self, payload: dict[str, Any]) -> None: + """Stash the summariser's event. Called during Tier 2 regardless of + whether Tier 2 goes on to win, so ``_selected`` reads it only for the + label that actually used it.""" + self._tier2_payload = payload + + def _selected( + self, + label: str, + messages: list[Message], + *, + before_tokens: int, + scale: float, + spill_refs: int, + ) -> list[Message]: + after_tokens = estimate_tokens(messages) + # Only a tier that summarises has a summary. Tier 2 can run and lose to + # a cheaper candidate, and its payload would still be in the stash — so + # key off the selected label, not off the stash being populated. + summary = "" + rollback_reason = "" + attempts = 0 + if label.startswith("tier2"): + attempts = int(self._tier2_payload.get("attempts") or 0) + if self._tier2_payload.get("rolled_back"): + # The slice a failed summariser falls back to can still be the + # smallest candidate and win. Recording only an empty summary + # would read as "no summariser ran". + rollback_reason = str( + self._tier2_payload.get("rollback_reason") or "unknown", + ) + else: + summary = str(self._tier2_payload.get("summary") or "") + self.last_event = CompactionEvent( + turn=0, # stamped by the loop, which owns the turn counter + seq=0, # stamped by the loop, which owns the sequence + selected=label, + tokens_before=before_tokens, + tokens_after=after_tokens, + relief_met=scale * after_tokens <= self._relief_target, + spill_refs=spill_refs, + attempts=attempts, + summary=summary, + rollback_reason=rollback_reason, + ) + # ``rollback_reason`` goes last: ``scripts/truncation_metrics.py`` reads + # ``selected=(\S+)`` off the front of this line, so appending is safe. + logger.info( + "TieredCompactor selected=%s tokens=%d->%d scaled_after=%d " + "relief_target=%d relief_met=%s spill_refs=%d attempts=%d " + "rollback_reason=%s", + label, + before_tokens, + after_tokens, + int(scale * after_tokens), + self._relief_target, + scale * after_tokens <= self._relief_target, + spill_refs, + attempts, + rollback_reason or "-", + ) + return messages + + async def compact(self, messages: list[Message], keep_recent: int) -> list[Message]: + # Calibrate the estimate to real tokens (see module docstring): scale by + # the ratio the gauge saw on the pre-compaction messages, so the relief + # threshold below means real ``max_len*0.6``, matching the trigger. + scale = self._real_token_scale() + # Per-compaction, so a Tier 2 summary from an earlier turn can never be + # reported as this turn's. + self._tier2_payload = {} + source_estimate = estimate_tokens(messages) + tier1 = self._tier1.compact(messages, keep_recent) + spill_refs = self._spill_refs(tier1) + selected_spill_ref_count = len(spill_refs) + tier1_tokens = estimate_tokens(tier1) + if scale * tier1_tokens <= self._relief_target: + self._last_no_relief_estimate = 0 + return self._selected( + "tier1", tier1, before_tokens=source_estimate, scale=scale, + spill_refs=len(spill_refs), + ) + + # Cheap fallback can reach large results inside the protected recent + # window. With a spill store that can name a readable path, a selected + # candidate persists every body it shortens below. Without one, keep the + # latest tool-call turn whole: compaction runs before the model has seen + # those results even once. + latest_result_ids: frozenset[str] = ( + frozenset() if self._spill_recoverable + else self._latest_tool_result_ids(messages) + ) + candidates: list[tuple[str, list[Message], int]] = [ + ("tier1", tier1, tier1_tokens), + ] + for width in (1_200, 600, 300): + # These candidates ARE the final history when they win — no summary + # stage follows, so protected results must survive them untouched. + candidate = compress_tool_results( + messages, + max_chars=width, + protect_tool_names=self._protect, + preserve_tool_result_ids=latest_result_ids, + ) + candidate_tokens = estimate_tokens(candidate) + if candidate_tokens < min(item[2] for item in candidates): + candidates.append((f"tool_compression_{width}", candidate, candidate_tokens)) + best_label, best, best_tokens = min(candidates, key=lambda item: item[2]) + if scale * best_tokens <= self._relief_target: + if best_label != "tier1": + # The manifest is this compaction's own recovery index, capped at + # ~3 KB. Charging it against the relief target would discard a + # candidate that already freed enough and fall through to a Tier 2 + # round-trip — slower, and under no obligation to come back + # smaller. The logged ``tokens=`` below still reports the real + # post-manifest size. + candidate_refs = self._spill_changed_tool_results(messages, best) + spill_refs = list(dict.fromkeys([*spill_refs, *candidate_refs])) + best = self._with_spill_manifest(best, spill_refs) + self._last_no_relief_estimate = 0 + return self._selected( + best_label, best, before_tokens=source_estimate, scale=scale, + spill_refs=len(spill_refs), + ) + + # A failed summary can otherwise run every turn without freeing a token. + if ( + self._last_no_relief_estimate + and source_estimate <= int(self._last_no_relief_estimate * 1.1) + ): + if best_label != "tier1": + candidate_refs = self._spill_changed_tool_results(messages, best) + spill_refs = list(dict.fromkeys([*spill_refs, *candidate_refs])) + best = self._with_spill_manifest(best, spill_refs) + return self._selected( + f"{best_label}_cached", best, before_tokens=source_estimate, + scale=scale, spill_refs=len(spill_refs), + ) + + # Summarize the real history, not Tier 1's placeholders. Ordinary tool + # bodies are bounded first, but protected fan-in is passed verbatim: a + # fixed head/tail cap on one collect_reports result can erase every + # middle sub-agent report before the summarizer ever sees it. + tier2_input = compress_tool_results( + messages, + protect_tool_names=self._protect, + preserve_tool_result_ids=latest_result_ids, + ) + tier2 = await self._tier2.compact( + tier2_input, keep_recent, preserve_tool_names=self._protect, + ) + tier2_tokens = estimate_tokens(tier2) + if tier2_tokens < best_tokens: + best_label = "tier2" + protected_refs = self._spill_protected_results(messages) + compressed_refs = self._spill_changed_tool_results( + messages, tier2_input, + ) + all_refs = list(dict.fromkeys([ + *spill_refs, *protected_refs, *compressed_refs, + ])) + best = self._with_spill_manifest(tier2, all_refs) + best_tokens = tier2_tokens + selected_spill_ref_count = len(all_refs) + elif best_label != "tier1": + candidate_refs = self._spill_changed_tool_results(messages, best) + spill_refs = list(dict.fromkeys([*spill_refs, *candidate_refs])) + best = self._with_spill_manifest(best, spill_refs) + selected_spill_ref_count = len(spill_refs) + # Compare the size BEFORE the manifest was attached, for the same reason + # the relief-target check above excludes it: the manifest is this + # compaction's own recovery index (~3 KB, ~750 tokens), not context the + # compactor failed to free. Charging it here recorded a candidate that + # genuinely beat Tier 1 — by less than the index costs — as "no relief", + # which then suppresses the Tier 2 attempt on every following pass until + # the estimate grows 10%, penalising a compactor that actually worked. + # ``_selected`` re-estimates the returned messages, so the real + # post-manifest size is still what gets logged. + if best_tokens >= tier1_tokens: + self._last_no_relief_estimate = source_estimate + else: + self._last_no_relief_estimate = 0 + return self._selected( + best_label, best, before_tokens=source_estimate, scale=scale, + spill_refs=selected_spill_ref_count, + ) + + def _real_token_scale(self) -> float: + """``real / estimate`` for the last request; 1.0 without a gauge or before + the first LLM response (falls back to the raw estimate). + + Unclamped, unlike the trigger's use of the same ratio: this one also gates + the relief target, where a scale below 1.0 is a real measurement (the + estimator over-stating) and flooring it would escalate to Tier 2 for + volume that is not there. + """ + return self._gauge.real_to_estimate_scale() if self._gauge is not None else 1.0 diff --git a/agent_core/runtime/spill.py b/agent_core/runtime/spill.py new file mode 100644 index 0000000..4a89eed --- /dev/null +++ b/agent_core/runtime/spill.py @@ -0,0 +1,452 @@ +"""Session-isolated, content-addressed storage for oversized tool results. + +The store owns persistence and preview shaping. Products own where the root is +mounted, which path is visible to an agent sandbox, and when a session ends. +No product execution context, tool registry, or sandbox implementation is +imported here. +""" + +from __future__ import annotations + +import contextlib +import hashlib +import json +import logging +import os +import time +import uuid +from collections.abc import Callable +from pathlib import Path +from typing import ClassVar, Literal + +logger = logging.getLogger(__name__) + +__all__ = [ + "DEFAULT_AGGREGATE_RESULT_CHARS", + "PreviewShape", + "SpillStore", + "TruncationMode", + "budgeted_preview", + "scope_component", + "truncate_preview", +] + +TruncationMode = Literal["middle", "head"] + +#: What the preview actually kept. ``head_and_tail`` is the only shape that +#: carries an in-band gap marker; promising one that is not there makes the +#: model read the truncation point as the end of the content. +PreviewShape = Literal["head", "head_and_tail"] + +DEFAULT_AGGREGATE_RESULT_CHARS = 200_000 +_MIN_AGGREGATE_KEEP = 2_000 +_MIN_PREVIEW_CHARS = 500 +_MIN_SIDE_CHARS = 200 +_SPILL_MIN_CHARS = 1_500 +_SPILL_SEPARATOR = "\n---\n\n" + + +def scope_component(scope: str) -> str: + """Map an arbitrary session identifier to one safe, stable component.""" + + return hashlib.sha256(scope.encode("utf-8")).hexdigest()[:16] if scope else "" + + +def _elision(removed: int) -> str: + return f"\n… {removed:,} chars elided …\n" + + +def _head_end(text: str, budget: int) -> int: + if budget >= len(text): + return len(text) + newline = text.rfind("\n", budget // 2, budget) + return newline if newline > 0 else budget + + +def _tail_start(text: str, budget: int) -> int: + start = max(len(text) - budget, 0) + if start == 0: + return 0 + newline = text.find("\n", start, start + max(budget // 2, 1)) + return newline + 1 if newline != -1 else start + + +def _truncate_with_shape( + text: str, + budget: int, + *, + mode: TruncationMode = "middle", +) -> tuple[str, PreviewShape]: + """Truncate and report which shape was produced. + + ``middle`` degrades to a head-only prefix whenever a marked gap does not + fit, so the shape has to come from here rather than from the requested mode: + the footer describes it to the model, and inspecting the returned text for a + marker would also match a body that merely quotes one. + """ + + if budget <= 0: + return "", "head" + if len(text) <= budget: + return text, "head_and_tail" + if mode == "head": + return text[: _head_end(text, budget)], "head" + + room = budget - len(_elision(len(text))) + if room < 2 * _MIN_SIDE_CHARS: + return text[: _head_end(text, budget)], "head" + head_end = _head_end(text, room // 2) + tail_start = _tail_start(text, room - room // 2) + if tail_start <= head_end: + return text[: _head_end(text, budget)], "head" + return ( + text[:head_end] + _elision(tail_start - head_end) + text[tail_start:], + "head_and_tail", + ) + + +def truncate_preview( + text: str, + budget: int, + *, + mode: TruncationMode = "middle", +) -> str: + """Return a bounded head-only or head-and-tail preview.""" + + return _truncate_with_shape(text, budget, mode=mode)[0] + + +def _shown_clause(shape: PreviewShape) -> str: + if shape == "head_and_tail": + return "head and tail; the gap is marked above" + return "the beginning only; the remainder is elided without a marker" + + +def _spill_footer( + ref: str, + *, + full_len: int, + note: str = "", + shape: PreviewShape = "head_and_tail", +) -> str: + suffix = f" {note}" if note else "" + if not ref: + return ( + f"\n\n[... only part of this {full_len:,}-char result is shown; the " + f"remainder is not readable from this backend.{suffix}]" + ) + directory = ref.rsplit("/", 1)[0] + return ( + f"\n\n[... only part of this {full_len:,}-char result is shown " + f"({_shown_clause(shape)}). Full content is saved read-only at " + f"{ref}. Only if the elided content is required, read that path with " + f"whichever tool you have — read_file, `cat` via bash, or " + f'grep_search(pattern="...", path="{directory}"). It is read-only; ' + f"do not write there.{suffix}]" + ) + + +def budgeted_preview( + body: str, + *, + cap: int, + ref: str, + full_len: int | None = None, + note: str = "", + mode: TruncationMode = "middle", +) -> str: + """Build a preview plus recovery pointer within one character budget. + + The footer is charged to ``cap``, but its wording depends on the truncation + that actually happened — so room is reserved for the longest variant and the + footer is rebuilt from the preview produced. + + ``body`` is returned unchanged when no bounded rewrite would be shorter than + it. ``_MIN_PREVIEW_CHARS`` is a floor on the preview, so for a small ``cap`` + the composed result can be longer than the body it replaces; a rewrite that + grows the payload while claiming to bound it is worse than the payload, and + a footer pointing at a recovery file is pointless when nothing was elided. + """ + + total = len(body) if full_len is None else full_len + reserve = max( + len(_spill_footer(ref, full_len=total, note=note, shape=shape)) + for shape in ("head", "head_and_tail") + ) + preview_budget = max(cap - reserve, _MIN_PREVIEW_CHARS) + if preview_budget >= len(body) and total <= len(body): + return body + preview, shape = _truncate_with_shape(body, preview_budget, mode=mode) + composed = preview + _spill_footer(ref, full_len=total, note=note, shape=shape) + if len(composed) >= len(body) and total <= len(body): + return body + return composed + + +class SpillStore: + """One conversation's read-only recovery store. + + ``root`` is the physical host directory containing all scoped stores. + ``visible_root`` is the path exposed inside the agent runtime, such as + ``/spill``. Passing ``None`` keeps persistence available to the host but + avoids advertising an unreadable path to the model. + """ + + _created_stores: ClassVar[set[Path]] = set() + + def __init__( + self, + root: str | Path, + scope: str, + *, + visible_root: str | None = None, + document_builder: Callable[[str, str, str], str] | None = None, + ) -> None: + if not scope: + raise ValueError("spill scope must not be empty") + self.root = Path(root).expanduser().resolve() + self.scope = scope + self.component = scope_component(scope) + self.directory = self.root / self.component + self.visible_root = visible_root.rstrip("/") if visible_root else None + self._document_builder = document_builder or self._document + + @property + def visible_directory(self) -> str: + if self.visible_root is None: + return "" + return f"{self.visible_root}/{self.component}" + + def _ensure_directory(self) -> None: + self.root.mkdir(parents=True, exist_ok=True) + if self.directory.is_symlink(): + raise OSError(f"spill scope directory must not be a symlink: {self.directory}") + self.directory.mkdir(parents=True, exist_ok=True) + try: + self.directory.resolve().relative_to(self.root) + except ValueError as exc: + raise OSError("spill scope escaped its configured root") from exc + with contextlib.suppress(OSError): + self.directory.chmod(0o755) + self._created_stores.add(self.directory) + + def ensure(self) -> Path: + """Create this scope's store and return its physical directory.""" + + self._ensure_directory() + return self.directory + + @staticmethod + def _document(tool_name: str, spill_id: str, result: str) -> str: + return ( + f"# {tool_name} — spilled tool result\n\n" + f"- id: `{spill_id}`\n" + f"- captured: {time.strftime('%Y-%m-%dT%H:%M:%S%z')}\n" + f"- length: {len(result):,} chars" + f"{_SPILL_SEPARATOR}{result}" + ) + + def write( + self, + tool_name: str, + body: str, + *, + require_visible: bool = False, + ) -> tuple[Path, str] | None: + """Persist one body atomically and return physical and visible paths.""" + + if require_visible and self.visible_root is None: + return None + self._ensure_directory() + digest = hashlib.sha256( + tool_name.encode("utf-8") + b"\x00" + body.encode("utf-8", "replace"), + ).hexdigest()[:16] + path = self.directory / f"{digest}.md" + if path.is_symlink(): + logger.warning("Refusing symlink at spill destination %s", path) + return None + if not path.exists(): + tmp = path.with_name(f".{digest}.{uuid.uuid4().hex[:8]}.tmp") + try: + tmp.write_text( + self._document_builder(tool_name, digest, body), + encoding="utf-8", + ) + os.replace(tmp, path) + with contextlib.suppress(OSError): + path.chmod(0o444) + except OSError as exc: + logger.warning("Failed to spill %s result: %s", tool_name, exc) + with contextlib.suppress(OSError): + tmp.unlink() + return None + ref = f"{self.visible_directory}/{path.name}" if self.visible_directory else "" + return path, ref + + def spill_compacted_body( + self, + tool_name: str, + body: str, + *, + min_chars: int = _SPILL_MIN_CHARS, + ) -> str | None: + """Persist a body before compaction removes its inline representation.""" + + if len(body) < min_chars: + return None + spilled = self.write(tool_name, body, require_visible=True) + return spilled[1] if spilled else None + + def overflow( + self, + tool_name: str, + result: str, + *, + cap: int, + mode: TruncationMode = "middle", + ) -> str: + """Spill and preview one result when it exceeds ``cap``.""" + + if cap <= 0 or len(result) <= cap: + return result + spilled = self.write(tool_name, result) + ref = spilled[1] if spilled else "" + return budgeted_preview(result, cap=cap, ref=ref, mode=mode) + + def enforce_aggregate_budget( + self, + results: list[str], + *, + tool_names: list[str] | None = None, + cap: int = DEFAULT_AGGREGATE_RESULT_CHARS, + mode: TruncationMode = "middle", + ) -> list[str]: + """Bound all results in one turn, spilling every shortened body first.""" + + total = sum(len(result) for result in results) + if cap <= 0 or total <= cap: + return results + names = list(tool_names or []) + adjusted = list(results) + for idx, result in sorted( + enumerate(results), + key=lambda item: len(item[1]), + reverse=True, + ): + if total <= cap: + break + excess = total - cap + result_cap = max(_MIN_AGGREGATE_KEEP, len(result) - excess) + if result_cap >= len(result): + continue + name = names[idx] if idx < len(names) else "tool" + spilled = ( + self.write(name, result, require_visible=True) + if len(result) >= _SPILL_MIN_CHARS + else None + ) + replacement = budgeted_preview( + result, + cap=result_cap, + ref=spilled[1] if spilled else "", + note="Cut further to fit the per-turn tool-result budget.", + mode=mode, + ) + total -= len(result) - len(replacement) + adjusted[idx] = replacement + return adjusted + + def _physical_path(self, ref: str | Path) -> Path | None: + raw = str(ref) + if self.visible_directory and ( + raw == self.visible_directory or raw.startswith(self.visible_directory + "/") + ): + relative = raw[len(self.visible_directory) :].lstrip("/") + candidate = self.directory / relative + else: + candidate = Path(raw) + try: + resolved = candidate.resolve() + resolved.relative_to(self.directory.resolve()) + except (OSError, ValueError): + return None + return resolved + + def contains_path(self, path: str | Path) -> bool: + """Whether a physical or agent-visible path belongs to this store.""" + + return self._physical_path(path) is not None + + def read(self, ref: str | Path) -> str | None: + """Read the captured body, rejecting paths outside this scope.""" + + path = self._physical_path(ref) + if path is None or not path.is_file(): + return None + return self._read_path(path) + + @staticmethod + def _read_path(path: Path) -> str | None: + try: + text = path.read_text(encoding="utf-8") + if _SPILL_SEPARATOR in text: + return text.split(_SPILL_SEPARATOR, 1)[1] + data = json.loads(text) + content = data.get("content") + return str(content) if content is not None else None + except (OSError, ValueError, TypeError) as exc: + logger.warning("Failed to read spill file %s: %s", path, exc) + return None + + @classmethod + def read_created(cls, path: str | Path) -> str | None: + """Read a physical file only when its store was created by this process.""" + + try: + resolved = Path(path).resolve() + except OSError: + return None + if resolved.parent not in cls._created_stores or not resolved.is_file(): + return None + return cls._read_path(resolved) + + def body_names_file(self, body: str) -> bool: + roots = [str(self.directory), self.visible_directory] + return any(root and f"{root.rstrip('/')}/" in body for root in roots) + + def cleanup(self) -> int: + """Remove files created for this scope, never traversing subdirectories.""" + + if not self.directory.is_dir(): + return 0 + count = 0 + for entry in self.directory.iterdir(): + if not entry.is_file() and not entry.is_symlink(): + continue + try: + entry.unlink() + count += 1 + except OSError: + pass + with contextlib.suppress(OSError): + self.directory.rmdir() + self._created_stores.discard(self.directory) + return count + + @classmethod + def cleanup_process(cls) -> int: + """Remove only stores created by this process.""" + + count = 0 + for directory in list(cls._created_stores): + if directory.is_dir(): + for entry in directory.iterdir(): + if not entry.is_file() and not entry.is_symlink(): + continue + with contextlib.suppress(OSError): + entry.unlink() + count += 1 + with contextlib.suppress(OSError): + directory.rmdir() + cls._created_stores.discard(directory) + return count diff --git a/docs/context-management-boundary.md b/docs/context-management-boundary.md new file mode 100644 index 0000000..ca9aa24 --- /dev/null +++ b/docs/context-management-boundary.md @@ -0,0 +1,38 @@ +# Context-management boundary + +AgentCore owns the portable lifecycle of prompt context: + +- token estimation and calibration against provider-reported input tokens; +- projection of the next request before it crosses the context ceiling; +- deterministic tool-result compression and LLM-backed summaries; +- tier selection, bounded retry/fallback, pairing invariants, and events; +- content-addressed spill persistence, bounded previews, recovery manifests, + aggregate result budgets, session isolation, and cleanup. + +Products configure policy rather than copying the implementation. They supply: + +- the summary LLM and optional prompt selector; +- the context window, trigger ratio, relief target, and protected tool names; +- the physical spill root and the path visible inside their sandbox; +- the session identifier and lifecycle point at which cleanup runs; +- sandbox read-only mounts and write/path authorization. + +`SpillStore` deliberately does not inspect a product execution scope or tool +registry. A product constructs one store per conversation and may pass it +directly to `TieredCompactor`. Tool adapters can also call `overflow()` and +`enforce_aggregate_budget()` on the same instance, so individual overflow, +per-turn aggregate trimming, and later compaction all point to one recovery +store. + +The physical store should be outside every agent-writable workspace. Products +that expose it inside a sandbox must mount it read-only. A backend that cannot +name host paths passes no `visible_root`; AgentCore may retain a host-side copy +for diagnostics, but generated previews do not advertise an unusable path. Such +a store recovers nothing the model can read, so tiered compaction treats it as +if no spill were configured: results from the latest tool-call turn — which the +model has not seen when compaction runs — are kept verbatim rather than +shortened against a path nobody can open. + +Checkpoint persistence and durable user memory are not context compaction. +They remain product concerns; AgentCore only preserves spill references in +messages so a product checkpoint can round-trip them. diff --git a/tests/test_compact_llm.py b/tests/test_compact_llm.py new file mode 100644 index 0000000..0182630 --- /dev/null +++ b/tests/test_compact_llm.py @@ -0,0 +1,257 @@ +"""``LLMSummaryCompactor`` unit tests. + +Cover the four flavours of behaviour: + +- Happy path: summary LLM returns text → output is + ``[system, user_msg(summary), recent]``. +- Rollback on LLM raise: output is + ``[system, placeholder, last_user_query]`` and the emitter receives + ``rolled_back=True`` with a reason. +- Rollback on empty summary: same shape; reason is ``empty_summary``. +- No-LLM fallthrough: when ``summary_llm=None`` the compactor delegates + to the string-slice path (no LLM call, no rollback). + +The agent-loop integration is covered separately by an integration test +in ``tests/test_agent_loop_pre_llm_compact.py`` (existing) once the +async-aware call site lands. +""" +from __future__ import annotations + +import pytest + +from agent_core.llm import LLMResponse +from agent_core.messages import ( + Message, + assistant_msg, + system_msg, + tool_msg, + user_msg, +) +from agent_core.runtime.loop.compact_llm import ( + LLMSummaryCompactor, + is_transient_summary_error, +) + + +class _StaticSummaryLLM: + """Fake LLM whose ``chat`` returns a configured string.""" + + def __init__(self, content: str = "## summary\n- entity X confirmed") -> None: + self.model = "static-summary" + self._content = content + self.calls: list[list] = [] + + async def chat(self, messages, **kwargs) -> LLMResponse: + self.calls.append(list(messages)) + return LLMResponse(content=self._content) + + +class _FailingSummaryLLM: + def __init__(self, exc: Exception | None = None) -> None: + self.model = "failing-summary" + self._exc = exc or RuntimeError("primary down") + + async def chat(self, messages, **kwargs): + raise self._exc + + +def _long_history(n: int = 8) -> list[Message]: + """Build a history with system + n turns of user/AI/tool messages.""" + msgs: list[Message] = [system_msg("be helpful")] + for i in range(n): + msgs.append(user_msg(f"user query #{i}")) + msgs.append(assistant_msg(f"thinking #{i}")) + msgs.append(tool_msg(f"tool result #{i}", f"t{i}")) + msgs.append(user_msg("final user query")) + return msgs + + +@pytest.mark.asyncio +async def test_happy_path_returns_system_summary_recent() -> None: + history = _long_history(n=6) + summary_llm = _StaticSummaryLLM(content="ROLLUP-CONTENT") + + captured: list[dict] = [] + + async def _emit(payload): + captured.append(payload) + + compactor = LLMSummaryCompactor( + summary_llm=summary_llm, emit_event=_emit, + ) + out = await compactor.compact(history, keep_recent=4) + + assert out[0].get("role") == "system" + # The single rollup user message carries the LLM-supplied summary. + assert out[1].get("role") == "user" + assert "ROLLUP-CONTENT" in out[1].get("content") + # ``keep_recent`` recent messages must follow verbatim. + assert out[-1] is history[-1] + + assert len(summary_llm.calls) == 1 + assert len(captured) == 1 + assert captured[0]["rolled_back"] is False + assert captured[0]["messages_before"] == len(history) + assert captured[0]["messages_after"] == len(out) + assert captured[0]["compactor"] == "llm" + assert captured[0]["attempts"] == 1 + assert captured[0]["summary"] == "ROLLUP-CONTENT" + + +@pytest.mark.asyncio +async def test_rollback_on_llm_failure_keeps_only_system_and_last_user() -> None: + history = _long_history(n=6) + last_user = history[-1] + + captured: list[dict] = [] + + async def _emit(payload): + captured.append(payload) + + compactor = LLMSummaryCompactor( + summary_llm=_FailingSummaryLLM(RuntimeError("502 backend")), + emit_event=_emit, + ) + out = await compactor.compact(history, keep_recent=4) + + assert out[0].get("role") == "system" + assert out[1].get("role") == "user" + assert "Compaction failed" in out[1].get("content") + assert "llm_error" in out[1].get("content") + # Last user query preserved so the next turn has something to answer. + assert out[-1] is last_user + + assert captured[0]["rolled_back"] is True + assert captured[0]["rollback_reason"] == "llm_error" + assert "502 backend" in captured[0]["error"] + + +@pytest.mark.asyncio +async def test_rollback_on_empty_summary() -> None: + history = _long_history(n=6) + captured: list[dict] = [] + + async def _emit(payload): + captured.append(payload) + + compactor = LLMSummaryCompactor( + summary_llm=_StaticSummaryLLM(content=" \n \t "), + emit_event=_emit, + ) + out = await compactor.compact(history, keep_recent=4) + + assert out[1].get("role") == "user" + assert "Compaction failed" in out[1].get("content") + assert captured[0]["rollback_reason"] == "empty_summary" + + +@pytest.mark.asyncio +async def test_no_llm_falls_through_to_string_slice() -> None: + """No summarizer → behaves as the deterministic string-slice path.""" + history = _long_history(n=6) + compactor = LLMSummaryCompactor(summary_llm=None) + out = await compactor.compact(history, keep_recent=4) + + # String-slice keeps system + ONE compact summary user message + recent. + assert out[0].get("role") == "system" + assert out[1].get("role") == "user" + assert "Compacted" in out[1].get("content") + # Importantly: no rollback placeholder text. + assert "Compaction failed" not in out[1].get("content") + + +@pytest.mark.asyncio +async def test_short_history_below_keep_recent_is_unchanged() -> None: + history = [ + system_msg("be helpful"), + user_msg("hi"), + assistant_msg("hello"), + ] + compactor = LLMSummaryCompactor(summary_llm=_StaticSummaryLLM()) + out = await compactor.compact(history, keep_recent=4) + + assert out == history + + +def test_partition_keeps_a_parallel_tool_turn_whole() -> None: + """A tail made only of tool results must not leave one of them orphaned. + + Walking the split point forward past the tool results runs off the end of the + history, and the last result is then kept while the ``AIMessage`` carrying + its ``tool_call_id`` is summarised away — exactly the HTTP 400 this guard + exists to prevent. + """ + history: list[Message] = [ + system_msg("S"), + user_msg("q"), + { + "role": "assistant", + "content": "", + "tool_calls": [ + {"id": "a", "function": {"name": "web_search", "arguments": "{}"}}, + {"id": "b", "function": {"name": "web_search", "arguments": "{}"}}, + ], + }, + tool_msg("result a", "a"), + tool_msg("result b", "b"), + ] + + _system, _middle, recent = LLMSummaryCompactor._partition(history, keep_recent=2) + + assert recent[0].get("role") == "assistant" + answered = { + str(call.get("id")) + for message in recent + if message.get("role") == "assistant" + for call in message.get("tool_calls") or [] + } + orphans = { + str(message.get("tool_call_id")) + for message in recent + if message.get("role") == "tool" + } - answered + assert not orphans + + +def test_partition_still_shrinks_the_window_when_it_can() -> None: + history: list[Message] = [ + system_msg("S"), + { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "a", "function": {"name": "web_search", "arguments": "{}"}}], + }, + tool_msg("result a", "a"), + user_msg("next"), + ] + + _system, middle, recent = LLMSummaryCompactor._partition(history, keep_recent=2) + + # The tool result at the head of the window is walked past, not kept. + assert recent == [history[-1]] + assert middle == history[1:3] + + +def test_rate_limit_naming_a_retry_delay_is_still_transient() -> None: + """A bare ``400`` inside a provider message is not a status code. + + OpenAI-compatible endpoints put the retry hint in the message text, so a + substring match read "try again in 400ms" as a permanent HTTP 400 and burned + the whole retry budget of a textbook retriable 429. + """ + assert is_transient_summary_error( + RuntimeError("Rate limit reached for model. Please try again in 400ms."), + ) + assert is_transient_summary_error(RuntimeError("request req_413abc failed: overloaded")) + assert not is_transient_summary_error(RuntimeError("HTTP 400: invalid_request_error")) + + +def test_a_structured_status_field_beats_the_message_text() -> None: + class Rejected(RuntimeError): + status_code = 422 + + class Throttled(RuntimeError): + status_code = 429 + + assert not is_transient_summary_error(Rejected("unhelpful message")) + assert is_transient_summary_error(Throttled("bad request, allegedly")) diff --git a/tests/test_context_management_phase2.py b/tests/test_context_management_phase2.py new file mode 100644 index 0000000..42018ac --- /dev/null +++ b/tests/test_context_management_phase2.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from agent_core.loop_types import TurnContext +from agent_core.messages import assistant_msg, tool_msg, user_msg +from agent_core.runtime.loop.budget_consistency import check_context_budget +from agent_core.runtime.loop.compact import INPUT_ESTIMATE_KEY +from agent_core.runtime.loop.compact_llm import LLMSummaryCompactor +from agent_core.runtime.loop.summary_prompt import ( + HANDOFF_COMPACTION_PROMPT, + RESEARCH_COMPACTION_PROMPT, + compaction_prompt, + format_conversation_for_summary, +) +from agent_core.runtime.loop.tiered_compact import ( + InputTokenGauge, + InputTokenThresholdPolicy, + TieredCompactor, + compaction_trigger_tokens, +) +from agent_core.runtime.spill import SpillStore + + +def _call(name: str, arguments: dict[str, str], call_id: str = "call-1") -> dict: + return { + "id": call_id, + "type": "function", + "function": {"name": name, "arguments": json.dumps(arguments)}, + } + + +def test_summary_renderer_preserves_tool_arguments() -> None: + rendered = format_conversation_for_summary([ + assistant_msg("", tool_calls=[_call("web_search", {"q": "exact query"})]), + tool_msg("result without the query", "call-1"), + ]) + + assert "web_search" in rendered + assert "exact query" in rendered + + +def test_prompt_style_is_explicit_and_auto_requires_host_metadata() -> None: + messages = [assistant_msg("", tool_calls=[_call("bash", {"cmd": "pytest"})])] + + assert compaction_prompt(messages, style="research") == RESEARCH_COMPACTION_PROMPT + assert compaction_prompt(messages, style="handoff") == HANDOFF_COMPACTION_PROMPT + assert compaction_prompt(messages, style="auto") == RESEARCH_COMPACTION_PROMPT + assert compaction_prompt( + messages, + style="auto", + tool_category=lambda _name: "compute", + ) == HANDOFF_COMPACTION_PROMPT + + +def test_auto_dispatch_is_the_documented_default() -> None: + """The default has to be ``auto``, or the routing never fires for any caller. + + Both compactors pass ``prompt_builder=compaction_prompt`` with no style, so a + ``research`` default made the measured handoff routing unreachable while the + docstring said it was on. It is behaviour-preserving: without a + tool-category callback ``auto`` returns the research prompt. + """ + machine_work = [assistant_msg("", tool_calls=[_call("bash", {"cmd": "pytest"})])] + + assert compaction_prompt(machine_work) == RESEARCH_COMPACTION_PROMPT + assert compaction_prompt( + machine_work, + tool_category=lambda _name: "compute", + ) == HANDOFF_COMPACTION_PROMPT + + +@pytest.mark.asyncio +async def test_summary_compactor_uses_host_prompt_builder() -> None: + class CaptureLLM: + def __init__(self) -> None: + self.prompt = "" + + async def chat(self, messages): + self.prompt = str(messages[0].get("content") or "") + return SimpleNamespace(content="summary") + + llm = CaptureLLM() + compactor = LLMSummaryCompactor( + summary_llm=llm, + prompt_builder=lambda _messages: "HOST POLICY\n{conversation}", + ) + history = [user_msg(f"old-{index}") for index in range(8)] + + await compactor.compact(history, keep_recent=2) + + assert llm.prompt.startswith("HOST POLICY") + assert "old-0" in llm.prompt + + +def test_projection_uses_real_to_estimated_ratio_for_the_next_request() -> None: + gauge = InputTokenGauge() + gauge.tokens = 207_803 + gauge.estimate = 182_283 + policy = InputTokenThresholdPolicy(gauge, 209_715) + + assert gauge.tokens < 209_715 + assert policy.should_compact( + turn=40, + messages=[], + estimated_tokens=232_231, + ) + + +def test_gauge_reads_usage_and_the_exact_sent_estimate() -> None: + gauge = InputTokenGauge() + asyncio.run(gauge.on_llm_response(TurnContext( + turn=1, + max_turns=2, + task_id="task", + role_id="role", + ai_text="", + thinking="", + tool_calls=[], + messages=[], + usage={"prompt_tokens": 12_000}, + metadata={INPUT_ESTIMATE_KEY: 10_000}, + ))) + + assert gauge.tokens == 12_000 + assert gauge.estimate == 10_000 + assert gauge.real_to_estimate_scale() == pytest.approx(1.2) + + +def test_trigger_ratio_is_an_explicit_validated_policy() -> None: + assert compaction_trigger_tokens(262_144) == 209_715 + assert compaction_trigger_tokens(262_144, 0.65) == 170_393 + with pytest.raises(ValueError, match="between 0 and 1"): + compaction_trigger_tokens(262_144, 1.0) + + +def test_budget_consistency_accepts_a_fitting_configuration() -> None: + assert check_context_budget( + max_len=262_144, + max_input_tokens=229_376, + max_tokens=32_768, + reasoning_only_max_tokens=16_384, + label="test", + ) == [] + + +def test_budget_consistency_reports_independent_violations() -> None: + problems = check_context_budget( + max_len=262_144, + max_input_tokens=229_376, + max_tokens=65_536, + reasoning_only_max_tokens=65_536, + label="test", + ) + + assert any("exceeds max_len" in problem for problem in problems) + assert any("reasoning watchdog" in problem for problem in problems) + + +@pytest.mark.asyncio +async def test_tiered_compactor_accepts_the_shared_spill_store(tmp_path: Path) -> None: + store = SpillStore(tmp_path, "session", visible_root="/spill") + history = [ + user_msg("research"), + assistant_msg("", tool_calls=[_call("web_fetch", {}, "old")]), + tool_msg("old body " * 1_000, "old"), + assistant_msg("", tool_calls=[_call("web_fetch", {}, "new")]), + tool_msg("new body " * 1_000, "new"), + ] + compactor = TieredCompactor( + keep_tool_result=1, + summary_llm=None, + relief_target=10**9, + spill_store=store, + ) + + result = await compactor.compact(history, keep_recent=1) + + assert any(message.get("spill_refs") for message in result) + assert list(store.directory.glob("*.md")) + + +def test_tiered_rejects_two_spill_owners(tmp_path: Path) -> None: + store = SpillStore(tmp_path, "session", visible_root="/spill") + with pytest.raises(ValueError, match="not both"): + TieredCompactor( + keep_tool_result=1, + summary_llm=None, + relief_target=1, + spill=lambda _name, _body: None, + spill_store=store, + ) + + +@pytest.mark.asyncio +async def test_unseen_results_stay_verbatim_when_the_store_cannot_name_a_path( + tmp_path: Path, +) -> None: + """A spill store with no ``visible_root`` recovers nothing, so it protects nothing. + + ``spill_compacted_body`` writes with ``require_visible=True`` and returns + ``None`` for every body here, so a candidate that shortened the latest + tool-call turn would discard results the model has not seen once, with no + path in the manifest to read them back from. + """ + store = SpillStore(tmp_path, "session") + latest = "new body " * 400 + history = [ + user_msg("research"), + assistant_msg("", tool_calls=[_call("web_fetch", {}, "old")]), + tool_msg("old body " * 400, "old"), + assistant_msg("", tool_calls=[_call("web_fetch", {}, "new")]), + tool_msg(latest, "new"), + ] + compactor = TieredCompactor( + keep_tool_result=1, + summary_llm=None, + relief_target=600, + spill_store=store, + ) + + result = await compactor.compact(history, keep_recent=1) + + bodies = { + message.get("tool_call_id"): str(message.get("content") or "") + for message in result + if message.get("role") == "tool" + } + assert bodies["new"] == latest + assert compactor.last_event is not None + assert not compactor.last_event.selected.startswith("tool_compression_") + + +@pytest.mark.asyncio +async def test_unseen_results_may_be_compressed_once_a_path_is_recoverable( + tmp_path: Path, +) -> None: + store = SpillStore(tmp_path, "session", visible_root="/spill") + latest = "new body " * 400 + history = [ + user_msg("research"), + assistant_msg("", tool_calls=[_call("web_fetch", {}, "old")]), + tool_msg("old body " * 400, "old"), + assistant_msg("", tool_calls=[_call("web_fetch", {}, "new")]), + tool_msg(latest, "new"), + ] + compactor = TieredCompactor( + keep_tool_result=1, + summary_llm=None, + relief_target=600, + spill_store=store, + ) + + result = await compactor.compact(history, keep_recent=1) + + bodies = { + message.get("tool_call_id"): str(message.get("content") or "") + for message in result + if message.get("role") == "tool" + } + assert len(bodies["new"]) < len(latest) + assert any(message.get("spill_refs") for message in result) + + +def test_budget_consistency_shares_the_trigger_ratio() -> None: + from agent_core.runtime.loop import budget_consistency + from agent_core.runtime.loop.tiered_compact import DEFAULT_TRIGGER_RATIO + + assert budget_consistency.COMPACTION_TRIGGER_RATIO == DEFAULT_TRIGGER_RATIO + + +def test_budget_consistency_is_reachable_from_the_package_api() -> None: + from agent_core.runtime import loop + + assert loop.check_context_budget( + max_len=262_144, + max_input_tokens=229_376, + max_tokens=32_768, + label="test", + ) == [] diff --git a/tests/test_spill.py b/tests/test_spill.py new file mode 100644 index 0000000..6ecb66f --- /dev/null +++ b/tests/test_spill.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from agent_core.runtime.spill import ( + SpillStore, + budgeted_preview, + scope_component, + truncate_preview, +) + + +def test_scope_is_hashed_and_session_isolated(tmp_path: Path) -> None: + first = SpillStore(tmp_path, "task:session-a", visible_root="/spill") + second = SpillStore(tmp_path, "task:session-b", visible_root="/spill") + + assert first.component == scope_component("task:session-a") + assert first.directory != second.directory + assert first.visible_directory == f"/spill/{first.component}" + + +def test_write_is_content_addressed_idempotent_and_readable(tmp_path: Path) -> None: + store = SpillStore(tmp_path, "session", visible_root="/spill") + body = "important body\n" * 500 + + first = store.write("bash", body, require_visible=True) + second = store.write("bash", body, require_visible=True) + + assert first is not None and second is not None + assert first == second + physical, visible = first + assert physical.is_file() + assert visible.startswith(store.visible_directory + "/") + assert store.read(physical) == body + assert store.read(visible) == body + assert physical.stat().st_mode & 0o222 == 0 + + +def test_read_rejects_paths_outside_the_scope(tmp_path: Path) -> None: + store = SpillStore(tmp_path / "spill", "session", visible_root="/spill") + outside = tmp_path / "outside.md" + outside.write_text("secret", encoding="utf-8") + + assert store.read(outside) is None + assert store.read(f"{store.visible_directory}/../outside.md") is None + assert not store.contains_path(outside) + + +def test_read_created_supports_lifecycle_reads_without_scope(tmp_path: Path) -> None: + store = SpillStore(tmp_path / "spill", "session", visible_root="/spill") + written = store.write("bash", "captured") + outside = tmp_path / "outside.md" + outside.write_text("secret", encoding="utf-8") + + assert written is not None + assert SpillStore.read_created(written[0]) == "captured" + assert SpillStore.read_created(outside) is None + + +def test_compaction_spill_skips_tiny_bodies(tmp_path: Path) -> None: + store = SpillStore(tmp_path, "session", visible_root="/spill") + + assert store.spill_compacted_body("tool", "small") is None + assert not store.directory.exists() + + +def test_overflow_keeps_head_tail_and_a_recovery_pointer(tmp_path: Path) -> None: + store = SpillStore(tmp_path, "session", visible_root="/spill") + body = "HEAD\n" + "middle\n" * 2_000 + "TAIL\n" + + preview = store.overflow("bash", body, cap=2_000) + + assert len(preview) <= 2_000 + assert "HEAD" in preview + assert "TAIL" in preview + assert "chars elided" in preview + assert store.visible_directory in preview + saved = next(store.directory.glob("*.md")) + assert store.read(saved) == body + + +def test_head_mode_uses_budget_on_ranked_prefix() -> None: + text = "first\n" + "x" * 3_000 + "\nlast" + + result = truncate_preview(text, 1_000, mode="head") + + assert len(result) <= 1_000 + assert result.startswith("first") + assert "last" not in result + + +def test_footer_is_charged_to_the_same_cap() -> None: + result = budgeted_preview( + "x" * 10_000, + cap=2_000, + ref="/spill/scope/body.md", + ) + + assert len(result) <= 2_000 + assert "/spill/scope/body.md" in result + + +def test_aggregate_budget_spills_before_shortening(tmp_path: Path) -> None: + store = SpillStore(tmp_path, "session", visible_root="/spill") + results = ["a" * 7_000, "b" * 7_000, "small"] + + adjusted = store.enforce_aggregate_budget( + results, + tool_names=["first", "second", "third"], + cap=9_000, + ) + + assert len(adjusted) == len(results) + assert sum(map(len, adjusted)) <= 9_000 + assert adjusted[2] == "small" + assert list(store.directory.glob("*.md")) + + +def test_unmounted_store_persists_but_does_not_advertise_a_path(tmp_path: Path) -> None: + store = SpillStore(tmp_path, "session") + body = "x" * 5_000 + + preview = store.overflow("tool", body, cap=1_500) + + assert "not readable from this backend" in preview + assert store.spill_compacted_body("tool", body) is None + assert list(store.directory.glob("*.md")) + + +def test_cleanup_removes_only_its_own_session(tmp_path: Path) -> None: + first = SpillStore(tmp_path, "first", visible_root="/spill") + second = SpillStore(tmp_path, "second", visible_root="/spill") + assert first.write("tool", "a" * 2_000) + assert second.write("tool", "b" * 2_000) + + assert first.cleanup() == 1 + assert not first.directory.exists() + assert second.directory.exists() + assert second.cleanup() == 1 + + +def test_body_pointer_detection_is_root_bounded(tmp_path: Path) -> None: + store = SpillStore(tmp_path, "session", visible_root="/spill") + written = store.write("tool", "x" * 2_000) + assert written is not None + + assert store.body_names_file(f"recover at {written[1]}") + assert not store.body_names_file("/spillover/not-a-pointer.md") + + +@pytest.mark.skipif(not hasattr(os, "symlink"), reason="symlinks unavailable") +def test_read_does_not_follow_a_spill_file_symlink_outside_scope(tmp_path: Path) -> None: + store = SpillStore(tmp_path / "root", "session", visible_root="/spill") + store.directory.mkdir(parents=True) + outside = tmp_path / "outside" + outside.write_text("secret", encoding="utf-8") + link = store.directory / "escape.md" + link.symlink_to(outside) + + assert store.read(link) is None + + +def test_budgeted_preview_never_grows_the_body_it_bounds() -> None: + """A cap too small for preview floor plus footer must not inflate the result. + + ``_MIN_PREVIEW_CHARS`` overrides the cap, and the footer is appended on top, + so a 700-char body under ``cap=600`` used to come back at ~950 characters — + the function that exists to bound a result growing it, and pointing at a + recovery file for content that is entirely inline. + """ + body = "x" * 700 + + result = budgeted_preview(body, cap=600, ref="/spill/scope/body.md") + + assert result == body + + +def test_head_only_preview_does_not_promise_a_gap_marker() -> None: + body = "\n".join(f"line {index}" for index in range(2_000)) + + result = budgeted_preview(body, cap=3_000, ref="/spill/scope/body.md", mode="head") + + assert "chars elided" not in result + assert "the gap is marked above" not in result + assert "the beginning only" in result + + +def test_head_and_tail_preview_reports_its_marked_gap() -> None: + body = "\n".join(f"line {index}" for index in range(2_000)) + + result = budgeted_preview(body, cap=3_000, ref="/spill/scope/body.md") + + assert "chars elided" in result + assert "the gap is marked above" in result diff --git a/tests/test_summary_retry_budget.py b/tests/test_summary_retry_budget.py new file mode 100644 index 0000000..71e90ff --- /dev/null +++ b/tests/test_summary_retry_budget.py @@ -0,0 +1,290 @@ +"""A failed summary must not be retried blindly, nor be unbounded in time. + +`_generate_summary` used to be a single call with no retry: any failure returned +the fallback layout, which on a long research run is the whole point of the +compaction being thrown away. The naive repair — retry twice — is worse than it +looks, because one summariser call can occupy the full LLM timeout (600s in the +shipped profiles). So the two properties that matter are inseparable: retry only +what can succeed on a second attempt, and bound the sequence in wall time. +""" + +from __future__ import annotations + +import asyncio +import time + +import pytest + +from agent_core.runtime.loop import compact_llm +from agent_core.runtime.loop.compact_llm import ( + LLMSummaryCompactor, + is_transient_summary_error, +) +from agent_core.runtime.loop.tiered_compact import TieredCompactor + + +@pytest.fixture(autouse=True) +def _fast_backoff(monkeypatch: pytest.MonkeyPatch) -> None: + """Shorten the real backoff. Every case here asserts on attempt counts or on + a ceiling being respected; none of them is about how long 2s is.""" + monkeypatch.setattr(compact_llm, "_RETRY_BACKOFF_BASE_S", 0.01) + monkeypatch.setattr(compact_llm, "_RETRY_BACKOFF_CAP_S", 0.02) + + +class _Resp: + def __init__(self, content: str) -> None: + self.content = content + + +class _LLM: + """Fails `fail_times` times with `exc`, then succeeds.""" + + def __init__(self, exc: Exception, fail_times: int = 99) -> None: + self._exc = exc + self._fail_times = fail_times + self.calls = 0 + + async def chat(self, _messages: object, **_kw: object) -> _Resp: + self.calls += 1 + if self.calls <= self._fail_times: + raise self._exc + return _Resp("a real summary of the earlier turns") + + +def _history(n: int = 12) -> list[dict]: + msgs: list[dict] = [{"role": "system", "content": "sys"}] + for i in range(n): + msgs.append({"role": "user", "content": f"question {i} " + "x" * 400}) + msgs.append({"role": "assistant", "content": f"answer {i} " + "y" * 400}) + return msgs + + +# ── classification ─────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("message", [ + "Error code: 400 - context length exceeded", + "maximum context length is 32768 tokens", + "invalid_request_error: input too large", + "413 Payload Too Large", + "401 unauthorized", +]) +def test_a_deterministic_failure_is_not_worth_retrying(message: str) -> None: + """These fail again identically, and the second failure costs the same + minutes as the first.""" + assert is_transient_summary_error(RuntimeError(message)) is False + + +@pytest.mark.parametrize("message", [ + "429 Too Many Requests", + "503 Service Unavailable", + "upstream connection reset", + "the request timed out", + "model is overloaded", +]) +def test_a_transient_failure_is_worth_one_more_attempt(message: str) -> None: + assert is_transient_summary_error(RuntimeError(message)) is True + + +@pytest.mark.parametrize("message", [ + "Error code: 400 - code=bad_response_status_code, upstream 503", + "Error code: 400 - type=new_api_error, upstream timeout", +]) +def test_a_proxy_wrapped_transient_400_is_retried(message: str) -> None: + """A gateway's outer 400 must not hide the recoverable upstream error.""" + assert is_transient_summary_error(RuntimeError(message)) is True + + +def test_a_proxy_wrapped_transient_in_a_structured_body_is_retried() -> None: + class _ProxyError(RuntimeError): + def __init__(self) -> None: + super().__init__("Bad request") + self.status_code = 400 + self.body = { + "error": { + "code": "bad_response_status_code", + "message": "upstream 503", + }, + } + + assert is_transient_summary_error(_ProxyError()) is True + + +def test_an_unrecognised_failure_is_treated_as_transient() -> None: + """The bias is deliberate: one extra attempt is cheaper than losing the + research history. `retry_total_timeout_s` is what keeps the bias bounded.""" + assert is_transient_summary_error(RuntimeError("something new")) is True + + +def test_cancellation_is_never_retried() -> None: + """Retrying a cancellation fights the caller that asked us to stop.""" + assert is_transient_summary_error(asyncio.CancelledError()) is False + + +# ── retry behaviour ────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_a_transient_failure_is_retried_and_can_succeed() -> None: + llm = _LLM(RuntimeError("503 Service Unavailable"), fail_times=1) + compactor = LLMSummaryCompactor( + summary_llm=llm, max_transient_retries=2, retry_total_timeout_s=30, + ) + + out = await compactor.compact(_history(), keep_recent=4) + + assert llm.calls == 2 + assert any("a real summary" in str(m.get("content")) for m in out), ( + "the retry succeeded, so its summary must be what lands in the history" + ) + + +@pytest.mark.asyncio +async def test_a_deterministic_failure_costs_exactly_one_attempt() -> None: + """The saving this classification exists for: two extra calls against an + endpoint that cannot answer, each able to run for the full LLM timeout.""" + llm = _LLM(RuntimeError("400 context length exceeded")) + events: list[dict] = [] + compactor = LLMSummaryCompactor( + summary_llm=llm, + max_transient_retries=2, + retry_total_timeout_s=30, + emit_event=events.append, + ) + + await compactor.compact(_history(), keep_recent=4) + + assert llm.calls == 1 + assert events[-1]["rollback_reason"] == "llm_error_permanent", ( + "a permanent failure must be distinguishable in the record from having " + "run out of attempts — only the former means the target is unreachable" + ) + assert events[-1]["attempts"] == 1 + + +@pytest.mark.asyncio +async def test_a_transient_failure_stops_at_the_retry_count() -> None: + llm = _LLM(RuntimeError("429 Too Many Requests")) + events: list[dict] = [] + compactor = LLMSummaryCompactor( + summary_llm=llm, + max_transient_retries=2, + retry_total_timeout_s=30, + emit_event=events.append, + ) + + await compactor.compact(_history(), keep_recent=4) + + assert llm.calls == 3, "one initial attempt plus two retries" + assert events[-1]["rollback_reason"] == "llm_error" + assert events[-1]["attempts"] == 3 + + +# ── the time ceiling is a ceiling ──────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_the_ceiling_bounds_a_single_hanging_attempt() -> None: + """The property a between-attempts check cannot give you. + + A summariser call carries its own transport timeout, which is far larger + than the retry budget. Checking the deadline only *after* an attempt returns + means one hanging call blows through the ceiling entirely, so each attempt + is wrapped in `asyncio.wait_for`. + """ + + class _Hangs: + calls = 0 + + async def chat(self, _messages: object, **_kw: object) -> _Resp: + _Hangs.calls += 1 + await asyncio.sleep(60) + return _Resp("never reached") + + compactor = LLMSummaryCompactor( + summary_llm=_Hangs(), max_transient_retries=2, retry_total_timeout_s=0.2, + ) + + started = time.monotonic() + out = await compactor.compact(_history(), keep_recent=4) + elapsed = time.monotonic() - started + + assert elapsed < 5, f"ceiling was 0.2s, sequence took {elapsed:.1f}s" + assert out, "a blown ceiling must still return a usable history" + + +@pytest.mark.asyncio +async def test_a_retry_whose_backoff_alone_exceeds_the_budget_is_refused( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Sleeping out the remaining budget and *then* returning the same failure + is strictly worse than returning now — the ceiling is meant to be strict, + not approximate.""" + monkeypatch.setattr(compact_llm, "_RETRY_BACKOFF_BASE_S", 2.0) + monkeypatch.setattr(compact_llm, "_RETRY_BACKOFF_CAP_S", 5.0) + llm = _LLM(RuntimeError("503 Service Unavailable")) + compactor = LLMSummaryCompactor( + summary_llm=llm, max_transient_retries=5, retry_total_timeout_s=0.05, + ) + + started = time.monotonic() + await compactor.compact(_history(), keep_recent=4) + elapsed = time.monotonic() - started + + assert elapsed < 2, f"backoff is 2s+, so it must not have been slept: {elapsed:.1f}s" + assert llm.calls < 6 + + +# ── the default must not change for existing callers ───────────────────────── + + +@pytest.mark.asyncio +async def test_a_bare_caller_still_gets_exactly_one_attempt() -> None: + """`apodex.session`, `apodex.task_runner` and `session_history` construct + this class without a time ceiling. A retry default of anything but zero + would hand them retries that nothing bounds in time. + """ + llm = _LLM(RuntimeError("503 Service Unavailable")) + compactor = LLMSummaryCompactor(summary_llm=llm) + + await compactor.compact(_history(), keep_recent=4) + + assert llm.calls == 1 + + +@pytest.mark.asyncio +async def test_tiered_without_a_ceiling_does_not_retry_either() -> None: + """The two knobs are not independent. A retry count with no time bound is + the same defect one layer up, so the half-specified pair resolves to the + conservative reading rather than to an unbounded one. + """ + llm = _LLM(RuntimeError("503 Service Unavailable")) + compactor = TieredCompactor( + keep_tool_result=2, + summary_llm=llm, + relief_target=1, + summary_retries=2, + summary_retry_timeout_s=None, + ) + + await compactor.compact(_history(40), keep_recent=4) + + assert llm.calls == 1 + + +@pytest.mark.asyncio +async def test_tiered_with_a_ceiling_does_retry() -> None: + llm = _LLM(RuntimeError("503 Service Unavailable"), fail_times=1) + compactor = TieredCompactor( + keep_tool_result=2, + summary_llm=llm, + relief_target=1, + summary_retries=2, + summary_retry_timeout_s=30, + ) + + await compactor.compact(_history(40), keep_recent=4) + + assert llm.calls == 2 + assert compactor.last_event is not None + assert compactor.last_event.attempts == 2 diff --git a/tests/test_tiered_compact.py b/tests/test_tiered_compact.py new file mode 100644 index 0000000..11ab4ce --- /dev/null +++ b/tests/test_tiered_compact.py @@ -0,0 +1,176 @@ +"""Tiered context compaction: protect-set, real-token gauge/policy, escalation.""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from agent_core.llm import LLMResponse +from agent_core.loop_types import LoopConfig, TurnContext +from agent_core.messages import system_msg, tool_msg +from agent_core.runtime.loop.agent_loop import run_agent_loop +from agent_core.runtime.loop.compact import ( + OMITTED_TOOL_RESULT_PLACEHOLDER, + KeepLastNToolResultsCompactor, + estimate_tokens, +) +from agent_core.runtime.loop.tiered_compact import ( + InputTokenGauge, + InputTokenThresholdPolicy, + TieredCompactor, +) + + +def _ai(tid: str, name: str) -> dict: + return {"role": "assistant", "content": "", + "tool_calls": [{"id": tid, "function": {"name": name, "arguments": "{}"}}]} + + +def _msgs() -> list[dict]: + m = [system_msg("S"), _ai("c1", "collect_reports"), tool_msg("REPORT-" + "x" * 500, "c1")] + for i in range(4): + m += [_ai(f"w{i}", "web_search"), tool_msg("SR-" + "y" * 500, f"w{i}")] + return m + + +def test_keeplastn_protects_named_tools(): + # keep last 1 tool result, but never blank collect_reports regardless of age. + k = KeepLastNToolResultsCompactor( + keep_tool_result=1, protect_tool_names=frozenset({"collect_reports"})) + out = k.compact(_msgs(), 0) + bodies = [m["content"] for m in out if m.get("role") == "tool"] + assert any(b.startswith("REPORT-") for b in bodies) # protected, kept full + blanked = sum(1 for b in bodies if b.startswith(OMITTED_TOOL_RESULT_PLACEHOLDER)) + assert blanked == 3 # 3 of 4 web results blanked + + +def test_keeplastn_default_has_no_protect(): + # Backward compatible: no protect set → old collect_reports IS blanked by age. + k = KeepLastNToolResultsCompactor(keep_tool_result=1) + out = k.compact(_msgs(), 0) + bodies = [m["content"] for m in out if m.get("role") == "tool"] + assert not any(b.startswith("REPORT-") for b in bodies) + + +def test_input_token_gauge_and_policy(): + g = InputTokenGauge() + ctx = TurnContext(turn=1, max_turns=9, task_id="t", role_id="r", ai_text="", + thinking="", tool_calls=[], messages=[], + usage={"prompt_tokens": 210_000}, metadata={}) + asyncio.run(g.on_llm_response(ctx)) + pol = InputTokenThresholdPolicy(g, int(262_144 * 0.8)) + assert pol.should_compact(1, [], 0) is True # 210k > 209k + g.tokens = 100_000 + assert pol.should_compact(1, [], 0) is False + + +class _FakeLLM: + def __init__(self) -> None: + self.calls = 0 + + async def chat(self, messages, **kw): + self.calls += 1 + class _R: + content = "SUMMARY" + return _R() + + +def test_tiered_escalates_only_when_tier1_insufficient(): + async def run(): + # relief huge → Tier1 alone suffices, Tier2 (LLM summary) must NOT fire. + t1 = TieredCompactor(keep_tool_result=1, summary_llm=_FakeLLM(), relief_target=10**9) + o1 = await t1.compact(_msgs(), 6) + assert not any("SUMMARY" in (m.get("content") or "") for m in o1) + # relief 0 → Tier1 never enough → escalate to Tier2 summary. + llm = _FakeLLM() + t2 = TieredCompactor(keep_tool_result=1, summary_llm=llm, relief_target=0) + await t2.compact(_msgs(), 6) + # Escalation means Tier 2 is evaluated. The compactor still returns the + # smallest candidate, which may be deterministic compression. + assert llm.calls == 1 + + asyncio.run(run()) + + +def test_tiered_calibration_escalates_via_real_token_scale(): + # The relief check is in REAL tokens: with a gauge reporting real >> estimate, + # a Tier1 result that the RAW estimate would clear still escalates to Tier2. + async def run(): + tier1_out = KeepLastNToolResultsCompactor(keep_tool_result=1).compact(_msgs(), 6) + relief = estimate_tokens(tier1_out) + 10 # raw est(out) clears this + # No gauge → raw estimate path → Tier1 alone, NO summary. + t_raw = TieredCompactor(keep_tool_result=1, summary_llm=_FakeLLM(), relief_target=relief) + o_raw = await t_raw.compact(_msgs(), 6) + assert not any("SUMMARY" in (m.get("content") or "") for m in o_raw) + # Gauge saw real = 50× the estimate → scale pushes over relief → Tier2 fires. + g = InputTokenGauge() + g.estimate = estimate_tokens(_msgs()) + g.tokens = g.estimate * 50 + llm = _FakeLLM() + t_cal = TieredCompactor( + keep_tool_result=1, summary_llm=llm, relief_target=relief, gauge=g) + await t_cal.compact(_msgs(), 6) + assert llm.calls == 1 + + asyncio.run(run()) + + +def _tool(name: str): + t = MagicMock() + t.name = name + t.ainvoke = AsyncMock(return_value="ok") + return t + + +def _resp(content: str, tool_calls: list[dict] | None, prompt_tokens: int) -> LLMResponse: + wire = [{"id": tc["id"], "type": "function", + "function": {"name": tc["name"], "arguments": "{}"}} + for tc in (tool_calls or [])] + return LLMResponse(content=content, tool_calls=wire, + usage={"prompt_tokens": prompt_tokens, "completion_tokens": 5}) + + +class _SeqUsageLLM: + """Fake LLM returning queued responses (last repeats) carrying usage.""" + model = "fake-usage" + + def __init__(self, responses: list[LLMResponse]) -> None: + self._r = responses + self._i = 0 + + async def chat(self, messages, **kw): + r = self._r[min(self._i, len(self._r) - 1)] + self._i += 1 + return r + + +class _SpyCompactor: + def __init__(self) -> None: + self.calls = 0 + + async def compact(self, messages, keep_recent): + self.calls += 1 + return messages + + +@pytest.mark.asyncio +async def test_gauge_feeds_policy_inside_run_agent_loop(): + # End-to-end wiring contract: the critical gauge updates in on_llm_response + # (loop step 6) BEFORE the policy reads it in the compaction step (step 10), + # so a real-token trigger actually drives compaction through the live loop. + gauge = InputTokenGauge() + spy = _SpyCompactor() + llm = _SeqUsageLLM([ + _resp("use tool", [{"id": "c1", "name": "web_search"}], prompt_tokens=5000), + _resp("done", None, prompt_tokens=5000), + ]) + cfg = LoopConfig(max_turns=5, no_tool_max_retries=1, + compactor=spy, compaction_policy=InputTokenThresholdPolicy(gauge, 1000)) + await run_agent_loop(system_prompt="s", user_message="q", llm=llm, + tools=[_tool("web_search")], config=cfg, observers=[gauge]) + assert gauge.tokens == 5000 # gauge captured the real prompt_tokens + # Only the threshold policy can fire → compaction running proves the gauge + # was fed (step 6) before the policy read it (step 10) within the live loop. + assert spy.calls >= 1