diff --git a/env.example b/env.example index 191a8ab9..2159bc6f 100644 --- a/env.example +++ b/env.example @@ -10,6 +10,10 @@ # BUB_MAX_STEPS=50 # BUB_MAX_TOKENS=16384 # BUB_MODEL_TIMEOUT_SECONDS=300 +# Estimated tokens (4 chars each) above which a tool result is spilled to the spill tape. 0 disables. +# BUB_TOOL_SPILL_THRESHOLD=4096 +# Hard cap on the serialized model request body in bytes; oversized tool messages are clamped. 0 disables. +# BUB_MAX_REQUEST_BYTES=262144 # BUB_HOME=~/.bub # --------------------------------------------------------------------------- diff --git a/src/bub/builtin/__init__.py b/src/bub/builtin/__init__.py index e69de29b..47c1844d 100644 --- a/src/bub/builtin/__init__.py +++ b/src/bub/builtin/__init__.py @@ -0,0 +1 @@ +"""Bub builtin runtime package.""" diff --git a/src/bub/builtin/agent.py b/src/bub/builtin/agent.py index 8406d111..c7baa5d0 100644 --- a/src/bub/builtin/agent.py +++ b/src/bub/builtin/agent.py @@ -102,6 +102,11 @@ async def run_stream( tape = self.tape.session_tape( session_id, workspace_from_state(state), context=replace(self.tape.context, state=state) ) + tape_store = self.framework.get_tape_store() + if tape_store is not None: + if not is_async_tape_store(tape_store): + tape_store = AsyncTapeStoreAdapter(tape_store) + state.setdefault("_runtime_spill_store", tape_store) merge_back = not session_id.startswith("temp/") stack = AsyncExitStack() # The fork_tape context manager must not be exited until the last chunk of the stream is consumed. diff --git a/src/bub/builtin/model_runner.py b/src/bub/builtin/model_runner.py index aae8e1c0..36b24f09 100644 --- a/src/bub/builtin/model_runner.py +++ b/src/bub/builtin/model_runner.py @@ -178,7 +178,7 @@ async def fire_after(error: Exception | None = None) -> None: async with asyncio.timeout(self.settings.model_timeout_seconds): completion = await self.completion_response( model=request.model, - messages=list(request.messages), + messages=self._clamp_oversized_messages(list(request.messages)), tools=tools, max_tokens=request.max_tokens, reasoning_effort=tape.context.state.get("reasoning_effort"), @@ -241,6 +241,51 @@ async def fire_after(error: Exception | None = None) -> None: def generate_run_id() -> str: return f"run-{datetime.now(UTC).strftime('%Y%m%dT%H%M%S%fZ')}" + def _clamp_oversized_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Clamp tool messages so the serialized request body stays under a hard cap. + + This is the last-resort fuse against provider/reverse-proxy ``413``: it + shrinks the largest oversized ``role: tool`` messages with a head+tail + clamp and an explicit marker, so even a plugin that bypasses spilling + can never send an unbounded request body. + """ + cap = self.settings.max_request_bytes + if cap <= 0 or not messages: + return messages + + import json as _json + + def size() -> int: + try: + return len(_json.dumps(messages, ensure_ascii=False, default=str)) + except TypeError: + return sum(len(str(m.get("content", ""))) for m in messages if isinstance(m, dict)) + + if size() <= cap: + return messages + + def clamp_tool_messages(messages: list[dict[str, Any]], budget: int) -> list[dict[str, Any]]: + """Head+tail clamp every oversized tool message to ``budget`` chars.""" + result: list[dict[str, Any]] = [] + for message in messages: + content = message.get("content") + if isinstance(content, str) and message.get("role") == "tool" and len(content) > budget: + head = content[: budget // 2] + tail = content[-(budget - budget // 2) :] + message = dict(message) + message["content"] = ( + f"{head}\n\n[clamped: {len(content) - budget:,} chars removed to keep the request body bounded]\n\n{tail}" + ) + result.append(message) + return result + + clamped = clamp_tool_messages(messages, budget=2000) + if size() <= cap: + return clamped + + # Still over: shrink every tool message to a per-message budget within the cap. + return clamp_tool_messages(clamped, budget=cap // max(1, len(clamped))) + async def _fire_after_llm_call( self, request: LlmCallRequest, diff --git a/src/bub/builtin/settings.py b/src/bub/builtin/settings.py index b440d864..8095bdd4 100644 --- a/src/bub/builtin/settings.py +++ b/src/bub/builtin/settings.py @@ -59,6 +59,14 @@ class AgentSettings(Settings): max_steps: int = 50 max_tokens: int = DEFAULT_MAX_TOKENS model_timeout_seconds: int | None = None + tool_spill_threshold: int = Field( + 4096, + description="Estimated tokens (4 chars each) above which a tool result is spilled to the spill tape. 0 disables spilling.", + ) + max_request_bytes: int = Field( + 262_144, + description="Hard cap on the serialized model request body in bytes; oversized tool messages are clamped before sending. 0 disables.", + ) client_args: dict[str, Any] = Field(default_factory=dict) completion_args: dict[str, Any] = Field(default_factory=dict) verbose: int = Field(default=0, description="Verbosity level for logging. Higher means more verbose.", ge=0, le=2) diff --git a/src/bub/builtin/spill.py b/src/bub/builtin/spill.py new file mode 100644 index 00000000..6cd19cb6 --- /dev/null +++ b/src/bub/builtin/spill.py @@ -0,0 +1,150 @@ +"""Spill oversized tool outputs into a dedicated spill tape. + +Large tool results are written once to a ``spill`` tape (the same store as the +session tapes) and the model-facing result is replaced with a short ref: a +handle, a shape sketch, and a bounded preview. The full payload stays in the +tape store — queryable, replayable, and deletable by the user — without ever +entering a model request. + +The spill tape is intentionally shared across sessions and has no built-in +cleanup: the user owns retention, exactly like the session tapes themselves. +""" + +from __future__ import annotations + +import uuid +from typing import Any, Protocol + +from bub.tape import AsyncTapeStore, TapeEntry, TapeQuery + +SPILL_TAPE = "spill" +"""Name of the tape that stores full tool outputs.""" + +READ_TOOL_RESULT_NAME = "read_tool_result" +"""Name of the bounded read-back tool. Its own returns are never spilled.""" + +PREVIEW_CHARS = 600 +"""Characters of head+tail preview kept inline in a spill ref.""" + +MAX_READ_LINES = 1000 +"""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.""" + + +class SpillStore(Protocol): + """The only capability spilling needs: append entries.""" + + async def append(self, tape: str, entry: TapeEntry) -> None: ... + + +def needs_spill(output: str, *, threshold: int) -> bool: + """Decide whether one stringified tool result should be spilled. + + ``threshold`` is measured in estimated tokens (4 chars per token), matching + the heuristic used by pydantic-ai-harness. ``0`` disables spilling. + """ + if threshold <= 0: + return False + return len(output) // 4 >= threshold + + +def handle_key(run_id: str, tool: str) -> str: + """Build a unique handle for one spill (the spill tape's lookup key).""" + return f"{run_id}/{tool}.{uuid.uuid4().hex[:8]}" + + +def spill_ref(handle: str, output: str) -> str: + """Build the model-visible ref that replaces an oversized tool result.""" + lines = output.splitlines() + total = len(output) + body = _head_tail_preview(output) + return ( + f"[tool output spilled: {total:,} chars in {len(lines):,} lines; " + f"handle: {handle}]\n" + f"[read it back: read_tool_result(handle={handle!r}, offset=0, limit=200, " + f"from_end=False, pattern=None)]\n" + f"{body}" + ) + + +def _head_tail_preview(text: str, preview_chars: int = PREVIEW_CHARS) -> str: + if len(text) <= preview_chars: + return text + head = preview_chars // 2 + tail = preview_chars - head + omitted = len(text) - head - tail + return f"{text[:head]}\n...[{omitted:,} chars omitted]...\n{text[-tail:]}" + + +async def maybe_spill( + *, + tool: str, + run_id: str | None, + result: Any, + store: SpillStore | None, +) -> Any: + """Rewrite an oversized string tool result into a spill ref; no-op otherwise. + + Every "can't spill" case — non-string result, the read-back tool itself, a + missing store, spilling disabled, or a failed write — keeps the original + result. Errors are never spilled (the model needs full error text to + recover) and spilling must never fail a turn. + """ + if not isinstance(result, str): + return result + if tool == READ_TOOL_RESULT_NAME: + return result + if store is None: + return result + + from bub.builtin.settings import load_settings + + threshold = load_settings().tool_spill_threshold + if not needs_spill(result, threshold=threshold): + return result + + handle = handle_key(run_id or "run", tool) + entry = TapeEntry.tool_result([result], spill_handle=handle) + try: + await store.append(SPILL_TAPE, entry) + except Exception: + return result + return spill_ref(handle, result) + + +async def read_spilled(*, store: AsyncTapeStore, handle: str) -> str | None: + """Return the full spilled payload for ``handle``, or None when unknown.""" + key = handle.strip().lstrip("/") + query = TapeQuery(tape=SPILL_TAPE, store=store).kinds("tool_result") + entries = await store.fetch_all(query) + for entry in reversed(list(entries)): # newest wins; handles are unique per spill + if entry.meta.get("spill_handle") == key: + payload = entry.payload.get("results") + if isinstance(payload, list) and payload: + value = payload[0] + return value if isinstance(value, str) else str(value) + return None + + +def read_slice(output: str, *, offset: int, limit: int, from_end: bool, pattern: str | None) -> str: + """Slice a spilled payload with hard bounds; ``pattern`` is a literal substring.""" + lines = output.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: {total:,} matching line(s); showing {len(window)}{capped}]" + return f"{header}\n{body}" if body else header diff --git a/src/bub/builtin/store.py b/src/bub/builtin/store.py index c5f33adf..7039fc30 100644 --- a/src/bub/builtin/store.py +++ b/src/bub/builtin/store.py @@ -82,7 +82,10 @@ def _redact_payload(payload: dict) -> None: async def append(self, tape: str, entry: TapeEntry) -> None: self._redact_payload(entry.payload) - self._store.append(tape, entry) + if tape == self._tape: + self._store.append(tape, entry) + return + await self._parent.append(tape, entry) async def merge_back(self) -> None: if self._tape_was_reset: diff --git a/src/bub/builtin/tools.py b/src/bub/builtin/tools.py index 3e4cb916..d7d71109 100644 --- a/src/bub/builtin/tools.py +++ b/src/bub/builtin/tools.py @@ -10,6 +10,11 @@ from pydantic import BaseModel, Field from bub.builtin.shell_manager import shell_manager +from bub.builtin.spill import ( + MAX_READ_LINES, + read_slice, + read_spilled, +) from bub.skills import discover_skills from bub.tools import REGISTRY, Tool, ToolContext, tool @@ -195,6 +200,44 @@ async def kill_bash(shell_id: str) -> str: return f"id: {shell.shell_id}\nstatus: {shell.status}\nexit_code: {shell.returncode}" +@tool(context=True, name="read_tool_result") +async def read_tool_result( + handle: str, + offset: int = 0, + limit: int = 200, + from_end: bool = False, + pattern: str | None = None, + *, + context: ToolContext, +) -> str: + """Read a bounded slice of a spilled tool result. + + Args: + handle: The handle from a `[tool output spilled ...]` ref. + offset: Number of lines to skip from the start (or end when `from_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. + """ + if offset < 0: + return "`offset` must be >= 0." + if limit < 1: + return "`limit` must be >= 1." + limit = min(limit, MAX_READ_LINES) + + store = context.state.get("_runtime_spill_store") + if store is None: + return "spill store unavailable in this context." + output = await read_spilled(store=store, handle=handle) + if output is None: + return ( + f"[No stored tool result for handle {handle!r}. Use the exact handle from a " + '"[tool output spilled ...]" marker; if the result is no longer available, ' + "re-run the original tool.]" + ) + return read_slice(output, offset=offset, limit=limit, from_end=from_end, pattern=pattern) + + @tool(context=True, name="fs.read") def fs_read(path: str, offset: int = 0, limit: int | None = None, *, context: ToolContext) -> str: """Read a text file and return its content. Supports optional pagination with offset and limit.""" diff --git a/src/bub/tools.py b/src/bub/tools.py index 9b3bc01e..b21d8e0e 100644 --- a/src/bub/tools.py +++ b/src/bub/tools.py @@ -13,6 +13,7 @@ from loguru import logger from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, validate_call +from bub.builtin.spill import maybe_spill from bub.builtin.tape import Tape from bub.errors import BubError, ErrorKind from bub.hooks.interception import ToolCall, ToolCallResult @@ -256,6 +257,13 @@ async def _handle_tool_response_async( raise else: await self._fire_after_tool_call(call, hook_state, started, result=result) + if context is not None and isinstance(result, str): + result = await maybe_spill( + tool=call.tool, + run_id=context.run_id, + result=result, + store=context.state.get("_runtime_spill_store"), + ) return result async def _invoke_normalized(self, tool_obj: Tool, call: ToolCall, context: ToolContext | None) -> Any: diff --git a/tests/test_spill.py b/tests/test_spill.py new file mode 100644 index 00000000..680fe73f --- /dev/null +++ b/tests/test_spill.py @@ -0,0 +1,247 @@ +"""Behavior and regression tests for tool-output spilling. + +These are acceptance tests for what a user actually observes: +- an oversized tool result never re-enters a model request in full (no 413), +- the full payload stays reachable through the spill tape and read_tool_result, +- small results and tool errors are untouched, +- a spill write failure degrades to the original result (never a failed turn). +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +import bub.builtin.tools as builtin_tools +from bub.builtin.model_runner import ModelRunner +from bub.builtin.settings import AgentSettings +from bub.builtin.spill import ( + MAX_READ_LINES, + SPILL_TAPE, + handle_key, + read_slice, + read_spilled, + spill_ref, +) +from bub.builtin.store import ForkTapeStore +from bub.builtin.tape import Tape +from bub.tape import AsyncTapeStoreAdapter, InMemoryTapeStore, TapeContext +from bub.tools import Tool, ToolContext, ToolExecutor + + +def _make_context(tmp_path: Path, store: Any) -> ToolContext: + tape = Tape(tmp_path, AsyncTapeStoreAdapter(InMemoryTapeStore()), TapeContext()).scoped("test-tape") + return ToolContext( + tape=tape, run_id="run-1", state={"_runtime_workspace": str(tmp_path), "_runtime_spill_store": store} + ) + + +@pytest.fixture +def spill_store() -> AsyncTapeStoreAdapter: + return AsyncTapeStoreAdapter(InMemoryTapeStore()) + + +@pytest.mark.asyncio +async def test_oversized_tool_result_is_spilled_and_readable(tmp_path: Path, spill_store: Any) -> None: + """An oversized bash result becomes a ref in the result, and the full payload reads back.""" + big = "line-%04d\n" * 3000 # ~30k chars -> well over the 4096-token threshold + executor = ToolExecutor() + context = _make_context(tmp_path, spill_store) + + tool_obj = Tool(name="bash", handler=lambda cmd: big, description="", parameters={}) + execution = await executor.execute_async([(tool_obj, {"cmd": "x"})], context=context) + + result = execution.tool_results[0] + assert isinstance(result, str) + assert "tool output spilled" in result + assert "handle:" in result + + handle = result.split("handle: ")[1].split("]")[0].strip() + full = await read_spilled(store=spill_store, handle=handle) + assert full == big + + +@pytest.mark.asyncio +async def test_spill_ref_is_small_and_self_describing(tmp_path: Path, spill_store: Any) -> None: + """The ref a model sees is a short marker with handle, shape, and preview — not the payload.""" + big = "x" * 200_000 + executor = ToolExecutor() + context = _make_context(tmp_path, spill_store) + tool_obj = Tool(name="bash", handler=lambda cmd: big, description="", parameters={}) + + execution = await executor.execute_async([(tool_obj, {"cmd": "x"})], context=context) + + ref = execution.tool_results[0] + assert "200,000 chars" in ref + assert "read_tool_result(" in ref + assert len(ref) < 2_000 # the ref itself is tiny + + +@pytest.mark.asyncio +async def test_small_tool_result_is_untouched(tmp_path: Path, spill_store: Any) -> None: + context = _make_context(tmp_path, spill_store) + executor = ToolExecutor() + tool_obj = Tool(name="bash", handler=lambda cmd: "tiny", description="", parameters={}) + + execution = await executor.execute_async([(tool_obj, {"cmd": "x"})], context=context) + + assert execution.tool_results == ["tiny"] + + +@pytest.mark.asyncio +async def test_tool_error_is_never_spilled(tmp_path: Path, spill_store: Any) -> None: + def boom(cmd: str) -> str: + raise ValueError("boom") + + executor = ToolExecutor() + context = _make_context(tmp_path, spill_store) + tool_obj = Tool(name="bash", handler=boom, description="", parameters={}) + + execution = await executor.execute_async([(tool_obj, {"cmd": "x"})], context=context) + + assert execution.error is not None + assert execution.error.details["error"] == "ValueError('boom')" + assert not any(entry.kind == "tool_result" for entry in (spill_store._store.read(SPILL_TAPE) or [])) + + +@pytest.mark.asyncio +async def test_spill_write_failure_keeps_original_result(tmp_path: Path) -> None: + """A failing spill store degrades to the original result — never a failed turn.""" + + class BrokenStore: + async def append(self, tape: str, entry: Any) -> None: + raise OSError("disk full") + + async def fetch_all(self, query: Any) -> Any: + return [] + + executor = ToolExecutor() + context = _make_context(tmp_path, BrokenStore()) + tool_obj = Tool(name="bash", handler=lambda cmd: "x" * 100_000, description="", parameters={}) + + execution = await executor.execute_async([(tool_obj, {"cmd": "x"})], context=context) + + assert execution.error is None + assert execution.tool_results == ["x" * 100_000] + + +@pytest.mark.asyncio +async def test_spill_goes_to_spill_tape_through_forked_session_tape(tmp_path: Path, spill_store: Any) -> None: + """Spill entries bypass the session-tape fork and land in the shared spill tape immediately.""" + parent = spill_store._store + fork = ForkTapeStore(spill_store, "session-tape") + executor = ToolExecutor() + context = _make_context(tmp_path, fork) + tool_obj = Tool(name="bash", handler=lambda cmd: "y" * 100_000, description="", parameters={}) + + await executor.execute_async([(tool_obj, {"cmd": "x"})], context=context) + + entries = list(parent.read(SPILL_TAPE) or []) + assert len(entries) == 1 + assert entries[0].payload["results"] == ["y" * 100_000] + assert entries[0].meta["spill_handle"] + + +@pytest.mark.asyncio +async def test_read_spilled_unknown_handle_returns_none(spill_store: Any) -> None: + assert await read_spilled(store=spill_store, handle="run-1/bash.deadbeef") is None + + +@pytest.mark.asyncio +async def test_read_tool_result_tool_is_bounded_and_literal(tmp_path: Path, spill_store: Any) -> None: + """read_tool_result enforces bounds and treats pattern as a literal substring.""" + big = "\n".join(f"line-{i}" for i in range(3000)) + handle = handle_key("run-1", "bash") + await spill_store.append( + SPILL_TAPE, __import__("bub.tape", fromlist=["TapeEntry"]).TapeEntry.tool_result([big], spill_handle=handle) + ) + + context = _make_context(tmp_path, spill_store) + result = await builtin_tools.read_tool_result.run( + handle=handle, offset=0, limit=3, from_end=False, pattern=None, context=context + ) + assert result.startswith("[handle: 3,000 matching line(s); showing 3]") + assert "line-0" in result + assert "line-2" in result + + tail = await builtin_tools.read_tool_result.run( + handle=handle, offset=0, limit=3, from_end=True, pattern=None, context=context + ) + assert "line-2999" in tail + + literal = await builtin_tools.read_tool_result.run( + handle=handle, offset=0, limit=2000, from_end=False, pattern="line-1999", context=context + ) + assert "line-1999" in literal + assert "line-1998" not in literal # literal substring, not a prefix match + + too_many = await builtin_tools.read_tool_result.run( + handle=handle, offset=0, limit=MAX_READ_LINES + 100, from_end=False, pattern=None, context=context + ) + assert "showing 1000" in too_many # limit clamped + + +@pytest.mark.asyncio +async def test_read_tool_result_unknown_handle_is_friendly(tmp_path: Path, spill_store: Any) -> None: + context = _make_context(tmp_path, spill_store) + result = await builtin_tools.read_tool_result.run( + handle="run-1/bash.nope", offset=0, limit=5, from_end=False, pattern=None, context=context + ) + assert "No stored tool result" in result + assert "re-run the original tool" in result + + +def test_read_slice_bounds_and_literal_pattern() -> None: + output = "\n".join(f"line-{i}" for i in range(100)) + window = read_slice(output, offset=10, limit=5, from_end=False, pattern=None) + assert "line-10" in window and "line-14" in window + assert "line-9" not in window + + tail = read_slice(output, offset=0, limit=5, from_end=True, pattern=None) + assert "line-99" in tail and "line-95" in tail + + filtered = read_slice( + "\n".join(["foo", "bar-baz", "qux", "bar"]), offset=0, limit=10, from_end=False, pattern="bar" + ) + assert filtered.count("bar") == 2 # literal substring matches every line containing it + regexish = read_slice("\n".join(["bar", "b.r"]), offset=0, limit=10, from_end=False, pattern="b.r") + assert regexish.count("b.r") == 1 # pattern is literal, not a regex + + +def test_spill_ref_is_self_describing() -> None: + ref = spill_ref("run-1/bash.abc", "a\nb\nc\n") + assert "handle: run-1/bash.abc" in ref + assert "read_tool_result(handle=" in ref + + +@pytest.mark.asyncio +async def test_next_model_request_never_contains_full_payload(tmp_path: Path, spill_store: Any) -> None: + """Regression: after a spill, the serialized request body sent to the model stays bounded.""" + big = "x" * 5_000_000 # multi-MB single-line output (the 413 scenario) + executor = ToolExecutor() + context = _make_context(tmp_path, spill_store) + tool_obj = Tool(name="bash", handler=lambda cmd: big, description="", parameters={}) + + execution = await executor.execute_async([(tool_obj, {"cmd": "grep -R foo"})], context=context) + ref = execution.tool_results[0] + assert "tool output spilled" in ref + + # What the next request would carry: the tool_result entry the model sees. + messages = [{"role": "tool", "content": ref}] + body = len(json.dumps(messages, ensure_ascii=False)) + assert body < 100_000 # orders of magnitude below any 413 threshold + assert "x" * 1000 not in ref # the raw payload is not inline + + +@pytest.mark.asyncio +async def test_hard_request_cap_clamps_oversized_messages() -> None: + """Regression: even a bypassed spill (huge inline tool message) never sends an unbounded body.""" + runner = ModelRunner(AgentSettings.model_construct(model="openai:gpt-test", max_request_bytes=2048, max_tokens=100)) + messages = [{"role": "tool", "content": "z" * 100_000}] + clamped = runner._clamp_oversized_messages(messages) + body = json.dumps(clamped, ensure_ascii=False) + assert len(body) < 4096 + assert "clamped" in clamped[0]["content"]