"""Lifecycle hook Mixins and the aggregate AgentExtension contract."""
from __future__ import annotations
from collections.abc import AsyncIterator, Awaitable, Callable
from dataclasses import dataclass, field
from threading import Lock
from typing import TYPE_CHECKING
from .._compat import TypeAliasType, aclosing_stream
from ..tools.base import ToolResult
if TYPE_CHECKING:
from ..agent import AgentRunConfig, AgentRunContext
from ..messages import AssistantMessage, ToolCall, ToolMessage
from ..model import ModelEvent, ModelRequest, ModelResponse
from .events import ExtensionEvent
from .external import ExternalEvent
@dataclass(slots=True)
class _PendingRequests:
"""Thread-safe active-request identities owned by one extension instance."""
contexts: set[AgentRunContext] = field(default_factory=set)
lock: Lock = field(default_factory=Lock)
_PENDING_TRACKER_CREATION_LOCK = Lock()
ModelRequestNext = TypeAliasType("ModelRequestNext", Callable[["ModelRequest"], AsyncIterator["ModelEvent"]])
ToolCallNext = TypeAliasType("ToolCallNext", Callable[[], Awaitable[ToolResult]])
[docs]
class MiddlewareHook:
"""Wrap provider requests and local tool handlers for one extension.
Every :class:`AgentExtension` inherits these hooks. Extensions form an
onion in their existing priority order: the first extension is outermost
and the last extension is closest to the provider or tool handler.
One extension instance may serve concurrent sessions and parallel tool
calls. Keep request-local data in local variables or on ``context``, not
mutable instance attributes without synchronization.
"""
[docs]
async def on_model_request(
self,
context: AgentRunContext,
request: ModelRequest,
call_next: ModelRequestNext,
) -> AsyncIterator[ModelEvent]:
"""Wrap one provider request and its complete streamed response.
Pass ``request`` (or a dataclass replacement) to ``call_next``. The
default closes the inner stream so cancellation reaches the provider.
"""
async with aclosing_stream(call_next(request)) as events:
async for event in events:
yield event
[docs]
class AgentSetupHooksMixin:
"""Hooks that prepare request-scoped tools and conversation context."""
[docs]
async def on_state(self, context: AgentRunContext) -> None:
"""Restore history and initialize state before any input conversion.
Examples:
Usage::
async def on_state(self, context):
context.add_message(SystemMessage(content="Use concise answers."), index=0)
"""
[docs]
async def on_message(self, context: AgentRunContext) -> None:
"""Transform context.input_message after state restoration, before append.
Replace the pending UserMessage to implement externally selected modes.
Do not append or persist it here: the runtime appends the final message
exactly once after all input hooks finish. before_run sees that message
through :meth:`AgentRunContext.add_message` or
:meth:`AgentRunContext.replace_messages`, and persistence subscribers
have already received it.
"""
[docs]
class AgentRunHooksMixin:
"""Hooks around one complete request and its terminal outcome."""
[docs]
async def before_run(self, context: AgentRunContext) -> None:
"""Run after the transformed user input is appended and published."""
[docs]
async def after_run(self, context: AgentRunContext, result: AssistantMessage) -> None:
"""Run after a successful final answer is appended."""
[docs]
async def on_success(self, context: AgentRunContext, result: AssistantMessage) -> None:
"""Run once after all after_run hooks succeed, before RUN_COMPLETED.
This callback observes a completed request, not an intermediate model
step. Failed or cancelled requests do not trigger it. Callback errors
propagate through on_error and prevent RUN_COMPLETED from being emitted.
Examples:
Usage::
async def on_success(self, context, result):
await save_answer(context.session_id, result.content)
"""
[docs]
async def on_error(self, context: AgentRunContext, error: Exception) -> None:
"""Run before an agent error is propagated to the caller."""
[docs]
class AgentTurnHooksMixin:
"""Hooks around one model/tool turn inside a request.
A turn is one primary model call plus every tool invocation that call
requested. One request runs one or more turns::
before_run() -> [before_turn() -> model step -> tools -> after_turn()] * n -> after_run()
Later turns follow steering or queued internal input that continues the same
request. A request that uses no tools, receives no further input, and is not
steered has exactly one turn, whose turn hooks sit between before_run and
after_run.
"""
[docs]
async def before_turn(self, context: AgentRunContext) -> None:
"""Run before this turn's model call.
Runs after before_run() on the first turn and after the previous turn's
after_turn() on later ones. Model request assembly, compaction, and
before_model() follow this hook, so context edited here reaches the
provider.
"""
[docs]
async def after_turn(self, context: AgentRunContext, result: AssistantMessage) -> None:
"""Run once the turn's assistant message and tool results are settled.
``result`` is the assistant message that opened the turn: either the
final answer, or the tool-calling message whose ToolMessages were
appended. Every requested call is settled first, including calls that
failed, and calls that steering skipped, which arrive as unsuccessful or
skipped ToolMessages.
Steering that interrupts a turn is accepted while the runtime settles
that batch, so its UserMessage is already appended when this hook runs
and the steering turn follows as the next turn.
A turn that never completed, because the request failed or was
cancelled, has no after_turn call; use on_error and RunCancelledEvent for
those paths.
"""
[docs]
class AgentModelHooksMixin:
"""Hooks immediately before and after each primary model invocation."""
[docs]
async def before_model(self, context: AgentRunContext, request: ModelRequest) -> None:
"""Inspect the current request before preprocessing a primary model call.
Request fields are frozen. Change messages through
:meth:`AgentRunContext.add_message` or
:meth:`AgentRunContext.replace_messages`, local tools through
context.tools, and provider tools through context.server_tools; the
runtime rebuilds these fields before the next hook and before calling
the provider. This is a shallow request view, not an immutable copy of
each message or tool definition.
Examples:
Inspect the tools offered in this iteration::
async def before_model(self, context, request):
offered = {definition.name for definition in request.tools}
if "write_file" in offered:
context.add_message(
SystemMessage(content="Read existing files before editing."),
)
"""
[docs]
async def after_model(self, context: AgentRunContext, response: ModelResponse) -> None:
"""Run after the complete assistant message is appended."""
class AgentMaintenanceHooksMixin:
"""Hooks for work a caller asks for outside a model turn."""
async def on_compact(self, context: AgentRunContext) -> None:
"""Summarize restored context when a caller asks for compaction.
``Agent.compact`` restores the session through ``on_tool`` and
``on_state`` exactly as a request does, appends no input message, and
makes no primary model call: the only work a compaction pass does is
what this hook performs. The default does nothing, so an extension with
nothing to compact pays nothing.
Examples:
Usage::
async def on_compact(self, context):
await self.compact(context)
"""
[docs]
class AgentEventHooksMixin:
"""Hooks for internal notifications and external input."""
[docs]
def accept(self, config: AgentRunConfig | None, event: ExternalEvent) -> bool:
"""Handle one external event and report whether it was accepted.
The Agent broadcasts each event to every registered extension. Override
this synchronous hook when an extension waits for input from another
thread, HTTP request, TUI, or Web UI. The default ignores the event.
Config is passed separately and can be None before initialization.
Each extension decides whether the supplied identity and payload match
its protocol. Return True to include this extension's name in the
broadcast result; returning False does not stop later receivers.
"""
return False
[docs]
async def on_event(self, context: AgentRunContext, event: ExtensionEvent) -> None:
"""Process a published notification; inspect its concrete type with match."""
[docs]
class AgentExtension(
MiddlewareHook,
AgentSetupHooksMixin,
AgentRunHooksMixin,
AgentTurnHooksMixin,
AgentModelHooksMixin,
AgentToolHooksMixin,
AgentEventHooksMixin,
AgentMaintenanceHooksMixin,
):
"""Combine all optional hooks for the model-tool request lifecycle.
.. note::
Override only the hooks needed by one concern. Hook groups are barriers:
every ``on_tool`` completes before any ``on_state`` starts, regardless
of extension priority. Lifecycle hooks use ``async def`` even with
SyncAgent; they receive the original AgentRunContext on the runtime loop.
External-event ``accept`` remains a synchronous method.
.. seealso::
:doc:`/extending/first-extension` builds a complete extension,
:doc:`/extending/hooks` lists every hook boundary, and
:doc:`/extending/state` covers concurrency and cleanup.
``priority`` controls the order in which the Agent invokes extensions.
Lower numbers run first; extensions with the same priority retain their
registration order. Override the class attribute for one extension type or
assign it on an instance when one registration needs a different order::
class RestoreHistory(AgentExtension):
priority = 10
metrics = MetricsExtension()
metrics.priority = 200
``[E]`` marks an internal event published to extensions. Agent-owned events
use the default broadcast; extension-authored events may target one unique
extension name.
.. zett-diagram:: extension-lifecycle
+-----------------------------------------+
| NEW REQUEST |
+-----------------------------------------+
|
v
+-----------------------------------------+
| SETUP |
| on_tool() -> on_state() |
| on_message() -> append UserMessage |
| before_run() |
+-----------------------------------------+
|
v
+-----------------------------------------+
| TURN START (PRE-MODEL) |<------------------------------------------------------+
| before_turn() | |
| before_model() / compaction | |
+-----------------------------------------+ |
| |
v |
+-----------------------------------------+ |
| MODEL STEP | |
| on_model_request() -> provider | |
| append AssistantMessage | |
| after_model() | |
+-----------------------------------------+ |
| |
v |
+-----------------------------------------+ +-------------------------------------+ |
| HAS TOOL CALLS? |-- yes --->| TOOL STEP | |
+-----------------------------------------+ | before_tool() | |
| | on_tool_call() -> handler | |
| | after_tool() |----+
| | append ToolMessage / TOOL_* | |
| | after_turn() | |
| +-------------------------------------+ |
| |
| |
v |
+-----------------------------------------+ |
| TURN END | |
| after_turn() | |
+-----------------------------------------+ |
| |
v |
+-----------------------------------------+ +-------------------------------------+ |
| STEERING OR INTERNAL INPUT? |-- yes --->| SELECT INPUT | |
+-----------------------------------------+ | append typed input |-----+
| | reset iteration budget |
| +-------------------------------------+
| empty
|
|
v
+-----------------------------------------+
| SUCCESS |
| after_run() -> on_success() |
| COMPLETED [E] -> RUN_COMPLETED |
+-----------------------------------------+
+-----------------------------+ +-------------------------------------+
| ANY ACTIVE STAGE |-- exception --------->| FAILED [E] |
+-----------------------------+ | on_error() -> re-raise |
| +-------------------------------------+
|
|
| +-------------------------------------+
+- cancellation ---------------------->| CANCELLED [E] |
| RunCancelledEvent -> re-raise |
+-------------------------------------+
+-----------------------------+ +-------------------------------------+
| ExternalEvent |-- emit_external_event | accept() on each extension |
+-----------------------------+ +-------------------------------------+
All on_tool() hooks finish before on_state() restores history and system
instructions. All on_state() hooks finish before on_message() transforms
context.input_message. The runtime then appends and publishes the transformed
UserMessage exactly once, followed by before_run(). Each turn then runs
before_turn(), one model step, every requested tool, and after_turn() before
the runtime looks for steering or queued internal input. Because the final
answer ends its own turn, after_turn() precedes after_run(). Each hook
runs by ascending priority; equal priorities retain registration order. The
right return line runs after every tool in the response has been processed,
or after one AgentMessage is appended with a fresh iteration budget.
Internal messages have a per-request count limit. Their processing starts
after the current model/tool loop completes. SteeringExtension is checked
after MODEL_COMPLETED and every TOOL_COMPLETED/TOOL_FAILED, in READY. Urgent
user input takes priority over internal input and the remaining tool calls.
Those calls receive skipped ToolMessages [E] and TOOL_SKIPPED events before
the steering UserMessage [E] and STEERING_STARTED event. The loop then returns
to MODEL STEP with a fresh iteration budget. An active internal/steering input
emits ``*_INTERRUPTED`` when superseded, or ``*_COMPLETED`` after a final answer.
Extensions emit outward events with ``await context.emit(event)`` from an
ordinary async lifecycle hook. Compaction events enter COMPACTING and return
to READY. Only Agent.stream consumes the request event queue and yields to
callers. [E] marks awaited, sequential delivery through the separate async
on_event() hook. AgentRunContext.publish broadcasts by default or routes to
one named extension when ``target`` is supplied.
Tool execution errors become failed ToolMessages and still run after_tool(),
which receives the original exception separately and can modify the message
before it is appended and published.
Unhandled model or hook errors take the exception exit. Cancellation skips
on_error(), after_run(), and on_success(). Errors raised by subscribers after
a terminal phase was committed do not change that terminal phase.
Setup, run, turn, model, tool, event, and middleware hooks are supplied by
their corresponding base classes. Subclasses override only the hooks they
need; lifecycle defaults are no-ops and middleware defaults call the next
layer. External events are independent from a request's internal [E]
notifications: callers emit them through the Agent instead of addressing an
extension.
"""
priority: int = 100
@property
def pending(self) -> bool:
"""Return whether this extension is participating in an active request.
The Agent updates this state for every registered extension, including
extensions whose hooks are currently idle. One extension instance may
be shared by multiple Agents and sessions; ``pending`` remains true
until all requests using that instance have finished, failed, or been
cancelled. Persistent data retained between requests does not count.
"""
tracker = self._pending_request_tracker()
with tracker.lock:
return bool(tracker.contexts)
def _request_started(self, context: AgentRunContext) -> None:
"""Record one active request; called only by the Agent runtime."""
tracker = self._pending_request_tracker()
with tracker.lock:
tracker.contexts.add(context)
def _request_finished(self, context: AgentRunContext) -> None:
"""Forget one terminal request without affecting concurrent sessions."""
tracker = self._pending_request_tracker()
with tracker.lock:
tracker.contexts.discard(context)
def _pending_request_tracker(self) -> _PendingRequests:
"""Lazily allocate tracking because subclasses need not call super().__init__."""
with _PENDING_TRACKER_CREATION_LOCK:
tracker = getattr(self, "_extension_pending_requests", None)
if tracker is None:
tracker = _PendingRequests()
self._extension_pending_requests = tracker
return tracker
@property
def name(self) -> str:
"""Unique registration identity, defaulting to the concrete class name.
Override with a class attribute or assign an instance name to register
independently configured instances of the same extension type.
"""
return getattr(self, "_extension_name", type(self).__name__)
@name.setter
def name(self, value: str) -> None:
self._extension_name = value
