diff --git a/env.example b/env.example index 2159bc6f..191a8ab9 100644 --- a/env.example +++ b/env.example @@ -10,10 +10,6 @@ # 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 47c1844d..e69de29b 100644 --- a/src/bub/builtin/__init__.py +++ b/src/bub/builtin/__init__.py @@ -1 +0,0 @@ -"""Bub builtin runtime package.""" diff --git a/src/bub/builtin/agent.py b/src/bub/builtin/agent.py index c7baa5d0..8406d111 100644 --- a/src/bub/builtin/agent.py +++ b/src/bub/builtin/agent.py @@ -102,11 +102,6 @@ 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 36b24f09..aae8e1c0 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=self._clamp_oversized_messages(list(request.messages)), + messages=list(request.messages), tools=tools, max_tokens=request.max_tokens, reasoning_effort=tape.context.state.get("reasoning_effort"), @@ -241,51 +241,6 @@ 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 8095bdd4..b440d864 100644 --- a/src/bub/builtin/settings.py +++ b/src/bub/builtin/settings.py @@ -59,14 +59,6 @@ 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 deleted file mode 100644 index 6cd19cb6..00000000 --- a/src/bub/builtin/spill.py +++ /dev/null @@ -1,150 +0,0 @@ -"""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 7039fc30..c5f33adf 100644 --- a/src/bub/builtin/store.py +++ b/src/bub/builtin/store.py @@ -82,10 +82,7 @@ def _redact_payload(payload: dict) -> None: async def append(self, tape: str, entry: TapeEntry) -> None: self._redact_payload(entry.payload) - if tape == self._tape: - self._store.append(tape, entry) - return - await self._parent.append(tape, entry) + self._store.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 d7d71109..3e4cb916 100644 --- a/src/bub/builtin/tools.py +++ b/src/bub/builtin/tools.py @@ -10,11 +10,6 @@ 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 @@ -200,44 +195,6 @@ 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 b21d8e0e..9b3bc01e 100644 --- a/src/bub/tools.py +++ b/src/bub/tools.py @@ -13,7 +13,6 @@ 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 @@ -257,13 +256,6 @@ 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 deleted file mode 100644 index 680fe73f..00000000 --- a/tests/test_spill.py +++ /dev/null @@ -1,247 +0,0 @@ -"""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"]