feat: add context budget plugin

README.md

@@ -16,6 +16,7 @@ Below is the list of packages currently included in this repository.
 | Package                                                                     | PyPI Status                                                                                               | Description                                                                                                                            |
 | --------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------- |
 | [`bub-codex`](./packages/bub-codex/README.md)                               |                                                                                                           | Provides a `run_model` hook that delegates model execution to the Codex CLI.                                                           |
+| [`bub-context-budget`](./packages/bub-context-budget/README.md)             |                                                                                                           | Estimates input context tokens before LLM calls and prompts the model to hand off when the configured budget is exceeded.              |
 | [`bub-cursor`](./packages/bub-cursor/README.md)                             |                                                                                                           | Provides a `run_model` hook that delegates model execution to the Cursor CLI, plus `bub login cursor`.                                 |
 | [`bub-acp-server`](./packages/bub-acp-server/README.md)                     | [![PyPI version](https://img.shields.io/pypi/v/bub-acp-server)](https://pypi.org/project/bub-acp-server/) | Exposes Bub as an Agent Client Protocol agent with `bub acp serve` for ACP-compatible editors.                                         |
 | [`bub-schedule`](./packages/bub-schedule/README.md)                         | [![PyPI version](https://img.shields.io/pypi/v/bub-schedule)](https://pypi.org/project/bub-schedule/)     | Provides scheduling channel/tools backed by APScheduler with a JSON job store.                                                         |

packages/bub-context-budget/README.md

@@ -0,0 +1,53 @@
+# bub-context-budget
+
+Proactive context-budget guard for Bub's builtin agent loop.
+
+## What It Provides
+
+- Bub plugin entry point: `context-budget`
+- A `before_llm_call` hook that runs after other request modifiers where possible
+- Provider-reported token usage combined with an incremental message estimate
+- A configurable `tiktoken` fallback when no matching usage baseline exists
+- A fixed 16,384-token reserve below the configured context limit
+- A system instruction asking the model to call `tape.handoff` before continuing when the budget is exceeded
+
+## Installation
+
+```bash
+uv pip install "git+https://github.com/bubbuild/bub-contrib.git#subdirectory=packages/bub-context-budget"
+```
+
+You can also install it with Bub:
+
+```bash
+bub install bub-context-budget@main
+```
+
+## Configuration
+
+Add the plugin section to the Bub config file:
+
+```yaml
+context-budget:
+  max_context_tokens: 200000
+  encoding_name: o200k_base
+```
+
+The defaults are equivalent to the example above. Environment variables are also supported:
+
+- `BUB_CONTEXT_BUDGET_MAX_CONTEXT_TOKENS`
+- `BUB_CONTEXT_BUDGET_ENCODING_NAME`
+
+Use an encoding that approximates the configured model's tokenizer. `o200k_base` is the default because Bub model identifiers may refer to providers that `tiktoken.encoding_for_model` cannot resolve.
+
+## Runtime Behavior
+
+After each successful builtin agent-loop LLM call, the plugin records the provider-reported total usage and a digest of the messages covered by that usage. On the next call, a matching model, tool set, and message prefix uses that real total plus a `tiktoken` estimate of only the newly appended messages. This follows Pi's usage-baseline approach without retaining a second copy of the conversation in memory.
+
+If usage is unavailable, the model or tool set changed, or the message prefix no longer matches, the plugin estimates the full message list. This fallback count covers the serialized messages but not provider-added tool schemas or protocol framing.
+
+The handoff threshold is `max_context_tokens - 16384`, leaving fixed headroom for the injected instruction, provider framing, and the handoff response. With the default 200,000-token limit, the threshold is 183,616 tokens. A count equal to the threshold is allowed. A count above it adds a high-priority, idempotent system instruction telling the model to call `tape.handoff` with a concise continuity summary before doing any other work.
+
+Bub normalizes streaming and non-streaming completion usage into `LlmCallResult.usage`; the same data backs `StreamState` and `AsyncStreamEvents` and is later persisted to tape. The plugin consumes the normalized hook result directly and stores its lightweight baseline in the current `TurnState`. A fresh state falls back to full estimation until its first successful LLM call establishes a baseline.
+
+The plugin only modifies the outgoing request. It does not call `tape.handoff` itself, and it requires that tool to be available to the model. Alternate `run_model` providers that bypass Bub's builtin agent-loop interception hooks are outside its scope.

packages/bub-context-budget/pyproject.toml

@@ -0,0 +1,26 @@
+[project]
+name = "bub-context-budget"
+version = "0.1.0"
+description = "Proactive context-budget guard for Bub LLM calls"
+readme = "README.md"
+authors = [
+    { name = "Frost Ming", email = "me@frostming.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "pydantic>=2.0.0",
+    "pydantic-settings>=2.10.1",
+    "tiktoken>=0.13.0,<0.14.0",
+]
+
+[project.entry-points.bub]
+context-budget = "bub_context_budget.plugin"
+
+[build-system]
+requires = ["uv_build>=0.9.7,<0.10.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "pytest>=9.0.3",
+]

packages/bub-context-budget/src/bub_context_budget/__init__.py

@@ -0,0 +1 @@
+"""Proactive context-budget guard for Bub."""

packages/bub-context-budget/src/bub_context_budget/plugin.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass, replace
+from functools import cache
+from typing import Any
+
+import bub
+import tiktoken
+from bub import hookimpl
+from bub.hooks.interception import LlmCallRequest, LlmCallResult
+from bub.turn import TurnState
+from pydantic import Field
+from pydantic_settings import SettingsConfigDict
+
+CONFIG_NAME = "context-budget"
+DEFAULT_MAX_CONTEXT_TOKENS = 200_000
+DEFAULT_ENCODING_NAME = "o200k_base"
+HANDOFF_TOOL_NAME = "tape.handoff"
+HANDOFF_INSTRUCTION_MARKER = "<context_budget_exceeded>"
+RESERVE_TOKENS = 16_384
+STATE_BASELINE_KEY = "_context_budget_usage_baseline"
+
+
+@dataclass(frozen=True)
+class _UsageBaseline:
+    model: str
+    tool_names: tuple[str, ...]
+    message_count: int
+    message_digest: str
+    total_tokens: int
+
+
+@bub.config(name=CONFIG_NAME)
+class ContextBudgetSettings(bub.Settings):
+    model_config = SettingsConfigDict(env_prefix="BUB_CONTEXT_BUDGET_", extra="ignore")
+    max_context_tokens: int = Field(default=DEFAULT_MAX_CONTEXT_TOKENS, gt=0)
+    encoding_name: str = DEFAULT_ENCODING_NAME
+
+
+@cache
+def _get_encoding(name: str) -> tiktoken.Encoding:
+    return tiktoken.get_encoding(name)
+
+
+def estimate_context_tokens(
+    messages: list[dict[str, Any]],
+    *,
+    encoding_name: str,
+) -> int:
+    if not messages:
+        return 0
+    encoding = _get_encoding(encoding_name)
+    return len(encoding.encode(_serialize_messages(messages), disallowed_special=()))
+
+
+def _serialize_messages(messages: list[dict[str, Any]]) -> str:
+    return json.dumps(
+        messages,
+        ensure_ascii=False,
+        separators=(",", ":"),
+        sort_keys=True,
+        default=str,
+    )
+
+
+def _message_digest(messages: list[dict[str, Any]]) -> str:
+    return hashlib.sha256(_serialize_messages(messages).encode()).hexdigest()
+
+
+def _token_count(value: object) -> int | None:
+    if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
+        return value
+    return None
+
+
+def _usage_total_tokens(usage: Mapping[str, Any] | None) -> int | None:
+    if usage is None:
+        return None
+    total = _token_count(usage.get("total_tokens"))
+    if total is not None:
+        return total
+
+    prompt = _token_count(usage.get("prompt_tokens"))
+    if prompt is None:
+        prompt = _token_count(usage.get("input_tokens"))
+    completion = _token_count(usage.get("completion_tokens"))
+    if completion is None:
+        completion = _token_count(usage.get("output_tokens"))
+    if prompt is None and completion is None:
+        return None
+    return (prompt or 0) + (completion or 0)
+
+
+def _assistant_message(result: LlmCallResult) -> dict[str, Any]:
+    if result.tool_calls:
+        return {"role": "assistant", "content": "", "tool_calls": result.tool_calls}
+    return {"role": "assistant", "content": result.text or ""}
+
+
+def _estimate_request_tokens(
+    request: LlmCallRequest,
+    state: TurnState,
+    *,
+    encoding_name: str,
+) -> int:
+    baseline = state.get(STATE_BASELINE_KEY)
+    if (
+        not isinstance(baseline, _UsageBaseline)
+        or baseline.model != request.model
+        or baseline.tool_names != request.tool_names
+        or len(request.messages) < baseline.message_count
+        or _message_digest(request.messages[: baseline.message_count])
+        != baseline.message_digest
+    ):
+        return estimate_context_tokens(
+            request.messages,
+            encoding_name=encoding_name,
+        )
+
+    trailing_messages = request.messages[baseline.message_count :]
+    return baseline.total_tokens + estimate_context_tokens(
+        trailing_messages,
+        encoding_name=encoding_name,
+    )
+
+
+def _handoff_instruction(
+    *,
+    estimated_tokens: int,
+    threshold: int,
+    limit: int,
+) -> str:
+    return (
+        f"{HANDOFF_INSTRUCTION_MARKER}\n"
+        f"The estimated input context is {estimated_tokens} tokens, above the handoff "
+        f"threshold of {threshold} tokens. The configured context limit is {limit} tokens, "
+        f"with {RESERVE_TOKENS} tokens reserved. Before doing any other work, call the "
+        "`tape.handoff` "
+        "(`tape_handoff`) tool with name `context-budget` and a concise summary that "
+        "preserves the current goal, progress, key decisions, modified files, and remaining "
+        "work. Do not continue the task in this call.\n"
+        "</context_budget_exceeded>"
+    )
+
+
+def _inject_system_instruction(
+    messages: list[dict[str, Any]],
+    instruction: str,
+) -> list[dict[str, Any]]:
+    if any(
+        message.get("role") == "system"
+        and isinstance(message.get("content"), str)
+        and HANDOFF_INSTRUCTION_MARKER in message["content"]
+        for message in messages
+    ):
+        return messages
+
+    updated = [dict(message) for message in messages]
+    for index, message in enumerate(updated):
+        if message.get("role") != "system" or not isinstance(
+            message.get("content"), str
+        ):
+            continue
+        content = message["content"]
+        updated[index] = {
+            **message,
+            "content": f"{content}\n\n{instruction}" if content else instruction,
+        }
+        return updated
+
+    return [{"role": "system", "content": instruction}, *updated]
+
+
+@hookimpl(trylast=True)
+def before_llm_call(
+    request: LlmCallRequest,
+    state: TurnState,
+) -> LlmCallRequest | None:
+    settings = bub.ensure_config(ContextBudgetSettings)
+    estimated_tokens = _estimate_request_tokens(
+        request,
+        state,
+        encoding_name=settings.encoding_name,
+    )
+    threshold = max(0, settings.max_context_tokens - RESERVE_TOKENS)
+    if estimated_tokens <= threshold:
+        return None
+
+    instruction = _handoff_instruction(
+        estimated_tokens=estimated_tokens,
+        threshold=threshold,
+        limit=settings.max_context_tokens,
+    )
+    messages = _inject_system_instruction(request.messages, instruction)
+    if messages is request.messages:
+        return None
+    return replace(
+        request,
+        messages=messages,
+    )
+
+
+@hookimpl
+def after_llm_call(
+    request: LlmCallRequest,
+    result: LlmCallResult,
+    state: TurnState,
+) -> None:
+    total_tokens = _usage_total_tokens(result.usage)
+    if total_tokens is None or result.error is not None:
+        return
+
+    covered_messages = [*request.messages, _assistant_message(result)]
+    state[STATE_BASELINE_KEY] = _UsageBaseline(
+        model=request.model,
+        tool_names=request.tool_names,
+        message_count=len(covered_messages),
+        message_digest=_message_digest(covered_messages),
+        total_tokens=total_tokens,
+    )