diff --git a/pydantic_ai_harness/experimental/overflow/README.md b/pydantic_ai_harness/experimental/overflow/README.md new file mode 100644 index 0000000..b00ed15 --- /dev/null +++ b/pydantic_ai_harness/experimental/overflow/README.md @@ -0,0 +1,201 @@ +# Overflow capability + +> [!WARNING] +> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may +> change or be removed in any release, without a deprecation period. Import it from the +> experimental path -- there is no top-level export: +> +> ```python +> from pydantic_ai_harness.experimental.overflow import OverflowingToolOutput +> ``` +> +> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all** +> harness experimental warnings with a single filter (no per-capability lines needed): +> +> ```python +> import warnings +> from pydantic_ai_harness.experimental import HarnessExperimentalWarning +> +> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning) +> ``` + +A tool can return a payload large enough to dominate the context window. Tool returns +persist in history as `ToolReturnPart`s, so an oversized one is re-sent on every later +model request -- paying its token cost for the rest of the run. `OverflowingToolOutput` +intercepts a return in the `after_tool_execute` hook, reduces it once, and lets the reduced +form persist. The reduction is not recomputed per request. + +This is the overflow-to-file follow-up the `compaction` README names as out of scope: it +moves large tool outputs *out* of the window at production time, rather than compressing or +dropping context already inside it. + +## The three modes + +| Mode | Cost | Lossy? | What the model gets | +|---|---|---|---| +| `Truncate` | zero-LLM | yes | A head / tail / head+tail clamp of the text | +| `Spill` | zero-LLM | no | A handle + preview + shape sketch; full payload read back on demand | +| `Summarize` | one LLM call | yes | A size-gated summary (inherits the run's model by default) | + +`Spill` is lossless: the full payload is persisted and the model reads slices of it through +the registered `read_tool_result(handle, offset, limit, from_end, pattern)` tool (the Claude +Code pattern, the core [#4352](https://github.com/pydantic/pydantic-ai/issues/4352) design). +That tool is bounded: `offset >= 0`, `limit` clamped to a built-in line cap, the joined output +capped, and `pattern` is a literal substring (not a regex), so a model-supplied value cannot +hang the host with catastrophic backtracking. + +### Both `return_value` and `content` are reduced + +A `ToolReturn` carries a `return_value` and an optional `content` that core renders as a +separate, model-visible part which also persists in history. This capability measures and +reduces both with the same band logic (they spill to distinct handles). Text `content` is +reduced in place; non-text `content` (multimodal parts) that overflows is left unreduced with +a `warnings.warn`, since it cannot be safely truncated. + +## Bands: combine the modes + +Configure an ordered list of size `bands`. Each band is a `(over, action)` pair: when a +return's measured size reaches `over`, its action runs. The band with the largest threshold +that fits wins; anything below the smallest threshold passes through. + +```python +from pydantic_ai import Agent +from pydantic_ai_harness.experimental.overflow import ( + Band, + OverflowingToolOutput, + Spill, + Summarize, + Truncate, +) + +agent = Agent( + 'openai:gpt-4o', + capabilities=[ + OverflowingToolOutput( + bands=[ + Band(over=100_000, action=Spill()), # huge: keep losslessly, read back on demand + Band(over=20_000, action=Summarize()), # large: compress with the run's model + Band(over=5_000, action=Truncate()), # medium: cheap clamp + ], + # below 5,000: passthrough + ) + ], +) +``` + +The default band, when you pass no `bands`, is `Spill(then=Truncate())`: lossless when a +store accepts the write, a bounded truncation otherwise -- zero LLM cost and no silent drop. + +### Fallbacks with `then` + +Every action takes an optional `then`, applied when the action cannot run: a `Spill` whose +store errors, a `Truncate` / `Summarize` on a binary payload, a `Summarize` whose model call +raises. `then` chains, so `Summarize(then=Spill(then=Truncate()))` degrades summarize -> +spill -> truncate. + +### Per-tool overrides and filtering + +`per_tool` replaces the global band list for named tools (file reads to `head`, logs to +`tail`); `tool_filter` (a `ToolSelector`) scopes which tools the capability touches at all. + +```python +OverflowingToolOutput( + per_tool={ + 'read_file': [Band(over=8_000, action=Truncate(strategy=TruncationStrategy.head))], + 'run_shell': [Band(over=8_000, action=Truncate(strategy=TruncationStrategy.tail))], + }, + tool_filter=['read_file', 'run_shell', 'search'], +) +``` + +## Size unit + +Thresholds are measured in characters by default. Set `over_tokens=True` to measure in +estimated tokens (the same ~4-chars-per-token heuristic as `compaction`); pass a `tokenizer` +callable for accuracy. `Truncate.max_chars` is always characters -- truncation is a +character operation regardless of the threshold unit. + +## Spill store + +Spilled payloads go through the narrow `OverflowStore` protocol. The default `LocalFileStore` +writes one file per `(run_id, tool_call_id, retry)` under a stable root directory and keeps it +after the run, so a later `read_tool_result` -- in this run or a subsequent agent/run -- can +still reach it. The handle is backend-addressable (a relative key), not an absolute local +path, so a durable backend (Temporal, a blob store, or the core `ExecutionEnvironment` +workspace once #4352 lands) can resolve the same handle in another process. Supply your own +backend with `store=...`. + +```python +class OverflowStore(Protocol): + async def write(self, key: str, data: bytes) -> str: ... # returns a handle + async def read(self, handle: str) -> bytes: ... +``` + +### Security model (shared root, not isolation) + +The store root is stable and shareable on purpose -- spilled files must be readable by a later +agent or run -- so security does not come from per-instance isolation. It comes from two +mechanisms: the root is created with `0700` (owner-only) permissions, and `read` resolves the +target (following symlinks) and rejects anything that escapes the root via symlink, `..`, or +an absolute path. Handle segments are also sanitized so a crafted handle cannot traverse out. + +### Cleanup: keep-forever by default, opt-in TTL pruning + +By default the store keeps spilled files forever -- deleting on run end would break a later +agent that still wants to read a spill. To bound disk use, opt into age-based pruning: + +```python +from datetime import timedelta + +store = LocalFileStore(cleanup_after=timedelta(hours=6)) # default: None = keep forever +``` + +When set, a `write` schedules a background prune (a daemon thread, off the hot path) that +deletes files whose modification time (`st_mtime`) is older than `cleanup_after`. Pruning is +non-blocking and non-erroring: any failure is caught and surfaced via `warnings.warn`, never +propagated into the agent run, so cleanup can never fail a run or block the hot path. +Last-read time (`st_atime`) is unreliable on `noatime`/`relatime` mounts and is not used. + +Prefer external cleanup (cron, a sweeper) over the in-process TTL? Point it at the store root +and delete by mtime: + +```python +import time +from pathlib import Path + +root = Path('/tmp/pyai_harness_overflow') # or your configured base_dir +cutoff = time.time() - 6 * 3600 +for path in root.rglob('*'): + if path.is_file() and path.stat().st_mtime < cutoff: + path.unlink(missing_ok=True) +``` + +## Usage accounting + +A `Summarize` call is a real request to the model, so its full usage -- tokens and the +request itself -- folds into the run's `ctx.usage`, exactly like `SummarizingCompaction`. No +token caps are imposed on the summary call. A `UsageLimits` request limit will see it. + +## Edge cases + +- Binary returns spill verbatim and are never stringify-truncated; `Truncate` / `Summarize` + on binary fall through to `then`. +- Structured / nested returns spill (or summarize) by preference -- truncating JSON produces + invalid JSON. `Spill` includes a one-line shape sketch of the top level. +- `ModelRetry` and tool errors never reach this hook (they are raised, not returned), so the + model always gets the full error it needs to recover. +- A large `ToolReturn.content` is reduced with the same bands as `return_value`; non-text + content that overflows is left unreduced with a warning. +- Multiple oversized returns in one step get distinct handles (keyed per `tool_call_id`); + retries get distinct handles too (keyed per `retry`), so a retried call never clobbers the + earlier attempt's spill. + +## Relationship to other capabilities + +- Supersedes the spill scope of PR #185 `ToolOutputManagement` (one-way truncate / spill with + no read-back); this capability's truncation and ANSI / binary handling are harvested from it. +- Consumes core [#4352](https://github.com/pydantic/pydantic-ai/issues/4352) (the canonical + queryable-file primitive) through the `OverflowStore` seam once it lands. +- Distinct from `compaction`, which compresses or drops context already inside the window, and + from `ClampOversizedMessages` (PR #286), which clamps runaway model responses, not tool + returns. diff --git a/pydantic_ai_harness/experimental/overflow/__init__.py b/pydantic_ai_harness/experimental/overflow/__init__.py new file mode 100644 index 0000000..fbdaf6b --- /dev/null +++ b/pydantic_ai_harness/experimental/overflow/__init__.py @@ -0,0 +1,42 @@ +"""Overflow capability: reduce oversized tool returns at production time. + +`OverflowingToolOutput` intercepts a tool return when it is produced and reduces it -- +truncating, spilling to a queryable file, or summarizing -- so an oversized payload does +not persist in history and get re-sent on every later model request. Combine the three +modes through an ordered list of size `bands`. + +Spilled payloads are read back on demand through the registered `read_tool_result` tool; +the `OverflowStore` protocol is the seam for a durable backend (the local-file default +ships for single-process runs). +""" + +from pydantic_ai_harness.experimental._warn import warn_experimental +from pydantic_ai_harness.experimental.overflow._bands import ( + Action, + Band, + Passthrough, + Spill, + Summarize, + SummarizeFunc, + Truncate, +) +from pydantic_ai_harness.experimental.overflow._capability import READ_TOOL_NAME, OverflowingToolOutput +from pydantic_ai_harness.experimental.overflow._payload import TruncationStrategy +from pydantic_ai_harness.experimental.overflow._store import LocalFileStore, OverflowStore + +warn_experimental('overflow') + +__all__ = [ + 'READ_TOOL_NAME', + 'Action', + 'Band', + 'LocalFileStore', + 'OverflowStore', + 'OverflowingToolOutput', + 'Passthrough', + 'Spill', + 'Summarize', + 'SummarizeFunc', + 'Truncate', + 'TruncationStrategy', +] diff --git a/pydantic_ai_harness/experimental/overflow/_bands.py b/pydantic_ai_harness/experimental/overflow/_bands.py new file mode 100644 index 0000000..1f820b1 --- /dev/null +++ b/pydantic_ai_harness/experimental/overflow/_bands.py @@ -0,0 +1,85 @@ +"""Size bands and the actions they trigger. + +A band is a `(over, action)` pair: when a tool return's measured size is at least `over`, +its `action` runs. `OverflowingToolOutput` holds an ordered band list and picks the first +match (largest threshold that fits), passing through anything below the smallest threshold. + +Every action carries an optional `then` fallback, applied when the action cannot run -- +`Spill` whose store errors, `Truncate`/`Summarize` on a binary payload, a `Summarize` +whose model call raises. `Spill(then=Truncate())` is the default: lossless when the store +works, a bounded truncation otherwise, never a silent drop. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from pydantic_ai_harness.experimental.overflow._payload import TruncationStrategy + +if TYPE_CHECKING: + from pydantic_ai.models import Model + +_DEFAULT_TRUNCATE_CHARS = 4_000 +_DEFAULT_PREVIEW_CHARS = 1_000 + +# A summarizer callable: `(tool_name, full_text) -> summary`, sync or async. +SummarizeFunc = Callable[[str, str], str | Awaitable[str]] + + +@dataclass(frozen=True) +class Passthrough: + """Leave the tool return untouched. Useful as an explicit no-op band.""" + + +@dataclass(frozen=True) +class Truncate: + """Clamp the stringified return to `max_chars`. Lossy, zero-cost, no read-back. + + `max_chars` is always characters, independent of the capability's `over_tokens` size + unit (truncation is a character operation). Falls back to `then` for binary payloads, + which cannot be stringify-truncated. + """ + + strategy: TruncationStrategy = TruncationStrategy.head_tail + max_chars: int = _DEFAULT_TRUNCATE_CHARS + then: Action | None = None + + +@dataclass(frozen=True) +class Spill: + """Persist the full return and replace it with a handle, preview, and shape sketch. + + Lossless: the model gets a `read_tool_result` handle to slice / grep / tail the full + payload on demand. Falls back to `then` when no store accepts the write. + """ + + preview_chars: int = _DEFAULT_PREVIEW_CHARS + then: Action | None = None + + +@dataclass(frozen=True) +class Summarize: + """Replace the return with a size-gated LLM summary. The expensive band. + + `model=None` inherits the running agent's model (`ctx.model`), mirroring + `SummarizingCompaction`. Pass a model id / instance to override, or a `summarize` + callable to bypass the built-in prompt entirely. Summary usage folds into `ctx.usage`; + no token caps are imposed. Falls back to `then` on a binary payload or a failed call. + """ + + model: str | Model | None = None + summarize: SummarizeFunc | None = None + then: Action | None = None + + +Action = Passthrough | Truncate | Spill | Summarize + + +@dataclass(frozen=True) +class Band: + """Trigger `action` once a return's measured size reaches `over` (chars or tokens).""" + + over: int + action: Action diff --git a/pydantic_ai_harness/experimental/overflow/_capability.py b/pydantic_ai_harness/experimental/overflow/_capability.py new file mode 100644 index 0000000..126ffee --- /dev/null +++ b/pydantic_ai_harness/experimental/overflow/_capability.py @@ -0,0 +1,552 @@ +"""`OverflowingToolOutput` -- reduce oversized tool returns at production time.""" + +from __future__ import annotations + +import warnings +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, TypeGuard + +from pydantic_ai import FunctionToolset +from pydantic_ai.capabilities import AbstractCapability +from pydantic_ai.exceptions import ModelRetry +from pydantic_ai.messages import ToolCallPart, ToolReturn, ToolReturnContent, UserContent +from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition, ToolSelector, matches_tool_selector +from pydantic_ai.toolsets import AgentToolset + +from pydantic_ai_harness.experimental.overflow._bands import ( + Action, + Band, + Passthrough, + Spill, + Summarize, + Truncate, +) +from pydantic_ai_harness.experimental.overflow._payload import ( + is_binary, + json_sketch, + measure, + strip_ansi, + to_bytes, + to_text, + truncate_text, +) +from pydantic_ai_harness.experimental.overflow._store import LocalFileStore, OverflowStore + +READ_TOOL_NAME = 'read_tool_result' +"""Name of the registered read-back tool. Its own returns are exempt from reduction.""" + +_DEFAULT_THRESHOLD = 10_000 +"""Default band threshold (characters) -- below this, returns pass through untouched.""" + +_DEFAULT_SUMMARY_PROMPT = """\ +The following output from the `{tool_name}` tool is too large to keep in full. Summarize it \ +so the summary carries everything needed to keep working: concrete values, identifiers, \ +errors, and structure. Respond ONLY with the summary, no preamble. + + +{output} +\ +""" + + +def _default_bands() -> list[Band]: + """Lossless spill with a bounded truncation fallback: zero LLM cost, no silent drop.""" + return [Band(over=_DEFAULT_THRESHOLD, action=Spill(then=Truncate()))] + + +@dataclass +class _Unit: + """One reducible piece of a tool return: its `return_value` or its `content`. + + `suffix` distinguishes the two so they spill to distinct handles for the same call. + """ + + binary: bool + text: str | None + data: bytes + value: ToolReturnContent + suffix: str + + +@dataclass +class OverflowingToolOutput(AbstractCapability[AgentDepsT]): + """Reduce oversized tool returns when they are produced, persisting the reduction. + + A tool can return a payload large enough to dominate the context window. Tool returns + persist in history, so an oversized one is re-sent on every later request. This + capability intercepts a return in `after_tool_execute`, reduces it once, and lets the + reduced form persist -- it is not recomputed per request. + + Three reduction modes, freely combined through an ordered list of size `bands`: + + - `Truncate`: clamp to a character budget. Lossy, zero-cost. + - `Spill`: persist the full payload, hand the model a `read_tool_result` handle plus a + preview. Lossless. + - `Summarize`: size-gated LLM summary. Inherits the run's model by default. + + The first band whose `over` threshold the measured size meets wins; smaller returns pass + through. `per_tool` replaces the band list for named tools; `tool_filter` scopes which + tools are touched at all. The default is `Spill(then=Truncate())`: lossless when a store + accepts the write, a bounded truncation otherwise. + + `ModelRetry` and other errors never reach this hook (they are raised, not returned), so + error payloads the model needs to recover are never spilled or summarized. + + Example: + ```python + from pydantic_ai import Agent + from pydantic_ai_harness.experimental.overflow import ( + Band, + OverflowingToolOutput, + Spill, + Summarize, + Truncate, + ) + + agent = Agent( + 'openai:gpt-4o', + capabilities=[ + OverflowingToolOutput( + bands=[ + Band(over=100_000, action=Spill()), + Band(over=20_000, action=Summarize()), + Band(over=5_000, action=Truncate()), + ], + ) + ], + ) + ``` + """ + + bands: Sequence[Band] = field(default_factory=_default_bands) + """Ordered size bands. The first band whose `over` threshold is met wins.""" + + per_tool: Mapping[str, Sequence[Band]] = field(default_factory=dict[str, Sequence[Band]]) + """Per-tool band lists that replace `bands` for the named tools.""" + + tool_filter: ToolSelector[AgentDepsT] = 'all' + """Which tools this capability touches. Non-matching tools always pass through.""" + + over_tokens: bool = False + """Measure band thresholds in estimated tokens instead of characters.""" + + tokenizer: Callable[[str], int] | None = None + """Optional `(str) -> int` tokenizer for `over_tokens`. Defaults to a ~4-char heuristic.""" + + store: OverflowStore | None = None + """Backend for spilled payloads. Defaults to a `LocalFileStore`.""" + + strip_ansi: bool = False + """Strip ANSI escape sequences from text returns before measuring and reducing.""" + + summary_prompt: str = _DEFAULT_SUMMARY_PROMPT + """Prompt template for `Summarize`. Must contain `{tool_name}` and `{output}`.""" + + _store: OverflowStore = field(init=False, repr=False) + _bands: list[Band] = field(init=False, repr=False) + _per_tool: dict[str, list[Band]] = field(init=False, repr=False) + + def __post_init__(self) -> None: + self._store = self.store if self.store is not None else LocalFileStore() + self._bands = self._prepare_bands(self.bands) + self._per_tool = {name: self._prepare_bands(bands) for name, bands in self.per_tool.items()} + + @staticmethod + def _prepare_bands(bands: Sequence[Band]) -> list[Band]: + """Validate thresholds and order bands largest-first so first-match means largest-fit.""" + for band in bands: + if band.over < 0: + raise ValueError('Band.over must be non-negative.') + return sorted(bands, key=lambda b: b.over, reverse=True) + + # --- toolset --- + + def get_toolset(self) -> AgentToolset[AgentDepsT] | None: + """Register the `read_tool_result` tool for reading spilled payloads on demand.""" + store = self._store + + async def read_tool_result( + ctx: RunContext[AgentDepsT], + handle: str, + offset: int = 0, + limit: int = 200, + from_end: bool = False, + pattern: str | None = None, + ) -> str: + """Read a slice of a spilled tool result. + + Args: + ctx: The run context (supplied by the agent). + handle: The handle from the overflowed tool return. + offset: Number of matching lines to skip from the start (or end). Must be >= 0. + limit: Maximum number of lines to return (>= 1; clamped to a built-in cap). + from_end: Count `offset`/`limit` from the end of the result. + pattern: Optional literal substring; only lines containing it are returned. + """ + return await _read_slice(store, handle, offset, limit, from_end, pattern) + + return FunctionToolset([read_tool_result]) + + # --- reduction --- + + async def after_tool_execute( + self, + ctx: RunContext[AgentDepsT], + *, + call: ToolCallPart, + tool_def: ToolDefinition, + args: dict[str, Any], + result: Any, + ) -> Any: + """Reduce the tool result -- both `return_value` and model-visible `content`.""" + original: object = result + if call.tool_name == READ_TOOL_NAME: + return original + if not await matches_tool_selector(self.tool_filter, ctx, tool_def): + return original + + metadata: object + if isinstance(result, ToolReturn): + return_value: ToolReturnContent = result.return_value + content = result.content + metadata = result.metadata + wrapped = True + else: + return_value = result + content = None + metadata = None + wrapped = False + + if isinstance(return_value, BaseException): + return original + + bands = self._per_tool.get(call.tool_name, self._bands) + value_unit = self._make_unit(return_value, suffix='') + value_text, value_handle = await self._reduce(ctx, call, bands, value_unit) + content_text, content_handle = await self._reduce_content(ctx, call, bands, content) + + if value_text is None and content_text is None: + return original + + return self._assemble( + wrapped=wrapped, + return_value=return_value, + content=content, + metadata=metadata, + value_unit=value_unit, + value_text=value_text, + value_handle=value_handle, + content_text=content_text, + content_handle=content_handle, + ) + + def _assemble( + self, + *, + wrapped: bool, + return_value: ToolReturnContent, + content: str | Sequence[UserContent] | None, + metadata: object, + value_unit: _Unit, + value_text: str | None, + value_handle: str | None, + content_text: str | None, + content_handle: str | None, + ) -> object: + """Rebuild the tool result from the reduced parts, preserving the envelope.""" + if wrapped: + new_metadata = metadata + if value_handle is not None or content_handle is not None: + new_metadata = _with_handles(metadata, value_handle, len(value_unit.data), content_handle) + wrapped_out: ToolReturn[object] = ToolReturn( + return_value=value_text if value_text is not None else return_value, + content=content_text if content_text is not None else content, + metadata=new_metadata, + ) + return wrapped_out + + # A plain (non-`ToolReturn`) result has no separate content part. + if value_handle is not None: + spilled_out: ToolReturn[object] = ToolReturn( + return_value=value_text, metadata=_with_handles(None, value_handle, len(value_unit.data)) + ) + return spilled_out + return value_text + + def _make_unit(self, value: ToolReturnContent, *, suffix: str) -> _Unit: + """Pre-render a value into the text / bytes the reduction pipeline needs.""" + if is_binary(value): + return _Unit(binary=True, text=None, data=to_bytes(value), value=value, suffix=suffix) + text = to_text(value) + if self.strip_ansi: + text = strip_ansi(text) + return _Unit(binary=False, text=text, data=text.encode('utf-8'), value=value, suffix=suffix) + + async def _reduce( + self, + ctx: RunContext[AgentDepsT], + call: ToolCallPart, + bands: Sequence[Band], + unit: _Unit, + ) -> tuple[str | None, str | None]: + """Select a band for `unit` and apply it. Returns `(replacement, handle)`. + + `replacement` is None when the unit passes through unchanged; `handle` is set only + when the unit was spilled. + """ + size = ( + len(unit.data) + if unit.binary + else measure(unit.text or '', over_tokens=self.over_tokens, tokenizer=self.tokenizer) + ) + action = _select_action(bands, size) + if action is None: + return None, None + return await self._apply(ctx, call, action, unit) + + async def _reduce_content( + self, + ctx: RunContext[AgentDepsT], + call: ToolCallPart, + bands: Sequence[Band], + content: str | Sequence[UserContent] | None, + ) -> tuple[str | None, str | None]: + """Reduce model-visible `content`. Text content is reduced; other content warns.""" + if content is None: + return None, None + if isinstance(content, str): + return await self._reduce(ctx, call, bands, self._make_unit(content, suffix='.content')) + + text = ''.join(part for part in content if isinstance(part, str)) + size = measure(text, over_tokens=self.over_tokens, tokenizer=self.tokenizer) + action = _select_action(bands, size) + if action is not None and not isinstance(action, Passthrough): + warnings.warn( + f'OverflowingToolOutput: tool {call.tool_name!r} returned large non-text ' + f'content ({len(content)} parts); leaving it unreduced.', + stacklevel=2, + ) + return None, None + + async def _apply( + self, + ctx: RunContext[AgentDepsT], + call: ToolCallPart, + action: Action, + unit: _Unit, + ) -> tuple[str | None, str | None]: + """Apply one action to a unit, falling back to its `then` when it cannot run.""" + if isinstance(action, Passthrough): + return None, None + + if isinstance(action, Truncate): + if unit.binary: + return await self._fallback(ctx, call, action.then, unit) + assert unit.text is not None + return truncate_text(unit.text, action.max_chars, action.strategy), None + + if isinstance(action, Spill): + return await self._spill(ctx, call, action, unit) + + return await self._summarize_action(ctx, call, action, unit) + + async def _fallback( + self, + ctx: RunContext[AgentDepsT], + call: ToolCallPart, + then: Action | None, + unit: _Unit, + ) -> tuple[str | None, str | None]: + """Run the fallback action, or keep the unit unchanged when there is none.""" + if then is None: + return None, None + return await self._apply(ctx, call, then, unit) + + async def _spill( + self, + ctx: RunContext[AgentDepsT], + call: ToolCallPart, + action: Spill, + unit: _Unit, + ) -> tuple[str | None, str | None]: + key = _handle_key(ctx, call, unit.suffix) + try: + handle = await self._store.write(key, unit.data) + except Exception: + return await self._fallback(ctx, call, action.then, unit) + + preview = _build_spill_preview(handle, unit, action.preview_chars, over_tokens=self.over_tokens) + return preview, handle + + async def _summarize_action( + self, + ctx: RunContext[AgentDepsT], + call: ToolCallPart, + action: Summarize, + unit: _Unit, + ) -> tuple[str | None, str | None]: + if unit.binary: + return await self._fallback(ctx, call, action.then, unit) + assert unit.text is not None + try: + summary = await self._summarize(ctx, call, action, unit.text) + except Exception: + return await self._fallback(ctx, call, action.then, unit) + return summary, None + + async def _summarize( + self, + ctx: RunContext[AgentDepsT], + call: ToolCallPart, + action: Summarize, + text: str, + ) -> str: + """Generate the summary via a custom callable or the inherited-model agent.""" + if action.summarize is not None: + outcome = action.summarize(call.tool_name, text) + if isinstance(outcome, Awaitable): + return await outcome + return outcome + + from pydantic_ai import Agent + + model = action.model if action.model is not None else ctx.model + prompt = self.summary_prompt.format(tool_name=call.tool_name, output=text) + agent: Agent[None, str] = Agent(model, instructions='You summarize oversized tool output.') + run = await agent.run(prompt, usage=ctx.usage) + return run.output.strip() + + +def _select_action(bands: Sequence[Band], size: int) -> Action | None: + """Return the first (largest-threshold) band action whose threshold `size` meets.""" + for band in bands: + if size >= band.over: + return band.action + return None + + +def _handle_key(ctx: RunContext[AgentDepsT], call: ToolCallPart, suffix: str = '') -> str: + """Build a per-run, per-call, per-retry key so concurrent and retried calls never clash. + + `suffix` keeps a return's `return_value` and `content` spills on distinct handles. + """ + run_id = ctx.run_id or 'run' + call_id = call.tool_call_id or 'call' + return f'{run_id}/{call_id}.{ctx.retry}{suffix}' + + +def _is_mapping(value: object) -> TypeGuard[Mapping[object, object]]: + """`TypeGuard` so a mapping narrows to a known element type, not `Unknown`.""" + return isinstance(value, Mapping) + + +def _with_handles( + existing: object, + value_handle: str | None, + value_bytes: int, + content_handle: str | None = None, +) -> dict[str, object]: + """Stash spill handle(s) in `ToolReturn.metadata` (app-only, costs no model tokens).""" + base: dict[str, object] = {} + if _is_mapping(existing): + base.update(_copy_mapping(existing)) + if value_handle is not None: + base['overflow_handle'] = value_handle + base['overflow_bytes'] = value_bytes + if content_handle is not None: + base['overflow_content_handle'] = content_handle + return base + + +def _copy_mapping(source: Mapping[object, object]) -> dict[str, object]: + """Copy an arbitrary mapping with stringified keys (tool metadata is app-defined).""" + return {str(key): source[key] for key in source} + + +def _build_spill_preview(handle: str, unit: _Unit, preview_chars: int, *, over_tokens: bool) -> str: + """Compose the model-visible spill stand-in: marker, sketch, and a head/tail preview.""" + if unit.binary: + size_desc = f'{len(unit.data):,} bytes (binary)' + body = f'<{len(unit.data):,} bytes of binary data>' + sketch = '' + else: + text = unit.text or '' + size_unit = 'tokens' if over_tokens else 'chars' + amount = measure(text, over_tokens=over_tokens, tokenizer=None) if over_tokens else len(text) + size_desc = f'{amount:,} {size_unit}' + body = _head_tail_preview(text, preview_chars) + sketch = json_sketch(unit.value) + + header = ( + f'[Tool output too large ({size_desc}); stored to handle {handle!r}. ' + f'Read it with read_tool_result(handle={handle!r}, offset=0, limit=200, ' + f'from_end=False, pattern=None).]' + ) + parts = [header] + if sketch: + parts.append(f'shape: {sketch}') + parts.append(body) + return '\n'.join(parts) + + +def _head_tail_preview(text: str, preview_chars: int) -> str: + """Return a head+tail slice of `text` with a middle-elision marker.""" + if len(text) <= preview_chars: + return text + head_chars = preview_chars // 2 + tail_chars = preview_chars - head_chars + omitted = len(text) - head_chars - tail_chars + return f'{text[:head_chars]}\n...[{omitted:,} chars omitted]...\n{text[-tail_chars:]}' + + +_MAX_READ_LINES = 1_000 +"""Hard cap on lines returned by one `read_tool_result` call.""" + +_MAX_READ_CHARS = 50_000 +"""Hard cap on characters returned by one `read_tool_result` call.""" + + +async def _read_slice( + store: OverflowStore, + handle: str, + offset: int, + limit: int, + from_end: bool, + pattern: str | None, +) -> str: + """Filter and slice a spilled payload for `read_tool_result`, bounded in both axes. + + `pattern` is a literal substring (not a regex), so a model-supplied value cannot hang + the host with catastrophic backtracking. `limit` is clamped and the joined output is + capped, so one call can never return an unbounded amount of text. + """ + if offset < 0: + raise ModelRetry('`offset` must be >= 0.') + if limit < 1: + raise ModelRetry('`limit` must be >= 1.') + limit = min(limit, _MAX_READ_LINES) + + try: + data = await store.read(handle) + except OSError as exc: + raise ModelRetry(f'No stored tool result for handle {handle!r}: {exc}.') from exc + + lines = data.decode('utf-8', errors='replace').splitlines() + if pattern is not None: + lines = [line for line in lines if pattern in line] + + total = len(lines) + if from_end: + end = max(0, total - offset) + window = lines[max(0, end - limit) : end] + else: + window = lines[offset : offset + limit] + + body = '\n'.join(window) + capped = '' + if len(body) > _MAX_READ_CHARS: + body = body[:_MAX_READ_CHARS] + capped = ', output capped' + header = f'[handle {handle!r}: {total:,} matching line(s); showing {len(window)}{capped}]' + return f'{header}\n{body}' if body else header diff --git a/pydantic_ai_harness/experimental/overflow/_payload.py b/pydantic_ai_harness/experimental/overflow/_payload.py new file mode 100644 index 0000000..a4b7e99 --- /dev/null +++ b/pydantic_ai_harness/experimental/overflow/_payload.py @@ -0,0 +1,138 @@ +"""Size measurement, stringification, truncation, and binary detection. + +Harvested from PR #185 (`ToolOutputManagement`) and adapted: character-based truncation +strategies, ANSI stripping, and binary detection. Token measurement reuses the compaction +heuristic via `estimate_token_count` so the two capabilities stay aligned. +""" + +from __future__ import annotations + +import re +from collections.abc import Callable, Mapping, Sequence +from enum import Enum +from typing import TypeGuard + +from pydantic_ai.messages import ModelMessage, ModelRequest, SystemPromptPart +from pydantic_core import to_json + +from pydantic_ai_harness.experimental.compaction._shared import estimate_token_count + + +class TruncationStrategy(str, Enum): + """Which end(s) of an oversized text to keep when truncating.""" + + head = 'head' + """Keep the first characters (good for headers / schemas).""" + + tail = 'tail' + """Keep the last characters (good for build / test output, where errors land last).""" + + head_tail = 'head_tail' + """Keep the first and last characters, eliding the middle.""" + + +# CSI sequences, OSC sequences, and simple escapes. Terminal tool output is full of color +# codes that waste tokens and can confuse models. +_ANSI_ESCAPE_RE = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07|\x1b[^[\]()]') + + +def strip_ansi(text: str) -> str: + """Remove ANSI escape sequences from `text`.""" + return _ANSI_ESCAPE_RE.sub('', text) + + +def is_binary(value: object) -> bool: + """Return True for raw byte payloads, which must never be stringify-truncated.""" + return isinstance(value, (bytes, bytearray, memoryview)) + + +def to_bytes(value: object) -> bytes: + """Serialize any tool return value to the bytes that get spilled. + + Strings spill as UTF-8 text; byte payloads spill verbatim; everything else spills as + JSON so the stored payload stays valid and grep-able. + """ + if isinstance(value, str): + return value.encode('utf-8') + if isinstance(value, memoryview): + return value.tobytes() + if isinstance(value, (bytes, bytearray)): + return bytes(value) + return to_json(value) + + +def to_text(value: object) -> str: + """Render a non-binary tool return value as the text used for measuring and truncating. + + Strings pass through; structured values become JSON (truncating JSON is lossy, so prefer + spill or summarize for them -- see the README). + """ + if isinstance(value, str): + return value + return to_json(value).decode('utf-8', errors='replace') + + +def measure(text: str, *, over_tokens: bool, tokenizer: Callable[[str], int] | None) -> int: + """Measure `text` in characters (default) or estimated tokens (`over_tokens=True`).""" + if not over_tokens: + return len(text) + message: ModelMessage = ModelRequest(parts=[SystemPromptPart(content=text)]) + return estimate_token_count([message], tokenizer) + + +def json_sketch(value: object) -> str: + """Build a one-line shape hint for a structured value, or '' for anything else. + + The `_is_*` guards are `TypeGuard`s, so a `Mapping`/`Sequence` value narrows to a known + element type (`object`) the strict type checker accepts -- no `Any` and no `Unknown`. + """ + if _is_mapping(value): + return _sketch_mapping(value) + if _is_text_sequence(value): + return _sketch_sequence(value) + return '' + + +def _is_mapping(value: object) -> TypeGuard[Mapping[object, object]]: + return isinstance(value, Mapping) + + +def _is_text_sequence(value: object) -> TypeGuard[Sequence[object]]: + return isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)) + + +def _sketch_mapping(mapping: Mapping[object, object]) -> str: + keys = list(mapping) + shown = ', '.join(f'{key!r}: {type(mapping[key]).__name__}' for key in keys[:10]) + more = '' if len(keys) <= 10 else f', ... ({len(keys)} keys)' + return f'{{{shown}{more}}}' + + +def _sketch_sequence(items: Sequence[object]) -> str: + elem = type(items[0]).__name__ if items else 'empty' + return f'[{len(items)} items of {elem}]' + + +def truncate_text(text: str, max_chars: int, strategy: TruncationStrategy) -> str: + """Cut `text` down to roughly `max_chars`, annotating what was removed. + + Returns `text` unchanged when it already fits. + """ + total = len(text) + if total <= max_chars: + return text + + if strategy is TruncationStrategy.head: + return f'{text[:max_chars]}\n\n[truncated: showing first {max_chars:,} of {total:,} chars]' + if strategy is TruncationStrategy.tail: + return f'[truncated: showing last {max_chars:,} of {total:,} chars]\n\n{text[-max_chars:]}' + + head_chars = max_chars * 2 // 5 + tail_chars = max_chars - head_chars + omitted = total - head_chars - tail_chars + return ( + f'{text[:head_chars]}\n\n' + f'[truncated: {omitted:,} chars omitted from the middle; ' + f'showing first {head_chars:,} + last {tail_chars:,} of {total:,} chars]\n\n' + f'{text[-tail_chars:]}' + ) diff --git a/pydantic_ai_harness/experimental/overflow/_store.py b/pydantic_ai_harness/experimental/overflow/_store.py new file mode 100644 index 0000000..705bbd9 --- /dev/null +++ b/pydantic_ai_harness/experimental/overflow/_store.py @@ -0,0 +1,153 @@ +"""Storage backend for spilled tool outputs. + +`OverflowStore` is a narrow protocol: persist a payload under a key, read it back by +handle. `LocalFileStore` is the dependency-free default -- it writes each payload to a +file under a stable root directory. The handle is backend-addressable (a relative key), +not an absolute local path, so a durable backend (Temporal, a blob store) can resolve the +same handle in another process. This is the seam for consuming the core queryable-file +primitive (pydantic-ai #4352 / `ExecutionEnvironment`) once it lands. +""" + +from __future__ import annotations + +import re +import tempfile +import threading +import time +import warnings +from dataclasses import dataclass, field +from datetime import timedelta +from pathlib import Path +from typing import Protocol, runtime_checkable + + +@runtime_checkable +class OverflowStore(Protocol): + """Persist and retrieve spilled tool-output payloads. + + `write` takes a caller-chosen `key` and returns a `handle`. The handle is the only + thing a later `read` needs, so it must be self-contained (a backend can encode the + run, the call, and the retry into it). Implementations may return the key unchanged. + """ + + async def write(self, key: str, data: bytes) -> str: + """Persist `data` under `key` and return a handle that `read` accepts.""" + ... # pragma: no cover + + async def read(self, handle: str) -> bytes: + """Return the payload previously stored for `handle`. + + Raise `FileNotFoundError` (or another `OSError`) when the handle is unknown. + """ + ... # pragma: no cover + + +_UNSAFE_SEGMENT = re.compile(r'[^A-Za-z0-9._-]+') + + +def _safe_segment(segment: str) -> str: + """Make one path segment filesystem-safe without collapsing distinct keys. + + Empty or dot-only segments are replaced so a handle can never escape the root via + `.`/`..`; the resolve-within-root check in `read` is the second line of defense. + """ + cleaned = _UNSAFE_SEGMENT.sub('_', segment) + if cleaned in ('', '.', '..'): + return '_' + return cleaned + + +@dataclass +class LocalFileStore: + """Dependency-free `OverflowStore` that writes each payload to a local file. + + The handle equals the key: a relative `run_id/tool_call_id.retry` path under + `base_dir`. The root is stable and shareable on purpose -- a later agent or run can + read a spill a previous run produced, so the store is not isolated per instance. + + Security comes from two mechanisms, not isolation: the root is created with `0700` + perms (owner-only), and `read` resolves the target (following symlinks) and rejects + anything that escapes the root via symlink, `..`, or an absolute path. Handle segments + are also sanitized by `_safe_segment`. + + Files are kept after the run by default (a later `read_tool_result` may need them). + Set `cleanup_after` to opt into age-based pruning; see that field. + """ + + base_dir: Path | None = None + """Root directory for spilled files. Defaults to a stable temp subdirectory.""" + + cleanup_after: timedelta | None = None + """Opt-in TTL for spilled files. `None` (default) keeps files forever. + + When set, a `write` schedules a background prune (a daemon thread, off the hot path) + that deletes files whose modification time is older than `cleanup_after`. Pruning is + best-effort: any failure is caught and surfaced via `warnings.warn`, never propagated + into the agent run. Modification time (`st_mtime`) is the age signal; last-read time + (`st_atime`) is unreliable on `noatime`/`relatime` mounts and is not used. + """ + + _root: Path = field(init=False, repr=False) + + def __post_init__(self) -> None: + self._root = ( + self.base_dir if self.base_dir is not None else Path(tempfile.gettempdir()) / 'pyai_harness_overflow' + ) + + def _path(self, key: str) -> Path: + segments = [_safe_segment(part) for part in key.split('/') if part] + if not segments: + segments = ['_'] + return self._root.joinpath(*segments) + + def _ensure_root(self) -> None: + """Create the root directory owned by the current user with `0700` perms.""" + self._root.mkdir(parents=True, exist_ok=True) + try: + self._root.chmod(0o700) + except OSError: # pragma: no cover - best effort on a root we do not own + pass + + async def write(self, key: str, data: bytes) -> str: + self._ensure_root() + path = self._path(key) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + self._schedule_cleanup() + return key + + async def read(self, handle: str) -> bytes: + target = self._path(handle).resolve() + root = self._root.resolve() + if not target.is_relative_to(root): + raise PermissionError(f'Handle {handle!r} resolves outside the store root.') + return target.read_bytes() + + # --- opt-in TTL pruning (non-blocking, non-erroring) --- + + def _schedule_cleanup(self) -> threading.Thread | None: + """Fire a background prune when `cleanup_after` is set. Never blocks `write`.""" + if self.cleanup_after is None: + return None + thread = threading.Thread(target=self._run_prune, name='overflow-prune', daemon=True) + thread.start() + return thread + + def _run_prune(self) -> None: + try: + self._prune_sync() + except Exception as exc: # never let cleanup fail a run or block the hot path + warnings.warn(f'LocalFileStore cleanup failed: {exc}', stacklevel=2) + + def _prune_sync(self) -> None: + """Delete files older than `cleanup_after` (by `st_mtime`).""" + assert self.cleanup_after is not None + cutoff = time.time() - self.cleanup_after.total_seconds() + for path in self._root.rglob('*'): + if not path.is_file(): + continue + try: + if path.stat().st_mtime < cutoff: + path.unlink() + except OSError: # pragma: no cover - file vanished mid-prune + continue diff --git a/tests/experimental/overflow/__init__.py b/tests/experimental/overflow/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/experimental/overflow/test_overflow.py b/tests/experimental/overflow/test_overflow.py new file mode 100644 index 0000000..231c8e1 --- /dev/null +++ b/tests/experimental/overflow/test_overflow.py @@ -0,0 +1,754 @@ +"""Tests for pydantic_ai_harness.experimental.overflow.""" + +from __future__ import annotations + +import dataclasses +import os +import time +from datetime import timedelta +from pathlib import Path +from typing import Any + +import pytest +from pydantic_ai import Agent +from pydantic_ai.exceptions import ModelRetry +from pydantic_ai.messages import ( + BinaryContent, + ModelMessage, + ModelResponse, + TextPart, + ToolCallPart, + ToolReturn, + ToolReturnPart, +) +from pydantic_ai.models.function import AgentInfo, FunctionModel +from pydantic_ai.models.test import TestModel +from pydantic_ai.tools import ToolDefinition +from pydantic_ai.usage import RunUsage + +from pydantic_ai_harness.experimental.overflow import ( + Band, + LocalFileStore, + OverflowingToolOutput, + Passthrough, + Spill, + Summarize, + Truncate, + TruncationStrategy, +) +from pydantic_ai_harness.experimental.overflow._capability import ( + READ_TOOL_NAME, + _build_spill_preview, + _handle_key, + _head_tail_preview, + _read_slice, + _select_action, + _Unit, + _with_handles, +) +from pydantic_ai_harness.experimental.overflow._payload import ( + is_binary, + json_sketch, + measure, + strip_ansi, + to_bytes, + to_text, + truncate_text, +) +from pydantic_ai_harness.experimental.overflow._store import _safe_segment + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_ctx(*, run_id: str | None = 'run-1', retry: int = 0, model: Any = None) -> Any: + """Build a minimal RunContext-like object for testing the hook directly.""" + + @dataclasses.dataclass + class _FakeModel: + model_id: str = 'test-model' + + @dataclasses.dataclass + class _FakeCtx: + usage: RunUsage + run_id: str | None + retry: int + tool_call_id: str | None = 'call-1' + model: Any = dataclasses.field(default_factory=_FakeModel) + deps: None = None + + ctx = _FakeCtx(usage=RunUsage(), run_id=run_id, retry=retry) + if model is not None: + ctx.model = model + return ctx + + +def _fixed_model(text: str) -> FunctionModel: + """A `FunctionModel` whose single text response is `text` (no tool calls).""" + + def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + return ModelResponse(parts=[TextPart(content=text)]) + + return FunctionModel(respond) + + +def _call(tool_name: str = 'big_tool', tool_call_id: str = 'call-1') -> ToolCallPart: + return ToolCallPart(tool_name=tool_name, args='{}', tool_call_id=tool_call_id) + + +def _tool_def(name: str = 'big_tool') -> ToolDefinition: + return ToolDefinition(name=name) + + +async def _run(cap: OverflowingToolOutput[None], result: Any, *, ctx: Any = None, tool_name: str = 'big_tool') -> Any: + return await cap.after_tool_execute( + ctx if ctx is not None else _make_ctx(), + call=_call(tool_name), + tool_def=_tool_def(tool_name), + args={}, + result=result, + ) + + +@pytest.fixture +def anyio_backend() -> str: + return 'asyncio' + + +# --------------------------------------------------------------------------- +# _payload helpers +# --------------------------------------------------------------------------- + + +class TestPayloadHelpers: + def test_strip_ansi(self): + assert strip_ansi('\x1b[31mred\x1b[0m') == 'red' + + def test_is_binary(self): + assert is_binary(b'x') is True + assert is_binary(bytearray(b'x')) is True + assert is_binary('x') is False + + def test_to_bytes_variants(self): + assert to_bytes('hi') == b'hi' + assert to_bytes(memoryview(b'mv')) == b'mv' + assert to_bytes(bytearray(b'ba')) == b'ba' + assert to_bytes({'a': 1}) == b'{"a":1}' + + def test_to_text_variants(self): + assert to_text('hi') == 'hi' + assert to_text({'a': 1}) == '{"a":1}' + + def test_measure_chars_and_tokens(self): + assert measure('x' * 100, over_tokens=False, tokenizer=None) == 100 + assert measure('x' * 100, over_tokens=True, tokenizer=None) == 25 + assert measure('abcd', over_tokens=True, tokenizer=lambda s: len(s)) == 4 + + def test_json_sketch_mapping(self): + assert json_sketch({'a': 1, 'b': 'x'}) == "{'a': int, 'b': str}" + + def test_json_sketch_mapping_truncated(self): + big = {f'k{i}': i for i in range(12)} + assert json_sketch(big).endswith('... (12 keys)}') + + def test_json_sketch_sequence(self): + assert json_sketch([1, 2, 3]) == '[3 items of int]' + + def test_json_sketch_empty_sequence(self): + assert json_sketch([]) == '[0 items of empty]' + + def test_json_sketch_scalar(self): + assert json_sketch(42) == '' + assert json_sketch('plain') == '' + + def test_truncate_under_limit(self): + assert truncate_text('short', 100, TruncationStrategy.head_tail) == 'short' + + def test_truncate_head(self): + out = truncate_text('a' * 100, 10, TruncationStrategy.head) + assert out.startswith('aaaaaaaaaa') + assert 'showing first 10' in out + + def test_truncate_tail(self): + out = truncate_text('a' * 100, 10, TruncationStrategy.tail) + assert out.endswith('aaaaaaaaaa') + assert 'showing last 10' in out + + def test_truncate_head_tail(self): + out = truncate_text('a' * 100, 10, TruncationStrategy.head_tail) + assert 'omitted from the middle' in out + + +# --------------------------------------------------------------------------- +# Store: write/read, S1 hardening +# --------------------------------------------------------------------------- + + +class TestStore: + def test_safe_segment(self): + assert _safe_segment('a b!@#') == 'a_b_' + assert _safe_segment('') == '_' + assert _safe_segment('..') == '_' + assert _safe_segment('.') == '_' + assert _safe_segment('ok-1.2') == 'ok-1.2' + + def test_default_root(self): + store = LocalFileStore() + assert store._root.name == 'pyai_harness_overflow' + + async def test_write_read_roundtrip(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path / 'store') + handle = await store.write('run-1/call-1.0', b'payload') + assert handle == 'run-1/call-1.0' + assert await store.read(handle) == b'payload' + + async def test_root_created_0700(self, tmp_path: Path): + root = tmp_path / 'store' + store = LocalFileStore(base_dir=root) + await store.write('run/c.0', b'x') + assert oct(root.stat().st_mode & 0o777) == '0o700' + + async def test_empty_key(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path / 'store') + handle = await store.write('', b'data') + assert await store.read(handle) == b'data' + + async def test_read_missing_raises(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path / 'store') + with pytest.raises(OSError): + await store.read('nope/x.0') + + async def test_dotdot_handle_stays_in_root(self, tmp_path: Path): + # `_safe_segment` neutralizes `..`, so the read resolves inside the root and 404s + # rather than escaping. + store = LocalFileStore(base_dir=tmp_path / 'store') + await store.write('run/c.0', b'inside') + with pytest.raises(OSError): + await store.read('../c.0') + + async def test_symlink_escape_rejected(self, tmp_path: Path): + secret = tmp_path / 'secret.txt' + secret.write_bytes(b'top secret') + root = tmp_path / 'store' + store = LocalFileStore(base_dir=root) + await store.write('run/c.0', b'inside') # creates the root + (root / 'evil').symlink_to(secret) + with pytest.raises(PermissionError, match='outside the store root'): + await store.read('evil') + + +# --------------------------------------------------------------------------- +# Store: opt-in TTL cleanup +# --------------------------------------------------------------------------- + + +class TestCleanup: + def test_prune_removes_old_keeps_new(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path, cleanup_after=timedelta(seconds=1)) + old = tmp_path / 'old.bin' + old.write_bytes(b'x') + new = tmp_path / 'new.bin' + new.write_bytes(b'y') + (tmp_path / 'sub').mkdir() # a directory rglob yields -- must be skipped + past = time.time() - 100 + os.utime(old, (past, past)) + + store._prune_sync() + + assert not old.exists() + assert new.exists() + + def test_run_prune_swallows_errors(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + store = LocalFileStore(base_dir=tmp_path, cleanup_after=timedelta(seconds=1)) + + def boom() -> None: + raise OSError('disk gone') + + monkeypatch.setattr(store, '_prune_sync', boom) + with pytest.warns(UserWarning, match='cleanup failed'): + store._run_prune() + + def test_schedule_none_when_disabled(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + assert store._schedule_cleanup() is None + + def test_schedule_starts_thread(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path, cleanup_after=timedelta(seconds=1)) + (tmp_path / 'f.bin').write_bytes(b'z') + thread = store._schedule_cleanup() + assert thread is not None + thread.join(timeout=5) + assert not thread.is_alive() + + async def test_write_schedules_cleanup(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + store = LocalFileStore(base_dir=tmp_path / 'store', cleanup_after=timedelta(seconds=1)) + scheduled: list[int] = [] + monkeypatch.setattr(store, '_schedule_cleanup', lambda: scheduled.append(1)) + await store.write('run/c.0', b'data') + assert scheduled == [1] + + +# --------------------------------------------------------------------------- +# Capability construction +# --------------------------------------------------------------------------- + + +class TestConstruction: + def test_default_band_is_spill_then_truncate(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput() + assert len(cap._bands) == 1 + action = cap._bands[0].action + assert isinstance(action, Spill) + assert isinstance(action.then, Truncate) + + def test_bands_sorted_descending(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=10, action=Truncate()), Band(over=100, action=Spill())] + ) + assert [b.over for b in cap._bands] == [100, 10] + + def test_negative_threshold_rejected(self): + with pytest.raises(ValueError, match='non-negative'): + OverflowingToolOutput(bands=[Band(over=-1, action=Passthrough())]) + + def test_provided_store_used(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + cap: OverflowingToolOutput[None] = OverflowingToolOutput(store=store) + assert cap._store is store + + def test_per_tool_prepared(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + per_tool={'read_file': [Band(over=5, action=Truncate())]} + ) + assert 'read_file' in cap._per_tool + + +# --------------------------------------------------------------------------- +# Passthrough / filtering / guards +# --------------------------------------------------------------------------- + + +class TestPassthrough: + async def test_read_tool_exempt(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Truncate(max_chars=2))]) + out = await _run(cap, 'x' * 100, tool_name=READ_TOOL_NAME) + assert out == 'x' * 100 + + async def test_tool_filter_skips_unmatched(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=1, action=Truncate(max_chars=2))], tool_filter=['other'] + ) + out = await _run(cap, 'x' * 100) + assert out == 'x' * 100 + + async def test_callable_filter(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=1, action=Truncate(max_chars=2))], + tool_filter=lambda ctx, td: td.name == 'big_tool', + ) + out = await _run(cap, 'x' * 100) + assert isinstance(out, str) and 'truncated' in out + + async def test_below_threshold_passthrough(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1000, action=Truncate())]) + out = await _run(cap, 'small') + assert out == 'small' + + async def test_exception_result_passthrough(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Truncate(max_chars=2))]) + err = ValueError('boom') + assert await _run(cap, err) is err + + +# --------------------------------------------------------------------------- +# Truncate +# --------------------------------------------------------------------------- + + +class TestTruncate: + async def test_truncates_text(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=10, action=Truncate(max_chars=20, strategy=TruncationStrategy.head))] + ) + out = await _run(cap, 'a' * 100) + assert isinstance(out, str) and out.startswith('a' * 20) + + async def test_strip_ansi_applied(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=5, action=Truncate(max_chars=1000))], strip_ansi=True + ) + out = await _run(cap, '\x1b[31m' + 'red text ' * 10 + '\x1b[0m') + assert isinstance(out, str) and '\x1b[' not in out + + async def test_binary_truncate_falls_back_to_passthrough(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Truncate())]) + data = b'\x00\x01' * 100 + assert await _run(cap, data) == data + + async def test_tool_return_envelope_preserved(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=10, action=Truncate(max_chars=20))]) + out = await _run(cap, ToolReturn(return_value='a' * 100, content='note', metadata={'k': 1})) + assert isinstance(out, ToolReturn) + assert out.content == 'note' + assert out.metadata == {'k': 1} + + +# --------------------------------------------------------------------------- +# Spill +# --------------------------------------------------------------------------- + + +class TestSpill: + async def test_spill_roundtrip(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=10, action=Spill(preview_chars=20))], store=store + ) + text = 'line\n' * 1000 + out = await _run(cap, text) + assert isinstance(out, ToolReturn) + assert isinstance(out.return_value, str) and 'too large' in out.return_value + handle = out.metadata['overflow_handle'] + assert handle == 'run-1/call-1.0' + assert await store.read(handle) == text.encode('utf-8') + + async def test_spill_binary_verbatim(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Spill())], store=store) + data = b'\x00\xff' * 100 + out = await _run(cap, data) + assert isinstance(out, ToolReturn) + assert 'binary' in out.return_value # type: ignore[operator] + assert await store.read(out.metadata['overflow_handle']) == data + + async def test_spill_structured_includes_sketch(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=5, action=Spill())], store=store) + out = await _run(cap, {'rows': list(range(1000)), 'ok': True}) + assert isinstance(out, ToolReturn) + assert 'shape:' in out.return_value # type: ignore[operator] + + async def test_spill_failure_falls_back_to_truncate(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=10, action=Spill(then=Truncate(max_chars=15)))], store=_BrokenStore() + ) + out = await _run(cap, 'a' * 100) + assert isinstance(out, str) and 'truncated' in out + + async def test_spill_failure_no_fallback_returns_original(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=10, action=Spill())], store=_BrokenStore() + ) + out = await _run(cap, 'a' * 100) + assert out == 'a' * 100 + + async def test_handle_distinct_per_retry(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=5, action=Spill())], store=store) + out0 = await _run(cap, 'a' * 100, ctx=_make_ctx(retry=0)) + out1 = await _run(cap, 'b' * 100, ctx=_make_ctx(retry=1)) + assert out0.metadata['overflow_handle'] != out1.metadata['overflow_handle'] # type: ignore[union-attr] + + async def test_spill_merges_existing_metadata(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=5, action=Spill())], store=store) + out = await _run(cap, ToolReturn(return_value='a' * 100, metadata={'orig': True})) + assert isinstance(out, ToolReturn) + assert out.metadata['orig'] is True + assert 'overflow_handle' in out.metadata + + +class _BrokenStore: + """An `OverflowStore` whose writes always fail (for fallback tests).""" + + async def write(self, key: str, data: bytes) -> str: + raise OSError('disk full') + + async def read(self, handle: str) -> bytes: # pragma: no cover - never reached + raise FileNotFoundError(handle) + + +# --------------------------------------------------------------------------- +# C1: model-visible ToolReturn.content is reduced too +# --------------------------------------------------------------------------- + + +class TestContentReduction: + async def test_large_content_spilled(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=100, action=Spill(preview_chars=20))], store=store + ) + out = await _run(cap, ToolReturn(return_value='small', content='C' * 5000)) + assert isinstance(out, ToolReturn) + assert out.return_value == 'small' # small return_value untouched + assert isinstance(out.content, str) and 'too large' in out.content + handle = out.metadata['overflow_content_handle'] + assert await store.read(handle) == ('C' * 5000).encode('utf-8') + + async def test_large_content_truncated(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=10, action=Truncate(max_chars=20))]) + out = await _run(cap, ToolReturn(return_value='small', content='C' * 200)) + assert isinstance(out, ToolReturn) + assert isinstance(out.content, str) and 'truncated' in out.content + + async def test_both_value_and_content_reduced(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=50, action=Spill())], store=store) + out = await _run(cap, ToolReturn(return_value='v' * 500, content='c' * 500)) + assert isinstance(out, ToolReturn) + assert out.metadata['overflow_handle'] != out.metadata['overflow_content_handle'] + + async def test_nontext_content_warns_and_passes_through(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=10, action=Truncate())]) + content = ['x' * 5000, BinaryContent(data=b'\x00', media_type='application/octet-stream')] + original = ToolReturn(return_value='small', content=content) + with pytest.warns(UserWarning, match='non-text content'): + out = await _run(cap, original) + assert out is original + + async def test_nontext_content_passthrough_action_no_warn(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Passthrough())]) + content = ['x' * 5000, BinaryContent(data=b'\x00', media_type='application/octet-stream')] + original = ToolReturn(return_value='small', content=content) + out = await _run(cap, original) # Passthrough action -> no warning, returned unchanged + assert out is original + + async def test_small_nontext_content_no_warn(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=10_000, action=Truncate())]) + content = ['tiny', BinaryContent(data=b'\x00', media_type='application/octet-stream')] + original = ToolReturn(return_value='small', content=content) + out = await _run(cap, original) + assert out is original + + +# --------------------------------------------------------------------------- +# Summarize (M1: assert model + usage, not just a wholesale mock) +# --------------------------------------------------------------------------- + + +class TestSummarize: + async def test_custom_sync_summarizer(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=5, action=Summarize(summarize=lambda name, text: f'{name}:{len(text)}'))] + ) + out = await _run(cap, 'x' * 100) + assert out == 'big_tool:100' + + async def test_custom_async_summarizer(self): + async def summ(name: str, text: str) -> str: + return f'async:{len(text)}' + + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=5, action=Summarize(summarize=summ))]) + out = await _run(cap, 'x' * 100) + assert out == 'async:100' + + async def test_inherited_model_and_usage(self): + # model=None resolves to ctx.model, and the call threads usage=ctx.usage. + ctx = _make_ctx(model=_fixed_model('THE SUMMARY')) + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=5, action=Summarize())]) + out = await _run(cap, 'x' * 100, ctx=ctx) + assert out == 'THE SUMMARY' + assert ctx.usage.requests == 1 + + async def test_explicit_model_overrides_ctx(self): + ctx = _make_ctx(model=_fixed_model('FROM CTX MODEL')) + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=5, action=Summarize(model=_fixed_model('FROM EXPLICIT MODEL')))] + ) + out = await _run(cap, 'x' * 100, ctx=ctx) + assert out == 'FROM EXPLICIT MODEL' + assert ctx.usage.requests == 1 + + async def test_binary_summarize_falls_back(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=1, action=Summarize(then=Passthrough()))] + ) + data = b'\x00' * 100 + assert await _run(cap, data) == data + + async def test_summarize_failure_falls_back(self): + def boom(name: str, text: str) -> str: + raise RuntimeError('model down') + + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=5, action=Summarize(summarize=boom, then=Truncate(max_chars=10)))] + ) + out = await _run(cap, 'a' * 100) + assert isinstance(out, str) and 'truncated' in out + + +# --------------------------------------------------------------------------- +# Passthrough action + per-tool + band selection +# --------------------------------------------------------------------------- + + +class TestActionsAndSelection: + async def test_passthrough_action(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput(bands=[Band(over=1, action=Passthrough())]) + assert await _run(cap, 'x' * 100) == 'x' * 100 + + async def test_per_tool_replaces_bands(self): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=1, action=Truncate(max_chars=5))], + per_tool={'big_tool': [Band(over=100_000, action=Truncate())]}, + ) + # global band would truncate, but per_tool threshold is huge -> passthrough + assert await _run(cap, 'x' * 100) == 'x' * 100 + + def test_select_action_no_match(self): + assert _select_action([Band(over=100, action=Passthrough())], 50) is None + + def test_select_action_first_match(self): + bands = [Band(over=100, action=Spill()), Band(over=10, action=Truncate())] + assert isinstance(_select_action(bands, 50), Truncate) + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +class TestInternals: + def test_handle_key_defaults(self): + ctx = _make_ctx(run_id=None, retry=2) + ctx.tool_call_id = None + key = _handle_key(ctx, ToolCallPart(tool_name='t', args='{}', tool_call_id='')) + assert key == 'run/call.2' + + def test_handle_key_suffix(self): + key = _handle_key(_make_ctx(), ToolCallPart(tool_name='t', args='{}', tool_call_id='c'), '.content') + assert key.endswith('.content') + + def test_with_handles_non_mapping(self): + meta = _with_handles('not-a-mapping', 'h/1.0', 42) + assert meta == {'overflow_handle': 'h/1.0', 'overflow_bytes': 42} + + def test_with_handles_content_only(self): + meta = _with_handles({'orig': 1}, None, 0, 'h/1.0.content') + assert meta == {'orig': 1, 'overflow_content_handle': 'h/1.0.content'} + + def test_head_tail_preview_under(self): + assert _head_tail_preview('short', 1000) == 'short' + + def test_head_tail_preview_over(self): + assert 'omitted' in _head_tail_preview('a' * 100, 10) + + def test_build_spill_preview_tokens_unit(self): + unit = _Unit(binary=False, text='x' * 100, data=b'x' * 100, value='x' * 100, suffix='') + assert 'tokens' in _build_spill_preview('h/1.0', unit, 20, over_tokens=True) + + +# --------------------------------------------------------------------------- +# read_tool_result / _read_slice (C2 bounds + literal pattern) +# --------------------------------------------------------------------------- + + +class TestReadBack: + async def test_read_slice_basic(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + await store.write('h/1.0', '\n'.join(f'line {i}' for i in range(50)).encode('utf-8')) + out = await _read_slice(store, 'h/1.0', offset=0, limit=3, from_end=False, pattern=None) + assert 'line 0' in out and 'line 2' in out and 'line 3' not in out + + async def test_read_slice_from_end(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + await store.write('h/1.0', '\n'.join(f'line {i}' for i in range(50)).encode('utf-8')) + out = await _read_slice(store, 'h/1.0', offset=0, limit=2, from_end=True, pattern=None) + assert 'line 49' in out and 'line 48' in out + + async def test_read_slice_literal_pattern(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + await store.write('h/1.0', b'apple\nbanana\navocado\ncherry') + out = await _read_slice(store, 'h/1.0', offset=0, limit=200, from_end=False, pattern='av') + assert 'avocado' in out and 'apple' not in out and 'banana' not in out + + async def test_read_slice_pattern_is_literal_not_regex(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + await store.write('h/1.0', b'plain line\n^anchored') + # A regex metacharacter is matched literally, so it cannot trigger backtracking. + out = await _read_slice(store, 'h/1.0', offset=0, limit=200, from_end=False, pattern='^a') + assert 'anchored' in out and 'plain line' not in out + + async def test_read_slice_offset_negative(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + await store.write('h/1.0', b'data') + with pytest.raises(ModelRetry, match='offset'): + await _read_slice(store, 'h/1.0', offset=-1, limit=10, from_end=False, pattern=None) + + async def test_read_slice_limit_too_small(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + await store.write('h/1.0', b'data') + with pytest.raises(ModelRetry, match='limit'): + await _read_slice(store, 'h/1.0', offset=0, limit=0, from_end=False, pattern=None) + + async def test_read_slice_limit_clamped(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + await store.write('h/1.0', '\n'.join(f'l{i}' for i in range(2000)).encode('utf-8')) + out = await _read_slice(store, 'h/1.0', offset=0, limit=10_000, from_end=False, pattern=None) + assert out.count('\n') <= 1_000 # clamped to the line cap + + async def test_read_slice_output_capped(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + await store.write('h/1.0', ('x' * 60_000).encode('utf-8')) + out = await _read_slice(store, 'h/1.0', offset=0, limit=10, from_end=False, pattern=None) + assert 'output capped' in out + assert len(out) < 60_000 + + async def test_read_slice_missing_handle(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + with pytest.raises(ModelRetry, match='No stored tool result'): + await _read_slice(store, 'missing/1.0', offset=0, limit=10, from_end=False, pattern=None) + + async def test_get_toolset_registers_read_tool(self, tmp_path: Path): + store = LocalFileStore(base_dir=tmp_path) + await store.write('h/1.0', b'hello\nworld') + cap: OverflowingToolOutput[None] = OverflowingToolOutput(store=store) + toolset = cap.get_toolset() + assert toolset is not None + tool = toolset.tools[READ_TOOL_NAME] # type: ignore[union-attr] + out = await tool.function(_make_ctx(), 'h/1.0') # type: ignore[attr-defined] + assert 'hello' in out + + +# --------------------------------------------------------------------------- +# Agent-path integration +# --------------------------------------------------------------------------- + + +class TestAgentIntegration: + async def test_spill_persists_in_history(self, tmp_path: Path, anyio_backend: str): + store = LocalFileStore(base_dir=tmp_path) + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=100, action=Spill(preview_chars=50))], store=store + ) + agent = Agent(TestModel(call_tools=['big_tool']), capabilities=[cap]) + + @agent.tool_plain + def big_tool() -> str: + return 'data line\n' * 500 + + result = await agent.run('go') + returns = [p for m in result.all_messages() for p in m.parts if isinstance(p, ToolReturnPart)] + spilled = [p for p in returns if p.tool_name == 'big_tool'] + assert spilled + part = spilled[0] + assert isinstance(part.content, str) and 'too large' in part.content + assert part.metadata is not None and 'overflow_handle' in part.metadata + assert await store.read(part.metadata['overflow_handle']) == ('data line\n' * 500).encode('utf-8') + + async def test_small_output_untouched(self, tmp_path: Path, anyio_backend: str): + cap: OverflowingToolOutput[None] = OverflowingToolOutput( + bands=[Band(over=10_000, action=Spill())], store=LocalFileStore(base_dir=tmp_path) + ) + agent = Agent(TestModel(call_tools=['small_tool']), capabilities=[cap]) + + @agent.tool_plain + def small_tool() -> str: + return 'tiny' + + result = await agent.run('go') + returns = [p for m in result.all_messages() for p in m.parts if isinstance(p, ToolReturnPart)] + small = [p for p in returns if p.tool_name == 'small_tool'] + assert small and small[0].content == 'tiny'