Improve hooks

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:
     ...
 

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]:

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 = [

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:
 

src/vercel_ai_sdk/core/hooks.py

@@ -59,20 +59,21 @@ 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.
     """
 
     _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

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,18 +394,18 @@ def run(
     root: Callable[..., Coroutine[Any, Any, Any]],
     *args: Any,
     checkpoint: checkpoint_.Checkpoint | None = None,
-    cancel_on_hooks: bool = False,
 ) -> RunResult:
     """
     Main entry point.
 
     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.

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)
 

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
 

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