Write your first extension#

Goal: implement /note extensions as a request-scoped capability. The model will receive a rewritten instruction, call a note tool, and emit progress events that a UI can display. Prerequisites: Define and register tools and The request lifecycle. No model API credentials are required.

When an extension is the right abstraction#

Use one for dynamic tools, context preparation, custom input commands, approvals, or persistence hooks. Do not put a text-rendering callback here unless it must participate in lifecycle decisions; AgentEventDispatcher is simpler for display.

Step 1: register a tool in on_tool#

    async def on_tool(self, context: AgentRunContext) -> None:
        @tool
        def lookup_note(topic: str) -> str:
            """Look up a small demonstration note.

            Args:
                topic: Exact note topic to look up.

            Snippet:
                lookup_note(topic="extensions")

            Guidelines:
                - Use when a user requests a project note.
            """
            return {"extensions": "Extensions customize lifecycle hooks."}.get(topic, "No note found.")

        context.register_tool(lookup_note)

context.register_tool rejects duplicate names and copies mutable schema metadata for this request. Registering through context avoids leaking tools into other sessions. The registered handler can close over context when it needs request-local information; do not store that handler on a global singleton.

All on_tool hooks finish before any on_state hook, so prompt extensions can see the complete tool registry even when they were registered earlier in the list.

Step 2: inject instructions and transform only the new input#

    async def on_state(self, context: AgentRunContext) -> None:
        context.add_message(SystemMessage(content="Use the note tool for project facts."), index=0)

    async def on_message(self, context: AgentRunContext) -> None:
        message = context.input_message
        if message is not None and message.text.startswith("/note "):
            attributes = deepcopy(message.attributes)
            attributes["original_content"] = deepcopy(message.content)
            context.input_message = UserMessage(
                content=f"Look up this note: {message.text.removeprefix('/note ')}",
                attributes=attributes,
            )
        # The Agent appends input after all on_message hooks. Do not append here.
        self.active[context] = 0

on_state is where an extension restores history or adds system instructions. on_message sees context.input_message after restoration but before append. The example keeps the original input in attributes, replaces the command with an instruction the model understands, and creates a request-local counter.

Do not append the user message here. The runtime appends the final transformed input exactly once after every on_message hook finishes. This matters for both in-memory history and append-only persistence.

Step 3: emit a UI progress event#

    async def before_model(self, context: AgentRunContext, request: ModelRequest) -> None:
        self.active[context] += 1
        await context.emit(
            AgentEvent(
                AgentEventType.CUSTOM,
                session_id=context.session_id,
                name="note.model_step",
                payload={"step": self.active[context]},
            )
        )

The hook is an ordinary async function. It runs before each model call, so a tool round trip produces two progress events. Use a namespaced custom name such as note.model_step and a small payload suitable for your UI transport. await context.emit(...) preserves hook ordering; CUSTOM does not change the Agent phase.

Step 4: clean up every terminal path#

    async def on_success(self, context: AgentRunContext, result: AssistantMessage) -> None:
        self.active.pop(context, None)

The full class below additionally handles on_error and RunCancelledEvent. An on_success-only cleanup leaks state when a user stops generation. Contexts are request identities, so two sessions can safely coexist in the same dictionary. Read State, concurrency, and cleanup for ownership and idempotent cleanup rules.

Step 5: compose with history and prompt guidance#

async def main() -> None:
    extension = NoteExtension()
    client = await create_agent(
        NoteModel(),
        extensions=[InMemoryMessageAccumulator(), extension, ToolGuidelinesExtension()],
        event_dispatcher=NoteEvents(),
    )
    reply = await client.run("/note extensions")
    assert not extension.active
    print(reply.content)
    print("Request state cleared")

NoteExtension has priority 50, ahead of default-priority 100 extensions. Hooks still execute stage by stage, not extension by extension. Explicit extension lists replace defaults, which is why memory and tool guidance appear in this list.

Run and inspect the result#

uv run python docs/_examples/custom_extension.py

Expected output:

Model step: 1
Model step: 2
Extensions customize lifecycle hooks.
Request state cleared

The scripted NoteModel checks the transformed input, preserved attributes, and injected instructions before requesting the real registered tool. The final assertion verifies request-local state was removed. To use an LLM, replace only NoteModel with a provider adapter; the extension and UI callbacks stay the same.

Complete implementation#

Download custom_extension.py.

  1"""An extension with tools, input rewriting, UI events, and request cleanup."""
  2
  3import asyncio
  4from collections.abc import AsyncIterator
  5from copy import deepcopy
  6
  7from zett_agent.agent import AgentRunContext
  8from zett_agent.client import create_agent
  9from zett_agent.dispatcher import AgentEventDispatcher
 10from zett_agent.events import (
 11    AgentEvent,
 12    AgentEventType,
 13)
 14from zett_agent.extensions.base import AgentExtension
 15from zett_agent.extensions.events import (
 16    ExtensionEvent,
 17    RunCancelledEvent,
 18)
 19from zett_agent.extensions.memory import InMemoryMessageAccumulator
 20from zett_agent.extensions.tool_guidelines import ToolGuidelinesExtension
 21from zett_agent.messages import (
 22    AssistantMessage,
 23    SystemMessage,
 24    ToolCall,
 25    ToolMessage,
 26    UserMessage,
 27)
 28from zett_agent.model import (
 29    ModelEvent,
 30    ModelRequest,
 31    ModelResponse,
 32    RetryOptions,
 33)
 34from zett_agent.tools.base import tool
 35
 36
 37class NoteExtension(AgentExtension):
 38    """Expose a request-bound note tool and translate an explicit /note command."""
 39
 40    priority = 50
 41
 42    def __init__(self) -> None:
 43        self.active: dict[AgentRunContext, int] = {}
 44
 45    async def on_tool(self, context: AgentRunContext) -> None:
 46        @tool
 47        def lookup_note(topic: str) -> str:
 48            """Look up a small demonstration note.
 49
 50            Args:
 51                topic: Exact note topic to look up.
 52
 53            Snippet:
 54                lookup_note(topic="extensions")
 55
 56            Guidelines:
 57                - Use when a user requests a project note.
 58            """
 59            return {"extensions": "Extensions customize lifecycle hooks."}.get(topic, "No note found.")
 60
 61        context.register_tool(lookup_note)
 62
 63    async def on_state(self, context: AgentRunContext) -> None:
 64        context.add_message(SystemMessage(content="Use the note tool for project facts."), index=0)
 65
 66    async def on_message(self, context: AgentRunContext) -> None:
 67        message = context.input_message
 68        if message is not None and message.text.startswith("/note "):
 69            attributes = deepcopy(message.attributes)
 70            attributes["original_content"] = deepcopy(message.content)
 71            context.input_message = UserMessage(
 72                content=f"Look up this note: {message.text.removeprefix('/note ')}",
 73                attributes=attributes,
 74            )
 75        # The Agent appends input after all on_message hooks. Do not append here.
 76        self.active[context] = 0
 77
 78    async def before_model(self, context: AgentRunContext, request: ModelRequest) -> None:
 79        self.active[context] += 1
 80        await context.emit(
 81            AgentEvent(
 82                AgentEventType.CUSTOM,
 83                session_id=context.session_id,
 84                name="note.model_step",
 85                payload={"step": self.active[context]},
 86            )
 87        )
 88
 89    async def on_success(self, context: AgentRunContext, result: AssistantMessage) -> None:
 90        self.active.pop(context, None)
 91
 92    async def on_error(self, context: AgentRunContext, error: Exception) -> None:
 93        self.active.pop(context, None)
 94
 95    async def on_event(self, context: AgentRunContext, event: ExtensionEvent) -> None:
 96        if isinstance(event, RunCancelledEvent):
 97            self.active.pop(context, None)
 98
 99
100class NoteModel:
101    retry = RetryOptions(max_retries=0)
102
103    async def stream(self, request: ModelRequest) -> AsyncIterator[ModelEvent]:
104        last = request.messages[-1]
105        if isinstance(last, ToolMessage):
106            message = AssistantMessage(content="Extensions customize lifecycle hooks.")
107        else:
108            assert isinstance(last, UserMessage)
109            assert last.text == "Look up this note: extensions"
110            assert last.attributes["original_content"] == "/note extensions"
111            assert any(isinstance(item, SystemMessage) and "note tool" in item.content for item in request.messages)
112            message = AssistantMessage(tool_calls=(ToolCall("note-1", "lookup_note", {"topic": "extensions"}),))
113        yield ModelEvent.completed(ModelResponse(message))
114
115
116class NoteEvents(AgentEventDispatcher):
117    async def on_custom_event(self, event: AgentEvent) -> None:
118        if event.name == "note.model_step" and event.payload is not None:
119            print(f"Model step: {event.payload['step']}")
120
121
122async def main() -> None:
123    extension = NoteExtension()
124    client = await create_agent(
125        NoteModel(),
126        extensions=[InMemoryMessageAccumulator(), extension, ToolGuidelinesExtension()],
127        event_dispatcher=NoteEvents(),
128    )
129    reply = await client.run("/note extensions")
130    assert not extension.active
131    print(reply.content)
132    print("Request state cleared")
133
134
135if __name__ == "__main__":
136    asyncio.run(main())

Next: Lifecycle hook reference, Emit events and accept user input, State, concurrency, and cleanup.

Related API: AgentExtension, AgentRunContext, register_tool(), register_server_tool(), and AgentEvent.