feat: agent-loop hooks — before/after LLM & tool calls + wrap_tool_call

src/bub/agent_hooks.py

@@ -0,0 +1,125 @@
+"""Agent-loop interception payload contracts (issue #253).
+
+These dataclasses are the generic shapes passed to the ``before_llm_call`` /
+``after_llm_call`` / ``before_tool_call`` / ``after_tool_call`` hookspecs.
+They carry no chat or provider business beyond what the hooks need to
+observe or modify; the execution semantics live in
+:class:`bub.hook_runtime.AgentHooks`.
+
+``run_id`` is threaded through every payload so plugins can correlate hook
+observations with tape entries (which carry the same ``run_id`` meta).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+
+@dataclass(frozen=True)
+class LlmCallRequest:
+    """Outgoing agent-loop LLM request, as seen by ``before_llm_call``.
+
+    Hooks may return a modified copy (``dataclasses.replace``) to change the
+    model or messages for this call. Tool *objects* are not exposed —
+    ``tool_names`` is observational; altering the toolset is out of scope.
+    """
+
+    run_id: str
+    model: str
+    messages: list[dict[str, Any]]
+    tool_names: tuple[str, ...] = ()
+    max_tokens: int | None = None
+
+
+@dataclass(frozen=True)
+class LlmCallResult:
+    """Terminal outcome of one LLM call, as seen by ``after_llm_call``.
+
+    For streaming completions this is the fully-accumulated final state,
+    not a per-chunk view. ``error`` is the original raised exception (and
+    other fields best-effort) when the call failed. Cancellation and
+    consumer close are not observed: after hooks fire only for real
+    completions and ``Exception`` failures.
+    """
+
+    run_id: str
+    text: str | None = None
+    tool_calls: list[dict[str, Any]] = field(default_factory=list)
+    usage: dict[str, Any] | None = None
+    error: Exception | None = None
+    duration_ms: int = 0
+
+
+@dataclass(frozen=True)
+class LlmCallDecision:
+    """Short-circuit verdict returned by ``before_llm_call``.
+
+    ``finish`` skips the provider call entirely and emits ``text`` as the
+    final assistant output for this call — the cost-guard / call-limit
+    primitive (cf. LangChain ``ModelCallLimitMiddleware``).
+    """
+
+    action: Literal["finish"] = "finish"
+    text: str = ""
+
+    @classmethod
+    def finish(cls, text: str) -> LlmCallDecision:
+        return cls(action="finish", text=text)
+
+
+@dataclass(frozen=True)
+class ToolCall:
+    """One tool invocation about to run, as seen by ``before_tool_call``."""
+
+    run_id: str
+    tool: str
+    arguments: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class ToolCallDecision:
+    """Verdict returned by a ``before_tool_call`` implementation.
+
+    - ``proceed``: continue (optionally with modified ``arguments``);
+      later hook implementations still run and see the updated call.
+    - ``replace``: skip the tool handler entirely and use ``result``.
+    - ``deny``: skip the tool handler and surface ``message`` as a tool
+      error result.
+
+    ``replace`` and ``deny`` short-circuit any remaining implementations.
+    """
+
+    action: Literal["proceed", "replace", "deny"] = "proceed"
+    arguments: dict[str, Any] | None = None
+    result: Any = None
+    message: str | None = None
+
+    @classmethod
+    def proceed(cls, arguments: dict[str, Any] | None = None) -> ToolCallDecision:
+        return cls(action="proceed", arguments=arguments)
+
+    @classmethod
+    def replace(cls, result: Any) -> ToolCallDecision:
+        return cls(action="replace", result=result)
+
+    @classmethod
+    def deny(cls, message: str) -> ToolCallDecision:
+        return cls(action="deny", message=message)
+
+
+@dataclass(frozen=True)
+class ToolCallResult:
+    """Terminal outcome of one tool invocation, as seen by ``after_tool_call``.
+
+    ``error`` is the original ``BubError`` when the invocation raised or was
+    denied (kind/message/details preserved); ``result`` is unset in that
+    case. Cancellation is not observed by after hooks.
+    """
+
+    run_id: str
+    tool: str
+    arguments: dict[str, Any]
+    result: Any = None
+    error: Exception | None = None
+    duration_ms: int = 0

src/bub/builtin/agent.py

@@ -47,7 +47,7 @@ class Agent:
     def __init__(self, framework: BubFramework) -> None:
         self.settings = load_settings()
         self.framework = framework
-        self.model_runner = ModelRunner(self.settings)
+        self.model_runner = ModelRunner(self.settings, hooks=framework.get_agent_hooks())
 
     @cached_property
     def tape(self) -> Tape:

src/bub/builtin/model_runner.py

@@ -25,9 +25,11 @@
 from loguru import logger
 from pydantic import TypeAdapter, ValidationError
 
+from bub.agent_hooks import LlmCallDecision, LlmCallRequest, LlmCallResult
 from bub.builtin.codex_provider import OpenaiCodexProvider, should_use_openai_codex_provider
 from bub.builtin.settings import AgentSettings, ModelCandidate
 from bub.builtin.tape import Tape
+from bub.hook_runtime import AgentHooks
 from bub.runtime import AsyncStreamEvents, BubError, ErrorKind, StreamEvent, StreamState
 from bub.tools import Tool, ToolContext, ToolExecutor
 
@@ -53,8 +55,9 @@ def _stream_usage_options(llm: AnyLLM, *, stream: bool) -> dict[str, Any] | None
 
 
 class ModelRunner:
-    def __init__(self, settings: AgentSettings) -> None:
+    def __init__(self, settings: AgentSettings, hooks: AgentHooks | None = None) -> None:
         self.settings = settings
+        self.hooks = hooks
 
     def iter_llm_clients(self, model: str) -> Iterator[tuple[ModelCandidate, AnyLLM]]:
         for candidate in self.settings.model_candidates(model):
@@ -76,7 +79,7 @@ def create_llm_client(candidate: ModelCandidate, client_kwargs: dict[str, Any])
         return AnyLLM.create(candidate.provider, **client_kwargs)
 
     async def completion_response(
-        self, *, model: str, messages: list[dict[str, Any]], tools: list[Tool]
+        self, *, model: str, messages: list[dict[str, Any]], tools: list[Tool], max_tokens: int | None = None
     ) -> CompletionResult:
         from bub.builtin.tools import completion_tools
 
@@ -91,7 +94,7 @@ async def completion_response(
                     model=candidate.model_id,
                     messages=completion_messages,
                     tools=tool_payloads,
-                    max_tokens=self.settings.max_tokens,
+                    max_tokens=max_tokens if max_tokens is not None else self.settings.max_tokens,
                     stream=streaming,
                     stream_options=_stream_usage_options(llm, stream=streaming),
                 )
@@ -127,10 +130,57 @@ async def iterator() -> AsyncGenerator[StreamEvent, None]:
                 steering_messages=steering_messages,
             )
             output = ModelOutputAccumulator()
-            async with asyncio.timeout(self.settings.model_timeout_seconds):
-                completion = await self.completion_response(model=model, messages=messages, tools=tools)
-                async for event in self._completion_events(completion, state, output):
-                    yield event
+            request = LlmCallRequest(
+                run_id=run_id,
+                model=model,
+                messages=messages,
+                tool_names=tuple(tool_item.name for tool_item in tools),
+                max_tokens=self.settings.max_tokens,
+            )
+            decision: LlmCallDecision | None = None
+            if self.hooks is not None:
+                request, decision = await self.hooks.before_llm_call(request, state=tape.context.state)
+            if decision is not None:
+                await self.record_chat(
+                    tape=tape,
+                    run_id=run_id,
+                    system_prompt=system_prompt,
+                    new_messages=new_messages,
+                    response_text=decision.text,
+                    model=request.model,
+                )
+                yield StreamEvent("text", {"delta": decision.text})
+                yield StreamEvent("final", {"ok": True, "text": decision.text})
+                return
+            llm_started = datetime.now(UTC)
+            after_fired = False
+
+            async def fire_after(error: Exception | None = None) -> None:
+                """Fire after_llm_call once per completed call (success or Exception failure); cancellation/consumer close bypasses it."""
+
+                nonlocal after_fired
+                if after_fired:
+                    return
+                after_fired = True
+                await self._fire_after_llm_call(request, output, state, llm_started, tape, error=error)
+
+            try:
+                async with asyncio.timeout(self.settings.model_timeout_seconds):
+                    completion = await self.completion_response(
+                        model=request.model,
+                        messages=list(request.messages),
+                        tools=tools,
+                        max_tokens=request.max_tokens,
+                    )
+                    async for event in self._completion_events(completion, state, output):
+                        yield event
+            except Exception as exc:
+                # Cancellation / consumer close (BaseException) intentionally
+                # bypasses after_llm_call: only real completions and failures
+                # are terminal observations.
+                await fire_after(exc)
+                raise
+            await fire_after()
 
             tool_calls = output.tool_calls
             if tool_calls:
@@ -139,7 +189,7 @@ async def iterator() -> AsyncGenerator[StreamEvent, None]:
                 tool_invocations = [tool_invocation_from_native(tool_call, tool_map) for tool_call in tool_calls]
                 yield StreamEvent("tool_call", {"tool_calls": serialized_tool_calls})
                 context = ToolContext(tape=tape, run_id=run_id, state=tape.context.state)
-                execution = await ToolExecutor().execute_async(
+                execution = await ToolExecutor(hooks=self.hooks).execute_async(
                     tool_invocations,
                     context=context,
                 )
@@ -152,7 +202,7 @@ async def iterator() -> AsyncGenerator[StreamEvent, None]:
                     tool_calls=serialized_tool_calls,
                     tool_results=execution.tool_results,
                     response=output.response,
-                    model=model,
+                    model=request.model,
                     usage=state.usage,
                 )
                 yield StreamEvent("tool_result", {"tool_results": execution.tool_results})
@@ -169,7 +219,7 @@ async def iterator() -> AsyncGenerator[StreamEvent, None]:
                 new_messages=new_messages,
                 response_text=text,
                 response=output.response,
-                model=model,
+                model=request.model,
                 usage=state.usage,
             )
             yield StreamEvent("final", {"ok": True, "text": text})
@@ -180,6 +230,28 @@ async def iterator() -> AsyncGenerator[StreamEvent, None]:
     def generate_run_id() -> str:
         return f"run-{datetime.now(UTC).strftime('%Y%m%dT%H%M%S%fZ')}"
 
+    async def _fire_after_llm_call(
+        self,
+        request: LlmCallRequest,
+        output: ModelOutputAccumulator,
+        state: StreamState,
+        started: datetime,
+        tape: Tape,
+        error: Exception | None = None,
+    ) -> None:
+        if self.hooks is None:
+            return
+        duration_ms = int((datetime.now(UTC) - started).total_seconds() * 1000)
+        result = LlmCallResult(
+            run_id=request.run_id,
+            text=output.text or None,
+            tool_calls=[call.model_dump(exclude_none=True) for call in output.tool_calls],
+            usage=state.usage,
+            error=error,
+            duration_ms=duration_ms,
+        )
+        await self.hooks.after_llm_call(request, result, state=tape.context.state)
+
     async def build_messages(
         self,
         *,

src/bub/framework.py

@@ -15,7 +15,7 @@
 
 from bub import configure
 from bub.envelope import content_of, field_of, unpack_batch
-from bub.hook_runtime import _SKIP_VALUE, HookRuntime
+from bub.hook_runtime import _SKIP_VALUE, AgentHooks, HookRuntime
 from bub.hookspecs import BUB_HOOK_NAMESPACE, BubHookSpecs
 from bub.runtime import BubError, ErrorKind
 from bub.runtime_options import RuntimeOptions
@@ -48,6 +48,7 @@ def __init__(self, config_file: Path = DEFAULT_CONFIG_FILE) -> None:
         self._plugin_manager = pluggy.PluginManager(BUB_HOOK_NAMESPACE)
         self._plugin_manager.add_hookspecs(BubHookSpecs)
         self._hook_runtime = HookRuntime(self._plugin_manager)
+        self._agent_hooks = AgentHooks(self._hook_runtime)
         self._plugin_status: dict[str, PluginStatus] = {}
         self._outbound_router: OutboundChannelRouter | None = None
         self._tape_store: TapeStore | AsyncTapeStore | None = None
@@ -360,6 +361,9 @@ def get_tape_store(self) -> TapeStore | AsyncTapeStore | None:
     def get_steering_inbox(self) -> SteeringInboxProtocol | None:
         return self._steering_inbox
 
+    def get_agent_hooks(self) -> AgentHooks:
+        return self._agent_hooks
+
     def get_system_prompt(self, prompt: str | list[dict], state: dict[str, Any]) -> str:
         return "\n\n".join(
             result