Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions env.example
Original file line number Diff line number Diff line change
Expand Up @@ -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

# ---------------------------------------------------------------------------
Expand Down
1 change: 1 addition & 0 deletions src/bub/builtin/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
"""Bub builtin runtime package."""
5 changes: 5 additions & 0 deletions src/bub/builtin/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
47 changes: 46 additions & 1 deletion src/bub/builtin/model_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 8 additions & 0 deletions src/bub/builtin/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
150 changes: 150 additions & 0 deletions src/bub/builtin/spill.py
Original file line number Diff line number Diff line change
@@ -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
5 changes: 4 additions & 1 deletion src/bub/builtin/store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
43 changes: 43 additions & 0 deletions src/bub/builtin/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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."""
Expand Down
8 changes: 8 additions & 0 deletions src/bub/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
Loading
Loading