Nicer errors under the capture_run_messages context (#2219)

Author: alonme

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -815,14 +815,21 @@ def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]:
         If you call `run`, `run_sync`, or `run_stream` more than once within a single `capture_run_messages` context,
         `messages` will represent the messages exchanged during the first call only.
     """
+    token = None
+    messages: list[_messages.ModelMessage] = []
+
+    # Try to reuse existing message context if available
     try:
-        yield _messages_ctx_var.get().messages
+        messages = _messages_ctx_var.get().messages
     except LookupError:
-        messages: list[_messages.ModelMessage] = []
+        # No existing context, create a new one
         token = _messages_ctx_var.set(_RunMessages(messages))
-        try:
-            yield messages
-        finally:
+
+    try:
+        yield messages
+    finally:
+        # Clean up context if we created it
+        if token is not None:
             _messages_ctx_var.reset(token)
 
 

tests/test_agent.py

@@ -2673,6 +2673,18 @@ def test_double_capture_run_messages() -> None:
     )
 
 
+def test_capture_run_messages_with_user_exception_does_not_contain_internal_errors() -> None:
+    """Test that user exceptions within capture_run_messages context have clean stack traces."""
+    agent = Agent('test')
+
+    try:
+        with capture_run_messages():
+            agent.run_sync('Hello')
+            raise ZeroDivisionError('division by zero')
+    except Exception as e:
+        assert e.__context__ is None
+
+
 def test_dynamic_false_no_reevaluate():
     """When dynamic is false (default), the system prompt is not reevaluated
     i.e: SystemPromptPart(