diff --git a/docs/ai-python/content/docs/basics/durable-execution.mdx b/docs/ai-python/content/docs/basics/durable-execution.mdx index 643dc9d8..6591eb78 100644 --- a/docs/ai-python/content/docs/basics/durable-execution.mdx +++ b/docs/ai-python/content/docs/basics/durable-execution.mdx @@ -104,5 +104,6 @@ output. Consider using the serverless hook flow for approvals. When a deferred `HookEvent` appears, call `ai.defer_hook(event.hook)`, persist `stream.messages`, and return the approval request to the client. On the next workflow turn, call -`ai.resolve_hook(...)` before `agent.run(...)`. The SDK replays the interrupted -assistant turn and continues when the hook reads the pre-registered resolution. +`ai.resolve_hook(...)` inside the `agent.run(...)` block before iterating the +stream. The SDK replays the interrupted assistant turn and continues when the +hook reads the pre-registered resolution. diff --git a/docs/ai-python/content/docs/basics/human-in-the-loop.mdx b/docs/ai-python/content/docs/basics/human-in-the-loop.mdx index 0424706b..931faa25 100644 --- a/docs/ai-python/content/docs/basics/human-in-the-loop.mdx +++ b/docs/ai-python/content/docs/basics/human-in-the-loop.mdx @@ -31,12 +31,10 @@ async with agent.run(model, messages) as stream: and event.hook.status == "pending" ): print(event.hook.hook_id, event.hook.metadata) - -# Later, from the approval path: -ai.resolve_hook( - "approve_tool_call_id_here", - ai.tools.ToolApproval(granted=True, reason="approved"), -) + ai.resolve_hook( + event.hook, + ai.tools.ToolApproval(granted=True, reason="approved"), + ) ``` Return a denial by resolving the hook with `granted=False`: @@ -54,6 +52,44 @@ Cancel a live hook when the waiting workflow should stop: await ai.cancel_hook("approve_tool_call_id_here", reason="client disconnected") ``` +Hook operations target the current `HookRegistry`, which is set inside the +`async with agent.run(...)` block. To resolve from a different task — a UI +callback, another endpoint — pass a registry explicitly. Create one up front +and hand it to the run (or use the run's own via `stream.hook_registry`): + +```python +registry = ai.HookRegistry() + +async with agent.run(model, messages, hook_registry=registry) as stream: + ... + +# elsewhere, e.g. in a UI callback: +ai.resolve_hook(hook_id, approval, registry=registry) +``` + +`ai.get_hook_registry()` returns the current registry (raising `LookupError` +when there is none), which is handy to capture the run's registry without +threading `stream` around: + +```python +async with agent.run(model, messages) as stream: + on_approval_needed(registry=ai.get_hook_registry()) + ... +``` + +To make a registry current for a block of code outside a run — so plain +`resolve_hook(...)` calls inside it need no `registry=` argument — use +`ai.agents.hooks.use_hook_registry`. It also works as a decorator on sync +and async functions: + +```python +from ai.agents import hooks + +@hooks.use_hook_registry(registry) +def on_decision(hook_id: str, approval: ai.tools.ToolApproval) -> None: + ai.resolve_hook(hook_id, approval) +``` + ## Build custom hooks Use `ai.hook` to define your own custom hooks to suspend execution @@ -89,15 +125,15 @@ Using the serverless flow only requires you to update the `agent.run` callsite (e.g. your serverless endpoint code). ```python -# pre-register tool approvals -for approval in approvals: - ai.resolve_hook( - approval.hook_id, - ai.tools.ToolApproval.model_validate(approval.data) - ) - # start (or replay pre-hook part of) a run async with agent.run(model, messages) as stream: + # pre-register tool approvals before iterating the stream + for approval in approvals: + ai.resolve_hook( + approval.hook_id, + ai.tools.ToolApproval.model_validate(approval.data) + ) + async for event in stream: if ( isinstance(event, ai.events.HookEvent) @@ -116,8 +152,9 @@ scheduled tasks to complete or get aborted. That way you can run approval-gated tools on serverless concurrently. Once the client has gathered payloads for all hooks, it will re-enter via the -same endpoint, which will then pre-register hook resolutions *before* calling -`agent.run` again. The loop will replay results of LLM calls and tool results +same endpoint, which will then pre-register hook resolutions inside the +`agent.run` block, before iterating the stream. The loop will replay results +of LLM calls and tool results from the first run without redoing non-deterministic and duplicating side-effects. When it hits the hook for which it has a pre-registered resolution, it will continue through without suspending. diff --git a/examples/agents/tool_approval.py b/examples/agents/tool_approval.py index fa483862..60278f33 100644 --- a/examples/agents/tool_approval.py +++ b/examples/agents/tool_approval.py @@ -45,13 +45,14 @@ async def stateless( history: list[ai.messages.Message], approvals: dict[str, ai.tools.ToolApproval] | None = None, ) -> tuple[list[ai.messages.Message], list[ai.messages.HookPart[Any]], str]: - for hook_id, approval in (approvals or {}).items(): - ai.resolve_hook(hook_id, approval) - deferred: list[ai.messages.HookPart[Any]] = [] text: list[str] = [] async with agent.run(model, history) as stream: + # Pre-register the approvals before iterating, so the replayed + # hooks find them immediately. + for hook_id, approval in (approvals or {}).items(): + ai.resolve_hook(hook_id, approval) async for event in stream: if isinstance(event, ai.events.TextDelta): text.append(event.chunk) diff --git a/examples/apps/coding_agent/main.py b/examples/apps/coding_agent/main.py index 3c977b26..df7a47e9 100644 --- a/examples/apps/coding_agent/main.py +++ b/examples/apps/coding_agent/main.py @@ -306,6 +306,10 @@ def __init__(self) -> None: # at a time so user/assistant alternation stays clean. self.pending: list[str] = [] self._worker: textual.worker.Worker[None] | None = None + # Approval prompts resolve hooks from the UI task, outside the + # run's context, so runs use this app-owned registry and the + # prompt handler passes it to resolve_hook explicitly. + self._hook_registry = ai.HookRegistry() def compose(self) -> textual.app.ComposeResult: yield ChatView() @@ -371,7 +375,9 @@ async def run_turns(self) -> None: async def _run_turn(self) -> None: """Run one agent turn, rendering events into the transcript.""" chat = self.chat - async with self.agent.run(self.model, self.messages) as stream: + async with self.agent.run( + self.model, self.messages, hook_registry=self._hook_registry + ) as stream: try: async for event in stream: if isinstance(event, ai.events.ReasoningDelta): @@ -426,6 +432,7 @@ def on_approval_prompt_decided(self, event: ApprovalPrompt.Decided) -> None: if event.granted else "denied by operator", ), + registry=self._hook_registry, ) self.chat.push("system", "approved" if event.granted else "denied") self._drop_prompt(event.hook_id) diff --git a/examples/apps/web_agent/backend/main.py b/examples/apps/web_agent/backend/main.py index 670faaf6..fb40d86b 100644 --- a/examples/apps/web_agent/backend/main.py +++ b/examples/apps/web_agent/backend/main.py @@ -63,12 +63,12 @@ async def chat(request: ChatRequest) -> fastapi.responses.StreamingResponse: """Handle chat requests and stream responses.""" messages, approvals = ai.agents.ui.ai_sdk.to_messages(request.messages) - # Pre-register hook resolutions so the agent loop's hooks find them - # immediately on the resume turn. - ai.agents.ui.ai_sdk.apply_approvals(approvals) - async def stream_response() -> AsyncGenerator[str]: async with agent_.chat_agent.run(agent_.MODEL, messages) as result: + # Pre-register hook resolutions so the agent loop's hooks + # find them immediately on the resume turn. + ai.agents.ui.ai_sdk.apply_approvals(approvals) + # We need to monitor the stream for HookEvents to abort; # since ui.ai_sdk.to_sse consumes a stream, we have a wrapper # async generator that does this check and yields the events. diff --git a/src/ai/__init__.py b/src/ai/__init__.py index 320af687..ab9b3e58 100644 --- a/src/ai/__init__.py +++ b/src/ai/__init__.py @@ -4,6 +4,7 @@ AgentTool, Context, HookDeferredException, + HookRegistry, StreamingStatusTool, StreamingTextTool, SubAgentTool, @@ -13,6 +14,7 @@ cancel_hook, defer_hook, deferred_tool_result, + get_hook_registry, hook, mcp, resolve_hook, @@ -120,6 +122,7 @@ "GeoRegion", "HTTPErrorContext", "HookDeferredException", + "HookRegistry", "ImageParams", "InferenceRequestParams", "InstallationError", @@ -185,6 +188,7 @@ "events", "file_part", "generate", + "get_hook_registry", "get_model", "get_provider", "hook", diff --git a/src/ai/agents/__init__.py b/src/ai/agents/__init__.py index b0c56367..273b83ce 100644 --- a/src/ai/agents/__init__.py +++ b/src/ai/agents/__init__.py @@ -26,8 +26,10 @@ from .hooks import ( TOOL_APPROVAL_HOOK_TYPE, HookDeferredException, + HookRegistry, cancel_hook, defer_hook, + get_hook_registry, hook, resolve_hook, ) @@ -42,6 +44,7 @@ "Context", "GatedToolCall", "HookDeferredException", + "HookRegistry", "LastAggregator", "MessageAggregator", "MessageBundle", @@ -56,6 +59,7 @@ "cancel_hook", "defer_hook", "deferred_tool_result", + "get_hook_registry", "hook", "mcp", "resolve_hook", diff --git a/src/ai/agents/_middleware.py b/src/ai/agents/_middleware.py index 93adfedb..197f5d09 100644 --- a/src/ai/agents/_middleware.py +++ b/src/ai/agents/_middleware.py @@ -48,6 +48,7 @@ from ..types import events as events_ from ..types.tools import Tool from .agent import Context + from .hooks import HookRegistry @dataclasses.dataclass(frozen=True) @@ -99,6 +100,7 @@ class HookContext: payload: type[pydantic.BaseModel] metadata: dict[str, Any] tool_call_id: str | None = None + registry: HookRegistry | None = None def __post_init__(self) -> None: object.__setattr__(self, "metadata", dict(self.metadata)) diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index d86b80e0..203762d4 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -998,9 +998,11 @@ def __init__( self, gen: AsyncGenerator[events_.AgentEvent], context: Context, + hook_registry: hooks_.HookRegistry, ) -> None: self._gen = gen self._context = context + self._hook_registry = hook_registry self._blocked = False self._deferred_hooks: dict[str, types.messages.HookPart[Any]] = {} @@ -1040,6 +1042,17 @@ async def aclose(self) -> None: def context(self) -> Context: return self._context + @property + def hook_registry(self) -> hooks_.HookRegistry: + """The run's :class:`~ai.agents.hooks.HookRegistry`. + + Inside the ``agent.run()`` block the registry is already + current, so hook operations find it implicitly. Pass this + explicitly (``resolve_hook(..., registry=...)``) when resolving + from a different task -- a UI callback, a server endpoint. + """ + return self._hook_registry + @property def blocked(self) -> bool: """Whether the run is blocked awaiting external hook resolution. @@ -1306,6 +1319,7 @@ def run( messages: list[types.messages.Message], *, params: models.InferenceRequestParams | None = None, + hook_registry: hooks_.HookRegistry | None = None, _middleware: list[middleware_._Middleware] | None = None, ) -> AbstractAsyncContextManager[AgentStream[str]]: ... @overload @@ -1316,6 +1330,7 @@ def run[T: pydantic.BaseModel]( *, output_type: type[T], params: models.InferenceRequestParams | None = None, + hook_registry: hooks_.HookRegistry | None = None, _middleware: list[middleware_._Middleware] | None = None, ) -> AbstractAsyncContextManager[AgentStream[T]]: ... def run( @@ -1325,6 +1340,7 @@ def run( *, output_type: type[pydantic.BaseModel] | None = None, params: models.InferenceRequestParams | None = None, + hook_registry: hooks_.HookRegistry | None = None, _middleware: list[middleware_._Middleware] | None = None, ) -> AbstractAsyncContextManager[AgentStream[Any]]: """Run the agent loop, yielding events to the consumer. @@ -1343,6 +1359,11 @@ def run( output_type: Optional Pydantic model the model's output must conform to. When set, ``stream.output`` validates the final assistant message's text against it. + hook_registry: The :class:`~ai.agents.hooks.HookRegistry` + for this run's hooks. Defaults to the current registry + (so a nested run joins the enclosing run's hooks), or a + fresh one at the top level. Pass a registry you own to + resolve hooks from outside the run's context. To attribute a sub-agent's events to a branch, wrap the run in ``yield_from(..., label=...)`` — the label flows via @@ -1354,6 +1375,7 @@ def run( messages, output_type=output_type, params=params, + hook_registry=hook_registry, _middleware=_middleware, ) @@ -1365,6 +1387,7 @@ async def _run( *, output_type: type[pydantic.BaseModel] | None, params: models.InferenceRequestParams | None, + hook_registry: hooks_.HookRegistry | None, _middleware: list[middleware_._Middleware] | None, ) -> AsyncIterator[AgentStream[Any]]: context = Context( @@ -1381,6 +1404,12 @@ async def _run( _populate_model_inputs(context.messages, context._agent_tools_by_name) _process_interrupted_hooks(context.messages) + registry = hook_registry + if registry is None: + # Join the current registry when there is one (a nested run + # shares the enclosing run's hooks); otherwise start fresh. + registry = hooks_._hook_registry.get(None) or hooks_.HookRegistry() + async def _real(call: Context) -> AsyncGenerator[events_.AgentEvent]: tracker = events_.RunStateTracker() source = self.loop(call) @@ -1429,9 +1458,10 @@ async def _stream() -> AsyncGenerator[events_.AgentEvent]: if mw_token is not None: middleware_.deactivate(mw_token) - # close the event generator on exit from the ``run()`` block. - astream = AgentStream(_stream(), context) - try: + async with ( + hooks_.use_hook_registry(registry), + contextlib.aclosing( + AgentStream(_stream(), context, registry) + ) as astream, + ): yield astream - finally: - await astream.aclose() diff --git a/src/ai/agents/hooks.py b/src/ai/agents/hooks.py index 57fd8f9a..4caeea64 100644 --- a/src/ai/agents/hooks.py +++ b/src/ai/agents/hooks.py @@ -21,39 +21,111 @@ from __future__ import annotations import asyncio -from typing import Any, cast +import contextvars +from typing import TYPE_CHECKING, Any, cast import pydantic -from .. import telemetry, types +from .. import telemetry, types, util from ..types import messages as messages_ from . import _middleware as middleware_ from . import runtime as runtime_ -# --------------------------------------------------------------------------- -# Module-level hook registries -# -# _live_hooks: -# Populated by hook() when it suspends inside a running agent. -# Maps hook label -> (future, metadata dict, Runtime). -# Consumed by resolve_hook() / cancel_hook() to unblock the awaiting -# coroutine. Entries are removed when the hook resolves, cancels, or -# the run completes. -# -# _pending_resolutions: -# Populated by resolve_hook() when no live hook exists yet (serverless -# re-entry: the user calls resolve_hook() *before* agent.run() replays). -# Maps hook label -> (payload type, validated resolution dict). -# Consumed by hook() at the start of execution — if a pre-registered -# resolution exists for the label, the hook returns immediately without -# suspending. Entries are removed on consumption. -# --------------------------------------------------------------------------- - -_live_hooks: dict[ - str, tuple[asyncio.Future[dict[str, Any]], dict[str, Any], runtime_.Runtime] -] = {} - -_pending_resolutions: dict[str, dict[str, Any] | BaseException] = {} +if TYPE_CHECKING: + from collections.abc import Iterator + + +class HookRegistry: + """Holds hook state: live suspensions and pre-registered resolutions. + + ``_live_hooks``: + Populated by ``hook()`` when it suspends inside a running agent. + Maps hook label -> (future, metadata dict, Runtime). + Consumed by ``resolve_hook()`` / ``cancel_hook()`` to unblock the + awaiting coroutine. Each entry is removed by the suspended hook + coroutine itself when it finishes -- resolved, cancelled, or torn + down with its run. + + ``_pending_resolutions``: + Populated by ``resolve_hook()`` when no live hook exists yet + (serverless re-entry: the user calls ``resolve_hook()`` inside + the ``agent.run()`` block, before iterating the stream). Maps + hook label -> resolution dict or exception. Consumed by + ``hook()`` at the start of execution -- if a pre-registered + resolution exists for the label, the hook returns immediately + without suspending. Entries are removed on consumption. + + Every hook operation takes an optional ``registry`` argument; when + omitted, the current registry is used. ``agent.run()`` makes its + registry current both for the agent loop and for the consumer's + ``async with`` block, and a nested run reuses the enclosing run's + registry, so hooks created anywhere in a run tree are resolvable + from the outermost consumer. Create a ``HookRegistry`` and pass + it explicitly to resolve from outside the run's context (e.g. a + UI callback on another task) or to isolate a run's hooks. + + Whoever owns a registry owns its contents: a pre-registered + resolution that no hook ever consumes stays until the registry is + dropped. The default per-run registry dies with the run; if you + keep a long-lived one, avoid pre-registering labels that may never + run. + """ + + def __init__(self) -> None: + self._live_hooks: dict[ + str, + tuple[ + asyncio.Future[dict[str, Any]], + dict[str, Any], + runtime_.Runtime, + ], + ] = {} + self._pending_resolutions: dict[ + str, dict[str, Any] | BaseException + ] = {} + + +_hook_registry: contextvars.ContextVar[HookRegistry] = contextvars.ContextVar( + "hook_registry" +) + + +def _registry(registry: HookRegistry | None) -> HookRegistry: + """Return the registry a hook operation should use.""" + if registry is not None: + return registry + try: + return _hook_registry.get() + except LookupError: + raise LookupError( + "no current HookRegistry: hook operations work inside an " + "agent.run() block, or pass an explicit registry" + ) from None + + +def get_hook_registry() -> HookRegistry: + """Return the current HookRegistry. + + Set inside an ``agent.run()`` block (and the run's loop) and within + a ``use_hook_registry()`` context. Raises LookupError when there + is none. + """ + return _registry(None) + + +@util.contextmanager_any_sync +def use_hook_registry(registry: HookRegistry) -> Iterator[None]: + """Make *registry* the current HookRegistry within this context. + + Scoped to the calling task and tasks spawned from it, and restored + on exit. ``agent.run()`` uses this to make the run's registry + current for the consumer's block. + """ + token = _hook_registry.set(registry) + try: + yield + finally: + _hook_registry.reset(token) class HookDeferredException(Exception): # noqa: N818 @@ -74,19 +146,13 @@ def _label(target: str | messages_.HookPart[Any]) -> str: return target.hook_id if isinstance(target, messages_.HookPart) else target -def cleanup_run(labels: set[str]) -> None: - """Remove all registry entries associated with a finished run.""" - for label in labels: - _live_hooks.pop(label, None) - _pending_resolutions.pop(label, None) - - async def hook[T: pydantic.BaseModel]( hook: str | messages_.HookPart[Any], *, payload: type[T], metadata: dict[str, Any] | None = None, tool_call_id: str | None = None, + registry: HookRegistry | None = None, ) -> T: """Create a hook suspension point and await its resolution. @@ -103,6 +169,8 @@ async def hook[T: pydantic.BaseModel]( (run-blocked tracking, UIs) can attribute the suspension to its tool call. Approval gating sets this automatically; custom gating wrappers should pass it too. + registry: The :class:`HookRegistry` to register the hook in. + Defaults to the current registry. """ call = middleware_.HookContext( @@ -110,6 +178,7 @@ async def hook[T: pydantic.BaseModel]( payload=payload, metadata=metadata or {}, tool_call_id=tool_call_id, + registry=registry, ) chain = middleware_._build_hook_chain(_hook_impl) @@ -120,6 +189,7 @@ async def hook[T: pydantic.BaseModel]( async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel: """Core hook logic — the innermost ``next`` in the middleware chain.""" rt = runtime_.get_runtime() + registry = _registry(call.registry) label = call.label payload = call.payload hook_metadata = call.metadata @@ -129,7 +199,7 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel: ) # Pre-registered resolution (serverless re-entry). - pre_registered = _pending_resolutions.pop(label, None) + pre_registered = registry._pending_resolutions.pop(label, None) if pre_registered is not None: async with telemetry.span(data, replay=True) as sp: if isinstance(pre_registered, BaseException): @@ -142,8 +212,7 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel: async with telemetry.span(data) as sp: future: asyncio.Future[dict[str, Any]] = asyncio.Future() - _live_hooks[label] = (future, hook_metadata, rt) - rt.track_hook_label(label) + registry._live_hooks[label] = (future, hook_metadata, rt) # Emit pending signal. hook_part: messages_.HookPart[Any] = messages_.HookPart( @@ -168,9 +237,10 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel: attrs["reason"] = exc.args[0] await sp.add_event(telemetry.HOOK_CANCELLED, **attrs) raise + finally: + # Clean up live registry. + registry._live_hooks.pop(label, None) - # Clean up live registry. - _live_hooks.pop(label, None) sp.data.status = "resolved" await sp.add_event(telemetry.HOOK_RESOLVED) @@ -194,6 +264,7 @@ def resolve_hook( data: pydantic.BaseModel | dict[str, Any] | BaseException, *, payload: type[pydantic.BaseModel] | None = None, + registry: HookRegistry | None = None, ) -> None: """Resolve a hook by label. @@ -206,7 +277,9 @@ def resolve_hook( 2. **No live hook yet** (serverless re-entry): stashes the resolution in the pre-registration registry. When ``hook()`` executes during replay, it finds the pre-registered value and returns without - suspending. + suspending. Pre-registration must target the run's registry: + either call this inside the ``async with agent.run(...)`` block + (before iterating the stream), or pass the registry explicitly. Passing an exception sends it to the awaiter (or stashes it for the next replay) so the awaiting ``ai.hook(...)`` call raises rather than @@ -220,6 +293,10 @@ def resolve_hook( exception to raise in the awaiter. payload: Optional pydantic model class for validation. Ignored when *data* is an exception. + registry: The :class:`HookRegistry` to resolve in. Defaults to + the current registry. Pass the run's registry explicitly + when resolving from a task outside the ``agent.run()`` block + (e.g. a UI callback). """ label = _label(hook) @@ -240,9 +317,11 @@ def resolve_hook( f"Expected dict or pydantic model, got {type(data).__name__}" ) + reg = _registry(registry) + # Path 1: live hook — resolve the future directly. - if label in _live_hooks: - future, _, _rt = _live_hooks[label] + if label in reg._live_hooks: + future, _, _rt = reg._live_hooks[label] if isinstance(resolution, BaseException): future.set_exception(resolution) else: @@ -250,10 +329,14 @@ def resolve_hook( return # Path 2: no live hook — pre-register for later consumption. - _pending_resolutions[label] = resolution + reg._pending_resolutions[label] = resolution -def defer_hook(hook_part: messages_.HookPart[Any]) -> None: +def defer_hook( + hook_part: messages_.HookPart[Any], + *, + registry: HookRegistry | None = None, +) -> None: """Defer the hook identified by ``hook_part.hook_id``. The deferred exception carries a :class:`HookDeferredException` wrapping @@ -264,23 +347,31 @@ def defer_hook(hook_part: messages_.HookPart[Any]) -> None: from inbound conversion) and needs to surface it back through the awaiting ``ai.hook(...)`` site as a structured suspension. """ - resolve_hook(hook_part.hook_id, HookDeferredException(hook_part)) + resolve_hook( + hook_part.hook_id, HookDeferredException(hook_part), registry=registry + ) async def cancel_hook( - hook: str | messages_.HookPart[Any], *, reason: str | None = None + hook: str | messages_.HookPart[Any], + *, + reason: str | None = None, + registry: HookRegistry | None = None, ) -> None: """Cancel a deferred hook. Only works for live hooks (long-running mode). Raises ValueError if the hook is not currently deferred. ``hook`` may be a label - string or a HookPart whose ``hook_id`` supplies it. + string or a HookPart whose ``hook_id`` supplies it. ``registry`` + selects the :class:`HookRegistry` to use, defaulting to the current + one. """ label = _label(hook) - if label not in _live_hooks: + reg = _registry(registry) + if label not in reg._live_hooks: raise ValueError(f"No deferred hook with label: {label!r}") - future, hook_metadata, rt = _live_hooks.pop(label) + future, hook_metadata, rt = reg._live_hooks.pop(label) future.cancel(reason) # Emit cancelled signal. diff --git a/src/ai/agents/runtime.py b/src/ai/agents/runtime.py index 5c5cc1e3..30ed2a63 100644 --- a/src/ai/agents/runtime.py +++ b/src/ai/agents/runtime.py @@ -9,7 +9,6 @@ from .. import util from ..types import events as events_ from ..types import messages as messages_ -from . import hooks as hooks_ from .mcp import client as mcp_client if TYPE_CHECKING: @@ -28,7 +27,6 @@ def __init__(self) -> None: self._event_queue: util.AsyncIterableQueue[events_.AgentEvent] = ( util.AsyncIterableQueue() ) - self._hook_labels: set[str] = set() async def put_event(self, event: events_.AgentEvent) -> None: await self._event_queue.put(event) @@ -40,15 +38,6 @@ async def put_hook(self, hook_part: messages_.HookPart[Any]) -> None: async def signal_done(self) -> None: await self._event_queue.astop() - def track_hook_label(self, label: str) -> None: - """Register a hook label for cleanup when the run ends.""" - self._hook_labels.add(label) - - def cleanup_hooks(self) -> None: - """Remove all hook registry entries for this run.""" - hooks_.cleanup_run(self._hook_labels) - self._hook_labels.clear() - _runtime: contextvars.ContextVar[Runtime] = contextvars.ContextVar("runtime") @@ -88,8 +77,6 @@ async def _drain() -> None: await rt.put_event(event) finally: - rt.cleanup_hooks() - await mcp_client.close_connections() mcp_client._pool.reset(mcp_token) diff --git a/src/ai/agents/ui/ai_sdk/approvals.py b/src/ai/agents/ui/ai_sdk/approvals.py index 4384f479..a4947cb0 100644 --- a/src/ai/agents/ui/ai_sdk/approvals.py +++ b/src/ai/agents/ui/ai_sdk/approvals.py @@ -5,7 +5,7 @@ from typing import Any, NamedTuple from ....types import messages as messages_ -from ...hooks import TOOL_APPROVAL_HOOK_TYPE, resolve_hook +from ...hooks import TOOL_APPROVAL_HOOK_TYPE, HookRegistry, resolve_hook from . import ui_messages ToolPart = ui_messages.UIToolPart | ui_messages.UIDynamicToolPart @@ -106,12 +106,22 @@ def extract_approvals( return approvals -def apply_approvals(approvals: list[ApprovalResponse]) -> None: - """Pre-register each approval resolution with the hooks registry.""" +def apply_approvals( + approvals: list[ApprovalResponse], + *, + registry: HookRegistry | None = None, +) -> None: + """Pre-register each approval resolution with a hook registry. + + ``registry`` defaults to the current one, so either call this + inside the ``agent.run()`` block (before iterating the stream), or + pass the registry the run will use. + """ for approval in approvals: resolve_hook( approval.hook_id, {"granted": approval.granted, "reason": approval.reason}, + registry=registry, ) diff --git a/src/ai/agents/ui/ai_sdk/inbound_messages.py b/src/ai/agents/ui/ai_sdk/inbound_messages.py index 6d96e476..0ca20cdd 100644 --- a/src/ai/agents/ui/ai_sdk/inbound_messages.py +++ b/src/ai/agents/ui/ai_sdk/inbound_messages.py @@ -538,8 +538,9 @@ def to_messages( :meth:`Agent.run` (which has the tool registry), not here. Returns ``(messages, approvals)``. The caller can pre-register - resolutions via :func:`apply_approvals` before calling - :meth:`Agent.run` if the run should resume from a hook. + resolutions via :func:`apply_approvals` inside the + :meth:`Agent.run` block (before iterating the stream) if the run + should resume from a hook. """ validate_context = ( messages_.tool_context(tools) if tools is not None else None diff --git a/tests/agents/test_hooks.py b/tests/agents/test_hooks.py index 726a2393..9710404d 100644 --- a/tests/agents/test_hooks.py +++ b/tests/agents/test_hooks.py @@ -95,9 +95,17 @@ async def loop( async def test_cancel_nonexistent_raises() -> None: - with pytest.raises(ValueError, match="No deferred hook"): + # Outside a run there is no current Runtime at all. + with pytest.raises(LookupError): await ai.cancel_hook("does_not_exist_xyz") + mock_llm([[text_msg("OK")]]) + async with ai.Agent().run(MOCK_MODEL, [ai.user_message("go")]) as stream: + with pytest.raises(ValueError, match="No deferred hook"): + await ai.cancel_hook("does_not_exist_xyz") + async for _msg in stream: + pass + # -- Pre-registration (serverless re-entry) -------------------------------- @@ -118,11 +126,79 @@ async def loop( my_agent = MyAgent() - # Pre-register BEFORE run. - ai.resolve_hook("pre_reg_1", {"approved": True}) + mock_llm([[text_msg("OK")]]) + async with my_agent.run(MOCK_MODEL, [ai.user_message("go")]) as stream: + # Pre-register before iterating: the loop hasn't started yet. + ai.resolve_hook("pre_reg_1", {"approved": True}) + async for _msg in stream: + pass + + assert resolved_value is not None + assert resolved_value.approved is True + + +# -- Explicit HookRegistry --------------------------------------------------- + + +async def test_explicit_registry_isolates_live_hook() -> None: + """A hook created in an explicit registry is invisible to the run's.""" + reg = ai.HookRegistry() + resolved_value: Confirmation | None = None + + class MyAgent(ai.Agent): + async def loop( + self, context: ai.Context + ) -> AsyncGenerator[ai.events.Event]: + nonlocal resolved_value + async with ai.models.stream(context=context) as stream: + async for event in stream: + yield event + resolved_value = await ai.hook( + "iso_1", payload=Confirmation, registry=reg + ) + + my_agent = MyAgent() mock_llm([[text_msg("OK")]]) + async with my_agent.run(MOCK_MODEL, [ai.user_message("go")]) as stream: + async for event in stream: + if not isinstance(event, agent_events_.HookEvent): + continue + if event.hook.status == "pending": + # The run's own registry doesn't know about this hook. + with pytest.raises(ValueError, match="No deferred hook"): + await ai.cancel_hook("iso_1") + ai.resolve_hook("iso_1", {"approved": True}, registry=reg) + + assert resolved_value is not None + assert resolved_value.approved is True + + +async def test_run_with_explicit_registry() -> None: + """A caller-owned registry can pre-register before the run starts.""" + reg = ai.HookRegistry() + resolved_value: Confirmation | None = None + + class MyAgent(ai.Agent): + async def loop( + self, context: ai.Context + ) -> AsyncGenerator[ai.events.Event]: + nonlocal resolved_value + async with ai.models.stream(context=context) as stream: + async for event in stream: + yield event + resolved_value = await ai.hook("pre_reg_rt", payload=Confirmation) + + my_agent = MyAgent() + + ai.resolve_hook("pre_reg_rt", {"approved": True}, registry=reg) + + mock_llm([[text_msg("OK")]]) + async with my_agent.run( + MOCK_MODEL, [ai.user_message("go")], hook_registry=reg + ) as stream: + assert stream.hook_registry is reg async for _msg in stream: pass @@ -130,6 +206,49 @@ async def loop( assert resolved_value.approved is True +# -- Nested runs ------------------------------------------------------------ + + +async def test_nested_run_joins_enclosing_registry() -> None: + """A nested run reuses the enclosing run's registry, so its hooks + are resolvable from the outermost consumer.""" + resolved_value: Confirmation | None = None + + class InnerAgent(ai.Agent): + async def loop( + self, context: ai.Context + ) -> AsyncGenerator[ai.events.Event]: + nonlocal resolved_value + async with ai.models.stream(context=context) as stream: + async for event in stream: + yield event + resolved_value = await ai.hook("nested_1", payload=Confirmation) + + class OuterAgent(ai.Agent): + async def loop( + self, context: ai.Context + ) -> AsyncGenerator[agent_events_.AgentEvent]: + inner = InnerAgent() + async with inner.run(MOCK_MODEL, [ai.user_message("sub")]) as s: + assert s.hook_registry is context_registry + async for event in s: + yield event + + mock_llm([[text_msg("OK")]]) + + async with OuterAgent().run(MOCK_MODEL, [ai.user_message("go")]) as stream: + context_registry = stream.hook_registry + async for event in stream: + if ( + isinstance(event, agent_events_.HookEvent) + and event.hook.status == "pending" + ): + ai.resolve_hook("nested_1", {"approved": True}) + + assert resolved_value is not None + assert resolved_value.approved is True + + # -- Schema validation on resolve ----------------------------------------- @@ -292,9 +411,11 @@ async def loop( async def test_pre_registered_hook_is_replay_span(recorder: Recorder) -> None: my_agent = _HookAgent() - ai.resolve_hook(my_agent.label, {"approved": True}, payload=Confirmation) mock_llm([[text_msg("OK")]]) async with my_agent.run(MOCK_MODEL, [ai.user_message("go")]) as stream: + ai.resolve_hook( + my_agent.label, {"approved": True}, payload=Confirmation + ) async for _ in stream: pass