Track hooks in a HookRegistry, rather than in globals

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.

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.

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)

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)

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.

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",

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",

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

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