AgentExtension#

Build extensions, choose hooks, and manage request-scoped state. See the guide.

class zett_agent.extensions.base.AgentExtension[source]#

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.

See also

Write your first extension builds a complete extension, Lifecycle hook reference lists every hook boundary, and State, concurrency, and cleanup 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.

        flowchart TD
    request[NEW REQUEST] --> setup["SETUP<br/>on_tool() / on_state()<br/>on_message() / append UserMessage<br/>before_run()"]
    setup --> premodel["TURN START<br/>before_turn() / before_model()<br/>optional compaction via context.emit()"]
    premodel --> model["MODEL STEP<br/>append AssistantMessage<br/>after_model() / MODEL_COMPLETED"]
    model --> calls[Has tool calls?]
    calls -- yes --> tools["TOOL STEP<br/>before_tool()<br/>execute / after_tool()<br/>append ToolMessage / TOOL_*<br/>after_turn()"]
    tools --> premodel
    calls -- no --> turnend["TURN END<br/>after_turn()"]
    turnend --> inbox[Steering or internal input?]
    inbox -- yes --> continuation[Append selected input]
    continuation --> premodel
    inbox -- no --> success["SUCCESS<br/>after_run() / on_success()<br/>COMPLETED / RUN_COMPLETED"]
    active[ANY ACTIVE STAGE] -. exception .-> failed["FAILED<br/>on_error() / re-raise"]
    active -. cancellation .-> cancelled["CANCELLED<br/>RunCancelledEvent / re-raise"]
    request -. active scope .-> active
    external[ExternalEvent] -. emit_external_event .-> accept[accept hooks]

Agent and extension lifecycle#

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.

property pending: 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.

accept(config, event)#

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.

Parameters:
Return type:

bool

async after_model(context, response)#

Run after the complete assistant message is appended.

Parameters:
Return type:

None

async after_run(context, result)#

Run after a successful final answer is appended.

Parameters:
Return type:

None

async after_tool(context, call, result, error)#

Inspect or modify a tool result before it is appended.

error is the original execution exception for a failed result and None after successful execution. Changes made to result are included in context, persistence, and the next model request. Raising prevents the result from being appended.

Parameters:
Return type:

None

async after_turn(context, result)#

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.

Parameters:
Return type:

None

async before_model(context, request)#

Inspect the current request before preprocessing a primary model call.

Request fields are frozen. Change messages through AgentRunContext.add_message() or 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."),
        )
Parameters:
Return type:

None

async before_run(context)#

Run after the transformed user input is appended and published.

Parameters:

context (AgentRunContext)

Return type:

None

async before_tool(context, call)#

Run before one requested tool is invoked.

Parameters:
Return type:

None

async before_turn(context)#

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.

Parameters:

context (AgentRunContext)

Return type:

None

property name: 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.

async on_compact(context)#

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)
Parameters:

context (AgentRunContext)

Return type:

None

async on_error(context, error)#

Run before an agent error is propagated to the caller.

Parameters:
Return type:

None

async on_event(context, event)#

Process a published notification; inspect its concrete type with match.

Parameters:
Return type:

None

async on_message(context)#

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 AgentRunContext.add_message() or AgentRunContext.replace_messages(), and persistence subscribers have already received it.

Parameters:

context (AgentRunContext)

Return type:

None

async on_model_request(context, request, call_next)#

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.

Parameters:
Return type:

AsyncIterator[ModelEvent]

async on_state(context)#

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)
Parameters:

context (AgentRunContext)

Return type:

None

async on_success(context, result)#

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)
Parameters:
Return type:

None

async on_tool(context)#

Register local or provider-hosted tools before context restoration.

All extensions finish this hook before the first on_state() call, so prompt extensions can reliably inspect context.tools and context.server_tools as complete request-scoped registries.

Examples

Usage:

async def on_tool(self, context):
    context.register_tool(read_file)
    context.register_server_tool(ServerToolDefinition(type="web_search"))
Parameters:

context (AgentRunContext)

Return type:

None

async on_tool_call(context, call, call_next)#

Wrap one registered local tool handler and return its ToolResult.

An extension may adjust call.arguments, return without calling the next layer, transform its result, or raise into normal tool-failure handling.

Parameters:
Return type:

ToolResult