From 7cfb60e7fd01cc231a85a15e7e331758a43ada58 Mon Sep 17 00:00:00 2001 From: CTW_CTWalk <100585900+CTWalk@users.noreply.github.com> Date: Sat, 1 Aug 2026 11:19:30 +0800 Subject: [PATCH] fix(openai): finalize completion spans on early stream close The Completions streaming wrappers end their spans after the iteration loop. That works when the stream is exhausted, but close() or aclose() exits the generator at its suspended yield, before span.end() is reached, so an early-closed stream exports no openai.completion span at all. Wrap both streaming generator bodies in try/finally and end a still-recording span in finally, so exhaustion, close(), and aclose() converge on one terminal span.end(). Full-consumption response attributes, token usage, events, and StatusCode.OK still run only on normal exhaustion. Because finally also runs on exceptions, a stream that raises mid-iteration now ends and exports its span instead of leaving one recording and never exported. Add three tests covering sync early close, async early close, and an exception raised during consumption, the last mirroring the chat-side coverage added in #3155. All three reuse the existing streaming cassettes via pytest-recording's default_cassette marker, so no fixture data is added. Most of the diff is re-indentation from the new try block; the semantic change is six lines, visible with git diff -w. Related to #3151 and #3155, which established this contract for chat streams. Co-Authored-By: Claude Opus 5 --- .../openai/shared/completion_wrappers.py | 54 +++++++------- .../tests/traces/test_completions.py | 71 +++++++++++++++++++ 2 files changed, 101 insertions(+), 24 deletions(-) diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py index b4461cd841..c77ba5c0b3 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py @@ -208,43 +208,49 @@ def _set_output_messages(span, choices): @dont_throw def _build_from_streaming_response(span, request_kwargs, response): complete_response = {"choices": [], "model": "", "id": ""} - for item in response: - yield item - _accumulate_streaming_response(complete_response, item) + try: + for item in response: + yield item + _accumulate_streaming_response(complete_response, item) - _set_response_attributes(span, complete_response) + _set_response_attributes(span, complete_response) - _set_token_usage(span, request_kwargs, complete_response) + _set_token_usage(span, request_kwargs, complete_response) - if should_emit_events(): - _emit_streaming_response_events(complete_response) - else: - if should_send_prompts(): - _set_completions(span, complete_response.get("choices")) + if should_emit_events(): + _emit_streaming_response_events(complete_response) + else: + if should_send_prompts(): + _set_completions(span, complete_response.get("choices")) - span.set_status(Status(StatusCode.OK)) - span.end() + span.set_status(Status(StatusCode.OK)) + finally: + if span.is_recording(): + span.end() @dont_throw async def _abuild_from_streaming_response(span, request_kwargs, response): complete_response = {"choices": [], "model": "", "id": ""} - async for item in response: - yield item - _accumulate_streaming_response(complete_response, item) + try: + async for item in response: + yield item + _accumulate_streaming_response(complete_response, item) - _set_response_attributes(span, complete_response) + _set_response_attributes(span, complete_response) - _set_token_usage(span, request_kwargs, complete_response) + _set_token_usage(span, request_kwargs, complete_response) - if should_emit_events(): - _emit_streaming_response_events(complete_response) - else: - if should_send_prompts(): - _set_completions(span, complete_response.get("choices")) + if should_emit_events(): + _emit_streaming_response_events(complete_response) + else: + if should_send_prompts(): + _set_completions(span, complete_response.get("choices")) - span.set_status(Status(StatusCode.OK)) - span.end() + span.set_status(Status(StatusCode.OK)) + finally: + if span.is_recording(): + span.end() def _emit_streaming_response_events(complete_response): diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py index 0574de076e..be44fe4a35 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py @@ -1,3 +1,4 @@ +import gc from unittest.mock import patch import httpx @@ -393,6 +394,55 @@ def test_completion_streaming( ), "Assert that it doesn't emit logs when use_legacy_attributes is True" +@pytest.mark.vcr +@pytest.mark.default_cassette("test_completion_streaming.yaml") +def test_completion_streaming_early_close( + instrument_legacy, span_exporter, log_exporter, mock_openai_client +): + response = mock_openai_client.completions.create( + model="davinci-002", + prompt="Tell me a joke about opentelemetry", + stream=True, + ) + + next(response) + response.close() + + spans = span_exporter.get_finished_spans() + assert [span.name for span in spans] == [ + "openai.completion", + ] + + +@pytest.mark.vcr +@pytest.mark.default_cassette("test_completion_streaming.yaml") +def test_completion_streaming_exception_during_consumption( + instrument_legacy, span_exporter, log_exporter, mock_openai_client +): + """The span is still ended when the caller aborts consumption by raising.""" + + response = mock_openai_client.completions.create( + model="davinci-002", + prompt="Tell me a joke about opentelemetry", + stream=True, + ) + + with pytest.raises(ValueError, match="simulated interruption"): + for _ in response: + raise ValueError("simulated interruption") + + del response + gc.collect() + + spans = span_exporter.get_finished_spans() + assert [span.name for span in spans] == [ + "openai.completion", + ] + # Status is deliberately not asserted: what an incompletely observed + # completion should report is a separate question from ending the span. + assert spans[0].end_time is not None + + @pytest.mark.vcr def test_completion_streaming_with_events_with_content( instrument_with_content, span_exporter, log_exporter, openai_client @@ -548,6 +598,27 @@ async def test_async_completion_streaming( ), "Assert that it doesn't emit logs when use_legacy_attributes is True" +@pytest.mark.vcr +@pytest.mark.default_cassette("test_async_completion_streaming.yaml") +@pytest.mark.asyncio +async def test_async_completion_streaming_early_close( + instrument_legacy, span_exporter, log_exporter, async_openai_client +): + response = await async_openai_client.completions.create( + model="davinci-002", + prompt="Tell me a joke about opentelemetry", + stream=True, + ) + + await anext(response) + await response.aclose() + + spans = span_exporter.get_finished_spans() + assert [span.name for span in spans] == [ + "openai.completion", + ] + + @pytest.mark.vcr @pytest.mark.asyncio async def test_async_completion_streaming_with_events_with_content(