Preserve Usage Tracking on Streaming Errors

src/agents/__init__.py

@@ -20,6 +20,7 @@
  MaxTurnsExceeded,
  ModelBehaviorError,
  OutputGuardrailTripwireTriggered,
+ RunError,
  RunErrorDetails,
  ToolInputGuardrailTripwireTriggered,
  ToolOutputGuardrailTripwireTriggered,
@@ -212,6 +213,8 @@ def enable_verbose_stdout_logging():
  "OutputGuardrailTripwireTriggered",
  "ToolInputGuardrailTripwireTriggered",
  "ToolOutputGuardrailTripwireTriggered",
+ "RunError",
+ "RunErrorDetails",
  "DynamicPromptFunction",
  "GenerateDynamicPromptData",
  "Prompt",

src/agents/exceptions.py

@@ -129,3 +129,20 @@ def __init__(self, guardrail: ToolOutputGuardrail[Any], output: ToolGuardrailFun
  self.guardrail = guardrail
  self.output = output
  super().__init__(f"Tool output guardrail {guardrail.__class__.__name__} triggered tripwire")
+
+
+class RunError(AgentsException):
+ """Wrapper exception for non-AgentsException errors that occur during agent runs.
+
+ This exception wraps external errors (API errors, connection failures, etc.) to ensure
+ that run data including usage information is preserved and accessible.
+ """
+
+ original_exception: Exception
+ """The original exception that was raised."""
+
+ def __init__(self, original_exception: Exception):
+ self.original_exception = original_exception
+ super().__init__(str(original_exception))
+ # Preserve the original exception as the cause
+ self.__cause__ = original_exception

src/agents/result.py

@@ -15,6 +15,7 @@
  AgentsException,
  InputGuardrailTripwireTriggered,
  MaxTurnsExceeded,
+ RunError,
  RunErrorDetails,
 )
 from .guardrail import InputGuardrailResult, OutputGuardrailResult
@@ -299,23 +300,40 @@ def _check_errors(self):
  if self._run_impl_task and self._run_impl_task.done():
  run_impl_exc = self._run_impl_task.exception()
  if run_impl_exc and isinstance(run_impl_exc, Exception):
- if isinstance(run_impl_exc, AgentsException) and run_impl_exc.run_data is None:
- run_impl_exc.run_data = self._create_error_details()
- self._stored_exception = run_impl_exc
+ if isinstance(run_impl_exc, AgentsException):
+ # For AgentsException, attach run_data if missing
+ if run_impl_exc.run_data is None:
+ run_impl_exc.run_data = self._create_error_details()
+ self._stored_exception = run_impl_exc
+ else:
+ # For non-AgentsException, wrap it to preserve run_data
+ wrapped_exc = RunError(run_impl_exc)
+ wrapped_exc.run_data = self._create_error_details()
+ self._stored_exception = wrapped_exc
 
  if self._input_guardrails_task and self._input_guardrails_task.done():
  in_guard_exc = self._input_guardrails_task.exception()
  if in_guard_exc and isinstance(in_guard_exc, Exception):
- if isinstance(in_guard_exc, AgentsException) and in_guard_exc.run_data is None:
- in_guard_exc.run_data = self._create_error_details()
- self._stored_exception = in_guard_exc
+ if isinstance(in_guard_exc, AgentsException):
+ if in_guard_exc.run_data is None:
+ in_guard_exc.run_data = self._create_error_details()
+ self._stored_exception = in_guard_exc
+ else:
+ wrapped_exc = RunError(in_guard_exc)
+ wrapped_exc.run_data = self._create_error_details()
+ self._stored_exception = wrapped_exc
 
  if self._output_guardrails_task and self._output_guardrails_task.done():
  out_guard_exc = self._output_guardrails_task.exception()
  if out_guard_exc and isinstance(out_guard_exc, Exception):
- if isinstance(out_guard_exc, AgentsException) and out_guard_exc.run_data is None:
- out_guard_exc.run_data = self._create_error_details()
- self._stored_exception = out_guard_exc
+ if isinstance(out_guard_exc, AgentsException):
+ if out_guard_exc.run_data is None:
+ out_guard_exc.run_data = self._create_error_details()
+ self._stored_exception = out_guard_exc
+ else:
+ wrapped_exc = RunError(out_guard_exc)
+ wrapped_exc.run_data = self._create_error_details()
+ self._stored_exception = wrapped_exc
 
  def _cleanup_tasks(self):
  if self._run_impl_task and not self._run_impl_task.done():

src/agents/run.py

@@ -35,6 +35,7 @@
  MaxTurnsExceeded,
  ModelBehaviorError,
  OutputGuardrailTripwireTriggered,
+ RunError,
  RunErrorDetails,
  UserError,
 )
@@ -702,6 +703,19 @@ async def run(
  output_guardrail_results=[],
  )
  raise
+ except Exception as exc:
+ # Wrap non-AgentsException to preserve run_data including usage
+ wrapped_exc = RunError(exc)
+ wrapped_exc.run_data = RunErrorDetails(
+ input=original_input,
+ new_items=generated_items,
+ raw_responses=model_responses,
+ last_agent=current_agent,
+ context_wrapper=context_wrapper,
+ input_guardrail_results=input_guardrail_results,
+ output_guardrail_results=[],
+ )
+ raise wrapped_exc from exc
  finally:
  if current_span:
  current_span.finish(reset_current=True)

