From f731ebbe6c2e869c801145714f12c6f0a87c7be6 Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Thu, 19 Mar 2026 19:28:23 -0700 Subject: [PATCH 1/2] Move future cancelling policy from runtime to per-hook --- README.md | 9 +++++---- examples/fastapi-vite/backend/main.py | 1 - skills/vercel-ai-sdk/SKILL.md | 7 ++++--- src/vercel_ai_sdk/core/hooks.py | 15 ++++++++++----- src/vercel_ai_sdk/core/runtime.py | 11 ++++++----- tests/ai_sdk_ui/test_adapter.py | 4 ++-- tests/core/test_checkpoint.py | 11 ++++++----- tests/core/test_hooks.py | 17 +++++++++++++---- 8 files changed, 46 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 301cbe3a..f2969e1d 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,7 @@ async for msg in ai.run(agent, llm, "When will the robots take over?"): ### Core Primitives -#### `ai.run(root, *args, checkpoint=None, cancel_on_hooks=False)` +#### `ai.run(root, *args, checkpoint=None)` Entry point. Starts `root` as a background task, processes the step/hook queue, yields `Message` objects. Returns a `RunResult`. @@ -97,6 +97,7 @@ Decorator that creates a suspension point from a pydantic model. The model defin ```python @ai.hook class Approval(pydantic.BaseModel): + cancels_future: ClassVar[bool] = True # cancel on suspend (serverless) granted: bool reason: str ``` @@ -124,12 +125,12 @@ if approval.granted: ... ``` -**Long-running mode** (`cancel_on_hooks=False`, the default): the `await` in `create()` blocks until `resolve()` or `cancel()` is called from external code. +**Long-running mode** (`cancels_future=False`, the default): the `await` in `create()` blocks until `resolve()` or `cancel()` is called from external code. -**Serverless mode** (`cancel_on_hooks=True`): if no resolution is available, the hook's future is cancelled and the branch dies. Inspect `result.pending_hooks` and `result.checkpoint` to resume later: +**Serverless mode** (`cancels_future=True`): if no resolution is available, the hook's future is cancelled and the branch dies. Inspect `result.pending_hooks` and `result.checkpoint` to resume later: ```python -result = ai.run(my_agent, llm, query, cancel_on_hooks=True) +result = ai.run(my_agent, llm, query) async for msg in result: ... diff --git a/examples/fastapi-vite/backend/main.py b/examples/fastapi-vite/backend/main.py index d34322f4..9107adb1 100644 --- a/examples/fastapi-vite/backend/main.py +++ b/examples/fastapi-vite/backend/main.py @@ -64,7 +64,6 @@ async def chat(request: ChatRequest) -> fastapi.responses.StreamingResponse: messages, agent.TOOLS, checkpoint=checkpoint, - cancel_on_hooks=True, ) async def stream_response() -> AsyncGenerator[str]: diff --git a/skills/vercel-ai-sdk/SKILL.md b/skills/vercel-ai-sdk/SKILL.md index da780636..e583bb95 100644 --- a/skills/vercel-ai-sdk/SKILL.md +++ b/skills/vercel-ai-sdk/SKILL.md @@ -15,7 +15,7 @@ import vercel_ai_sdk as ai ## Core workflow -`ai.run(root, *args, checkpoint=None, cancel_on_hooks=False)` is the entry point. It creates a `Runtime` (stored in a context var), starts `root` as a background task, processes an internal step queue, and yields `Message` objects. All SDK functions (`stream_step`, `execute_tool`, hooks) require this Runtime context -- they must be called within `ai.run()`. +`ai.run(root, *args, checkpoint=None)` is the entry point. It creates a `Runtime` (stored in a context var), starts `root` as a background task, processes an internal step queue, and yields `Message` objects. All SDK functions (`stream_step`, `execute_tool`, hooks) require this Runtime context -- they must be called within `ai.run()`. The root function is any async function. If it declares a param typed `ai.Runtime`, it's auto-injected. @@ -136,6 +136,7 @@ Hooks are typed suspension points for human-in-the-loop. Decorate a Pydantic mod ```python @ai.hook class Approval(pydantic.BaseModel): + cancels_future: ClassVar[bool] = True # cancel on suspend (serverless) granted: bool reason: str ``` @@ -155,9 +156,9 @@ Approval.resolve("approve_send_email", {"granted": True, "reason": "User approve Approval.cancel("approve_send_email") ``` -**Long-running mode** (`cancel_on_hooks=False`, default): `create()` blocks until `resolve()` or `cancel()` is called externally. Use for websocket/interactive UIs. +**Long-running mode** (`cancels_future=False`, default): `create()` blocks until `resolve()` or `cancel()` is called externally. Use for websocket/interactive UIs. -**Serverless mode** (`cancel_on_hooks=True`): unresolved hooks are cancelled, the run ends. Inspect `result.pending_hooks` and `result.checkpoint` to resume later. +**Serverless mode** (`cancels_future=True`): unresolved hooks are cancelled, the run ends. Inspect `result.pending_hooks` and `result.checkpoint` to resume later. Consuming hooks in the iterator: diff --git a/src/vercel_ai_sdk/core/hooks.py b/src/vercel_ai_sdk/core/hooks.py index f4001d5e..1d449e85 100644 --- a/src/vercel_ai_sdk/core/hooks.py +++ b/src/vercel_ai_sdk/core/hooks.py @@ -59,13 +59,13 @@ class Hook[T: pydantic.BaseModel]: ToolApproval.resolve("approve_delete", {"granted": True, ...}) - Behavior depends on the cancel_on_hooks flag passed to ai.run(): + Behavior depends on the ``cancels_future`` class variable: - cancel_on_hooks=False (default, long-running): the await blocks until + cancels_future=False (default, long-running): the await blocks until Hook.resolve() is called from outside the graph (e.g., websocket handler, API endpoint). - cancel_on_hooks=True (serverless): if no resolution is available, the + cancels_future=True (serverless): if no resolution is available, the hook's future is cancelled by run(). The branch receives CancelledError and dies cleanly. On re-entry, call Hook.resolve() before ai.run() to pre-register the resolution, then pass checkpoint= to replay. @@ -73,6 +73,7 @@ class Hook[T: pydantic.BaseModel]: _schema: ClassVar[type[pydantic.BaseModel]] hook_type: ClassVar[str] + cancels_future: ClassVar[bool] = False @classmethod async def create(cls, label: str, metadata: dict[str, Any] | None = None) -> T: @@ -82,8 +83,8 @@ async def create(cls, label: str, metadata: dict[str, Any] | None = None) -> T: The hook is submitted to the Runtime's step queue. run() will either: - Resolve immediately (if a resolution is available from checkpoint or pre-registered via Hook.resolve()) - - Cancel the future (cancel_on_hooks=True, serverless mode) - - Hold the future (cancel_on_hooks=False, long-running mode) + - Cancel the future (cancels_future=True, serverless mode) + - Hold the future (cancels_future=False, long-running mode) Args: label: Stable identifier for this hook. Used to match resolutions @@ -116,6 +117,7 @@ async def create(cls, label: str, metadata: dict[str, Any] | None = None) -> T: hook_type=cls.hook_type, metadata=metadata or {}, future=future, + cancels_future=cls.cancels_future, ) await rt.put_hook_suspension(suspension) @@ -236,6 +238,7 @@ def hook[T: pydantic.BaseModel](cls: type[T]) -> type[Hook[T]]: { "_schema": cls, "hook_type": cls.__name__, + "cancels_future": cls.__dict__.get("cancels_future", False), "__doc__": cls.__doc__, }, ) @@ -252,5 +255,7 @@ class ToolApproval(pydantic.BaseModel): hook system. """ + cancels_future: ClassVar[bool] = True + granted: bool reason: str | None = None diff --git a/src/vercel_ai_sdk/core/runtime.py b/src/vercel_ai_sdk/core/runtime.py index 2f6ac3e0..cc3af33d 100644 --- a/src/vercel_ai_sdk/core/runtime.py +++ b/src/vercel_ai_sdk/core/runtime.py @@ -32,6 +32,7 @@ class HookSuspension: hook_type: str metadata: dict[str, Any] future: asyncio.Future[Any] + cancels_future: bool = False # ── Runtime ─────────────────────────────────────────────────────── @@ -393,7 +394,6 @@ def run( root: Callable[..., Coroutine[Any, Any, Any]], *args: Any, checkpoint: checkpoint_.Checkpoint | None = None, - cancel_on_hooks: bool = False, ) -> RunResult: """ Main entry point. @@ -401,10 +401,11 @@ def run( 1. Starts the root function as a background task 2. Pulls steps and hook suspensions from the Runtime queue 3. Executes each step, yielding messages - 4. Resolves or suspends hooks depending on mode: - - cancel_on_hooks=True (serverless): cancel the future, branch dies, + 4. Resolves or suspends hooks depending on the hook's cancels_future + class variable: + - cancels_future=True (serverless): cancel the future, branch dies, caller inspects result.pending_hooks and result.checkpoint to resume - - cancel_on_hooks=False (long-running, default): future stays alive, + - cancels_future=False (long-running, default): future stays alive, external code calls Hook.resolve() / Hook.cancel() to unblock 5. Returns RunResult with .checkpoint and .pending_hooks """ @@ -485,7 +486,7 @@ async def _generate() -> AsyncGenerator[messages_.Message]: else: # No resolution available runtime._pending_hooks[step_item.label] = step_item - if cancel_on_hooks: + if step_item.cancels_future: # Serverless: cancel the future so the branch # dies with CancelledError. Caller inspects # result.pending_hooks to resume later. diff --git a/tests/ai_sdk_ui/test_adapter.py b/tests/ai_sdk_ui/test_adapter.py index 8a767800..267ef7cb 100644 --- a/tests/ai_sdk_ui/test_adapter.py +++ b/tests/ai_sdk_ui/test_adapter.py @@ -622,7 +622,7 @@ def test_approval_responded_resolves_hook() -> None: async def test_runtime_tool_approval_same_step() -> None: """E2E: tool-approval-request must land in the same SSE step as the tool call. - Runs a graph with ToolApproval through ai.run(cancel_on_hooks=True), + Runs a graph with ToolApproval (cancels_future=True) through ai.run(), collects runtime messages, streams through the adapter, and asserts that no spurious step boundary appears between tool-input-available and tool-approval-request. @@ -674,7 +674,7 @@ async def approve_and_execute(tc: ai.ToolPart) -> None: ) runtime_messages: list[messages.Message] = [] - result = ai.run(graph, mock_llm, cancel_on_hooks=True) + result = ai.run(graph, mock_llm) async for msg in result: runtime_messages.append(msg) diff --git a/tests/core/test_checkpoint.py b/tests/core/test_checkpoint.py index c3c209f1..4385a11b 100644 --- a/tests/core/test_checkpoint.py +++ b/tests/core/test_checkpoint.py @@ -1,7 +1,7 @@ """Checkpoint replay, hook cancellation/resolution, serialization.""" import asyncio -from typing import Any +from typing import Any, ClassVar import pydantic import pytest @@ -14,6 +14,7 @@ @ai.hook class Approval(pydantic.BaseModel): + cancels_future: ClassVar[bool] = True granted: bool @@ -82,7 +83,7 @@ async def graph(llm: ai.LanguageModel) -> Any: await ai.stream_step(llm, ai.make_messages(system="t", user="go")) return await Approval.create("my_approval", metadata={"tool": "test"}) # type: ignore[attr-defined] - result = ai.run(graph, MockLLM([[text_msg("OK")]]), cancel_on_hooks=True) + result = ai.run(graph, MockLLM([[text_msg("OK")]])) msgs = [msg async for msg in result] assert "my_approval" in result.pending_hooks hook_msgs = [m for m in msgs if any(isinstance(p, ai.HookPart) for p in m.parts)] @@ -96,7 +97,7 @@ async def graph(llm: ai.LanguageModel) -> Any: return await Approval.create("my_approval") # type: ignore[attr-defined] resp = [text_msg("OK")] - result1 = ai.run(graph, MockLLM([resp]), cancel_on_hooks=True) + result1 = ai.run(graph, MockLLM([resp])) [msg async for msg in result1] cp = result1.checkpoint @@ -122,7 +123,7 @@ async def b() -> Any: tg.create_task(a()) tg.create_task(b()) - result = ai.run(graph, MockLLM([[text_msg("OK")]]), cancel_on_hooks=True) + result = ai.run(graph, MockLLM([[text_msg("OK")]])) [msg async for msg in result] assert {"hook_a", "hook_b"} <= set(result.pending_hooks) @@ -144,7 +145,7 @@ async def b() -> Any: return ta.result(), tb.result() resp = [text_msg("OK")] - result1 = ai.run(graph, MockLLM([resp]), cancel_on_hooks=True) + result1 = ai.run(graph, MockLLM([resp])) [msg async for msg in result1] cp = result1.checkpoint diff --git a/tests/core/test_hooks.py b/tests/core/test_hooks.py index 87966a92..a51c2b55 100644 --- a/tests/core/test_hooks.py +++ b/tests/core/test_hooks.py @@ -1,7 +1,7 @@ """Hooks: live resolution, cancellation, pre-registration, schema validation.""" import asyncio -from typing import Any +from typing import Any, ClassVar import pydantic import pytest @@ -17,6 +17,13 @@ class Confirmation(pydantic.BaseModel): reason: str = "" +@ai.hook +class CancellingConfirmation(pydantic.BaseModel): + cancels_future: ClassVar[bool] = True + approved: bool + reason: str = "" + + # -- Hook.resolve() with live future (long-running mode) ------------------- @@ -32,7 +39,7 @@ async def graph(llm: ai.LanguageModel) -> None: resolved_value = result llm = MockLLM([[text_msg("OK")]]) - # Default cancel_on_hooks=False -> long-running mode + # Confirmation.cancels_future=False -> long-running mode run_result = ai.run(graph, llm) collected = [] @@ -159,9 +166,11 @@ async def graph(llm: ai.LanguageModel) -> None: async def test_hook_metadata_in_pending() -> None: async def graph(llm: ai.LanguageModel) -> None: await ai.stream_step(llm, ai.make_messages(user="go")) - await Confirmation.create("meta_test", metadata={"tool": "rm -rf", "path": "/"}) # type: ignore[attr-defined] + await CancellingConfirmation.create( # type: ignore[attr-defined] + "meta_test", metadata={"tool": "rm -rf", "path": "/"} + ) - run_result = ai.run(graph, MockLLM([[text_msg("OK")]]), cancel_on_hooks=True) + run_result = ai.run(graph, MockLLM([[text_msg("OK")]])) [m async for m in run_result] info = run_result.pending_hooks["meta_test"] From 8d750fba1d51109cb5a04bc489cfc6ca987246de Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Mon, 23 Mar 2026 10:07:00 -0700 Subject: [PATCH 2/2] Bump package version --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ba389767..ac5a0edd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "vercel-ai-sdk" -version = "0.0.1.dev7" +version = "0.0.1.dev8" description = "The AI Toolkit for Python" readme = "README.md" authors = [ diff --git a/uv.lock b/uv.lock index 4e171024..e811f862 100644 --- a/uv.lock +++ b/uv.lock @@ -1049,7 +1049,7 @@ wheels = [ [[package]] name = "vercel-ai-sdk" -version = "0.0.1.dev6" +version = "0.0.1.dev8" source = { editable = "." } dependencies = [ { name = "anthropic" },