---
title: "Trace streamed responses"
description: "Trace handlers that return a response stream before the model finishes."
seo:
  title: "Streaming AI Agent Tracing | Judgment Docs"
  description: "Trace streaming AI handlers in Judgment without buffering responses or ending the app root before generation finishes."
---

Streaming handlers return a response before generation finishes, so the tracing
boundary must follow the stream instead of the handler. For callback-based
frameworks, end the app root from the finish, error, and abort callbacks. In
Python, observe the async generator itself so the root stays active while the
framework consumes it.

If the app uses the Vercel AI SDK, follow its
[integration guide](/documentation/integrations/agent-frameworks/vercel-ai-sdk)
before adding spans.

The handler must return the stream immediately, but the root must remain open
until generation and its callbacks finish. Buffering changes the app's
behavior. Ending the root in the handler closes the trace too early.

## Callback-based streams

1. Pass `Tracer.getOTELTracer()` when the framework's telemetry setup accepts a
   tracer. Enabling telemetry without it can drop the framework spans.
2. Start the root before generation and set the session ID inside it.
3. Return the stream as soon as generation starts. Collecting it first makes the
   user wait for the full response instead of watching it arrive.
4. End the root from the framework's finish, error, and abort callbacks. Wait
   until the final output is known, until save steps and other callbacks end,
   and until all model and tool spans end.
5. Record the final output on the root from the finish callback.
6. Mark an aborted stream as an error and save a short, useful part
   of the output.
7. Flush after the root ends if the runtime can freeze
   (see [Bundled and serverless runtimes](/documentation/tracing/serverless-runtimes)).

A regular observed handler's span ends when the handler returns. For a
callback-based stream, keep the app root active until the framework ends the
stream. An observed async generator stays active while the framework consumes
it. `Tracer.set_session_id()` and `Tracer.set_customer_id()` in Python, or
`Tracer.setSessionId()` and `Tracer.setCustomerId()` in TypeScript, write to the
active span. Call them while the streaming root is active.

## Examples

**Python**

Put the observed root around the async generator itself, not around the
handler that returns `StreamingResponse`. The root remains active while
the framework consumes the generator. Setting
`disable_generator_yield_span=True` avoids a child span for every
streamed chunk.

```python
from fastapi.responses import StreamingResponse
from judgeval import Tracer

@Tracer.observe(
    span_type="agent",
    span_name="chat.stream",
    disable_generator_yield_span=True,
)
async def traced_stream(chat_id: str, messages: list[dict]):
    Tracer.set_session_id(chat_id)

    async for chunk in generate_chunks(messages):
        yield chunk

def make_response(chat_id: str, messages: list[dict]) -> StreamingResponse:
    return StreamingResponse(traced_stream(chat_id, messages))
```

**TypeScript**

Start the root before `streamText`, then end it from every callback that
can finish the stream. This example uses the AI SDK 6 per-call telemetry
setting. AI SDK 7 registers the tracer once at startup, as shown in the
[integration guide](/documentation/integrations/agent-frameworks/vercel-ai-sdk).

```typescript
import { streamText } from "ai";
import { Tracer } from "judgeval";

return Tracer.getOTELTracer().startActiveSpan("chat.stream", (root) => {
  Tracer.setSpanKind("agent", root);
  Tracer.setSessionId(chatId);
  Tracer.setInput({ messages }, root);

  let ended = false;
  const endRoot = () => {
    if (!ended) {
      ended = true;
      root.end();
    }
  };

  const result = streamText({
    model,
    messages,
    experimental_telemetry: {
      isEnabled: true,
      tracer: Tracer.getOTELTracer(),
    },
    onFinish({ text }) {
      Tracer.setOutput(text, root);
      endRoot();
    },
    onError({ error }) {
      Tracer.setError(error, root);
      endRoot();
    },
    onAbort({ steps }) {
      const partialOutput = steps.map((step) => step.text).join("");
      if (partialOutput) {
        Tracer.setOutput(partialOutput, root);
      }
      Tracer.setError(new Error("Stream aborted"), root);
      endRoot();
    },
  });

  return result.toUIMessageStreamResponse();
});
```

## Two shapes that fail quietly

- **Buffering.** Reading the whole stream before responding produces a clean
  trace and a broken app, because the user loses the live reply. This violates
  the no-behavior-change rule even when every span looks right.
- **The floating root.** Returning the stream from outside the traced callback,
  with the traced promise left un-awaited, drops the root span silently. The
  framework's own spans then export as sessionless roots, and the session ID
  lands nowhere. The platform still shows traces, so the loss is easy to miss.
  Check for the named root even when traffic appears.

## Verify

First run the
[shared verification checklist](/documentation/tracing/instrumentation#verify-before-you-claim-success).
Then test a full stream and an aborted stream. The full stream should save the
full output on your named root. The aborted stream should mark that root as an
error. If your named root is missing while framework spans appear, you have the
floating-root shape above.

## Related pages

- [Instrument your agent](/documentation/tracing/instrumentation)
- [Bundled and serverless runtimes](/documentation/tracing/serverless-runtimes)
- [Add attributes and context](/documentation/tracing/attributes)