Skip to content

CCM-33801: capture LiteLLM streaming responses in spans#12

Merged
t-santoshsahu merged 1 commit into
harness:mainfrom
shreyas70:CCM-33801-litellm-streaming-spans
Jul 10, 2026
Merged

CCM-33801: capture LiteLLM streaming responses in spans#12
t-santoshsahu merged 1 commit into
harness:mainfrom
shreyas70:CCM-33801-litellm-streaming-spans

Conversation

@shreyas70

Copy link
Copy Markdown
Contributor

Summary

The Harness SDK LiteLLM instrumentation did not capture response telemetry for streaming (stream=True) calls. The litellm_request span was ended as soon as completion() / acompletion() returned — but for streaming, the return value is a not-yet-consumed CustomStreamWrapper. Usage, response id, finish reasons, and content only exist after the stream is iterated, so streaming spans lost all response-side attributes.

This mirrors the deferred-span pattern already used by the OpenAI/Anthropic instrumentors in this SDK.

What changed

  • _StreamSpanWrapper (a transparent wrapt.ObjectProxy) now wraps the returned stream. It forwards every chunk to the caller unchanged while accumulating them, and on stream exhaustion (or error) rebuilds the aggregated response via litellm.stream_chunk_builder, copies response metadata onto the span, and ends it.
  • _LiteLLMSpanRun.wrap_stream() hands span ownership to the proxy: it detaches the active context and clears the re-dispatch guard without ending the span, so span lifecycle is owned by the stream.
  • The sync and async wrappers detect a streaming response (_is_stream_response) and route through the proxy; non-streaming behavior is unchanged.
  • A __del__ safety net ends the span if a consumer abandons the stream early.

Why the fix is at the SDK level

The consuming service already injects stream_options={"include_usage": True} and fully consumes the stream. The gap was purely that the SDK closed the span before those chunks were read.

Test plan

  • test_litellm_streaming_span_defers_until_consumed — sync: span is not finished before consumption; after consumption the span carries gen_ai.request.streaming=True and aggregated gen_ai.response.finish_reasons.
  • test_litellm_async_streaming_span_defers_until_consumed — async equivalent.
  • Full LiteLLM instrumentation suite passes (17 tests) — existing non-streaming behavior unchanged.

Made with Cursor

The LiteLLM instrumentation ended the litellm_request span as soon as
completion()/acompletion() returned, which for stream=True is the
not-yet-consumed CustomStreamWrapper. Response-side telemetry (usage,
response id, finish reasons, content) is only available after the stream
is iterated, so streaming spans were missing all of it.

Defer span completion for streaming: wrap the returned stream in a
transparent proxy (_StreamSpanWrapper) that forwards every chunk while
accumulating them, then on exhaustion (or error) rebuilds the aggregated
response via litellm.stream_chunk_builder, copies response metadata onto
the span, and ends it. Mirrors the deferred-span pattern already used by
the OpenAI/Anthropic instrumentors.

Adds sync + async streaming tests asserting the span is not finished
until the stream is consumed and that response metadata is captured.

Co-authored-by: Cursor <cursoragent@cursor.com>
@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

@t-santoshsahu
t-santoshsahu merged commit ee6a72f into harness:main Jul 10, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants