编写你的第一个扩展#

目标：把/note extensions实现为请求级能力。模型会收到重写后的指令，调用 note 工具，并发出可供 UI 显示的进度事件。前置条件：阅读定义并注册工具和请求生命周期。无需模型 API 凭据。

什么时候该用扩展#

在需要动态工具、上下文准备、自定义输入命令、审批或持久化钩子时使用它。不要在这里放文本渲染回调，除非它必须参与生命周期决策；只是用于展示时，AgentEventDispatcher 更简单。

第 1 步：在 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 会拒绝重名，并为本次请求复制可变的模式元数据。通过 context 注册可以避免把工具泄漏到其他会话。注册的处理函数在需要请求本地信息时可以闭包捕获 context；不要把该处理函数保存到全局单例上。

所有 on_tool 钩子都在任何 on_state 钩子之前结束，因此提示词扩展即使更早注册在列表中，也能看到完整的工具注册表。

第 2 步：注入指令，只转换新输入#

    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 是扩展恢复历史记录或添加系统指令的地方。on_message 在恢复之后、追加之前看到 context.input_message。示例把原始输入保存在属性中，用模型能理解的指令替换命令，并创建一个请求本地计数器。

不要在这里追加用户消息。运行时会等所有 on_message 钩子结束后，把最终转换过的输入恰好追加一次。这一点对内存中的历史记录和只追加的持久化都很重要。

第 3 步：发出 UI 进度事件#

    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]},
            )
        )

该钩子是一个普通的异步函数。它在每次模型调用之前运行，因此一次工具往返会产生两个进度事件。请使用带命名空间的自定义名称，例如 note.model_step，以及适合你的 UI 传输方式的少量载荷。await context.emit(...) 会保持钩子顺序；CUSTOM 不会改变 Agent 阶段。

第 4 步：清理每一条终止路径#

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

下面完整的类还额外处理 on_error 和 RunCancelledEvent。只在 on_success 中清理，会在用户停止生成时泄漏状态。context 就是请求的身份标识，因此两个会话可以安全地共存于同一个字典中。关于归属和幂等清理规则，请阅读状态、并发与清理。

第 5 步：与历史记录和提示词 guidelines 组合#

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 的优先级是 50，排在默认优先级 100 的扩展之前。钩子仍然按阶段执行，而不是按扩展逐个执行。显式给出的扩展列表会替换默认值，因此这个列表里出现了记忆和工具 guidelines。

运行并检查结果#

uv run python docs/_examples/custom_extension.py

预期输出:

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

预设脚本的 NoteModel 会先检查转换后的输入、保留下来的属性和注入的指令，然后再请求真正注册的工具。最后的断言验证请求本地状态已被清除。要使用 LLM，只需把 NoteModel 换成模型适配器；扩展和 UI 回调保持不变。

完整实现#

下载 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())

下一步：生命周期钩子参考、发出事件并接收用户输入、状态、并发与清理。

相关 API：AgentExtension、AgentRunContext、register_tool()、register_server_tool()和AgentEvent。