feat(core): Add AgentCancel exception for control-flow in callbacks (#894)

* feat(core): Add AgentCancel exception for control-flow in callbacks

Callbacks can now raise exceptions inheriting from AgentCancel to stop
agent execution. These exceptions propagate directly to the caller
without being wrapped in AgentRunError, allowing users to catch them
by their specific type.

- Add AgentCancel ABC with trace property for accessing spans
- Modify run_async to preserve AgentCancel exceptions
- Update AgentRunError docstring for clarity

* test(unit): Add tests for AgentCancel exception

- Test AgentCancel ABC cannot be instantiated directly
- Test subclasses can be instantiated and caught
- Test trace property and message preservation
- Test run_async preserves AgentCancel without wrapping
- Test regular exceptions are wrapped in AgentRunError

* test(integration): Add AgentCancel integration test

Verify AgentCancel subclasses propagate without being wrapped in
AgentRunError across all agent frameworks.

* docs: Document AgentCancel for stopping agent execution

- Add "Stopping Execution" section to callbacks documentation
- Document AgentCancel vs regular exception patterns
- Add examples for both exception types
- Add AgentCancel to API reference

* fix(tools): Update type: ignore for isinstance tuple type narrowing

The isinstance check using a runtime-constructed tuple doesn't allow
mypy to narrow the type, causing a no-any-return error instead of
the original return-value error.

* fix(openai): Remove type: ignore now that any-llm-sdk 1.7.0 supports xhigh

Author: peteski22

docs/agents/callbacks.md

@@ -48,20 +48,90 @@ class CountSearchWeb(Callback):
         return context
 ```
 
-Callbacks can raise exceptions to stop agent execution. This is useful for implementing safety guardrails or validation logic:
+## Stopping Execution
+
+Callbacks can raise exceptions to stop agent execution. This is useful for implementing safety guardrails or validation logic.
+
+!!! warning "Exceptions act as a circuit breaker"
+
+    Raising any exception from a callback immediately halts the agent loop. Use this intentionally to enforce limits or abort on invalid states.
+
+### Using `AgentCancel` (Recommended)
+
+For intentional cancellation (rate limits, guardrails, validation), subclass [`AgentCancel`][any_agent.AgentCancel]. These exceptions propagate directly to your code, allowing you to catch them by their specific type:
 
 ```python
+from any_agent import AgentCancel, AgentConfig, AnyAgent
+from any_agent.callbacks import Callback
+from any_agent.callbacks.context import Context
+
+class SearchLimitReached(AgentCancel):
+    """Raised when the search limit is exceeded."""
+
 class LimitSearchWeb(Callback):
     def __init__(self, max_calls: int):
         self.max_calls = max_calls
 
     def before_tool_execution(self, context: Context, *args, **kwargs) -> Context:
-        if context.shared["search_web_count"] > self.max_calls:
-            raise RuntimeError("Reached limit of `search_web` calls.")
+        if context.shared.get("search_web_count", 0) > self.max_calls:
+            raise SearchLimitReached(f"Exceeded {self.max_calls} search calls")
+        return context
+
+# In your application code:
+agent = AnyAgent.create(
+    "tinyagent",
+    AgentConfig(
+        model_id="gpt-4.1-nano",
+        callbacks=[LimitSearchWeb(max_calls=3)],
+    ),
+)
+try:
+    trace = agent.run("Find information about Python")
+except SearchLimitReached as e:
+    print(f"Search limit reached: {e}")
+    print(f"Trace: {e.trace}")  # Access spans collected before cancellation
+```
+
+### Using Regular Exceptions
+
+Regular exceptions (like `RuntimeError`) are automatically wrapped in [`AgentRunError`][any_agent.AgentRunError] by the framework, which provides access to the execution trace but requires you to inspect the wrapped exception:
+
+```python
+from any_agent import AgentConfig, AgentRunError, AnyAgent
+from any_agent.callbacks import Callback
+from any_agent.callbacks.context import Context
+
+class LimitSearchWeb(Callback):
+    def __init__(self, max_calls: int):
+        self.max_calls = max_calls
+
+    def before_tool_execution(self, context: Context, *args, **kwargs) -> Context:
+        if context.shared.get("search_web_count", 0) > self.max_calls:
+            msg = "Reached limit of `search_web` calls."
+            raise RuntimeError(msg)
+        return context
+
+# In your application code:
+agent = AnyAgent.create(
+    "tinyagent",
+    AgentConfig(
+        model_id="gpt-4.1-nano",
+        callbacks=[LimitSearchWeb(max_calls=3)],
+    ),
+)
+try:
+    trace = agent.run("Find information about Python")
+except AgentRunError as e:
+    print(f"Error: {e.original_exception}")
+    print(f"Trace: {e.trace}")
 ```
-!!! warning
 
-    Raising an exception is the standard way to halt execution. This effectively acts as a 'circuit breaker' for your agent.
+!!! tip "Choosing the right exception type"
+
+    - **`AgentCancel`**: Use when cancellation is expected behavior and you want to handle it distinctly (e.g., rate limits, safety guardrails).
+    - **Regular exceptions**: Use when something unexpected goes wrong and you want consistent error handling via `AgentRunError`.
+
+    Both expose the execution trace via `.trace` for debugging and inspection.
 
 ## Inspecting Data (`Context.current_span`)
 
@@ -264,46 +334,57 @@ You can find a working example in the [Callbacks Cookbook](../cookbook/callbacks
 
 ### Limit the number of steps
 
-Some agent frameworks allow to limit how many steps an agent can take and some don't. In addition,
-each framework defines a `step` differently: some count the llm calls, some the tool executions,
+Some agent frameworks allow you to limit how many steps an agent can take and some don't. In addition,
+each framework defines a `step` differently: some count the LLM calls, some the tool executions,
 and some the sum of both.
 
 You can use callbacks to limit how many steps an agent can take, and you can decide what to count
 as a `step`:
 
 ```python
+from any_agent import AgentCancel
 from any_agent.callbacks.base import Callback
 from any_agent.callbacks.context import Context
 
+
+class LLMCallLimitReached(AgentCancel):
+    """Raised when the LLM call limit is exceeded."""
+
+
+class ToolExecutionLimitReached(AgentCancel):
+    """Raised when the tool execution limit is exceeded."""
+
+
 class LimitLLMCalls(Callback):
     def __init__(self, max_llm_calls: int) -> None:
         self.max_llm_calls = max_llm_calls
 
     def before_llm_call(self, context: Context, *args, **kwargs) -> Context:
-
         if "n_llm_calls" not in context.shared:
             context.shared["n_llm_calls"] = 0
 
         context.shared["n_llm_calls"] += 1
 
         if context.shared["n_llm_calls"] > self.max_llm_calls:
-            raise RuntimeError("Reached limit of LLM Calls")
+            raise LLMCallLimitReached(f"Exceeded {self.max_llm_calls} LLM calls")
 
         return context
 
+
 class LimitToolExecutions(Callback):
     def __init__(self, max_tool_executions: int) -> None:
         self.max_tool_executions = max_tool_executions
 
     def before_tool_execution(self, context: Context, *args, **kwargs) -> Context:
-
         if "n_tool_executions" not in context.shared:
             context.shared["n_tool_executions"] = 0
 
         context.shared["n_tool_executions"] += 1
 
         if context.shared["n_tool_executions"] > self.max_tool_executions:
-            raise RuntimeError("Reached limit of Tool Executions")
+            raise ToolExecutionLimitReached(
+                f"Exceeded {self.max_tool_executions} tool executions"
+            )
 
         return context
 ```

docs/api/agent.md

@@ -2,4 +2,6 @@
 
 ::: any_agent.AnyAgent
 
+::: any_agent.AgentCancel
+
 ::: any_agent.AgentRunError

src/any_agent/__init__.py

@@ -1,7 +1,7 @@
 from importlib.metadata import PackageNotFoundError, version
 
 from .config import AgentConfig, AgentFramework
-from .frameworks.any_agent import AgentRunError, AnyAgent
+from .frameworks.any_agent import AgentCancel, AgentRunError, AnyAgent
 from .tracing.agent_trace import AgentTrace
 
 try:
@@ -12,6 +12,7 @@
     __version__ = "0.0.0-dev"
 
 __all__ = [
+    "AgentCancel",
     "AgentConfig",
     "AgentFramework",
     "AgentRunError",

src/any_agent/frameworks/any_agent.py

@@ -35,24 +35,111 @@
 INSIDE_NOTEBOOK = hasattr(builtins, "__IPYTHON__")
 
 
+class AgentCancel(ABC, Exception):  # noqa: N818
+    """Abstract base class for control-flow exceptions raised in callbacks.
+
+    Within a callback, raise an exception inherited from AgentCancel when you
+    want to intentionally stop agent execution and handle that specific case in
+    your application code.
+
+    Unlike regular exceptions (which are wrapped in AgentRunError), AgentCancel
+    subclasses propagate directly to the caller, allowing you to catch them by
+    their specific type.
+
+    When to use AgentCancel vs regular exceptions:
+        - Use AgentCancel: When stopping execution is expected behavior
+          (rate limits, safety guardrails, validation failures) and you
+          want to handle it distinctly in your application.
+        - Use regular exceptions: When something unexpected goes wrong,
+          and you want consistent error handling via AgentRunError.
+
+    Example:
+        class StopOnLimit(AgentCancel):
+            pass
+
+        class LimitCallsCallback(Callback):
+            def before_tool_execution(self, context, *args, **kwargs):
+                if context.shared.get("call_count", 0) > 10:
+                    raise StopOnLimit("Exceeded call limit")
+                return context
+
+        try:
+            agent.run("prompt")
+        except StopOnLimit as e:
+            # Handle the expected cancellation.
+            print(f"Canceled: {e}")
+            print(f"Collected {len(e.trace.spans)} spans")
+        except AgentRunError as e:
+            # Handle unexpected errors.
+            print(f"Unexpected error: {e.original_exception}")
+
+    """
+
+    _trace: AgentTrace | None
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> Self:
+        if cls is AgentCancel:
+            msg = "AgentCancel cannot be instantiated directly; subclass it instead"
+            raise TypeError(msg)
+        return super().__new__(cls)
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._trace = None
+
+    @property
+    def trace(self) -> AgentTrace | None:
+        """Execution trace collected before cancellation.
+
+        Returns None if accessed before the framework processes the exception.
+        """
+        return self._trace
+
+
 class AgentRunError(Exception):
-    """Error that wraps underlying framework specific errors and carries spans."""
+    """Wrapper for unexpected exceptions that occur during agent execution.
+
+    When an unexpected exception is raised during agent execution (from
+    callbacks, tools, or the underlying framework), it is caught and
+    wrapped in AgentRunError.
+
+    Note: Exceptions that inherit from AgentCancel are not wrapped,
+        they propagate directly to the caller.
+
+    AgentRunError ensures:
+
+    * The execution trace is preserved - you can inspect what happened
+       before the error via the `trace` property.
+    * Consistent error handling - all unexpected errors are wrapped in
+       the same type, regardless of the underlying framework.
+    * Original exception access - the wrapped exception is available
+       via `original_exception` for debugging.
+
+    Example:
+        try:
+            agent.run("prompt")
+        except AgentRunError as e:
+            print(f"Error: {e.original_exception}")
+            print(f"Trace had {len(e.trace.spans)} spans before failure")
+
+    """
 
     _trace: AgentTrace
     _original_exception: Exception
 
     def __init__(self, trace: AgentTrace, original_exception: Exception):
         self._trace = trace
         self._original_exception = original_exception
-        # Set the exception message to be the original exception's message
         super().__init__(str(original_exception))
 
     @property
     def trace(self) -> AgentTrace:
+        """The execution trace collected up to the point of failure."""
         return self._trace
 
     @property
     def original_exception(self) -> Exception:
+        """The underlying exception that was caught."""
         return self._original_exception
 
     def __str__(self) -> str:
@@ -262,6 +349,12 @@ async def run_async(self, prompt: str, **kwargs: Any) -> AgentTrace:
                         )
 
             trace.add_span(invoke_span)
+
+            # Preserve control-flow exceptions without wrapping.
+            if isinstance(e, AgentCancel):
+                e._trace = trace
+                raise
+
             raise AgentRunError(trace, e) from e
 
         async with self._lock:

src/any_agent/frameworks/openai.py

@@ -333,8 +333,7 @@ async def _fetch_response(
             parallel_tool_calls=parallel_tool_calls,
             stream=stream,
             stream_options=stream_options,
-            # TODO: Remove type: ignore after any-llm adds xhigh support.
-            reasoning_effort=reasoning_effort,  # type: ignore[arg-type]
+            reasoning_effort=reasoning_effort,
             top_logprobs=model_settings.top_logprobs,
             **extra_kwargs,  # type: ignore[arg-type]
         )

src/any_agent/tools/wrappers.py

@@ -50,7 +50,7 @@ def _wrap_tool_openai(tool: "Tool | AgentTool") -> "AgentTool":
     )
 
     if isinstance(tool, agent_tool_types):
-        return tool  # type: ignore[return-value]
+        return tool  # type: ignore[no-any-return]
 
     # Enabling strict mode required else
     # throws error "Only strict function tools can be auto-parsed"

tests/integration/frameworks/test_error_handling.py

@@ -4,6 +4,7 @@
 import pytest
 
 from any_agent import (
+    AgentCancel,
     AgentConfig,
     AgentFramework,
     AgentRunError,
@@ -122,3 +123,40 @@ def search_web(query: str) -> str:
         and exception_reason in getattr(span.status, "description", "")
         for span in agent_trace.spans
     )
+
+
+class StopExecution(AgentCancel):
+    """Test exception for cancelling agent execution."""
+
+
+class StopBeforeFirstLLMCall(Callback):
+    """Callback that raises StopExecution before the first LLM call."""
+
+    def before_llm_call(self, context: Context, *args: Any, **kwargs: Any) -> Context:
+        msg = "Stopped by callback"
+        raise StopExecution(msg)
+
+
+def test_agent_cancel_not_wrapped(
+    agent_framework: AgentFramework,
+) -> None:
+    """AgentCancel subclasses should propagate without being wrapped in AgentRunError.
+
+    When a callback raises an exception that inherits from AgentCancel,
+    the exception should propagate directly to the caller without being
+    wrapped in AgentRunError, and the trace should be attached.
+    """
+    agent_config = AgentConfig(
+        model_id=DEFAULT_SMALL_MODEL_ID,
+        tools=[],
+        callbacks=[StopBeforeFirstLLMCall()],
+        model_args=get_default_agent_model_args(agent_framework),
+    )
+    agent = AnyAgent.create(agent_framework, agent_config)
+
+    with pytest.raises(StopExecution) as exc_info:
+        agent.run("test")
+
+    # The exception should have a trace attached.
+    assert exc_info.value.trace is not None
+    assert len(exc_info.value.trace.spans) > 0