Fix retry error wrapping

Author: DouweM

docs/output.md

@@ -200,8 +200,8 @@ async def hand_off_to_sql_agent(ctx: RunContext, query: str) -> list[Row]:
         return output
     except UnexpectedModelBehavior as e:
         # Bubble up potentially retryable errors to the router agent
-        if (cause := e.__cause__) and hasattr(cause, 'tool_retry'):
-            raise ModelRetry(f'SQL agent failed: {cause.tool_retry.content}') from e
+        if (cause := e.__cause__) and isinstance(cause, ModelRetry):
+            raise ModelRetry(f'SQL agent failed: {cause.message}') from e
         else:
             raise
 

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -79,11 +79,13 @@ class GraphAgentState:
     retries: int
     run_step: int
 
-    def increment_retries(self, max_result_retries: int, error: Exception | None = None) -> None:
+    def increment_retries(self, max_result_retries: int, error: BaseException | None = None) -> None:
         self.retries += 1
         if self.retries > max_result_retries:
             message = f'Exceeded maximum retries ({max_result_retries}) for result validation'
             if error:
+                if isinstance(error, exceptions.UnexpectedModelBehavior) and error.__cause__ is not None:
+                    error = error.__cause__
                 raise exceptions.UnexpectedModelBehavior(message) from error
             else:
                 raise exceptions.UnexpectedModelBehavior(message)

pydantic_ai_slim/pydantic_ai/toolset.py

@@ -7,7 +7,7 @@
 from dataclasses import dataclass, field, replace
 from functools import partial
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Protocol, assert_never, overload
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Protocol, overload
 
 from pydantic import ValidationError
 from pydantic.json_schema import GenerateJsonSchema
@@ -689,38 +689,33 @@ def _with_retry(self, name: str, ctx: RunContext[AgentDepsT]) -> Iterator[RunCon
             ctx = replace(ctx, tool_name=name, retry=self._retries.get(name, 0), retries={})
             yield ctx
         except (ValidationError, ModelRetry, UnexpectedModelBehavior, ToolRetryError) as e:
-            if isinstance(e, ToolRetryError):
-                pass
-            elif isinstance(e, ValidationError):
-                if ctx.tool_call_id:
-                    m = _messages.RetryPromptPart(
-                        tool_name=name,
-                        content=e.errors(include_url=False, include_context=False),
-                        tool_call_id=ctx.tool_call_id,
-                    )
-                    e = ToolRetryError(m)
-            elif isinstance(e, ModelRetry):
-                if ctx.tool_call_id:
-                    m = _messages.RetryPromptPart(
-                        tool_name=name,
-                        content=e.message,
-                        tool_call_id=ctx.tool_call_id,
-                    )
-                    e = ToolRetryError(m)
-            elif isinstance(e, UnexpectedModelBehavior):
-                if e.__cause__ is not None:
-                    e = e.__cause__
-            else:
-                assert_never(e)
-
             try:
                 max_retries = self._max_retries_for_tool(name)
             except Exception:
                 max_retries = 1
             current_retry = self._retries.get(name, 0)
 
+            if isinstance(e, UnexpectedModelBehavior) and e.__cause__ is not None:
+                e = e.__cause__
+
             if current_retry == max_retries:
                 raise UnexpectedModelBehavior(f'Tool {name!r} exceeded max retries count of {max_retries}') from e
             else:
+                if ctx.tool_call_id:
+                    if isinstance(e, ValidationError):
+                        m = _messages.RetryPromptPart(
+                            tool_name=name,
+                            content=e.errors(include_url=False, include_context=False),
+                            tool_call_id=ctx.tool_call_id,
+                        )
+                        e = ToolRetryError(m)
+                    elif isinstance(e, ModelRetry):
+                        m = _messages.RetryPromptPart(
+                            tool_name=name,
+                            content=e.message,
+                            tool_call_id=ctx.tool_call_id,
+                        )
+                        e = ToolRetryError(m)
+
                 self._retries[name] = current_retry + 1
                 raise e