tests/test_run_hooks.py

@@ -172,6 +172,8 @@ async def test_async_run_hooks_with_agent_hooks_with_llm():
 
 @pytest.mark.asyncio
 async def test_run_hooks_llm_error_non_streaming(monkeypatch):
+ from agents import RunError
+
  hooks = RunHooksForTests()
  model = FakeModel()
  agent = Agent(name="A", model=model, tools=[get_function_tool("f", "res")], handoffs=[])
@@ -181,9 +183,16 @@ async def boom(*args, **kwargs):
 
  monkeypatch.setattr(FakeModel, "get_response", boom, raising=True)
 
- with pytest.raises(RuntimeError, match="boom"):
+ with pytest.raises(RunError) as exc_info:
  await Runner.run(agent, input="hello", hooks=hooks)
 
+ # Verify the original exception is preserved
+ assert isinstance(exc_info.value.original_exception, RuntimeError)
+ assert str(exc_info.value.original_exception) == "boom"
+ # Verify run_data is attached
+ assert exc_info.value.run_data is not None
+ assert exc_info.value.run_data.context_wrapper is not None
+
  # Current behavior is that hooks will not fire on LLM failure
  assert hooks.events["on_agent_start"] == 1
  assert hooks.events["on_llm_start"] == 1
@@ -229,16 +238,26 @@ async def test_streamed_run_hooks_llm_error(monkeypatch):
  Verify that when the streaming path raises, we still emit on_llm_start
  but do NOT emit on_llm_end (current behavior), and the exception propagates.
  """
+ from agents import RunError
+
  hooks = RunHooksForTests()
  agent = Agent(name="A", model=BoomModel(), tools=[get_function_tool("f", "res")], handoffs=[])
 
  stream = Runner.run_streamed(agent, input="hello", hooks=hooks)
 
- # Consuming the stream should surface the exception
- with pytest.raises(RuntimeError, match="stream blew up"):
+ # Consuming the stream should surface the exception (wrapped in RunError to preserve usage data)
+ with pytest.raises(RunError) as exc_info:
  async for _ in stream.stream_events():
  pass
 
+ # Verify the original exception is preserved and accessible
+ assert isinstance(exc_info.value.original_exception, RuntimeError)
+ assert str(exc_info.value.original_exception) == "stream blew up"
+ # Verify run_data is attached with usage information
+ assert exc_info.value.run_data is not None
+ assert exc_info.value.run_data.context_wrapper is not None
+ assert exc_info.value.run_data.context_wrapper.usage is not None
+
  # Current behavior: success-only on_llm_end; ensure starts fired but ends did not.
  assert hooks.events["on_agent_start"] == 1
  assert hooks.events["on_llm_start"] == 1

tests/test_tracing_errors.py

@@ -15,6 +15,7 @@
  MaxTurnsExceeded,
  ModelBehaviorError,
  RunContextWrapper,
+ RunError,
  Runner,
  TResponseInputItem,
 )
@@ -39,9 +40,12 @@ async def test_single_turn_model_error():
  name="test_agent",
  model=model,
  )
- with pytest.raises(ValueError):
+ with pytest.raises(RunError) as exc_info:
  await Runner.run(agent, input="first_test")
 
+ # Verify the original exception is preserved
+ assert isinstance(exc_info.value.original_exception, ValueError)
+
  assert fetch_normalized_spans() == snapshot(
  [
  {
@@ -92,9 +96,12 @@ async def test_multi_turn_no_handoffs():
  ]
  )
 
- with pytest.raises(ValueError):
+ with pytest.raises(RunError) as exc_info:
  await Runner.run(agent, input="first_test")
 
+ # Verify the original exception is preserved
+ assert isinstance(exc_info.value.original_exception, ValueError)
+
  assert fetch_normalized_spans() == snapshot(
  [
  {

tests/test_tracing_errors_streamed.py

@@ -18,6 +18,7 @@
  OutputGuardrail,
  OutputGuardrailTripwireTriggered,
  RunContextWrapper,
+ RunError,
  Runner,
  TResponseInputItem,
 )
@@ -42,11 +43,14 @@ async def test_single_turn_model_error():
  name="test_agent",
  model=model,
  )
- with pytest.raises(ValueError):
+ with pytest.raises(RunError) as exc_info:
  result = Runner.run_streamed(agent, input="first_test")
  async for _ in result.stream_events():
  pass
 
+ # Verify the original exception is preserved
+ assert isinstance(exc_info.value.original_exception, ValueError)
+
  assert fetch_normalized_spans() == snapshot(
  [
  {
@@ -98,11 +102,14 @@ async def test_multi_turn_no_handoffs():
  ]
  )
 
- with pytest.raises(ValueError):
+ with pytest.raises(RunError) as exc_info:
  result = Runner.run_streamed(agent, input="first_test")
  async for _ in result.stream_events():
  pass
 
+ # Verify the original exception is preserved
+ assert isinstance(exc_info.value.original_exception, ValueError)
+
  assert fetch_normalized_spans() == snapshot(
  [
  {