Start and finish streaming trace in impl metod (#540)

Closes #435 and closes #538. Unit tests.

Author: rm-openai

.gitignore

@@ -141,4 +141,5 @@ cython_debug/
 .ruff_cache/
 
 # PyPI configuration file
-.pypirc
+.pypirc
+.aider*

pyproject.toml

@@ -61,6 +61,7 @@ dev = [
  "graphviz",
  "mkdocs-static-i18n>=1.3.0",
  "eval-type-backport>=0.2.2",
+ "fastapi >= 0.110.0, <1",
 ]
 
 [tool.uv.workspace]

src/agents/result.py

@@ -126,7 +126,7 @@ class RunResultStreaming(RunResultBase):
 
  _current_agent_output_schema: AgentOutputSchemaBase | None = field(repr=False)
 
- _trace: Trace | None = field(repr=False)
+ trace: Trace | None = field(repr=False)
 
  is_complete: bool = False
  """Whether the agent has finished running."""
@@ -185,9 +185,6 @@ async def stream_events(self) -> AsyncIterator[StreamEvent]:
  yield item
  self._event_queue.task_done()
 
- if self._trace:
- self._trace.finish(reset_current=True)
-
  self._cleanup_tasks()
 
  if self._stored_exception:

src/agents/run.py

@@ -404,10 +404,6 @@ def run_streamed(
  disabled=run_config.tracing_disabled,
  )
  )
- # Need to start the trace here, because the current trace contextvar is captured at
- # asyncio.create_task time
- if new_trace:
- new_trace.start(mark_as_current=True)
 
  output_schema = cls._get_output_schema(starting_agent)
  context_wrapper: RunContextWrapper[TContext] = RunContextWrapper(
@@ -426,7 +422,7 @@ def run_streamed(
  input_guardrail_results=[],
  output_guardrail_results=[],
  _current_agent_output_schema=output_schema,
- _trace=new_trace,
+ trace=new_trace,
  )
 
  # Kick off the actual agent loop in the background and return the streamed result object.
@@ -499,6 +495,9 @@ async def _run_streamed_impl(
  run_config: RunConfig,
  previous_response_id: str | None,
  ):
+ if streamed_result.trace:
+ streamed_result.trace.start(mark_as_current=True)
+
  current_span: Span[AgentSpanData] | None = None
  current_agent = starting_agent
  current_turn = 0
@@ -625,6 +624,8 @@ async def _run_streamed_impl(
  finally:
  if current_span:
  current_span.finish(reset_current=True)
+ if streamed_result.trace:
+ streamed_result.trace.finish(reset_current=True)
 
  @classmethod
  async def _run_single_turn_streamed(

tests/fastapi/__init__.py

@@ -0,0 +1,30 @@
+from collections.abc import AsyncIterator
+
+from fastapi import FastAPI
+from starlette.responses import StreamingResponse
+
+from agents import Agent, Runner, RunResultStreaming
+
+agent = Agent(
+ name="Assistant",
+ instructions="You are a helpful assistant.",
+)
+
+
+app = FastAPI()
+
+
+@app.post("/stream")
+async def stream():
+ result = Runner.run_streamed(agent, input="Tell me a joke")
+ stream_handler = StreamHandler(result)
+ return StreamingResponse(stream_handler.stream_events(), media_type="application/x-ndjson")
+
+
+class StreamHandler:
+ def __init__(self, result: RunResultStreaming):
+ self.result = result
+
+ async def stream_events(self) -> AsyncIterator[str]:
+ async for event in self.result.stream_events():
+ yield f"{event.type}\n\n"

tests/fastapi/streaming_app.py

@@ -0,0 +1,29 @@
+import pytest
+from httpx import ASGITransport, AsyncClient
+from inline_snapshot import snapshot
+
+from ..fake_model import FakeModel
+from ..test_responses import get_text_message
+from .streaming_app import agent, app
+
+
+@pytest.mark.asyncio
+async def test_streaming_context():
+ """This ensures that FastAPI streaming works. The context for this test is that the Runner
+ method was called in one async context, and the streaming was ended in another context,
+ leading to a tracing error because the context was closed in the wrong context. This test
+ ensures that this actually works.
+ """
+ model = FakeModel()
+ agent.model = model
+ model.set_next_output([get_text_message("done")])
+
+ transport = ASGITransport(app)
+ async with AsyncClient(transport=transport, base_url="http://test") as ac:
+ async with ac.stream("POST", "/stream") as r:
+ assert r.status_code == 200
+ body = (await r.aread()).decode("utf-8")
+ lines = [line for line in body.splitlines() if line]
+ assert lines == snapshot(
+ ["agent_updated_stream_event", "raw_response_event", "run_item_stream_event"]
+ )