State, concurrency, and cleanup#

State has an owner and a lifetime#

Choose storage by lifetime#

Lifetime

Suitable location

Example

One callback

Local variable

A computed payload

One request

Dictionary keyed by AgentRunContext

Tool approval or model-step counter

One conversation

Map keyed by session ID or persistent storage

Accumulated message history

Application process

Immutable shared configuration or bounded sink

Model client pool or metrics exporter

One Agent can run different sessions concurrently. Same-session overlap is rejected. The agent.state convenience attribute points to the most recently prepared state, so it is not a safe way for a shared extension to select its request. Always use context.state passed to that hook.

AgentRunContext uses identity semantics, making it suitable as a request-local map key. A normal dictionary retains it strongly; remove entries on every terminal path. Do not rely on garbage collection as your cleanup protocol.

Complete cleanup pattern#

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)

discard/pop(..., None) make cleanup tolerant of partially initialized or already-released state. Cancellation is not an Exception subclass handled by on_error; it is delivered internally as RunCancelledEvent. Subclasses of the external-response base must preserve its cleanup by calling super().

Check early stream closure#

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

After the context manager closes the stream, the phase is CANCELLED and the extension’s active set is empty. A later request succeeds on the same Agent with a fresh state. Persisted prior messages are not automatically deleted or edited.

Testing an extension before shipping#

At minimum, exercise:

  1. Success and normal tool error.

  2. Failure before your state is fully initialized.

  3. Cancellation during model streaming and while waiting for external input.

  4. Two sessions using the same extension instance.

  5. Duplicate and late external replies.

  6. A second request after success, failure, and cancellation.

  7. Registration name/tool-name collisions and equal-priority ordering.

The documentation tests run these tutorial classes, not a separate copy. The core runtime suite additionally exercises state-machine and external-event races.

Avoid blocking async hooks#

Do not call synchronous terminal input or slow filesystem/network operations on the event-loop thread. Use an async library or asyncio.to_thread where appropriate. Do not add your own thread lock around an await. ExternalEventExtension already handles cross-thread response delivery safely.