取消并复用智能体

关闭被中断的流，清理扩展状态，并完成一次新的请求。

运行

在仓库根目录下执行：

uv run python docs/_examples/cancellation.py

输出包含：

Cancelled request; extension state cleared

这个程序在需要时使用确定性的离线模型和临时存储。它不需要 API 密钥，并由文档测试套件执行。

离开时关闭流

对于可能提前停止的流，请用aclosing包起来。该示例在第一个文本增量处跳出，退出上下文管理器，然后复用客户端完成一次完整的请求：

async def main() -> None:
    tracker = WorkTracker()
    client = await create_agent(StreamingModel(), extensions=[tracker])
    # stream() is annotated AsyncIterator but is an async generator, which aclosing needs.
    async with aclosing(client.stream("Start")) as events:  # type: ignore[type-var]
        async for event in events:
            if event.type == AgentEventType.TEXT_DELTA:
                break
    assert client.agent.state.phase == AgentPhase.CANCELLED and not tracker.active
    print("Cancelled request; extension state cleared")
    print((await client.run("Try again")).content)
    assert client.agent.state.phase == AgentPhase.COMPLETED and not tracker.active

如果扩展保留了请求级状态，那么在取消时也要像成功和出错时一样清理它。WorkTracker 演示了这三种路径。关闭流不会回滚已经完成的工具操作；设计有副作用的工具时请留意这一区别。

完整源码

下载 cancellation.py。

"""Close a streaming request and then safely reuse the same Agent."""

import asyncio
from collections.abc import AsyncIterator
from contextlib import aclosing

from zett_agent.agent import AgentRunContext
from zett_agent.client import create_agent
from zett_agent.events import (
    AgentEventType,
    AgentPhase,
)
from zett_agent.extensions.base import AgentExtension
from zett_agent.extensions.events import (
    ExtensionEvent,
    RunCancelledEvent,
)
from zett_agent.messages import AssistantMessage
from zett_agent.model import (
    ModelEvent,
    ModelRequest,
    ModelResponse,
    RetryOptions,
)


class WorkTracker(AgentExtension):
    def __init__(self) -> None:
        self.active: set[AgentRunContext] = set()

    async def before_run(self, context: AgentRunContext) -> None:
        self.active.add(context)

    async def on_event(self, context: AgentRunContext, event: ExtensionEvent) -> None:
        if isinstance(event, RunCancelledEvent):
            self.active.discard(context)

    async def on_error(self, context: AgentRunContext, error: Exception) -> None:
        self.active.discard(context)

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


class StreamingModel:
    retry = RetryOptions(max_retries=0)

    async def stream(self, request: ModelRequest) -> AsyncIterator[ModelEvent]:
        yield ModelEvent.text("Partial")
        yield ModelEvent.text(" answer")
        yield ModelEvent.completed(ModelResponse(AssistantMessage(content="Partial answer")))


async def main() -> None:
    tracker = WorkTracker()
    client = await create_agent(StreamingModel(), extensions=[tracker])
    # stream() is annotated AsyncIterator but is an async generator, which aclosing needs.
    async with aclosing(client.stream("Start")) as events:  # type: ignore[type-var]
        async for event in events:
            if event.type == AgentEventType.TEXT_DELTA:
                break
    assert client.agent.state.phase == AgentPhase.CANCELLED and not tracker.active
    print("Cancelled request; extension state cleared")
    print((await client.run("Try again")).content)
    assert client.agent.state.phase == AgentPhase.COMPLETED and not tracker.active


if __name__ == "__main__":
    asyncio.run(main())

下一步：消费端的处理模式见把流式输出接入终端或界面，扩展的清理见状态、并发与清理。