fix(util-genai,openai): scope stream invocation context

.codespellrc

@@ -4,4 +4,4 @@
 check-hidden = true
 # skipping auto generated folders
 skip = ./.git,./.tox,./.venv,./.mypy_cache,./docs/_build,./target,*/LICENSE,./venv,*/cassettes
-ignore-words-list = ot
+ignore-words-list = ot,asend

.github/instructions/instrumentation.instructions.md

@@ -55,7 +55,11 @@ prefer opt-in or additive. Breaking changes need explicit justification in the P
   capture path — never as unconditional span/log attributes.
 - Adding attributes to invocations produced by the util is fine.
 - Streaming responses must be instrumented by subclassing the util's `SyncStreamWrapper` /
-  `AsyncStreamWrapper` (`opentelemetry.util.genai.stream`). Flag hand-rolled stream wrappers.
+  `AsyncStreamWrapper` (`opentelemetry.util.genai.stream`). Flag hand-rolled stream wrappers, and
+  invocation-backed wrappers that do not pass the invocation to `super().__init__(stream, invocation)`
+  or do not `invocation.suspend()` before the stream is returned. Flag an `_execution_context()`
+  override on an invocation-backed wrapper: the base wrapper already activates a suspended
+  invocation during reads, and an override bypasses that guard.
 - Instrumentation should not change what a call returns or when its work happens. Flag: work the SDK
   didn't do (materializing a result early to build telemetry — stay lazy); a changed return type
   (`isinstance`/`__class__` should still resolve to the original; `wrapt.ObjectProxy` is the usual

AGENTS.md

@@ -229,13 +229,19 @@ as the reference:
 
 A streamed response only finishes once the caller has drained the stream, so the invocation must
 stay open until then. Do **not** call `invocation.stop()` when the SDK returns the stream — the
-span would close before any chunks arrive.
+span would close before any chunks arrive. The invocation's span must also not stay current in
+the caller's context while the stream is unconsumed: call `invocation.suspend()` before returning
+the stream, and re-activate the invocation only while a chunk is being read.
 
 Instrument streams by subclassing `SyncStreamWrapper` / `AsyncStreamWrapper` from
 `opentelemetry.util.genai.stream` (the public, supported helpers). The base class proxies the
 underlying SDK stream, drives iteration, and finalizes telemetry exactly once on success, error,
-or `close()`. Subclasses pass the SDK stream to `super().__init__(stream)` and implement three
-hooks:
+or `close()`. Subclasses pass the SDK stream and the invocation to
+`super().__init__(stream, invocation)`, which also marks the invocation as a streamed request
+(`gen_ai.request.stream`), call `invocation.suspend()` before returning the stream, and implement
+three hooks. The base `_execution_context()` makes a suspended invocation current again for each
+read and cleanup operation and restores the context before the chunk is returned, so
+invocation-backed wrappers don't override it.
 
 - `_process_chunk(chunk)` — accumulate per-chunk state (e.g. response model, finish reasons,
   token usage, streamed content) onto the invocation.
@@ -246,8 +252,8 @@ hooks:
 ```python
 class MyStreamWrapper(SyncStreamWrapper[Chunk]):
     def __init__(self, stream, invocation, capture_content):
-        super().__init__(stream)
-        self._self_invocation = invocation
+        super().__init__(stream, invocation)
+        invocation.suspend()
         ...
 
     def _process_chunk(self, chunk): ...  # accumulate state

instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/817.fixed

@@ -0,0 +1 @@
+Suspend the chat completion and Responses API invocations while a streamed response is unconsumed and re-activate them only while a chunk is read or the stream is cleaned up: the inference span is no longer current in the caller's context between reads, so spans a caller creates between chunks are siblings of the inference span, not children.

instrumentation/opentelemetry-instrumentation-genai-openai/pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
   "opentelemetry-api ~= 1.43",
   "opentelemetry-instrumentation >= 0.64b0, <1",
   "opentelemetry-semantic-conventions >= 0.64b0, <1",
-  "opentelemetry-util-genai >= 1.2b0, <2",
+  "opentelemetry-util-genai >= 1.3b0.dev, <2",
 ]
 
 [project.optional-dependencies]

instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/_raw_response.py

@@ -239,6 +239,13 @@ def __call__(
     ) -> object: ...
 
 
+def _stop_inside(invocation: _StreamingInvocation) -> None:
+    # A response closed without parse() ends the span inside it, as a parsed
+    # stream's cleanup does.
+    with invocation.activate():
+        invocation.stop()
+
+
 def wrap_stream_result(
     wrapper_cls: StreamWrapperFactory,
     result: RawResponseLike | AnyStream,
@@ -256,9 +263,13 @@ def wrap_stream_result(
         if served_model:
             if hasattr(invocation, "response_model_name"):
                 setattr(invocation, "response_model_name", served_model)
+        # The stream wrapper is built only on ``parse()``, possibly in another
+        # context; suspend now so the span isn't current before then and the
+        # context it detaches is the one ``create()`` ran in.
+        invocation.suspend()
         return RawResponseStreamProxy(
             result,
             lambda stream: wrapper_cls(stream, invocation, capture_content),
-            finalize=invocation.stop,
+            finalize=functools.partial(_stop_inside, invocation),
         )
     return wrapper_cls(result, invocation, capture_content)

instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/chat_wrappers.py

@@ -212,6 +212,7 @@ def __init__(
     ) -> None:
         super().__init__(stream, invocation=invocation)
         self._self_invocation = invocation
+        invocation.suspend()
         self._self_choice_buffers = []
         self._self_capture_content = capture_content
         self._self_response_id = None
@@ -234,6 +235,7 @@ def __init__(
     ) -> None:
         super().__init__(stream, invocation=invocation)
         self._self_invocation = invocation
+        invocation.suspend()
         self._self_choice_buffers = []
         self._self_capture_content = capture_content
         self._self_response_id = None

instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/response_wrappers.py

@@ -7,6 +7,7 @@
 
 import logging
 from collections.abc import Callable
+from contextlib import AbstractContextManager
 from contextvars import ContextVar
 from types import TracebackType
 from typing import TYPE_CHECKING, Generic, TypeVar, cast
@@ -123,6 +124,8 @@ class _ResponseStreamMixin(Generic[TextFormatT]):
     _self_invocation: _ResponseInvocation
     _self_capture_content: bool
     _self_response_telemetry_finalized: bool
+    # provided by the stream wrapper base class the mixin is combined with
+    _execution_context: Callable[[], AbstractContextManager[None]]
 
     def __init__(
         self,
@@ -132,6 +135,9 @@ def __init__(
         self._self_invocation = invocation
         self._self_capture_content = capture_content
         self._self_response_telemetry_finalized = False
+        # The stream returns to the caller undrained: leave the caller's
+        # context as it was and make the span current only while reading.
+        invocation.suspend()
 
     def _stop(
         self, result: ParsedResponse[TextFormatT] | Response | None
@@ -199,7 +205,13 @@ def response(self):
         response = _get_stream_response(self.stream)
         if response is None:
             return None
-        return finalize_on_close(response, lambda: self._stop(None))
+        # A close through the HTTP response is stream cleanup too: run it and
+        # the finalizer in the same context as the wrapper's own close.
+        return finalize_on_close(
+            response,
+            lambda: self._stop(None),
+            execution_context=self._execution_context,
+        )
 
     def process_event(self, event: ResponseStreamEvent[TextFormatT]) -> None:
         # raw-response stream can be parsed into a caller-defined event type.
@@ -412,7 +424,12 @@ def response(self):
         response = _get_stream_response(self.stream)
         if response is None:
             return None
-        return finalize_on_aclose(response, lambda: self._stop(None))
+        # See _ResponseStreamMixin.response.
+        return finalize_on_aclose(
+            response,
+            lambda: self._stop(None),
+            execution_context=self._execution_context,
+        )
 
 
 class AsyncFetchResponseStreamWrapper(

instrumentation/opentelemetry-instrumentation-genai-openai/tests/requirements.oldest.txt

@@ -30,3 +30,6 @@ httpx==0.27.2
 Deprecated==1.2.14
 importlib-metadata==6.11.0
 packaging==24.0
+
+# The stream execution hook is not released yet.
+-e util/opentelemetry-util-genai

instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_stream_context.py

@@ -0,0 +1,91 @@
+# Copyright The OpenTelemetry Authors
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator, Iterator
+
+import pytest
+from openai.types.chat import ChatCompletionChunk
+
+from opentelemetry import context, trace
+from opentelemetry.instrumentation.genai.openai.chat_wrappers import (
+    AsyncChatStreamWrapper,
+    ChatStreamWrapper,
+)
+from opentelemetry.util.genai.handler import TelemetryHandler
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("ending", ["drain", "close", "error"])
+async def test_chat_stream_activates_only_during_reads(
+    tracer_provider, span_exporter, caplog, asynchronous: bool, ending: str
+) -> None:
+    before = context.get_current()
+    invocation = TelemetryHandler(
+        tracer_provider=tracer_provider,
+    ).inference("openai", request_model="test-model")
+    error = ConnectionError("stream failed")
+    chunk = ChatCompletionChunk(
+        id="test",
+        created=1,
+        model="test-model",
+        object="chat.completion.chunk",
+        choices=[],
+    )
+
+    def produce() -> Iterator[ChatCompletionChunk]:
+        assert (
+            trace.get_current_span().get_span_context()
+            == trace.get_current_span(invocation.context).get_span_context()
+        )
+        yield chunk
+        assert (
+            trace.get_current_span().get_span_context()
+            == trace.get_current_span(invocation.context).get_span_context()
+        )
+        if ending == "error":
+            raise error
+
+    async def aproduce() -> AsyncIterator[ChatCompletionChunk]:
+        for item in produce():
+            await asyncio.sleep(0)
+            yield item
+
+    if asynchronous:
+        stream = AsyncChatStreamWrapper(aproduce(), invocation, False)
+    else:
+        stream = ChatStreamWrapper(produce(), invocation, False)
+    assert context.get_current() is before
+    if asynchronous:
+        assert await asyncio.create_task(anext(stream)) is chunk
+    else:
+        assert next(stream) is chunk
+    assert context.get_current() is before
+    assert span_exporter.get_finished_spans() == ()
+
+    if ending == "close":
+        if asynchronous:
+            await stream.aclose()
+        else:
+            stream.close()
+    elif ending == "error":
+        with pytest.raises(ConnectionError) as raised:
+            if asynchronous:
+                await anext(stream)
+            else:
+                next(stream)
+        assert raised.value is error
+    elif asynchronous:
+        assert [item async for item in stream] == []
+    else:
+        assert list(stream) == []
+    assert context.get_current() is before
+    assert len(span_exporter.get_finished_spans()) == 1
+    assert not [
+        r
+        for r in caplog.records
+        if r.name == "opentelemetry.context" and r.levelno >= 40
+    ]

instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_raw_response_stream_suspend.py

@@ -0,0 +1,120 @@
+# Copyright The OpenTelemetry Authors
+# SPDX-License-Identifier: Apache-2.0
+
+import asyncio
+
+from openai import AsyncStream, Stream
+
+from opentelemetry import context, trace
+from opentelemetry.instrumentation.genai.openai._raw_response import (
+    wrap_stream_result,
+)
+from opentelemetry.instrumentation.genai.openai.chat_wrappers import (
+    AsyncChatStreamWrapper,
+    ChatStreamWrapper,
+)
+from opentelemetry.util.genai.handler import TelemetryHandler
+
+from .test_raw_response_proxy import _AsyncRawResponse, _RawResponse
+
+
+class _EmptyStream(Stream):
+    def __init__(self):
+        pass
+
+    def __iter__(self):
+        return iter(())
+
+    def close(self):
+        pass
+
+
+class _EmptyAsyncStream(AsyncStream):
+    def __init__(self):
+        pass
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        raise StopAsyncIteration
+
+    async def close(self):
+        pass
+
+
+def _start(tracer_provider):
+    return TelemetryHandler(tracer_provider=tracer_provider).inference(
+        "openai", request_model="m"
+    )
+
+
+def test_raw_response_span_not_current_after_create(tracer_provider):
+    before = context.get_current()
+    raw = wrap_stream_result(
+        ChatStreamWrapper,
+        _RawResponse(_EmptyStream()),
+        _start(tracer_provider),
+        False,
+    )
+    assert context.get_current() is before
+    raw.http_response.close()
+
+
+def test_unparsed_raw_response_stops_inside_its_span(tracer_provider):
+    invocation = _start(tracer_provider)
+    stop = invocation.stop
+    current_at_stop = []
+
+    def recording_stop():
+        current_at_stop.append(trace.get_current_span())
+        stop()
+
+    invocation.stop = recording_stop
+    raw = wrap_stream_result(
+        ChatStreamWrapper, _RawResponse(_EmptyStream()), invocation, False
+    )
+    before = context.get_current()
+    with tracer_provider.get_tracer("t").start_as_current_span("caller"):
+        raw.http_response.close()
+
+    assert current_at_stop == [invocation.span]
+    assert context.get_current() is before
+
+
+def test_raw_response_parse_under_caller_span_then_drain(
+    tracer_provider, span_exporter
+):
+    before = context.get_current()
+    raw = wrap_stream_result(
+        ChatStreamWrapper,
+        _RawResponse(_EmptyStream()),
+        _start(tracer_provider),
+        False,
+    )
+    with tracer_provider.get_tracer("t").start_as_current_span("caller"):
+        stream = raw.parse()
+    assert list(stream) == []
+    assert len(span_exporter.get_finished_spans()) == 2
+    assert context.get_current() is before
+
+
+def test_async_raw_response_parse_under_caller_span_then_drain(
+    tracer_provider, span_exporter
+):
+    async def run():
+        before = context.get_current()
+        raw = wrap_stream_result(
+            AsyncChatStreamWrapper,
+            _AsyncRawResponse(_EmptyAsyncStream()),
+            _start(tracer_provider),
+            False,
+        )
+        assert context.get_current() is before
+        with tracer_provider.get_tracer("t").start_as_current_span("caller"):
+            stream = await raw.parse()
+        assert [chunk async for chunk in stream] == []
+        assert len(span_exporter.get_finished_spans()) == 2
+        assert context.get_current() is before
+
+    asyncio.run(run())