From 73e104c213764760584a4e6b0bd73471389bd7ab Mon Sep 17 00:00:00 2001 From: Rubick Date: Mon, 6 Jul 2026 18:24:49 +0800 Subject: [PATCH 01/10] =?UTF-8?q?feat:=20agent-loop=20hooks=20=E2=80=94=20?= =?UTF-8?q?before/after=20llm=20&=20tool=20calls=20+=20wrap=5Ftool=5Fcall?= =?UTF-8?q?=20(#253)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the missing interception points inside the agent loop: - before_llm_call: chain-modifiable request (model/messages), or LlmCallDecision.finish(text) to short-circuit (cost guards / call limits) - after_llm_call: observe-only terminal outcome (text/tool_calls/usage/ error/duration), fires once per call incl. error paths; streaming observed at accumulated final state, not per chunk - before_tool_call: per-invocation ToolCallDecision — proceed(args) folds argument changes, replace(result) / deny(message) short-circuit - after_tool_call: observe-only terminal outcome incl. deny/replace - wrap_tool_call: middleware-style continuation (call_next zero/one/many times) enabling retry, caching, human-in-the-loop — cf. LangChain middleware wrap_tool_call Design notes: - payload contracts in bub/agent_hooks.py; execution semantics in AgentHooks facade (bub/hook_runtime.py); wired via BubFramework.get_agent_hooks() -> ModelRunner/ToolExecutor - observer/chain hooks are fault-isolated (raising plugin logged+skipped, never fatal; veto only via decision objects); wrappers own the execution path so their exceptions surface as tool errors - ordering follows pluggy LIFO: last registered runs first / outermost - tool errors are normalized to BubError inside the continuation so wrappers catch a stable type; ToolExecutor() without hooks is unchanged Co-Authored-By: Claude Fable 5 --- src/bub/agent_hooks.py | 127 +++++++++++++ src/bub/builtin/agent.py | 2 +- src/bub/builtin/model_runner.py | 67 ++++++- src/bub/framework.py | 6 +- src/bub/hook_runtime.py | 138 ++++++++++++++ src/bub/hookspecs.py | 64 +++++++ src/bub/tools.py | 112 +++++++++++- tests/test_agent_hooks.py | 314 ++++++++++++++++++++++++++++++++ 8 files changed, 815 insertions(+), 15 deletions(-) create mode 100644 src/bub/agent_hooks.py create mode 100644 tests/test_agent_hooks.py diff --git a/src/bub/agent_hooks.py b/src/bub/agent_hooks.py new file mode 100644 index 00000000..eeca144a --- /dev/null +++ b/src/bub/agent_hooks.py @@ -0,0 +1,127 @@ +"""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 collections.abc import Awaitable, Callable +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 set (and other fields best-effort) + when the call raised. + """ + + run_id: str + text: str | None = None + tool_calls: list[dict[str, Any]] = field(default_factory=list) + usage: dict[str, Any] | None = None + error: str | 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) + + +ToolCallHandler = Callable[["ToolCall"], Awaitable[Any]] +"""Continuation passed to ``wrap_tool_call`` implementations.""" + + +@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 stringified tool failure when the invocation raised or + was denied; ``result`` is unset in that case. + """ + + run_id: str + tool: str + arguments: dict[str, Any] + result: Any = None + error: str | None = None + duration_ms: int = 0 diff --git a/src/bub/builtin/agent.py b/src/bub/builtin/agent.py index 879cf1a6..4404bae8 100644 --- a/src/bub/builtin/agent.py +++ b/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: diff --git a/src/bub/builtin/model_runner.py b/src/bub/builtin/model_runner.py index 00a8c06e..400df194 100644 --- a/src/bub/builtin/model_runner.py +++ b/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): @@ -127,10 +130,40 @@ 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) + try: + async with asyncio.timeout(self.settings.model_timeout_seconds): + completion = await self.completion_response( + model=request.model, messages=list(request.messages), tools=tools + ) + async for event in self._completion_events(completion, state, output): + yield event + except Exception as exc: + await self._fire_after_llm_call(request, output, state, llm_started, tape, error=exc) + raise + await self._fire_after_llm_call(request, output, state, llm_started, tape) tool_calls = output.tool_calls if tool_calls: @@ -139,7 +172,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, ) @@ -180,6 +213,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=repr(error) if error is not None else None, + duration_ms=duration_ms, + ) + await self.hooks.after_llm_call(request, result, state=tape.context.state) + async def build_messages( self, *, diff --git a/src/bub/framework.py b/src/bub/framework.py index f8de717b..cc8588cf 100644 --- a/src/bub/framework.py +++ b/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 diff --git a/src/bub/hook_runtime.py b/src/bub/hook_runtime.py index bdb0eaf7..54a0f3a4 100644 --- a/src/bub/hook_runtime.py +++ b/src/bub/hook_runtime.py @@ -9,6 +9,15 @@ import pluggy from loguru import logger +from bub.agent_hooks import ( + LlmCallDecision, + LlmCallRequest, + LlmCallResult, + ToolCall, + ToolCallDecision, + ToolCallHandler, + ToolCallResult, +) from bub.runtime import AsyncStreamEvents, StreamEvent, StreamState from bub.types import Envelope @@ -197,4 +206,133 @@ async def iterator() -> AsyncGenerator[StreamEvent, None]: return None +class AgentHooks: + """Narrow facade for agent-loop interception hooks (issue #253). + + Unlike ``call_first``/``call_many``, every implementation call here is + fault-isolated: a raising plugin is logged and skipped, never fatal to + the turn. Blocking a tool call is only possible through a returned + :class:`~bub.agent_hooks.ToolCallDecision` (``deny``/``replace``) — an + exception inside ``before_tool_call`` is treated as a broken plugin, + not as a veto. + """ + + def __init__(self, runtime: HookRuntime) -> None: + self._runtime = runtime + + async def before_llm_call( + self, request: LlmCallRequest, state: dict[str, Any] + ) -> tuple[LlmCallRequest, LlmCallDecision | None]: + """Chain ``before_llm_call`` impls; each sees the previous impl's request. + + The first ``LlmCallDecision`` (``finish``) short-circuits remaining + implementations and the provider call itself. + """ + + for impl in self._runtime._iter_hookimpls("before_llm_call"): + value = await self._safe_call_one("before_llm_call", impl, {"request": request, "state": state}) + if value is None or value is _SKIP_VALUE: + continue + if isinstance(value, LlmCallRequest): + request = value + elif isinstance(value, LlmCallDecision): + return request, value + else: + self._warn_bad_return( + "before_llm_call", impl, value, expected="LlmCallRequest | LlmCallDecision | None" + ) + return request, None + + async def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: dict[str, Any]) -> None: + """Observe-only; return values are ignored.""" + + await self._safe_calls("after_llm_call", lambda: {"request": request, "result": result, "state": state}) + + async def before_tool_call(self, call: ToolCall, state: dict[str, Any]) -> tuple[ToolCall, ToolCallDecision]: + """Chain ``before_tool_call`` impls. + + ``proceed`` decisions fold argument changes into the call visible to + later impls; the first ``replace``/``deny`` decision short-circuits. + """ + + from dataclasses import replace as dc_replace + + for impl in self._runtime._iter_hookimpls("before_tool_call"): + value = await self._safe_call_one("before_tool_call", impl, {"call": call, "state": state}) + if value is None or value is _SKIP_VALUE: + continue + if not isinstance(value, ToolCallDecision): + self._warn_bad_return("before_tool_call", impl, value, expected="ToolCallDecision | None") + continue + if value.action == "proceed": + if value.arguments is not None: + call = dc_replace(call, arguments=dict(value.arguments)) + continue + return call, value + return call, ToolCallDecision.proceed() + + async def after_tool_call(self, call: ToolCall, result: ToolCallResult, state: dict[str, Any]) -> None: + """Observe-only; return values are ignored.""" + + await self._safe_calls("after_tool_call", lambda: {"call": call, "state": state, "result": result}) + + def wrap_tool_call_chain(self, base: ToolCallHandler, state: dict[str, Any]) -> ToolCallHandler: + """Compose ``wrap_tool_call`` impls into an onion around ``base``. + + Consistent with pluggy's LIFO convention, the LAST-registered + implementation becomes the outermost layer (it runs first). + Unlike observer hooks, wrappers own the execution path and are NOT + fault-isolated: a raising wrapper surfaces as a tool error. + """ + + handler = base + for impl in reversed(self._runtime._iter_hookimpls("wrap_tool_call")): + handler = self._wrap_one(impl, handler, state) + return handler + + def _wrap_one(self, impl: Any, call_next: ToolCallHandler, state: dict[str, Any]) -> ToolCallHandler: + async def wrapped(call: ToolCall) -> Any: + kwargs = {"call": call, "call_next": call_next, "state": state} + call_kwargs = self._runtime._kwargs_for_impl(impl, kwargs) + value = impl.function(**call_kwargs) + if inspect.isawaitable(value): + value = await value + return value + + return wrapped + + async def _safe_calls(self, hook_name: str, kwargs_factory: Any) -> list[tuple[Any, Any]]: + outcomes: list[tuple[Any, Any]] = [] + for impl in self._runtime._iter_hookimpls(hook_name): + value = await self._safe_call_one(hook_name, impl, kwargs_factory()) + if value is _SKIP_VALUE: + continue + outcomes.append((impl, value)) + return outcomes + + async def _safe_call_one(self, hook_name: str, impl: Any, kwargs: dict[str, Any]) -> Any: + call_kwargs = self._runtime._kwargs_for_impl(impl, kwargs) + try: + return await self._runtime._invoke_impl_async( + hook_name=hook_name, impl=impl, call_kwargs=call_kwargs, kwargs=kwargs + ) + except Exception: + logger.opt(exception=True).warning( + "hook.agent_hook_failed hook={} adapter={}", + hook_name, + impl.plugin_name or "", + ) + return _SKIP_VALUE + + @staticmethod + def _warn_bad_return(hook_name: str, impl: Any, value: Any, *, expected: str) -> None: + logger.warning( + "hook.agent_hook_bad_return hook={} adapter={} got={} expected={}", + hook_name, + impl.plugin_name or "", + type(value).__name__, + expected, + ) + + _SKIP_VALUE = object() diff --git a/src/bub/hookspecs.py b/src/bub/hookspecs.py index f410ed81..d5ddc56f 100644 --- a/src/bub/hookspecs.py +++ b/src/bub/hookspecs.py @@ -7,6 +7,15 @@ import pluggy +from bub.agent_hooks import ( + LlmCallDecision, + LlmCallRequest, + LlmCallResult, + ToolCall, + ToolCallDecision, + ToolCallHandler, + ToolCallResult, +) from bub.runtime import AsyncStreamEvents from bub.runtime_options import RuntimeOptions from bub.tape import AsyncTapeStore, TapeContext, TapeStore @@ -105,6 +114,61 @@ def provide_runtime_options( def on_error(self, stage: str, error: Exception, message: Envelope | None) -> None: """Observe framework errors from any stage.""" + @hookspec + def before_llm_call(self, request: LlmCallRequest, state: State) -> LlmCallRequest | LlmCallDecision | None: + """Observe, modify or short-circuit an outgoing agent-loop LLM request. + + Implementations are chained in pluggy's LIFO order (last registered\n runs first): each receives the + request as modified by earlier implementations, and may return a + modified copy (``dataclasses.replace``), ``None`` to leave it + unchanged, or ``LlmCallDecision.finish(text)`` to skip the provider + call and emit ``text`` as the final output (cost guards / call + limits). Exceptions are logged and skipped, never fatal. + """ + + @hookspec + def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: State) -> None: + """Observe the terminal outcome of one agent-loop LLM call. + + Fires exactly once per call — for streaming completions after the + stream is fully consumed, and also on error (``result.error`` set). + Return values are ignored; exceptions are logged and skipped. + """ + + @hookspec + def before_tool_call(self, call: ToolCall, state: State) -> ToolCallDecision | None: + """Decide whether/how one tool invocation runs. + + Return ``None`` or ``ToolCallDecision.proceed(...)`` to continue + (optionally with modified arguments, visible to later + implementations), ``ToolCallDecision.replace(result)`` to skip the + tool and use ``result``, or ``ToolCallDecision.deny(message)`` to + surface a tool error instead. ``replace``/``deny`` short-circuit + remaining implementations. Blocking is only possible via the + decision object — exceptions are logged and skipped. + """ + + @hookspec + def after_tool_call(self, call: ToolCall, result: ToolCallResult, state: State) -> None: + """Observe the terminal outcome of one tool invocation. + + Fires for success, failure (``result.error`` set), denial and + replacement. Return values are ignored; exceptions are logged. + """ + + @hookspec + def wrap_tool_call(self, call: ToolCall, call_next: ToolCallHandler, state: State) -> Any: + """Wrap one tool invocation with full control over its execution. + + Middleware-style continuation hook: implementations may call + ``call_next(call)`` zero, one or multiple times (caching, retry with + backoff, human-in-the-loop pauses). Implementations nest in + registration order (first registered = outermost). Runs inside the + ``before_tool_call`` decision (a denied call never reaches wrappers) + and inside ``after_tool_call`` observation. Unlike observer hooks, + wrapper exceptions surface as tool errors. + """ + @hookspec def system_prompt(self, prompt: str | list[dict], state: State) -> str: """Provide a system prompt to be prepended to all model prompts.""" diff --git a/src/bub/tools.py b/src/bub/tools.py index 5735deee..00c10200 100644 --- a/src/bub/tools.py +++ b/src/bub/tools.py @@ -8,14 +8,18 @@ import time from collections.abc import Callable, Sequence from dataclasses import dataclass, field, replace -from typing import Any, Protocol, overload +from typing import TYPE_CHECKING, Any, Protocol, overload from loguru import logger from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, validate_call +from bub.agent_hooks import ToolCall, ToolCallHandler, ToolCallResult from bub.builtin.tape import Tape from bub.runtime import BubError, ErrorKind +if TYPE_CHECKING: + from bub.hook_runtime import AgentHooks + @dataclass(frozen=True) class ToolContext: @@ -159,6 +163,9 @@ def tool_call_reporter(reporter: ToolCallReporter): class ToolExecutor: """Execute already-resolved Bub tool invocations.""" + def __init__(self, hooks: AgentHooks | None = None) -> None: + self._hooks = hooks + async def execute_async( self, invocations: Sequence[tuple[Tool, dict[str, Any]]], @@ -205,16 +212,59 @@ async def _handle_tool_response_async( tool_args: dict[str, Any], context: ToolContext | None, ) -> Any: + tool_name = tool_obj.name + call = ToolCall( + run_id=(context.run_id if context is not None else None) or "", + tool=tool_name, + arguments=dict(tool_args), + ) + hook_state = context.state if context is not None else {} + started = time.monotonic() + if self._hooks is not None: + call, short_circuit = await self._apply_before_tool_call(call, hook_state, started) + if short_circuit is not None: + return short_circuit() + + async def invoke(current: ToolCall) -> Any: + return await self._invoke_normalized(tool_obj, current, context) + + handler: ToolCallHandler = invoke + if self._hooks is not None: + handler = self._hooks.wrap_tool_call_chain(invoke, state=hook_state) + try: + result = await handler(call) + except BubError as exc: + await self._fire_after_tool_call(call, hook_state, started, error=exc) + raise + except Exception as exc: + error = BubError( + ErrorKind.TOOL, + f"Tool '{tool_name}' hook wrapper failed.", + details={"error": repr(exc)}, + ) + await self._fire_after_tool_call(call, hook_state, started, error=error) + raise error from exc + else: + await self._fire_after_tool_call(call, hook_state, started, result=result) + return result + + async def _invoke_normalized(self, tool_obj: Tool, call: ToolCall, context: ToolContext | None) -> Any: + """Run the tool with errors normalized to BubError. + + Normalizing inside the continuation gives ``wrap_tool_call`` + middleware a stable exception type to catch (retry/fallback). + """ + tool_name = tool_obj.name try: - result = self._invoke_tool( + value = self._invoke_tool( tool_name=tool_name, tool_obj=tool_obj, - tool_args=tool_args, + tool_args=call.arguments, context=context, ) - if inspect.isawaitable(result): - return await result + if inspect.isawaitable(value): + value = await value except BubError: raise except ValidationError as exc: @@ -229,8 +279,56 @@ async def _handle_tool_response_async( f"Tool '{tool_name}' execution failed.", details={"error": repr(exc)}, ) from exc - else: - return result + return value + + async def _apply_before_tool_call( + self, + call: ToolCall, + hook_state: dict[str, Any], + started: float, + ) -> tuple[ToolCall, Callable[[], Any] | None]: + """Run before_tool_call and translate deny/replace into a short-circuit thunk.""" + + if self._hooks is None: + return call, None + call, decision = await self._hooks.before_tool_call(call, state=hook_state) + if decision.action == "deny": + error = BubError( + ErrorKind.TOOL, + decision.message or f"Tool '{call.tool}' call denied by policy hook.", + ) + await self._fire_after_tool_call(call, hook_state, started, error=error) + + def raise_denied() -> Any: + raise error + + return call, raise_denied + if decision.action == "replace": + await self._fire_after_tool_call(call, hook_state, started, result=decision.result) + return call, lambda: decision.result + return call, None + + async def _fire_after_tool_call( + self, + call: ToolCall, + state: dict[str, Any], + started: float, + *, + result: Any = None, + error: BubError | None = None, + ) -> None: + if self._hooks is None: + return + duration_ms = int((time.monotonic() - started) * 1000) + outcome = ToolCallResult( + run_id=call.run_id, + tool=call.tool, + arguments=call.arguments, + result=None if error is not None else result, + error=error.message if error is not None else None, + duration_ms=duration_ms, + ) + await self._hooks.after_tool_call(call, outcome, state=state) # Central registry for tools. Tools defined with the @tool decorator are automatically added here. diff --git a/tests/test_agent_hooks.py b/tests/test_agent_hooks.py new file mode 100644 index 00000000..24c6c0d6 --- /dev/null +++ b/tests/test_agent_hooks.py @@ -0,0 +1,314 @@ +"""Agent-loop hook semantics: chaining, short-circuit, isolation (issue #253).""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Any + +import pluggy +import pytest + +from bub.agent_hooks import ( + LlmCallRequest, + LlmCallResult, + ToolCall, + ToolCallDecision, + ToolCallResult, +) +from bub.hook_runtime import AgentHooks, HookRuntime +from bub.hookspecs import BUB_HOOK_NAMESPACE, BubHookSpecs, hookimpl +from bub.runtime import BubError +from bub.tools import Tool, ToolExecutor + + +def make_hooks(*plugins: Any) -> AgentHooks: + pm = pluggy.PluginManager(BUB_HOOK_NAMESPACE) + pm.add_hookspecs(BubHookSpecs) + for plugin in plugins: + pm.register(plugin) + return AgentHooks(HookRuntime(pm)) + + +def request() -> LlmCallRequest: + return LlmCallRequest(run_id="run-1", model="openai:gpt-x", messages=[{"role": "user", "content": "hi"}]) + + +class TestBeforeLlmCall: + @pytest.mark.asyncio + async def test_chain_folds_modifications_in_order(self) -> None: + class SwapModel: + @hookimpl + def before_llm_call(self, request: LlmCallRequest, state: dict) -> LlmCallRequest: + return replace(request, model="anthropic:claude") + + class AppendMessage: + @hookimpl + def before_llm_call(self, request: LlmCallRequest, state: dict) -> LlmCallRequest: + # must see SwapModel's change (registration-order chaining) + assert request.model == "anthropic:claude" + return replace(request, messages=[*request.messages, {"role": "user", "content": "extra"}]) + + # pluggy LIFO: last registered runs first -> register AppendMessage first + hooks = make_hooks(AppendMessage(), SwapModel()) + result, decision = await hooks.before_llm_call(request(), state={}) + assert decision is None + assert result.model == "anthropic:claude" + assert result.messages[-1]["content"] == "extra" + + @pytest.mark.asyncio + async def test_none_and_bad_returns_leave_request_unchanged(self) -> None: + class Noop: + @hookimpl + def before_llm_call(self, request: LlmCallRequest, state: dict) -> None: + return None + + class BadReturn: + @hookimpl + def before_llm_call(self, request: LlmCallRequest, state: dict) -> Any: + return "not-a-request" + + hooks = make_hooks(Noop(), BadReturn()) + original = request() + assert await hooks.before_llm_call(original, state={}) == (original, None) + + @pytest.mark.asyncio + async def test_raising_impl_is_isolated(self) -> None: + class Boom: + @hookimpl + def before_llm_call(self, request: LlmCallRequest, state: dict) -> LlmCallRequest: + raise RuntimeError("plugin broke") + + class After: + @hookimpl + def before_llm_call(self, request: LlmCallRequest, state: dict) -> LlmCallRequest: + return replace(request, model="fallback:model") + + hooks = make_hooks(Boom(), After()) + result, _ = await hooks.before_llm_call(request(), state={}) + assert result.model == "fallback:model" + + +class TestBeforeToolCall: + @pytest.mark.asyncio + async def test_proceed_folds_arguments_and_later_impl_sees_them(self) -> None: + class Rewrite: + @hookimpl + def before_tool_call(self, call: ToolCall, state: dict) -> ToolCallDecision: + return ToolCallDecision.proceed(arguments={**call.arguments, "safe": True}) + + class Verify: + @hookimpl + def before_tool_call(self, call: ToolCall, state: dict) -> None: + assert call.arguments["safe"] is True + return None + + hooks = make_hooks(Verify(), Rewrite()) # LIFO: Rewrite runs first + call, decision = await hooks.before_tool_call( + ToolCall(run_id="run-1", tool="shell", arguments={"cmd": "ls"}), state={} + ) + assert decision.action == "proceed" + assert call.arguments == {"cmd": "ls", "safe": True} + + @pytest.mark.asyncio + async def test_deny_short_circuits_remaining_impls(self) -> None: + seen: list[str] = [] + + class Deny: + @hookimpl + def before_tool_call(self, call: ToolCall, state: dict) -> ToolCallDecision: + return ToolCallDecision.deny("dangerous command") + + class Later: + @hookimpl + def before_tool_call(self, call: ToolCall, state: dict) -> None: + seen.append(call.tool) + return None + + hooks = make_hooks(Later(), Deny()) # LIFO: Deny runs first + _, decision = await hooks.before_tool_call(ToolCall(run_id="run-1", tool="shell", arguments={}), state={}) + assert decision.action == "deny" + assert decision.message == "dangerous command" + assert seen == [] + + +class TestToolExecutorIntegration: + def tool(self) -> Tool: + def handler(cmd: str) -> str: + return f"ran:{cmd}" + + return Tool(name="shell", handler=handler, description="", parameters={}) + + @pytest.mark.asyncio + async def test_deny_surfaces_tool_error_result(self) -> None: + class Deny: + @hookimpl + def before_tool_call(self, call: ToolCall, state: dict) -> ToolCallDecision: + return ToolCallDecision.deny("blocked by policy") + + executor = ToolExecutor(hooks=make_hooks(Deny())) + execution = await executor.execute_async([(self.tool(), {"cmd": "rm -rf /"})]) + assert execution.error is not None + assert "blocked by policy" in execution.error.message + assert execution.tool_results[0]["message"] == "blocked by policy" + + @pytest.mark.asyncio + async def test_replace_skips_handler(self) -> None: + class Replace: + @hookimpl + def before_tool_call(self, call: ToolCall, state: dict) -> ToolCallDecision: + return ToolCallDecision.replace("cached result") + + executor = ToolExecutor(hooks=make_hooks(Replace())) + execution = await executor.execute_async([(self.tool(), {"cmd": "ls"})]) + assert execution.error is None + assert execution.tool_results == ["cached result"] + + @pytest.mark.asyncio + async def test_after_tool_call_observes_success_and_error(self) -> None: + observed: list[ToolCallResult] = [] + + class Observe: + @hookimpl + def after_tool_call(self, call: ToolCall, result: ToolCallResult, state: dict) -> None: + observed.append(result) + + def failing(cmd: str) -> str: + raise ValueError("nope") + + executor = ToolExecutor(hooks=make_hooks(Observe())) + await executor.execute_async([(self.tool(), {"cmd": "ls"})]) + await executor.execute_async([(Tool(name="bad", handler=failing, description="", parameters={}), {"cmd": "x"})]) + assert observed[0].result == "ran:ls" + assert observed[0].error is None + assert observed[1].error is not None + assert "bad" in observed[1].tool + + @pytest.mark.asyncio + async def test_modified_arguments_reach_handler(self) -> None: + class Rewrite: + @hookimpl + def before_tool_call(self, call: ToolCall, state: dict) -> ToolCallDecision: + return ToolCallDecision.proceed(arguments={"cmd": "safe-ls"}) + + executor = ToolExecutor(hooks=make_hooks(Rewrite())) + execution = await executor.execute_async([(self.tool(), {"cmd": "rm"})]) + assert execution.tool_results == ["ran:safe-ls"] + + @pytest.mark.asyncio + async def test_no_hooks_keeps_current_behavior(self) -> None: + execution = await ToolExecutor().execute_async([(self.tool(), {"cmd": "ls"})]) + assert execution.tool_results == ["ran:ls"] + + +class TestAfterLlmCall: + @pytest.mark.asyncio + async def test_observe_only_and_isolated(self) -> None: + observed: list[LlmCallResult] = [] + + class Boom: + @hookimpl + def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: dict) -> None: + raise RuntimeError("metrics plugin broke") + + class Observe: + @hookimpl + def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: dict) -> None: + observed.append(result) + + hooks = make_hooks(Boom(), Observe()) + result = LlmCallResult(run_id="run-1", text="hello", usage={"total_tokens": 5}, duration_ms=12) + await hooks.after_llm_call(request(), result, state={}) + assert observed == [result] + + +class TestWrapToolCall: + def tool(self) -> Tool: + calls = {"n": 0} + + def handler(cmd: str) -> str: + calls["n"] += 1 + if calls["n"] == 1: + raise ValueError("transient") + return f"ok:{cmd}:attempt{calls['n']}" + + return Tool(name="flaky", handler=handler, description="", parameters={}) + + @pytest.mark.asyncio + async def test_retry_wrapper_can_call_next_twice(self) -> None: + class Retry: + @hookimpl + async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: + try: + return await call_next(call) + except BubError: + return await call_next(call) + + executor = ToolExecutor(hooks=make_hooks(Retry())) + execution = await executor.execute_async([(self.tool(), {"cmd": "x"})]) + assert execution.error is None + assert execution.tool_results == ["ok:x:attempt2"] + + @pytest.mark.asyncio + async def test_cache_wrapper_can_skip_call_next(self) -> None: + class Cache: + @hookimpl + async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: + return "cached" + + def never(cmd: str) -> str: + raise AssertionError("handler must not run") + + executor = ToolExecutor(hooks=make_hooks(Cache())) + execution = await executor.execute_async([ + (Tool(name="t", handler=never, description="", parameters={}), {"cmd": "x"}) + ]) + assert execution.tool_results == ["cached"] + + @pytest.mark.asyncio + async def test_wrappers_nest_first_registered_outermost(self) -> None: + order: list[str] = [] + + class Outer: + @hookimpl + async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: + order.append("outer-in") + value = await call_next(call) + order.append("outer-out") + return value + + class Inner: + @hookimpl + async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: + order.append("inner-in") + value = await call_next(call) + order.append("inner-out") + return value + + def handler(cmd: str) -> str: + order.append("handler") + return "done" + + executor = ToolExecutor(hooks=make_hooks(Inner(), Outer())) # LIFO: Outer outermost + await executor.execute_async([(Tool(name="t", handler=handler, description="", parameters={}), {"cmd": "x"})]) + assert order == ["outer-in", "inner-in", "handler", "inner-out", "outer-out"] + + +class TestBeforeLlmCallFinish: + @pytest.mark.asyncio + async def test_finish_decision_short_circuits(self) -> None: + from bub.agent_hooks import LlmCallDecision + + class Limit: + @hookimpl + def before_llm_call(self, request: LlmCallRequest, state: dict) -> LlmCallDecision: + return LlmCallDecision.finish("call budget exhausted") + + class Later: + @hookimpl + def before_llm_call(self, request: LlmCallRequest, state: dict) -> None: + raise AssertionError("must not run after finish") + + hooks = make_hooks(Later(), Limit()) # LIFO: Limit runs first + _req, decision = await hooks.before_llm_call(request(), state={}) + assert decision is not None + assert decision.text == "call budget exhausted" From 62249b69efb75eb79901549057d252faea9d11ec Mon Sep 17 00:00:00 2001 From: Rubick Date: Mon, 6 Jul 2026 20:10:04 +0800 Subject: [PATCH 02/10] =?UTF-8?q?fix:=20address=20review=20=E2=80=94=20hon?= =?UTF-8?q?or=20rewritten=20request=20end-to-end,=20exactly-once=20after?= =?UTF-8?q?=5Fllm=5Fcall,=20effective-call=20observation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review findings from PR #255 round 1: 1. before_llm_call rewrites now reach the provider AND the tape: completion_response accepts the effective max_tokens; both success-path record_chat calls record request.model instead of the pre-hook model. 2. after_llm_call is exactly-once on every exit path: the guard now catches BaseException so consumer aclose()/cancellation (GeneratorExit, CancelledError) still produce a terminal observation. 3. after_tool_call observes the effective invocation: the most recent ToolCall passed to call_next (wrappers may rewrite arguments or retry); wrappers that never invoke call_next are observed with the pre-wrap call. Contract documented in the wrap_tool_call hookspec. 4. Doc consistency: wrap nesting is pluggy LIFO (last registered = outermost) everywhere; fixed a literal backslash-n typo in the before_llm_call docstring and the misnamed nesting test. Adds 5 regression tests (rewritten model/max_tokens reach provider+tape, exactly-once on success and on early aclose, effective-call observation with and without call_next). Co-Authored-By: Claude Fable 5 --- src/bub/builtin/model_runner.py | 34 ++++++-- src/bub/hookspecs.py | 10 ++- src/bub/tools.py | 12 ++- tests/test_agent_hooks.py | 142 +++++++++++++++++++++++++++++++- 4 files changed, 182 insertions(+), 16 deletions(-) diff --git a/src/bub/builtin/model_runner.py b/src/bub/builtin/model_runner.py index 400df194..b3eda01d 100644 --- a/src/bub/builtin/model_runner.py +++ b/src/bub/builtin/model_runner.py @@ -79,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 @@ -94,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), ) @@ -153,17 +153,33 @@ async def iterator() -> AsyncGenerator[StreamEvent, None]: yield StreamEvent("final", {"ok": True, "text": decision.text}) return llm_started = datetime.now(UTC) + after_fired = False + + async def fire_after(error: BaseException | None = None) -> None: + """Fire after_llm_call exactly once per call, on every exit path.""" + + 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 + 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: - await self._fire_after_llm_call(request, output, state, llm_started, tape, error=exc) + except BaseException as exc: + # BaseException so consumer aclose()/cancellation (GeneratorExit, + # CancelledError) still produce a terminal after_llm_call. + await fire_after(exc) raise - await self._fire_after_llm_call(request, output, state, llm_started, tape) + await fire_after() tool_calls = output.tool_calls if tool_calls: @@ -185,7 +201,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}) @@ -202,7 +218,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}) @@ -220,7 +236,7 @@ async def _fire_after_llm_call( state: StreamState, started: datetime, tape: Tape, - error: Exception | None = None, + error: BaseException | None = None, ) -> None: if self.hooks is None: return diff --git a/src/bub/hookspecs.py b/src/bub/hookspecs.py index d5ddc56f..c3bec977 100644 --- a/src/bub/hookspecs.py +++ b/src/bub/hookspecs.py @@ -118,7 +118,8 @@ def on_error(self, stage: str, error: Exception, message: Envelope | None) -> No def before_llm_call(self, request: LlmCallRequest, state: State) -> LlmCallRequest | LlmCallDecision | None: """Observe, modify or short-circuit an outgoing agent-loop LLM request. - Implementations are chained in pluggy's LIFO order (last registered\n runs first): each receives the + Implementations are chained in pluggy's LIFO order (last registered + runs first): each receives the request as modified by earlier implementations, and may return a modified copy (``dataclasses.replace``), ``None`` to leave it unchanged, or ``LlmCallDecision.finish(text)`` to skip the provider @@ -162,8 +163,11 @@ def wrap_tool_call(self, call: ToolCall, call_next: ToolCallHandler, state: Stat Middleware-style continuation hook: implementations may call ``call_next(call)`` zero, one or multiple times (caching, retry with - backoff, human-in-the-loop pauses). Implementations nest in - registration order (first registered = outermost). Runs inside the + backoff, human-in-the-loop pauses). Implementations nest per pluggy's + LIFO convention (last registered = outermost, runs first). + ``after_tool_call`` observes the most recent call passed to + ``call_next`` (the effective invocation); wrappers that never invoke + ``call_next`` are observed with the pre-wrap call. Runs inside the ``before_tool_call`` decision (a denied call never reaches wrappers) and inside ``after_tool_call`` observation. Unlike observer hooks, wrapper exceptions surface as tool errors. diff --git a/src/bub/tools.py b/src/bub/tools.py index 00c10200..260dc2bd 100644 --- a/src/bub/tools.py +++ b/src/bub/tools.py @@ -225,7 +225,13 @@ async def _handle_tool_response_async( if short_circuit is not None: return short_circuit() + # Track the "effective call": the most recent ToolCall a wrapper passed + # to call_next. after_tool_call must observe what actually ran, not the + # pre-wrap call (wrappers may rewrite arguments or retry). + effective = {"call": call} + async def invoke(current: ToolCall) -> Any: + effective["call"] = current return await self._invoke_normalized(tool_obj, current, context) handler: ToolCallHandler = invoke @@ -234,7 +240,7 @@ async def invoke(current: ToolCall) -> Any: try: result = await handler(call) except BubError as exc: - await self._fire_after_tool_call(call, hook_state, started, error=exc) + await self._fire_after_tool_call(effective["call"], hook_state, started, error=exc) raise except Exception as exc: error = BubError( @@ -242,10 +248,10 @@ async def invoke(current: ToolCall) -> Any: f"Tool '{tool_name}' hook wrapper failed.", details={"error": repr(exc)}, ) - await self._fire_after_tool_call(call, hook_state, started, error=error) + await self._fire_after_tool_call(effective["call"], hook_state, started, error=error) raise error from exc else: - await self._fire_after_tool_call(call, hook_state, started, result=result) + await self._fire_after_tool_call(effective["call"], hook_state, started, result=result) return result async def _invoke_normalized(self, tool_obj: Tool, call: ToolCall, context: ToolContext | None) -> Any: diff --git a/tests/test_agent_hooks.py b/tests/test_agent_hooks.py index 24c6c0d6..d9ba3dc9 100644 --- a/tests/test_agent_hooks.py +++ b/tests/test_agent_hooks.py @@ -265,7 +265,7 @@ def never(cmd: str) -> str: assert execution.tool_results == ["cached"] @pytest.mark.asyncio - async def test_wrappers_nest_first_registered_outermost(self) -> None: + async def test_wrappers_nest_last_registered_outermost(self) -> None: order: list[str] = [] class Outer: @@ -312,3 +312,143 @@ def before_llm_call(self, request: LlmCallRequest, state: dict) -> None: _req, decision = await hooks.before_llm_call(request(), state={}) assert decision is not None assert decision.text == "call budget exhausted" + + +class TestModelRunnerHookIntegration: + """Regression tests for PR #255 review findings (effective request, exactly-once).""" + + def _runner_and_tape(self, hooks: AgentHooks, captured: dict): + import bub + from bub.builtin.model_runner import ModelRunner + from bub.builtin.settings import AgentSettings + from bub.builtin.tape import Tape + from bub.tape import AsyncTapeStoreAdapter, InMemoryTapeStore, TapeContext + + class FakeRunner(ModelRunner): + async def completion_response(self, *, model, messages, tools, max_tokens=None): + captured.update(model=model, max_tokens=max_tokens) + + async def chunks(): + return + yield # pragma: no cover + + return chunks() + + settings = AgentSettings.model_construct(model="openai:orig", max_tokens=100, model_timeout_seconds=None) + runner = FakeRunner(settings, hooks=hooks) + store = AsyncTapeStoreAdapter(InMemoryTapeStore()) + tape = Tape(bub.home / "tapes", store, TapeContext(anchor=None)).scoped("t1") + return runner, tape + + @pytest.mark.asyncio + async def test_rewritten_model_and_max_tokens_reach_provider_and_tape(self) -> None: + class Reroute: + @hookimpl + def before_llm_call(self, request: LlmCallRequest, state: dict) -> LlmCallRequest: + return replace(request, model="anthropic:new", max_tokens=42) + + captured: dict = {} + runner, tape = self._runner_and_tape(make_hooks(Reroute()), captured) + events = runner.run(tape=tape, model="openai:orig", tools=[], system_prompt=None, prompt="hi") + async for _ in events: + pass + assert captured == {"model": "anthropic:new", "max_tokens": 42} + entries = list(await tape.store.fetch_all(tape.query().kinds("event"))) + run_events = [e for e in entries if e.payload.get("name") == "run"] + assert run_events[-1].payload["data"]["model"] == "anthropic:new" + + @pytest.mark.asyncio + async def test_after_llm_call_fires_exactly_once_on_early_close(self) -> None: + observed: list[LlmCallResult] = [] + + class Observe: + @hookimpl + def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: dict) -> None: + observed.append(result) + + captured: dict = {} + runner, tape = self._runner_and_tape(make_hooks(Observe()), captured) + + from bub.builtin.model_runner import ModelRunner # noqa: F401 + + async def fake_events(completion, state, output): + from bub.runtime import StreamEvent + + yield StreamEvent("text", {"delta": "a"}) + yield StreamEvent("text", {"delta": "b"}) + + runner._completion_events = fake_events # type: ignore[method-assign] + events = runner.run(tape=tape, model="openai:orig", tools=[], system_prompt=None, prompt="hi") + iterator = events.__aiter__() + await iterator.__anext__() + await iterator.aclose() + assert len(observed) == 1 + assert observed[0].error is not None # closed before completion = abnormal terminal + + @pytest.mark.asyncio + async def test_after_llm_call_fires_exactly_once_on_success(self) -> None: + observed: list[LlmCallResult] = [] + + class Observe: + @hookimpl + def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: dict) -> None: + observed.append(result) + + captured: dict = {} + runner, tape = self._runner_and_tape(make_hooks(Observe()), captured) + events = runner.run(tape=tape, model="openai:orig", tools=[], system_prompt=None, prompt="hi") + async for _ in events: + pass + assert len(observed) == 1 + assert observed[0].error is None + + +class TestWrapperEffectiveCall: + @pytest.mark.asyncio + async def test_after_tool_call_observes_wrapper_rewritten_arguments(self) -> None: + observed: list[ToolCallResult] = [] + + class Rewrite: + @hookimpl + async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: + return await call_next(replace(call, arguments={"cmd": "rewritten"})) + + class Observe: + @hookimpl + def after_tool_call(self, call: ToolCall, result: ToolCallResult, state: dict) -> None: + observed.append(result) + + def handler(cmd: str) -> str: + return f"ran:{cmd}" + + executor = ToolExecutor(hooks=make_hooks(Observe(), Rewrite())) + execution = await executor.execute_async([ + (Tool(name="t", handler=handler, description="", parameters={}), {"cmd": "original"}) + ]) + assert execution.tool_results == ["ran:rewritten"] + assert observed[0].arguments == {"cmd": "rewritten"} + assert observed[0].result == "ran:rewritten" + + @pytest.mark.asyncio + async def test_after_tool_call_uses_prewrap_call_when_next_never_invoked(self) -> None: + observed: list[ToolCallResult] = [] + + class Cache: + @hookimpl + async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: + return "cached" + + class Observe: + @hookimpl + def after_tool_call(self, call: ToolCall, result: ToolCallResult, state: dict) -> None: + observed.append(result) + + def handler(cmd: str) -> str: + raise AssertionError("must not run") + + executor = ToolExecutor(hooks=make_hooks(Observe(), Cache())) + await executor.execute_async([ + (Tool(name="t", handler=handler, description="", parameters={}), {"cmd": "original"}) + ]) + assert observed[0].arguments == {"cmd": "original"} + assert observed[0].result == "cached" From c395e518910cc8324d330090f8450dd26851a473 Mon Sep 17 00:00:00 2001 From: Rubick Date: Mon, 6 Jul 2026 21:50:23 +0800 Subject: [PATCH 03/10] refactor: drop wrap_tool_call hook per review direction Keeps the four before/after hooks (incl. LlmCallDecision.finish and the exactly-once / effective-request fixes). The continuation-style wrapper can return later as a separate proposal if retry/HITL use cases press. Co-Authored-By: Claude Fable 5 --- src/bub/agent_hooks.py | 5 -- src/bub/hook_runtime.py | 26 -------- src/bub/hookspecs.py | 17 ------ src/bub/tools.py | 34 ++--------- tests/test_agent_hooks.py | 124 -------------------------------------- 5 files changed, 5 insertions(+), 201 deletions(-) diff --git a/src/bub/agent_hooks.py b/src/bub/agent_hooks.py index eeca144a..0cdd4963 100644 --- a/src/bub/agent_hooks.py +++ b/src/bub/agent_hooks.py @@ -12,7 +12,6 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from typing import Any, Literal @@ -67,10 +66,6 @@ def finish(cls, text: str) -> LlmCallDecision: return cls(action="finish", text=text) -ToolCallHandler = Callable[["ToolCall"], Awaitable[Any]] -"""Continuation passed to ``wrap_tool_call`` implementations.""" - - @dataclass(frozen=True) class ToolCall: """One tool invocation about to run, as seen by ``before_tool_call``.""" diff --git a/src/bub/hook_runtime.py b/src/bub/hook_runtime.py index 54a0f3a4..bd2c9795 100644 --- a/src/bub/hook_runtime.py +++ b/src/bub/hook_runtime.py @@ -15,7 +15,6 @@ LlmCallResult, ToolCall, ToolCallDecision, - ToolCallHandler, ToolCallResult, ) from bub.runtime import AsyncStreamEvents, StreamEvent, StreamState @@ -276,31 +275,6 @@ async def after_tool_call(self, call: ToolCall, result: ToolCallResult, state: d await self._safe_calls("after_tool_call", lambda: {"call": call, "state": state, "result": result}) - def wrap_tool_call_chain(self, base: ToolCallHandler, state: dict[str, Any]) -> ToolCallHandler: - """Compose ``wrap_tool_call`` impls into an onion around ``base``. - - Consistent with pluggy's LIFO convention, the LAST-registered - implementation becomes the outermost layer (it runs first). - Unlike observer hooks, wrappers own the execution path and are NOT - fault-isolated: a raising wrapper surfaces as a tool error. - """ - - handler = base - for impl in reversed(self._runtime._iter_hookimpls("wrap_tool_call")): - handler = self._wrap_one(impl, handler, state) - return handler - - def _wrap_one(self, impl: Any, call_next: ToolCallHandler, state: dict[str, Any]) -> ToolCallHandler: - async def wrapped(call: ToolCall) -> Any: - kwargs = {"call": call, "call_next": call_next, "state": state} - call_kwargs = self._runtime._kwargs_for_impl(impl, kwargs) - value = impl.function(**call_kwargs) - if inspect.isawaitable(value): - value = await value - return value - - return wrapped - async def _safe_calls(self, hook_name: str, kwargs_factory: Any) -> list[tuple[Any, Any]]: outcomes: list[tuple[Any, Any]] = [] for impl in self._runtime._iter_hookimpls(hook_name): diff --git a/src/bub/hookspecs.py b/src/bub/hookspecs.py index c3bec977..fd1cf479 100644 --- a/src/bub/hookspecs.py +++ b/src/bub/hookspecs.py @@ -13,7 +13,6 @@ LlmCallResult, ToolCall, ToolCallDecision, - ToolCallHandler, ToolCallResult, ) from bub.runtime import AsyncStreamEvents @@ -157,22 +156,6 @@ def after_tool_call(self, call: ToolCall, result: ToolCallResult, state: State) replacement. Return values are ignored; exceptions are logged. """ - @hookspec - def wrap_tool_call(self, call: ToolCall, call_next: ToolCallHandler, state: State) -> Any: - """Wrap one tool invocation with full control over its execution. - - Middleware-style continuation hook: implementations may call - ``call_next(call)`` zero, one or multiple times (caching, retry with - backoff, human-in-the-loop pauses). Implementations nest per pluggy's - LIFO convention (last registered = outermost, runs first). - ``after_tool_call`` observes the most recent call passed to - ``call_next`` (the effective invocation); wrappers that never invoke - ``call_next`` are observed with the pre-wrap call. Runs inside the - ``before_tool_call`` decision (a denied call never reaches wrappers) - and inside ``after_tool_call`` observation. Unlike observer hooks, - wrapper exceptions surface as tool errors. - """ - @hookspec def system_prompt(self, prompt: str | list[dict], state: State) -> str: """Provide a system prompt to be prepended to all model prompts.""" diff --git a/src/bub/tools.py b/src/bub/tools.py index 260dc2bd..162462f0 100644 --- a/src/bub/tools.py +++ b/src/bub/tools.py @@ -13,7 +13,7 @@ from loguru import logger from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, validate_call -from bub.agent_hooks import ToolCall, ToolCallHandler, ToolCallResult +from bub.agent_hooks import ToolCall, ToolCallResult from bub.builtin.tape import Tape from bub.runtime import BubError, ErrorKind @@ -225,41 +225,17 @@ async def _handle_tool_response_async( if short_circuit is not None: return short_circuit() - # Track the "effective call": the most recent ToolCall a wrapper passed - # to call_next. after_tool_call must observe what actually ran, not the - # pre-wrap call (wrappers may rewrite arguments or retry). - effective = {"call": call} - - async def invoke(current: ToolCall) -> Any: - effective["call"] = current - return await self._invoke_normalized(tool_obj, current, context) - - handler: ToolCallHandler = invoke - if self._hooks is not None: - handler = self._hooks.wrap_tool_call_chain(invoke, state=hook_state) try: - result = await handler(call) + result = await self._invoke_normalized(tool_obj, call, context) except BubError as exc: - await self._fire_after_tool_call(effective["call"], hook_state, started, error=exc) + await self._fire_after_tool_call(call, hook_state, started, error=exc) raise - except Exception as exc: - error = BubError( - ErrorKind.TOOL, - f"Tool '{tool_name}' hook wrapper failed.", - details={"error": repr(exc)}, - ) - await self._fire_after_tool_call(effective["call"], hook_state, started, error=error) - raise error from exc else: - await self._fire_after_tool_call(effective["call"], hook_state, started, result=result) + await self._fire_after_tool_call(call, hook_state, started, result=result) return result async def _invoke_normalized(self, tool_obj: Tool, call: ToolCall, context: ToolContext | None) -> Any: - """Run the tool with errors normalized to BubError. - - Normalizing inside the continuation gives ``wrap_tool_call`` - middleware a stable exception type to catch (retry/fallback). - """ + """Run the tool with errors normalized to BubError.""" tool_name = tool_obj.name try: diff --git a/tests/test_agent_hooks.py b/tests/test_agent_hooks.py index d9ba3dc9..98523374 100644 --- a/tests/test_agent_hooks.py +++ b/tests/test_agent_hooks.py @@ -17,7 +17,6 @@ ) from bub.hook_runtime import AgentHooks, HookRuntime from bub.hookspecs import BUB_HOOK_NAMESPACE, BubHookSpecs, hookimpl -from bub.runtime import BubError from bub.tools import Tool, ToolExecutor @@ -221,78 +220,6 @@ def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: assert observed == [result] -class TestWrapToolCall: - def tool(self) -> Tool: - calls = {"n": 0} - - def handler(cmd: str) -> str: - calls["n"] += 1 - if calls["n"] == 1: - raise ValueError("transient") - return f"ok:{cmd}:attempt{calls['n']}" - - return Tool(name="flaky", handler=handler, description="", parameters={}) - - @pytest.mark.asyncio - async def test_retry_wrapper_can_call_next_twice(self) -> None: - class Retry: - @hookimpl - async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: - try: - return await call_next(call) - except BubError: - return await call_next(call) - - executor = ToolExecutor(hooks=make_hooks(Retry())) - execution = await executor.execute_async([(self.tool(), {"cmd": "x"})]) - assert execution.error is None - assert execution.tool_results == ["ok:x:attempt2"] - - @pytest.mark.asyncio - async def test_cache_wrapper_can_skip_call_next(self) -> None: - class Cache: - @hookimpl - async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: - return "cached" - - def never(cmd: str) -> str: - raise AssertionError("handler must not run") - - executor = ToolExecutor(hooks=make_hooks(Cache())) - execution = await executor.execute_async([ - (Tool(name="t", handler=never, description="", parameters={}), {"cmd": "x"}) - ]) - assert execution.tool_results == ["cached"] - - @pytest.mark.asyncio - async def test_wrappers_nest_last_registered_outermost(self) -> None: - order: list[str] = [] - - class Outer: - @hookimpl - async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: - order.append("outer-in") - value = await call_next(call) - order.append("outer-out") - return value - - class Inner: - @hookimpl - async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: - order.append("inner-in") - value = await call_next(call) - order.append("inner-out") - return value - - def handler(cmd: str) -> str: - order.append("handler") - return "done" - - executor = ToolExecutor(hooks=make_hooks(Inner(), Outer())) # LIFO: Outer outermost - await executor.execute_async([(Tool(name="t", handler=handler, description="", parameters={}), {"cmd": "x"})]) - assert order == ["outer-in", "inner-in", "handler", "inner-out", "outer-out"] - - class TestBeforeLlmCallFinish: @pytest.mark.asyncio async def test_finish_decision_short_circuits(self) -> None: @@ -401,54 +328,3 @@ def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: pass assert len(observed) == 1 assert observed[0].error is None - - -class TestWrapperEffectiveCall: - @pytest.mark.asyncio - async def test_after_tool_call_observes_wrapper_rewritten_arguments(self) -> None: - observed: list[ToolCallResult] = [] - - class Rewrite: - @hookimpl - async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: - return await call_next(replace(call, arguments={"cmd": "rewritten"})) - - class Observe: - @hookimpl - def after_tool_call(self, call: ToolCall, result: ToolCallResult, state: dict) -> None: - observed.append(result) - - def handler(cmd: str) -> str: - return f"ran:{cmd}" - - executor = ToolExecutor(hooks=make_hooks(Observe(), Rewrite())) - execution = await executor.execute_async([ - (Tool(name="t", handler=handler, description="", parameters={}), {"cmd": "original"}) - ]) - assert execution.tool_results == ["ran:rewritten"] - assert observed[0].arguments == {"cmd": "rewritten"} - assert observed[0].result == "ran:rewritten" - - @pytest.mark.asyncio - async def test_after_tool_call_uses_prewrap_call_when_next_never_invoked(self) -> None: - observed: list[ToolCallResult] = [] - - class Cache: - @hookimpl - async def wrap_tool_call(self, call: ToolCall, call_next, state: dict) -> Any: - return "cached" - - class Observe: - @hookimpl - def after_tool_call(self, call: ToolCall, result: ToolCallResult, state: dict) -> None: - observed.append(result) - - def handler(cmd: str) -> str: - raise AssertionError("must not run") - - executor = ToolExecutor(hooks=make_hooks(Observe(), Cache())) - await executor.execute_async([ - (Tool(name="t", handler=handler, description="", parameters={}), {"cmd": "original"}) - ]) - assert observed[0].arguments == {"cmd": "original"} - assert observed[0].result == "cached" From 7cb26a0befb4805efa42b0a0d3c0e57d6a742998 Mon Sep 17 00:00:00 2001 From: Rubick Date: Mon, 6 Jul 2026 21:56:13 +0800 Subject: [PATCH 04/10] feat: preserve original exceptions in LlmCallResult/ToolCallResult MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review direction: observers get the raised exception object instead of a stringified form — ToolCallResult.error is the original BubError (kind/message/details intact), LlmCallResult.error is the original BaseException (incl. GeneratorExit/CancelledError on close/cancel). Tests pin the preserved types. Co-Authored-By: Claude Fable 5 --- src/bub/agent_hooks.py | 13 +++++++------ src/bub/builtin/model_runner.py | 2 +- src/bub/tools.py | 2 +- tests/test_agent_hooks.py | 6 ++++-- 4 files changed, 13 insertions(+), 10 deletions(-) diff --git a/src/bub/agent_hooks.py b/src/bub/agent_hooks.py index 0cdd4963..7aec410d 100644 --- a/src/bub/agent_hooks.py +++ b/src/bub/agent_hooks.py @@ -37,15 +37,16 @@ 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 set (and other fields best-effort) - when the call raised. + not a per-chunk view. ``error`` is the original raised exception (and + other fields best-effort) when the call failed — including + ``GeneratorExit`` / ``CancelledError`` for consumer close/cancellation. """ run_id: str text: str | None = None tool_calls: list[dict[str, Any]] = field(default_factory=list) usage: dict[str, Any] | None = None - error: str | None = None + error: BaseException | None = None duration_ms: int = 0 @@ -110,13 +111,13 @@ def deny(cls, message: str) -> ToolCallDecision: class ToolCallResult: """Terminal outcome of one tool invocation, as seen by ``after_tool_call``. - ``error`` is the stringified tool failure when the invocation raised or - was denied; ``result`` is unset in that case. + ``error`` is the original ``BubError`` when the invocation raised or was + denied (kind/message/details preserved); ``result`` is unset in that case. """ run_id: str tool: str arguments: dict[str, Any] result: Any = None - error: str | None = None + error: Exception | None = None duration_ms: int = 0 diff --git a/src/bub/builtin/model_runner.py b/src/bub/builtin/model_runner.py index b3eda01d..08b647cf 100644 --- a/src/bub/builtin/model_runner.py +++ b/src/bub/builtin/model_runner.py @@ -246,7 +246,7 @@ async def _fire_after_llm_call( text=output.text or None, tool_calls=[call.model_dump(exclude_none=True) for call in output.tool_calls], usage=state.usage, - error=repr(error) if error is not None else None, + error=error, duration_ms=duration_ms, ) await self.hooks.after_llm_call(request, result, state=tape.context.state) diff --git a/src/bub/tools.py b/src/bub/tools.py index 162462f0..b01a4c14 100644 --- a/src/bub/tools.py +++ b/src/bub/tools.py @@ -307,7 +307,7 @@ async def _fire_after_tool_call( tool=call.tool, arguments=call.arguments, result=None if error is not None else result, - error=error.message if error is not None else None, + error=error, duration_ms=duration_ms, ) await self._hooks.after_tool_call(call, outcome, state=state) diff --git a/tests/test_agent_hooks.py b/tests/test_agent_hooks.py index 98523374..07a5c064 100644 --- a/tests/test_agent_hooks.py +++ b/tests/test_agent_hooks.py @@ -179,7 +179,8 @@ def failing(cmd: str) -> str: await executor.execute_async([(Tool(name="bad", handler=failing, description="", parameters={}), {"cmd": "x"})]) assert observed[0].result == "ran:ls" assert observed[0].error is None - assert observed[1].error is not None + assert isinstance(observed[1].error, BubError) # original error object, kind/details preserved + assert observed[1].error.kind is not None assert "bad" in observed[1].tool @pytest.mark.asyncio @@ -310,7 +311,8 @@ async def fake_events(completion, state, output): await iterator.__anext__() await iterator.aclose() assert len(observed) == 1 - assert observed[0].error is not None # closed before completion = abnormal terminal + # original exception object preserved: consumer close surfaces as GeneratorExit + assert isinstance(observed[0].error, GeneratorExit) @pytest.mark.asyncio async def test_after_llm_call_fires_exactly_once_on_success(self) -> None: From 032c492774f19c0b00bab2374e85f654fc89504e Mon Sep 17 00:00:00 2001 From: Rubick Date: Mon, 6 Jul 2026 21:57:00 +0800 Subject: [PATCH 05/10] fix: restore BubError import dropped by unused-import autofix Co-Authored-By: Claude Fable 5 --- tests/test_agent_hooks.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_agent_hooks.py b/tests/test_agent_hooks.py index 07a5c064..2f775342 100644 --- a/tests/test_agent_hooks.py +++ b/tests/test_agent_hooks.py @@ -17,6 +17,7 @@ ) from bub.hook_runtime import AgentHooks, HookRuntime from bub.hookspecs import BUB_HOOK_NAMESPACE, BubHookSpecs, hookimpl +from bub.runtime import BubError from bub.tools import Tool, ToolExecutor From aa53057d18649b1d743aa93eb18d281c573b9a9b Mon Sep 17 00:00:00 2001 From: Rubick Date: Mon, 6 Jul 2026 22:01:50 +0800 Subject: [PATCH 06/10] fix: after_tool_call fires on cancellation; ToolCallResult.error widened to BaseException Mirrors the LLM-side exactly-once guarantee: a cancelled tool invocation (CancelledError, BaseException) now produces a terminal after_tool_call observation with the original exception, then re-raises unwrapped. Regression: real Task.cancel() on a blocking async tool. Co-Authored-By: Claude Fable 5 --- src/bub/agent_hooks.py | 8 +++++--- src/bub/tools.py | 8 +++++++- tests/test_agent_hooks.py | 34 ++++++++++++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 4 deletions(-) diff --git a/src/bub/agent_hooks.py b/src/bub/agent_hooks.py index 7aec410d..74d847d3 100644 --- a/src/bub/agent_hooks.py +++ b/src/bub/agent_hooks.py @@ -111,13 +111,15 @@ def deny(cls, message: str) -> ToolCallDecision: 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. + ``error`` is the original exception when the invocation failed: a + ``BubError`` for tool failures/denials (kind/message/details preserved), + or the raw ``BaseException`` for cancellation (``CancelledError``), + which is re-raised unwrapped. ``result`` is unset in that case. """ run_id: str tool: str arguments: dict[str, Any] result: Any = None - error: Exception | None = None + error: BaseException | None = None duration_ms: int = 0 diff --git a/src/bub/tools.py b/src/bub/tools.py index b01a4c14..446eeb34 100644 --- a/src/bub/tools.py +++ b/src/bub/tools.py @@ -230,6 +230,12 @@ async def _handle_tool_response_async( except BubError as exc: await self._fire_after_tool_call(call, hook_state, started, error=exc) raise + except BaseException as exc: + # Cancellation (CancelledError) must still produce a terminal + # after_tool_call observation — mirror the LLM-side guarantee — + # and is re-raised unwrapped. + await self._fire_after_tool_call(call, hook_state, started, error=exc) + raise else: await self._fire_after_tool_call(call, hook_state, started, result=result) return result @@ -297,7 +303,7 @@ async def _fire_after_tool_call( started: float, *, result: Any = None, - error: BubError | None = None, + error: BaseException | None = None, ) -> None: if self._hooks is None: return diff --git a/tests/test_agent_hooks.py b/tests/test_agent_hooks.py index 2f775342..220e4778 100644 --- a/tests/test_agent_hooks.py +++ b/tests/test_agent_hooks.py @@ -331,3 +331,37 @@ def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: pass assert len(observed) == 1 assert observed[0].error is None + + +class TestToolCancellation: + @pytest.mark.asyncio + async def test_after_tool_call_fires_exactly_once_on_cancel(self) -> None: + import asyncio + + observed: list[ToolCallResult] = [] + + class Observe: + @hookimpl + def after_tool_call(self, call: ToolCall, result: ToolCallResult, state: dict) -> None: + observed.append(result) + + started = asyncio.Event() + + async def blocking(cmd: str) -> str: + started.set() + await asyncio.Event().wait() # blocks until cancelled + return "unreachable" + + executor = ToolExecutor(hooks=make_hooks(Observe())) + task = asyncio.create_task( + executor.execute_async([ + (Tool(name="block", handler=blocking, description="", parameters={}), {"cmd": "x"}) + ]) + ) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert len(observed) == 1 + # original exception preserved, re-raised unwrapped + assert isinstance(observed[0].error, asyncio.CancelledError) From 8fb324d152d4d851fb4e01246272aa265c22e9b5 Mon Sep 17 00:00:00 2001 From: Rubick Date: Tue, 7 Jul 2026 08:06:40 +0800 Subject: [PATCH 07/10] =?UTF-8?q?docs:=20agent-loop=20hooks=20=E2=80=94=20?= =?UTF-8?q?reference=20table=20+=20AgentHooks=20semantics=20+=20build=20re?= =?UTF-8?q?cipes=20(en/zh)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - reference/hooks: four new table rows; 'Agent-loop hooks (AgentHooks)' section covering LIFO chaining, decision-only veto + fault isolation, exactly-once terminal observation (incl. aclose/cancellation), original exception preservation, run_id<->tape correlation, end-to-end request rewriting, and the removed wrap_tool_call boundary note - build/hooks: recipe section with model-routing/cost-guard, tool deny/replace policy, and terminal-observation metrics examples - zh-cn mirrors for both pages; website build verified Co-Authored-By: Claude Fable 5 --- website/src/content/docs/docs/build/hooks.mdx | 64 +++++++++++++++++++ .../src/content/docs/docs/reference/hooks.mdx | 16 +++++ .../content/docs/zh-cn/docs/build/hooks.mdx | 64 +++++++++++++++++++ .../docs/zh-cn/docs/reference/hooks.mdx | 16 +++++ 4 files changed, 160 insertions(+) diff --git a/website/src/content/docs/docs/build/hooks.mdx b/website/src/content/docs/docs/build/hooks.mdx index e9473e0b..0d64e5ee 100644 --- a/website/src/content/docs/docs/build/hooks.mdx +++ b/website/src/content/docs/docs/build/hooks.mdx @@ -218,6 +218,70 @@ class WeatherPlugin: Return `None` to skip without changing accumulated config. Returning a non-dict value raises `TypeError` and aborts onboarding. Verify by running `uv run bub onboard` and inspecting the saved config file. +## 9. Intercept agent-loop LLM and tool calls + +The agent loop exposes four interception hooks (see [Reference › Hooks](/docs/reference/hooks/#agent-loop-hooks-agenthooks) for the full contract). `before_*` hooks chain in LIFO order and can modify or short-circuit; `after_*` hooks observe the terminal outcome exactly once, including error and cancellation paths. + +Route models, cap spend, and redact — `before_llm_call`: + +```python +from dataclasses import replace + +from bub import hookimpl +from bub.agent_hooks import LlmCallDecision, LlmCallRequest + + +class RouteAndBudget: + @hookimpl + def before_llm_call(self, request: LlmCallRequest, state): + calls = state.get("llm_calls", 0) + state["llm_calls"] = calls + 1 + if calls >= 10: + # Skip the provider entirely; this text becomes the final output. + return LlmCallDecision.finish("Call budget exhausted for this turn.") + # Rewrites are honored end-to-end: the provider receives this model + # and max_tokens, and the tape records the effective model. + return replace(request, model="openai:gpt-4o-mini", max_tokens=512) +``` + +Guard tool calls — `before_tool_call` returns a decision, never raises to veto: + +```python +from bub import hookimpl +from bub.agent_hooks import ToolCall, ToolCallDecision + + +class ShellPolicy: + @hookimpl + def before_tool_call(self, call: ToolCall, state): + if call.tool == "shell" and "rm -rf" in str(call.arguments): + return ToolCallDecision.deny("Destructive shell commands are blocked.") + if call.tool == "web_search": + # Serve from cache without running the tool. + if cached := state.get("search_cache", {}).get(str(call.arguments)): + return ToolCallDecision.replace(cached) + return None # proceed unchanged +``` + +Observe terminal outcomes — `after_llm_call` / `after_tool_call` receive the original exception object on failure (`result.error`), including `CancelledError` when a call is cancelled mid-flight: + +```python +from bub import hookimpl + + +class Metrics: + @hookimpl + def after_llm_call(self, request, result, state): + usage = result.usage or {} + print(f"llm {request.model} {result.duration_ms}ms tokens={usage.get('total_tokens')} error={result.error!r}") + + @hookimpl + def after_tool_call(self, call, result, state): + print(f"tool {call.tool} {result.duration_ms}ms error={result.error!r}") +``` + +All payloads carry `run_id`, matching the tape entry meta, so metrics can be joined against the recorded conversation. + ## Next steps - [Hook reference](/docs/reference/hooks/) — full signature table, precedence, and sync/async rules diff --git a/website/src/content/docs/docs/reference/hooks.mdx b/website/src/content/docs/docs/reference/hooks.mdx index b8981795..4f64ede3 100644 --- a/website/src/content/docs/docs/reference/hooks.mdx +++ b/website/src/content/docs/docs/reference/hooks.mdx @@ -32,6 +32,10 @@ For the *why* and *how* of each stage see [Turn pipeline](/docs/concepts/turn-pi | `provide_channels` | sync-only consumer (deduped) | `(message_handler: MessageHandler) -> list[Channel]` | channels | `BubFramework.get_channels` (`call_many_sync`) | Channels are deduplicated by `Channel.name`; the first channel seen in hook priority order wins. | | `build_tape_context` | firstresult | `() -> TapeContext` | tape context | `BubFramework.build_tape_context` (`call_first_sync`) | Sync-only; awaitable returns are skipped. | | `admit_message` | firstresult | `(session_id, message, turn) -> AdmitDecision \| None` | turn admission decision | `ChannelManager` | Runs before channel scheduling. `None` keeps default concurrent scheduling; decision types are listed in [Types](/docs/reference/types/#turn-admission-types). | +| `before_llm_call` | chained | `(request: LlmCallRequest, state: State) -> LlmCallRequest \| LlmCallDecision \| None` | modified request or finish decision | `ModelRunner.run` via `AgentHooks` | Impls chain in LIFO order; each sees the previous impl's request. `LlmCallDecision.finish(text)` skips the provider call. Raising impls are logged and skipped. | +| `after_llm_call` | observer | `(request: LlmCallRequest, result: LlmCallResult, state: State) -> None` | none | `ModelRunner.run` via `AgentHooks` | Fires exactly once per call: success, error, consumer `aclose()` and cancellation. `result.error` is the original `BaseException`. | +| `before_tool_call` | chained | `(call: ToolCall, state: State) -> ToolCallDecision \| None` | decision | `ToolExecutor` via `AgentHooks` | Per tool invocation. `proceed(arguments=…)` folds argument changes; `replace(result)` / `deny(message)` short-circuit. Veto only via the decision object — exceptions are logged and skipped. | +| `after_tool_call` | observer | `(call: ToolCall, result: ToolCallResult, state: State) -> None` | none | `ToolExecutor` via `AgentHooks` | Fires for success, failure, deny/replace and cancellation. `result.error` is the original exception (`BubError`, or `CancelledError` re-raised unwrapped). | ## How hooks are invoked @@ -93,4 +97,16 @@ Each impl receives only the kwargs it declares. You can omit unused parameters f `notify_error` and `notify_error_sync` wrap each impl in a `try`/`except`; observer failures are logged (`hook.on_error_failed stage=… adapter=…`) but never propagate. This guarantees one broken observer does not block the others, and prevents an `on_error` from masking the original exception. +### Agent-loop hooks (`AgentHooks`) + +The four `*_llm_call` / `*_tool_call` hooks are dispatched through the `AgentHooks` facade (`src/bub/hook_runtime.py`), not `call_first`/`call_many`, with three deliberate differences: + +- **Chaining** — `before_llm_call` and `before_tool_call` run every implementation in LIFO order and *fold* modifications: each impl receives the request/call as modified by the impls before it. The first short-circuiting decision (`LlmCallDecision.finish`, `ToolCallDecision.replace`/`deny`) stops the chain. +- **Fault isolation** — a raising implementation is logged (`hook.agent_hook_failed`) and skipped, never fatal to the turn. Blocking is only expressible through decision objects, so a broken policy plugin cannot veto by crashing. +- **Exactly-once terminal observation** — `after_llm_call` and `after_tool_call` fire exactly once per call on *every* exit path, including consumer `aclose()` (`GeneratorExit`) and task cancellation (`CancelledError`). `result.error` carries the **original exception object** (a `BubError` for tool failures with `kind`/`details` intact); cancellation is re-raised unwrapped after the observation. + +Payload dataclasses (`LlmCallRequest`, `LlmCallResult`, `ToolCall`, `ToolCallDecision`, `ToolCallResult`, `LlmCallDecision`) live in `src/bub/agent_hooks.py`. Every payload carries `run_id`, matching the `run_id` meta on tape entries, so observers can correlate hook events with the tape. Rewritten `model`/`max_tokens` from `before_llm_call` are honored end-to-end: the provider receives them and the tape records the effective model. + +A continuation-style `wrap_tool_call` (retry / human-in-the-loop / caching) was prototyped in the original PR and removed by review direction; if those use cases materialize it will return as a separate proposal. + For the runtime contract behind every signature, read `src/bub/hookspecs.py`. To verify what is registered in your environment, run `bub hooks` (see [CLI › hooks](/docs/reference/cli/#hooks)). diff --git a/website/src/content/docs/zh-cn/docs/build/hooks.mdx b/website/src/content/docs/zh-cn/docs/build/hooks.mdx index dab1f61b..7376a2dd 100644 --- a/website/src/content/docs/zh-cn/docs/build/hooks.mdx +++ b/website/src/content/docs/zh-cn/docs/build/hooks.mdx @@ -218,6 +218,70 @@ class WeatherPlugin: 返回 `None` 表示跳过且不修改累积配置。返回非 dict 会抛 `TypeError` 并中止引导。运行 `uv run bub onboard` 并查看保存的配置文件来验证。 +## 9. 拦截 agent loop 的 LLM 与工具调用 + +agent loop 暴露四个拦截钩子(完整契约见 [参考 › 钩子](/zh-cn/docs/reference/hooks/))。`before_*` 按 LIFO 链式执行,可修改或短路;`after_*` 观察终态,每次调用恰好一次,覆盖异常与取消路径。 + +模型路由、限次、脱敏 —— `before_llm_call`: + +```python +from dataclasses import replace + +from bub import hookimpl +from bub.agent_hooks import LlmCallDecision, LlmCallRequest + + +class RouteAndBudget: + @hookimpl + def before_llm_call(self, request: LlmCallRequest, state): + calls = state.get("llm_calls", 0) + state["llm_calls"] = calls + 1 + if calls >= 10: + # Skip the provider entirely; this text becomes the final output. + return LlmCallDecision.finish("Call budget exhausted for this turn.") + # Rewrites are honored end-to-end: the provider receives this model + # and max_tokens, and the tape records the effective model. + return replace(request, model="openai:gpt-4o-mini", max_tokens=512) +``` + +工具策略 —— `before_tool_call` 用 decision 表达否决,不要抛异常: + +```python +from bub import hookimpl +from bub.agent_hooks import ToolCall, ToolCallDecision + + +class ShellPolicy: + @hookimpl + def before_tool_call(self, call: ToolCall, state): + if call.tool == "shell" and "rm -rf" in str(call.arguments): + return ToolCallDecision.deny("Destructive shell commands are blocked.") + if call.tool == "web_search": + # Serve from cache without running the tool. + if cached := state.get("search_cache", {}).get(str(call.arguments)): + return ToolCallDecision.replace(cached) + return None # proceed unchanged +``` + +观察终态 —— `after_llm_call` / `after_tool_call` 在失败时拿到**原始异常对象**(`result.error`),调用中途被取消时是 `CancelledError`: + +```python +from bub import hookimpl + + +class Metrics: + @hookimpl + def after_llm_call(self, request, result, state): + usage = result.usage or {} + print(f"llm {request.model} {result.duration_ms}ms tokens={usage.get('total_tokens')} error={result.error!r}") + + @hookimpl + def after_tool_call(self, call, result, state): + print(f"tool {call.tool} {result.duration_ms}ms error={result.error!r}") +``` + +所有载荷携带 `run_id`,与 tape 条目 meta 对应,指标可与记录的会话对齐。 + ## 下一步 - [钩子参考](/zh-cn/docs/reference/hooks/) —— 完整签名表、优先级与同步/异步规则 diff --git a/website/src/content/docs/zh-cn/docs/reference/hooks.mdx b/website/src/content/docs/zh-cn/docs/reference/hooks.mdx index 319af589..45416e74 100644 --- a/website/src/content/docs/zh-cn/docs/reference/hooks.mdx +++ b/website/src/content/docs/zh-cn/docs/reference/hooks.mdx @@ -32,6 +32,10 @@ description: BubHookSpecs 中每个钩子的类型、签名、返回值与调用 | `provide_channels` | sync-only consumer (deduped) | `(message_handler: MessageHandler) -> list[Channel]` | channels | `BubFramework.get_channels` (`call_many_sync`) | 按 `Channel.name` 去重；在钩子优先级顺序中最先出现的 channel 胜出。 | | `build_tape_context` | firstresult | `() -> TapeContext` | tape context | `BubFramework.build_tape_context` (`call_first_sync`) | 仅同步;awaitable 返回会被跳过。 | | `admit_message` | firstresult | `(session_id, message, turn) -> AdmitDecision \| None` | turn admission decision | `ChannelManager` | 调度 channel message 前调用。返回 `None` 保持默认并发调度；decision 类型见 [类型](/zh-cn/docs/reference/types/#turn-admission-类型)。 | +| `before_llm_call` | chained | `(request: LlmCallRequest, state: State) -> LlmCallRequest \| LlmCallDecision \| None` | 修改后的 request 或 finish 决定 | `ModelRunner.run` 经 `AgentHooks` | 实现按 LIFO 链式执行,每个实现看到的是前一个实现修改后的 request。返回 `LlmCallDecision.finish(text)` 跳过 provider 调用。实现抛异常仅记日志并跳过。 | +| `after_llm_call` | observer | `(request: LlmCallRequest, result: LlmCallResult, state: State) -> None` | 无 | `ModelRunner.run` 经 `AgentHooks` | 每次调用恰好触发一次:成功、异常、消费方 `aclose()`、取消。`result.error` 是原始 `BaseException`。 | +| `before_tool_call` | chained | `(call: ToolCall, state: State) -> ToolCallDecision \| None` | decision | `ToolExecutor` 经 `AgentHooks` | 逐次工具调用。`proceed(arguments=…)` 折叠参数修改;`replace(result)` / `deny(message)` 短路。否决只能通过 decision 对象——异常仅记日志并跳过。 | +| `after_tool_call` | observer | `(call: ToolCall, result: ToolCallResult, state: State) -> None` | 无 | `ToolExecutor` 经 `AgentHooks` | 成功、失败、deny/replace、取消都会触发。`result.error` 是原始异常(`BubError`,取消则为原样上抛的 `CancelledError`)。 | ## 钩子如何被调用 @@ -93,4 +97,16 @@ def _kwargs_for_impl(impl: Any, kwargs: dict[str, Any]) -> dict[str, Any]: `notify_error` 与 `notify_error_sync` 把每个实现包在 `try`/`except` 中;观察者失败会写入日志 (`hook.on_error_failed stage=… adapter=…`) 但不会向上传递。这样可以保证某个观察者出错不会阻塞其他观察者,也不会让 `on_error` 掩盖原始异常。 +### Agent-loop 钩子(`AgentHooks`) + +四个 `*_llm_call` / `*_tool_call` 钩子经由 `AgentHooks` 门面(`src/bub/hook_runtime.py`)分发,而不是 `call_first`/`call_many`,有三处刻意差异: + +- **链式** — `before_llm_call` 与 `before_tool_call` 按 LIFO 顺序执行全部实现并**折叠**修改:每个实现收到的是前序实现修改后的 request/call。第一个短路决定(`LlmCallDecision.finish`、`ToolCallDecision.replace`/`deny`)终止链。 +- **错误隔离** — 实现抛异常只记日志(`hook.agent_hook_failed`)并跳过,绝不炸掉回合。阻断只能通过 decision 对象表达,坏掉的策略插件无法靠 crash 否决。 +- **终态恰好一次** — `after_llm_call` 与 `after_tool_call` 在**每条**退出路径上恰好触发一次,包括消费方 `aclose()`(`GeneratorExit`)与任务取消(`CancelledError`)。`result.error` 携带**原始异常对象**(工具失败为完整 `BubError`,`kind`/`details` 保留);取消在观察后原样上抛。 + +载荷数据类(`LlmCallRequest`、`LlmCallResult`、`ToolCall`、`ToolCallDecision`、`ToolCallResult`、`LlmCallDecision`)位于 `src/bub/agent_hooks.py`。所有载荷携带 `run_id`,与 tape 条目 meta 中的 `run_id` 对应,观察者可将钩子事件与 tape 记录对齐。`before_llm_call` 改写的 `model`/`max_tokens` 全链路生效:provider 收到改写值,tape 记录 effective model。 + +continuation 风格的 `wrap_tool_call`(重试 / human-in-the-loop / 缓存)曾在原 PR 中实现,按 review 方向移除;若相关用例成熟,将作为独立提案回归。 + 每个签名的运行时契约见 `src/bub/hookspecs.py`。要核对当前环境注册的实现,运行 `bub hooks`(详见 [CLI › hooks](/zh-cn/docs/reference/cli/#hooks))。 From 63551337d3fee463ac9aa2471fcc365c356d8f0c Mon Sep 17 00:00:00 2001 From: Rubick Date: Tue, 7 Jul 2026 08:09:11 +0800 Subject: [PATCH 08/10] =?UTF-8?q?docs:=20budget=20example=20is=20session-s?= =?UTF-8?q?coped=20=E2=80=94=20say=20so?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit state persists across turns, so the counter is a session budget; the example now names it session_llm_calls, says so in the finish text, and notes how to build an actual per-turn budget. Co-Authored-By: Claude Fable 5 --- website/src/content/docs/docs/build/hooks.mdx | 9 ++++++--- website/src/content/docs/zh-cn/docs/build/hooks.mdx | 9 ++++++--- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/website/src/content/docs/docs/build/hooks.mdx b/website/src/content/docs/docs/build/hooks.mdx index 0d64e5ee..edf3bc63 100644 --- a/website/src/content/docs/docs/build/hooks.mdx +++ b/website/src/content/docs/docs/build/hooks.mdx @@ -234,11 +234,14 @@ from bub.agent_hooks import LlmCallDecision, LlmCallRequest class RouteAndBudget: @hookimpl def before_llm_call(self, request: LlmCallRequest, state): - calls = state.get("llm_calls", 0) - state["llm_calls"] = calls + 1 + # `state` persists for the session (it is not reset per turn), so + # this is a session-wide budget. For a per-turn budget, reset the + # counter yourself at a turn boundary (e.g. in build_prompt). + calls = state.get("session_llm_calls", 0) + state["session_llm_calls"] = calls + 1 if calls >= 10: # Skip the provider entirely; this text becomes the final output. - return LlmCallDecision.finish("Call budget exhausted for this turn.") + return LlmCallDecision.finish("LLM call budget exhausted for this session.") # Rewrites are honored end-to-end: the provider receives this model # and max_tokens, and the tape records the effective model. return replace(request, model="openai:gpt-4o-mini", max_tokens=512) diff --git a/website/src/content/docs/zh-cn/docs/build/hooks.mdx b/website/src/content/docs/zh-cn/docs/build/hooks.mdx index 7376a2dd..e01750d1 100644 --- a/website/src/content/docs/zh-cn/docs/build/hooks.mdx +++ b/website/src/content/docs/zh-cn/docs/build/hooks.mdx @@ -234,11 +234,14 @@ from bub.agent_hooks import LlmCallDecision, LlmCallRequest class RouteAndBudget: @hookimpl def before_llm_call(self, request: LlmCallRequest, state): - calls = state.get("llm_calls", 0) - state["llm_calls"] = calls + 1 + # `state` persists for the session (it is not reset per turn), so + # this is a session-wide budget. For a per-turn budget, reset the + # counter yourself at a turn boundary (e.g. in build_prompt). + calls = state.get("session_llm_calls", 0) + state["session_llm_calls"] = calls + 1 if calls >= 10: # Skip the provider entirely; this text becomes the final output. - return LlmCallDecision.finish("Call budget exhausted for this turn.") + return LlmCallDecision.finish("LLM call budget exhausted for this session.") # Rewrites are honored end-to-end: the provider receives this model # and max_tokens, and the tape records the effective model. return replace(request, model="openai:gpt-4o-mini", max_tokens=512) From 326f3b034957bec9b3c6ec1f20f860964de6a3f5 Mon Sep 17 00:00:00 2001 From: Rubick Date: Tue, 7 Jul 2026 08:42:13 +0800 Subject: [PATCH 09/10] =?UTF-8?q?refactor:=20after=20hooks=20observe=20Exc?= =?UTF-8?q?eption=20failures=20only=20=E2=80=94=20cancellation=20bypasses?= =?UTF-8?q?=20by=20design?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review direction: catching BaseException for terminal observation is dropped on both the LLM and tool paths. Cancellation / consumer close (CancelledError, GeneratorExit) intentionally does not fire after hooks; success, Exception failure and deny/replace still observe exactly once. Result error types narrowed back to Exception; tests inverted to pin the not-observed semantics; en/zh reference and build docs updated to match. Co-Authored-By: Claude Fable 5 --- src/bub/agent_hooks.py | 16 ++++++++-------- src/bub/builtin/model_runner.py | 11 ++++++----- src/bub/tools.py | 8 +------- tests/test_agent_hooks.py | 16 ++++++++-------- website/src/content/docs/docs/build/hooks.mdx | 4 ++-- .../src/content/docs/docs/reference/hooks.mdx | 6 +++--- .../src/content/docs/zh-cn/docs/build/hooks.mdx | 4 ++-- .../content/docs/zh-cn/docs/reference/hooks.mdx | 6 +++--- 8 files changed, 33 insertions(+), 38 deletions(-) diff --git a/src/bub/agent_hooks.py b/src/bub/agent_hooks.py index 74d847d3..6edcc1b8 100644 --- a/src/bub/agent_hooks.py +++ b/src/bub/agent_hooks.py @@ -38,15 +38,16 @@ class LlmCallResult: 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 — including - ``GeneratorExit`` / ``CancelledError`` for consumer close/cancellation. + 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: BaseException | None = None + error: Exception | None = None duration_ms: int = 0 @@ -111,15 +112,14 @@ def deny(cls, message: str) -> ToolCallDecision: class ToolCallResult: """Terminal outcome of one tool invocation, as seen by ``after_tool_call``. - ``error`` is the original exception when the invocation failed: a - ``BubError`` for tool failures/denials (kind/message/details preserved), - or the raw ``BaseException`` for cancellation (``CancelledError``), - which is re-raised unwrapped. ``result`` is unset in that case. + ``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: BaseException | None = None + error: Exception | None = None duration_ms: int = 0 diff --git a/src/bub/builtin/model_runner.py b/src/bub/builtin/model_runner.py index 08b647cf..b5da815f 100644 --- a/src/bub/builtin/model_runner.py +++ b/src/bub/builtin/model_runner.py @@ -155,7 +155,7 @@ async def iterator() -> AsyncGenerator[StreamEvent, None]: llm_started = datetime.now(UTC) after_fired = False - async def fire_after(error: BaseException | None = None) -> None: + async def fire_after(error: Exception | None = None) -> None: """Fire after_llm_call exactly once per call, on every exit path.""" nonlocal after_fired @@ -174,9 +174,10 @@ async def fire_after(error: BaseException | None = None) -> None: ) async for event in self._completion_events(completion, state, output): yield event - except BaseException as exc: - # BaseException so consumer aclose()/cancellation (GeneratorExit, - # CancelledError) still produce a terminal after_llm_call. + 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() @@ -236,7 +237,7 @@ async def _fire_after_llm_call( state: StreamState, started: datetime, tape: Tape, - error: BaseException | None = None, + error: Exception | None = None, ) -> None: if self.hooks is None: return diff --git a/src/bub/tools.py b/src/bub/tools.py index 446eeb34..c5b98e37 100644 --- a/src/bub/tools.py +++ b/src/bub/tools.py @@ -230,12 +230,6 @@ async def _handle_tool_response_async( except BubError as exc: await self._fire_after_tool_call(call, hook_state, started, error=exc) raise - except BaseException as exc: - # Cancellation (CancelledError) must still produce a terminal - # after_tool_call observation — mirror the LLM-side guarantee — - # and is re-raised unwrapped. - await self._fire_after_tool_call(call, hook_state, started, error=exc) - raise else: await self._fire_after_tool_call(call, hook_state, started, result=result) return result @@ -303,7 +297,7 @@ async def _fire_after_tool_call( started: float, *, result: Any = None, - error: BaseException | None = None, + error: Exception | None = None, ) -> None: if self._hooks is None: return diff --git a/tests/test_agent_hooks.py b/tests/test_agent_hooks.py index 220e4778..9319552d 100644 --- a/tests/test_agent_hooks.py +++ b/tests/test_agent_hooks.py @@ -287,7 +287,7 @@ def before_llm_call(self, request: LlmCallRequest, state: dict) -> LlmCallReques assert run_events[-1].payload["data"]["model"] == "anthropic:new" @pytest.mark.asyncio - async def test_after_llm_call_fires_exactly_once_on_early_close(self) -> None: + async def test_after_llm_call_not_fired_on_early_close(self) -> None: observed: list[LlmCallResult] = [] class Observe: @@ -311,9 +311,9 @@ async def fake_events(completion, state, output): iterator = events.__aiter__() await iterator.__anext__() await iterator.aclose() - assert len(observed) == 1 - # original exception object preserved: consumer close surfaces as GeneratorExit - assert isinstance(observed[0].error, GeneratorExit) + # Consumer close is intentionally NOT a terminal observation: + # after_llm_call fires only for real completions and Exception failures. + assert observed == [] @pytest.mark.asyncio async def test_after_llm_call_fires_exactly_once_on_success(self) -> None: @@ -335,7 +335,7 @@ def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: class TestToolCancellation: @pytest.mark.asyncio - async def test_after_tool_call_fires_exactly_once_on_cancel(self) -> None: + async def test_after_tool_call_not_fired_on_cancel(self) -> None: import asyncio observed: list[ToolCallResult] = [] @@ -362,6 +362,6 @@ async def blocking(cmd: str) -> str: task.cancel() with pytest.raises(asyncio.CancelledError): await task - assert len(observed) == 1 - # original exception preserved, re-raised unwrapped - assert isinstance(observed[0].error, asyncio.CancelledError) + # Cancellation is intentionally NOT a terminal observation: + # after_tool_call fires only for success, failure and deny/replace. + assert observed == [] diff --git a/website/src/content/docs/docs/build/hooks.mdx b/website/src/content/docs/docs/build/hooks.mdx index edf3bc63..7e408b9a 100644 --- a/website/src/content/docs/docs/build/hooks.mdx +++ b/website/src/content/docs/docs/build/hooks.mdx @@ -220,7 +220,7 @@ Return `None` to skip without changing accumulated config. Returning a non-dict ## 9. Intercept agent-loop LLM and tool calls -The agent loop exposes four interception hooks (see [Reference › Hooks](/docs/reference/hooks/#agent-loop-hooks-agenthooks) for the full contract). `before_*` hooks chain in LIFO order and can modify or short-circuit; `after_*` hooks observe the terminal outcome exactly once, including error and cancellation paths. +The agent loop exposes four interception hooks (see [Reference › Hooks](/docs/reference/hooks/#agent-loop-hooks-agenthooks) for the full contract). `before_*` hooks chain in LIFO order and can modify or short-circuit; `after_*` hooks observe the terminal outcome exactly once on success and `Exception` failure (cancellation is not observed). Route models, cap spend, and redact — `before_llm_call`: @@ -266,7 +266,7 @@ class ShellPolicy: return None # proceed unchanged ``` -Observe terminal outcomes — `after_llm_call` / `after_tool_call` receive the original exception object on failure (`result.error`), including `CancelledError` when a call is cancelled mid-flight: +Observe terminal outcomes — `after_llm_call` / `after_tool_call` receive the original exception object on failure (`result.error`); cancelled calls do not produce an observation: ```python from bub import hookimpl diff --git a/website/src/content/docs/docs/reference/hooks.mdx b/website/src/content/docs/docs/reference/hooks.mdx index 4f64ede3..9c1dfdf1 100644 --- a/website/src/content/docs/docs/reference/hooks.mdx +++ b/website/src/content/docs/docs/reference/hooks.mdx @@ -33,9 +33,9 @@ For the *why* and *how* of each stage see [Turn pipeline](/docs/concepts/turn-pi | `build_tape_context` | firstresult | `() -> TapeContext` | tape context | `BubFramework.build_tape_context` (`call_first_sync`) | Sync-only; awaitable returns are skipped. | | `admit_message` | firstresult | `(session_id, message, turn) -> AdmitDecision \| None` | turn admission decision | `ChannelManager` | Runs before channel scheduling. `None` keeps default concurrent scheduling; decision types are listed in [Types](/docs/reference/types/#turn-admission-types). | | `before_llm_call` | chained | `(request: LlmCallRequest, state: State) -> LlmCallRequest \| LlmCallDecision \| None` | modified request or finish decision | `ModelRunner.run` via `AgentHooks` | Impls chain in LIFO order; each sees the previous impl's request. `LlmCallDecision.finish(text)` skips the provider call. Raising impls are logged and skipped. | -| `after_llm_call` | observer | `(request: LlmCallRequest, result: LlmCallResult, state: State) -> None` | none | `ModelRunner.run` via `AgentHooks` | Fires exactly once per call: success, error, consumer `aclose()` and cancellation. `result.error` is the original `BaseException`. | +| `after_llm_call` | observer | `(request: LlmCallRequest, result: LlmCallResult, state: State) -> None` | none | `ModelRunner.run` via `AgentHooks` | Fires exactly once per completed call: success or `Exception` failure. Cancellation / consumer `aclose()` is not observed. `result.error` is the original exception. | | `before_tool_call` | chained | `(call: ToolCall, state: State) -> ToolCallDecision \| None` | decision | `ToolExecutor` via `AgentHooks` | Per tool invocation. `proceed(arguments=…)` folds argument changes; `replace(result)` / `deny(message)` short-circuit. Veto only via the decision object — exceptions are logged and skipped. | -| `after_tool_call` | observer | `(call: ToolCall, result: ToolCallResult, state: State) -> None` | none | `ToolExecutor` via `AgentHooks` | Fires for success, failure, deny/replace and cancellation. `result.error` is the original exception (`BubError`, or `CancelledError` re-raised unwrapped). | +| `after_tool_call` | observer | `(call: ToolCall, result: ToolCallResult, state: State) -> None` | none | `ToolExecutor` via `AgentHooks` | Fires for success, failure and deny/replace. Cancellation is not observed. `result.error` is the original `BubError`. | ## How hooks are invoked @@ -103,7 +103,7 @@ The four `*_llm_call` / `*_tool_call` hooks are dispatched through the `AgentHoo - **Chaining** — `before_llm_call` and `before_tool_call` run every implementation in LIFO order and *fold* modifications: each impl receives the request/call as modified by the impls before it. The first short-circuiting decision (`LlmCallDecision.finish`, `ToolCallDecision.replace`/`deny`) stops the chain. - **Fault isolation** — a raising implementation is logged (`hook.agent_hook_failed`) and skipped, never fatal to the turn. Blocking is only expressible through decision objects, so a broken policy plugin cannot veto by crashing. -- **Exactly-once terminal observation** — `after_llm_call` and `after_tool_call` fire exactly once per call on *every* exit path, including consumer `aclose()` (`GeneratorExit`) and task cancellation (`CancelledError`). `result.error` carries the **original exception object** (a `BubError` for tool failures with `kind`/`details` intact); cancellation is re-raised unwrapped after the observation. +- **Exactly-once terminal observation** — `after_llm_call` and `after_tool_call` fire exactly once per call for real completions: success or `Exception` failure. Cancellation and consumer close (`BaseException`) intentionally bypass after hooks. `result.error` carries the **original exception object** (a `BubError` for tool failures with `kind`/`details` intact). Payload dataclasses (`LlmCallRequest`, `LlmCallResult`, `ToolCall`, `ToolCallDecision`, `ToolCallResult`, `LlmCallDecision`) live in `src/bub/agent_hooks.py`. Every payload carries `run_id`, matching the `run_id` meta on tape entries, so observers can correlate hook events with the tape. Rewritten `model`/`max_tokens` from `before_llm_call` are honored end-to-end: the provider receives them and the tape records the effective model. diff --git a/website/src/content/docs/zh-cn/docs/build/hooks.mdx b/website/src/content/docs/zh-cn/docs/build/hooks.mdx index e01750d1..a62876c1 100644 --- a/website/src/content/docs/zh-cn/docs/build/hooks.mdx +++ b/website/src/content/docs/zh-cn/docs/build/hooks.mdx @@ -220,7 +220,7 @@ class WeatherPlugin: ## 9. 拦截 agent loop 的 LLM 与工具调用 -agent loop 暴露四个拦截钩子(完整契约见 [参考 › 钩子](/zh-cn/docs/reference/hooks/))。`before_*` 按 LIFO 链式执行,可修改或短路;`after_*` 观察终态,每次调用恰好一次,覆盖异常与取消路径。 +agent loop 暴露四个拦截钩子(完整契约见 [参考 › 钩子](/zh-cn/docs/reference/hooks/))。`before_*` 按 LIFO 链式执行,可修改或短路;`after_*` 观察终态,成功与 `Exception` 失败各恰好一次(取消不观察)。 模型路由、限次、脱敏 —— `before_llm_call`: @@ -266,7 +266,7 @@ class ShellPolicy: return None # proceed unchanged ``` -观察终态 —— `after_llm_call` / `after_tool_call` 在失败时拿到**原始异常对象**(`result.error`),调用中途被取消时是 `CancelledError`: +观察终态 —— `after_llm_call` / `after_tool_call` 在失败时拿到**原始异常对象**(`result.error`);被取消的调用不产生观察: ```python from bub import hookimpl diff --git a/website/src/content/docs/zh-cn/docs/reference/hooks.mdx b/website/src/content/docs/zh-cn/docs/reference/hooks.mdx index 45416e74..2fc0053b 100644 --- a/website/src/content/docs/zh-cn/docs/reference/hooks.mdx +++ b/website/src/content/docs/zh-cn/docs/reference/hooks.mdx @@ -33,9 +33,9 @@ description: BubHookSpecs 中每个钩子的类型、签名、返回值与调用 | `build_tape_context` | firstresult | `() -> TapeContext` | tape context | `BubFramework.build_tape_context` (`call_first_sync`) | 仅同步;awaitable 返回会被跳过。 | | `admit_message` | firstresult | `(session_id, message, turn) -> AdmitDecision \| None` | turn admission decision | `ChannelManager` | 调度 channel message 前调用。返回 `None` 保持默认并发调度；decision 类型见 [类型](/zh-cn/docs/reference/types/#turn-admission-类型)。 | | `before_llm_call` | chained | `(request: LlmCallRequest, state: State) -> LlmCallRequest \| LlmCallDecision \| None` | 修改后的 request 或 finish 决定 | `ModelRunner.run` 经 `AgentHooks` | 实现按 LIFO 链式执行,每个实现看到的是前一个实现修改后的 request。返回 `LlmCallDecision.finish(text)` 跳过 provider 调用。实现抛异常仅记日志并跳过。 | -| `after_llm_call` | observer | `(request: LlmCallRequest, result: LlmCallResult, state: State) -> None` | 无 | `ModelRunner.run` 经 `AgentHooks` | 每次调用恰好触发一次:成功、异常、消费方 `aclose()`、取消。`result.error` 是原始 `BaseException`。 | +| `after_llm_call` | observer | `(request: LlmCallRequest, result: LlmCallResult, state: State) -> None` | 无 | `ModelRunner.run` 经 `AgentHooks` | 每次完成的调用恰好触发一次:成功或 `Exception` 失败。取消 / 消费方 `aclose()` 不观察。`result.error` 是原始异常。 | | `before_tool_call` | chained | `(call: ToolCall, state: State) -> ToolCallDecision \| None` | decision | `ToolExecutor` 经 `AgentHooks` | 逐次工具调用。`proceed(arguments=…)` 折叠参数修改;`replace(result)` / `deny(message)` 短路。否决只能通过 decision 对象——异常仅记日志并跳过。 | -| `after_tool_call` | observer | `(call: ToolCall, result: ToolCallResult, state: State) -> None` | 无 | `ToolExecutor` 经 `AgentHooks` | 成功、失败、deny/replace、取消都会触发。`result.error` 是原始异常(`BubError`,取消则为原样上抛的 `CancelledError`)。 | +| `after_tool_call` | observer | `(call: ToolCall, result: ToolCallResult, state: State) -> None` | 无 | `ToolExecutor` 经 `AgentHooks` | 成功、失败、deny/replace 会触发。取消不观察。`result.error` 是原始 `BubError`。 | ## 钩子如何被调用 @@ -103,7 +103,7 @@ def _kwargs_for_impl(impl: Any, kwargs: dict[str, Any]) -> dict[str, Any]: - **链式** — `before_llm_call` 与 `before_tool_call` 按 LIFO 顺序执行全部实现并**折叠**修改:每个实现收到的是前序实现修改后的 request/call。第一个短路决定(`LlmCallDecision.finish`、`ToolCallDecision.replace`/`deny`)终止链。 - **错误隔离** — 实现抛异常只记日志(`hook.agent_hook_failed`)并跳过,绝不炸掉回合。阻断只能通过 decision 对象表达,坏掉的策略插件无法靠 crash 否决。 -- **终态恰好一次** — `after_llm_call` 与 `after_tool_call` 在**每条**退出路径上恰好触发一次,包括消费方 `aclose()`(`GeneratorExit`)与任务取消(`CancelledError`)。`result.error` 携带**原始异常对象**(工具失败为完整 `BubError`,`kind`/`details` 保留);取消在观察后原样上抛。 +- **终态恰好一次** — `after_llm_call` 与 `after_tool_call` 对真实完成的调用恰好触发一次:成功或 `Exception` 失败。取消与消费方关闭(`BaseException`)刻意不进入 after 钩子。`result.error` 携带**原始异常对象**(工具失败为完整 `BubError`,`kind`/`details` 保留)。 载荷数据类(`LlmCallRequest`、`LlmCallResult`、`ToolCall`、`ToolCallDecision`、`ToolCallResult`、`LlmCallDecision`)位于 `src/bub/agent_hooks.py`。所有载荷携带 `run_id`,与 tape 条目 meta 中的 `run_id` 对应,观察者可将钩子事件与 tape 记录对齐。`before_llm_call` 改写的 `model`/`max_tokens` 全链路生效:provider 收到改写值,tape 记录 effective model。 From d795a84efce077f506b93f0bd5340dd1a56207cf Mon Sep 17 00:00:00 2001 From: Rubick Date: Tue, 7 Jul 2026 08:44:51 +0800 Subject: [PATCH 10/10] docs: sync two stale docstrings with Exception-only terminal semantics Co-Authored-By: Claude Fable 5 --- src/bub/builtin/model_runner.py | 2 +- src/bub/hookspecs.py | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/bub/builtin/model_runner.py b/src/bub/builtin/model_runner.py index b5da815f..598e1211 100644 --- a/src/bub/builtin/model_runner.py +++ b/src/bub/builtin/model_runner.py @@ -156,7 +156,7 @@ async def iterator() -> AsyncGenerator[StreamEvent, None]: after_fired = False async def fire_after(error: Exception | None = None) -> None: - """Fire after_llm_call exactly once per call, on every exit path.""" + """Fire after_llm_call once per completed call (success or Exception failure); cancellation/consumer close bypasses it.""" nonlocal after_fired if after_fired: diff --git a/src/bub/hookspecs.py b/src/bub/hookspecs.py index fd1cf479..87307f6d 100644 --- a/src/bub/hookspecs.py +++ b/src/bub/hookspecs.py @@ -130,9 +130,11 @@ def before_llm_call(self, request: LlmCallRequest, state: State) -> LlmCallReque def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: State) -> None: """Observe the terminal outcome of one agent-loop LLM call. - Fires exactly once per call — for streaming completions after the - stream is fully consumed, and also on error (``result.error`` set). - Return values are ignored; exceptions are logged and skipped. + Fires exactly once per completed call — success (for streaming + completions, after the stream is fully consumed) or ``Exception`` + failure (``result.error`` set). Cancellation and consumer close + (``BaseException``) intentionally bypass this hook. Return values + are ignored; exceptions are logged and skipped. """ @hookspec