Cancel and reuse an agent

Close an interrupted stream, clean up extension state, and complete a new request.

Run it

From the repository root:

uv run python docs/_examples/cancellation.py

The output includes:

Cancelled request; extension state cleared

This program uses deterministic offline models and temporary storage where needed. It requires no API key and is executed by the documentation test suite.

Close the stream when you leave it

Use aclosing around a stream you might stop early. The example breaks on the first text delta, exits the context manager, and then reuses the client for a complete request:

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 an extension retains request-local state, clean it on cancellation as well as success and error. WorkTracker demonstrates all three paths. Closing a stream does not roll back a tool action that already completed; design side-effecting tools with that distinction in mind.

Complete source

Download 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())

Next: Stream to a terminal or UI for consumer patterns, or State, concurrency, and cleanup for extension cleanup.