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,9 @@ 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 `invocation.suspend()` before the stream is returned and
+  return `invocation.activate()` from `_execution_context()`.
 - 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,14 +229,20 @@ 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
+or `close()`. Subclasses pass the SDK stream to `super().__init__(stream)` and implement four
 hooks:
 
+- `_execution_context()` — return a fresh context manager for each stream read and cleanup
+  operation; for invocation-backed streams return `invocation.activate()`. Context is restored
+  before the chunk is returned to the consumer, so the token is created and released in the
+  same frame.
 - `_process_chunk(chunk)` — accumulate per-chunk state (e.g. response model, finish reasons,
   token usage, streamed content) onto the invocation.
 - `_on_stream_end()` — finalize on success; set the accumulated response attributes and call
@@ -248,8 +254,12 @@ class MyStreamWrapper(SyncStreamWrapper[Chunk]):
     def __init__(self, stream, invocation, capture_content):
         super().__init__(stream)
         self._self_invocation = invocation
+        invocation.suspend()
         ...
 
+    def _execution_context(self):
+        return self._self_invocation.activate()
+
     def _process_chunk(self, chunk): ...  # accumulate state
     def _on_stream_end(self):
         self._self_invocation.stop()

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: 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/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_response_wrappers.py

@@ -1,6 +1,7 @@
 # Copyright The OpenTelemetry Authors
 # SPDX-License-Identifier: Apache-2.0
 
+from contextlib import nullcontext
 from types import SimpleNamespace
 
 import pytest
@@ -78,13 +79,22 @@ def _noop_on_stream_chunk(chunk_at):
     del chunk_at
 
 
+def _fake_invocation(**overrides):
+    """A stand-in for an invocation with the stream wrapper's context hooks."""
+    fields = {
+        "request_model": None,
+        "stop": _noop_stop,
+        "fail": _noop_fail,
+        "_on_stream_chunk": _noop_on_stream_chunk,
+        "suspend": _noop_stop,
+        "activate": nullcontext,
+    }
+    fields.update(overrides)
+    return SimpleNamespace(**fields)
+
+
 def _make_wrapper(manager):
-    invocation = SimpleNamespace(
-        request_model=None,
-        stop=_noop_stop,
-        _on_stream_chunk=_noop_on_stream_chunk,
-        fail=_noop_fail,
-    )
+    invocation = _fake_invocation()
     return ResponseStreamManagerWrapper(
         manager=manager,
         invocation_factory=lambda: invocation,
@@ -94,12 +104,7 @@ def _make_wrapper(manager):
 
 def _make_stream_wrapper(stream, invocation=None):
     if invocation is None:
-        invocation = SimpleNamespace(
-            request_model=None,
-            stop=_noop_stop,
-            fail=_noop_fail,
-            _on_stream_chunk=_noop_on_stream_chunk,
-        )
+        invocation = _fake_invocation()
     return ResponseStreamWrapper(
         stream=stream,
         invocation=invocation,
@@ -108,12 +113,7 @@ def _make_stream_wrapper(stream, invocation=None):
 
 
 def _make_async_manager_wrapper(manager):
-    invocation = SimpleNamespace(
-        request_model=None,
-        stop=_noop_stop,
-        _on_stream_chunk=_noop_on_stream_chunk,
-        fail=_noop_fail,
-    )
+    invocation = _fake_invocation()
     return AsyncResponseStreamManagerWrapper(
         manager=manager,
         invocation_factory=lambda: invocation,
@@ -123,12 +123,7 @@ def _make_async_manager_wrapper(manager):
 
 def _make_async_stream_wrapper(stream, invocation=None):
     if invocation is None:
-        invocation = SimpleNamespace(
-            request_model=None,
-            stop=_noop_stop,
-            fail=_noop_fail,
-            _on_stream_chunk=_noop_on_stream_chunk,
-        )
+        invocation = _fake_invocation()
     return AsyncResponseStreamWrapper(
         stream=stream,
         invocation=invocation,
@@ -205,12 +200,7 @@ def test_manager_enter_failure_fails_invocation_and_reraises():
     error = RuntimeError("enter failure")
     manager = _FakeManager(stream=SimpleNamespace(), enter_error=error)
     failures = []
-    invocation = SimpleNamespace(
-        request_model=None,
-        stop=_noop_stop,
-        _on_stream_chunk=_noop_on_stream_chunk,
-        fail=failures.append,
-    )
+    invocation = _fake_invocation(fail=failures.append)
     wrapper = ResponseStreamManagerWrapper(
         manager=manager,
         invocation_factory=lambda: invocation,
@@ -272,12 +262,7 @@ async def test_async_manager_enter_failure_fails_invocation_and_reraises():
     error = RuntimeError("enter failure")
     manager = _FakeAsyncManager(stream=SimpleNamespace(), enter_error=error)
     failures = []
-    invocation = SimpleNamespace(
-        request_model=None,
-        stop=_noop_stop,
-        _on_stream_chunk=_noop_on_stream_chunk,
-        fail=failures.append,
-    )
+    invocation = _fake_invocation(fail=failures.append)
     wrapper = AsyncResponseStreamManagerWrapper(
         manager=manager,
         invocation_factory=lambda: invocation,
@@ -310,9 +295,7 @@ async def test_async_stream_wrapper_exit_fails_and_closes_on_exception():
     stream = _FakeAsyncStream()
     stopped = []
     failures = []
-    invocation = SimpleNamespace(
-        request_model=None, stop=_noop_stop, fail=failures.append
-    )
+    invocation = _fake_invocation(fail=failures.append)
     wrapper = _make_async_stream_wrapper(stream, invocation=invocation)
     wrapper._stop = stopped.append
 
@@ -387,9 +370,7 @@ async def test_async_stream_wrapper_fails_and_reraises_stream_errors():
     error = ValueError("boom")
     stream = _FakeAsyncStream(error=error)
     failures = []
-    invocation = SimpleNamespace(
-        request_model=None, stop=_noop_stop, fail=failures.append
-    )
+    invocation = _fake_invocation(fail=failures.append)
     wrapper = _make_async_stream_wrapper(stream, invocation=invocation)
 
     with pytest.raises(ValueError, match="boom"):
@@ -458,9 +439,7 @@ def _make_response(**overrides):
 
 def _capturing_invocation():
     calls = {"stop": 0, "fail": []}
-    invocation = SimpleNamespace(
-        request_model=None, response_model_name=None, attributes={}
-    )
+    invocation = _fake_invocation(response_model_name=None, attributes={})
     invocation.stop = lambda: calls.__setitem__("stop", calls["stop"] + 1)
     invocation.fail = lambda error: calls["fail"].append(error)
     return invocation, calls
@@ -562,7 +541,7 @@ def _make_fetch_stream_wrapper(stream, invocation):
 
 def _capturing_fetch_invocation():
     calls = {"stop": 0, "fail": []}
-    invocation = SimpleNamespace(
+    invocation = _fake_invocation(
         response_model_name=None,
         response_status=None,
         finish_reasons=None,