Repository navigation
Conversation
…enAI part) Keep callback-managed spans detached and activate their context around framework execution boundaries, including custom tools and retrievers, streaming, and graph persistence. Add inference and HTTP correlation tests covering concurrency, cancellation, and context restoration. Assisted-by: GPT-6
Pull request dashboard statusMerged · refreshed 2026-10-11 05:05 UTC Status above doesn't look right?
|
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Stream-manager exits and Responses close-proxy finalizers still perform cleanup outside the invocation context.
Review effort: Balanced
Findings: 2
Open (2)
What changed in this PR
Scopes streamed OpenAI invocations to reads and cleanup rather than leaving inference spans active between chunks.
Changes:
- Adds stream execution-context hooks and generator forwarding.
- Applies scoped contexts to OpenAI chat and Responses streams.
- Adds tests, documentation, dependency floor, and changelogs.
| File | Description |
|---|---|
util/.../tests/test_stream.py |
Tests scoped stream operations. |
util/.../genai/stream.py |
Adds execution-context handling. |
util/.../README.rst |
Documents the new hook. |
util/.../AGENTS.md |
Updates stream guidance. |
util/.../.changelog/817.added |
Records the util feature. |
openai/tests/test_responses_stream_context.py |
Tests Responses context behavior. |
openai/tests/test_response_wrappers.py |
Updates invocation test doubles. |
openai/tests/test_chat_stream_context.py |
Tests chat context behavior. |
openai/tests/requirements.oldest.txt |
Uses the unreleased local util. |
openai/.../response_wrappers.py |
Scopes Responses stream reads. |
openai/.../chat_wrappers.py |
Scopes chat stream reads. |
openai/pyproject.toml |
Raises the util dependency floor. |
openai/.changelog/817.fixed |
Records the context fix. |
AGENTS.md |
Updates repository stream guidance. |
.github/instructions/instrumentation.instructions.md |
Updates instrumentation review rules. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
… inside the stream execution context
| def _execution_context(self) -> AbstractContextManager[None]: | ||
| """Scope stream reads and cleanup, restoring context before returning a chunk.""" | ||
| return nullcontext() |
There was a problem hiding this comment.
Could the base wrappers provide the activate() part by default? Every wrapper that gets an invocation repeats the same override. With suspend/activate on _StreamTimingInvocation and the invocation remembering _attach_to_context:
def _execution_context(self) -> AbstractContextManager[None]:
invocation = self._self_invocation
if invocation is None or not invocation._attach_to_context:
return nullcontext()
return invocation.activate()This mirrors suspend(), which is already a no-op for detached invocations. suspend() can stay explicit in instrumentations since it has to run where the invocation was started. The OpenAI and tool-wrapper overrides can then go, and other instrumentations only need the suspend() call to opt in.
There was a problem hiding this comment.
Done in bbcdef5. GenAIInvocation remembers _attach_to_context, and the base _execution_context() returns invocation.activate() for a suspended invocation and nullcontext() otherwise. The OpenAI chat and Responses wrappers and both tool wrappers lose their overrides; the tool wrappers now hand the invocation to the base class so it reaches the default. suspend() stays where it is.
Two adjustments to the shape you sketched, both from things the first cut broke. The gate is the invocation's _suspended state (attached, then suspended) rather than _attach_to_context alone: wrappers that never call suspend() (Anthropic, Bedrock, Google GenAI, Portkey, smolagents) keep the invocation attached, and activating it around the read that finalizes reinstated the ended span after stop() detached the original token. With the gate in the hook, activate() itself keeps its meaning, and those instrumentations behave as before. And handing a ToolInvocation to the base class recorded chunk-timing histograms for tools, which the conventions don't define, so ToolInvocation._on_stream_chunk is a no-op and a test asserts a drained tool stream exports only gen_ai.execute_tool.duration. One consequence worth naming: a tool invocation created with _attach_to_context=False is no longer made current by the tool wrapper (its old override activated regardless of the flag); no in-repo caller does that, and a test documents it.
Regression tests: a never-suspended invocation ends with the parent current through exhaustion, error and close, sync and async; a suspended one is current during reads only. util 538 passed, openai latest 394 passed (the two bad_endpoint cassette failures are there without this change too), oldest 193, langchain 483, pyright on util 0.
| @@ -89,8 +95,12 @@ class MyStreamWrapper(SyncStreamWrapper[Chunk]): | |||
| def __init__(self, stream, invocation, capture_content): | |||
| super().__init__(stream) | |||
There was a problem hiding this comment.
Please pass invocation to the base wrapper here so it sets gen_ai.request.stream. The base wrapper now handles activation, so the _execution_context override can be removed.
There was a problem hiding this comment.
Updated in 82c4177, along with the copy of this section in the root AGENTS.md.
| - 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 |
There was a problem hiding this comment.
This rule still requires an _execution_context override, but the resolved implementation deliberately moved activation into the base wrapper. Could we instead require passing the invocation to super and calling invocation.suspend()? An explicit override would bypass the new suspended-state guard and could activate invocations created with _attach_to_context=False
| consumer. Generator ``send``/``throw`` and ``asend``/``athrow`` also use this | ||
| scope when the underlying stream supports them. | ||
|
|
||
| For invocation-backed streams, call ``invocation.suspend()`` before returning the |
There was a problem hiding this comment.
Could we update this guidance to reflect the new default? Invocation-backed wrappers should pass the invocation to the base constructor and call invocation.suspend(); the base wrapper then activates it during reads and cleanup. Recommending a custom _execution_context override no longer matches the implementation.
… overriding _execution_context
…-execution-context # Conflicts: # util/opentelemetry-util-genai/src/opentelemetry/util/genai/_invocation.py
| ) -> None: | ||
| super().__init__(stream, invocation=invocation) | ||
| self._self_invocation = invocation | ||
| invocation.suspend() |
There was a problem hiding this comment.
On the with_raw_response path this wrapper is only built on parse(), so suspend() runs wherever the caller parses. If that's inside a caller span, the inference span is left current after the stream is drained and the span has ended. On main the same flow ends with the context restored, since stop() detaches. The span is also still current between create() and parse().
Would it work to suspend in wrap_stream_result (or right before traced_method returns) so both paths are covered? Same applies to _ResponseStreamMixin. Both tests below fail on this branch; the second passes on main:
from openai import Stream
from opentelemetry import context
from opentelemetry.instrumentation.genai.openai._raw_response import (
wrap_stream_result,
)
from opentelemetry.instrumentation.genai.openai.chat_wrappers import (
ChatStreamWrapper,
)
from opentelemetry.util.genai.handler import TelemetryHandler
from .test_raw_response_proxy import _RawResponse
class _EmptyStream(Stream):
def __init__(self):
pass
def __iter__(self):
return iter(())
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()
wrap_stream_result(
ChatStreamWrapper, _RawResponse(_EmptyStream()), _start(tracer_provider), False
)
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 beforeThere was a problem hiding this comment.
Fixed in 76a7383: wrap_stream_result now suspends before returning the raw-response proxy, so create() leaves the context as it found it and parse() under a caller span no longer leaves the inference span current. The Responses API streams go through the same function, so _ResponseStreamMixin is covered too. Your two tests are in test_raw_response_stream_suspend.py, plus an async variant and one for a response closed without parse().
| @@ -0,0 +1 @@ | |||
| Add the `_execution_context()` hook on `SyncStreamWrapper`/`AsyncStreamWrapper` to scope each stream read and cleanup, including generator `send`/`throw` and `asend`/`athrow`. `finalize_on_close`/`finalize_on_aclose` accept an optional `execution_context` so a close through the HTTP response runs inside the same scope, and the stream manager wrappers run the wrapped manager's exit inside it. The hook defaults to making a suspended invocation current again for each read and for finalizing an abandoned stream, so an instrumentation opts in by passing the invocation to `super().__init__(stream, invocation)` and calling `invocation.suspend()` when it hands the stream back. | |||
There was a problem hiding this comment.
nit: this documents a private hook and implementation details; could be shorter and user-facing:
| Add the `_execution_context()` hook on `SyncStreamWrapper`/`AsyncStreamWrapper` to scope each stream read and cleanup, including generator `send`/`throw` and `asend`/`athrow`. `finalize_on_close`/`finalize_on_aclose` accept an optional `execution_context` so a close through the HTTP response runs inside the same scope, and the stream manager wrappers run the wrapped manager's exit inside it. The hook defaults to making a suspended invocation current again for each read and for finalizing an abandoned stream, so an instrumentation opts in by passing the invocation to `super().__init__(stream, invocation)` and calling `invocation.suspend()` when it hands the stream back. | |
| Stream wrappers make a suspended invocation current only while a chunk is read or the stream is cleaned up. |
…onsumed (open-telemetry#890) * fix(portkey): suspend the invocation while a completion stream is unconsumed PortkeyStreamWrapper and AsyncPortkeyStreamWrapper hand the invocation to the base stream wrapper but never call invocation.suspend(), so the inference span stays current in the caller's context between reads. The execution-scoped context landed in open-telemetry#817 for the OpenAI wrappers; Portkey is one entry of open-telemetry#872's list. Measured on this branch with the two suspend() calls removed, driving tests/test_stream_context.py (sync and async x drain, close, close before any read, error): E assert {'current-span-...': _Span(name="chat test-model", ...)} is {} E .../tests/test_stream_context.py:78 and pytest logs: Failed to detach context ... Token ... was created in a different Context before: 8 failed in 0.07s after: 8 passed Verification: tests/test_stream_context.py 8 passed; tests/test_streaming.py 11 passed; portkey tests/ 211 passed, 3 skipped; ruff check, ruff format --check, codespell and uv lock --check clean. Assisted-by: Doubao * docs(portkey): add the towncrier fragment for open-telemetry#890 The fragment records the behaviour change for the released package, as towncrier requires for a fix. * test(portkey): assert the inference span is current while the stream is closed The portkey wrapper closes the underlying stream inside the invocation context it restores, but test_stream_context.py only checked that the caller's context came back and that a span was exported, so nothing failed if that close path stopped being scoped. A fake stream whose close/aclose records the current span context now covers sync and async, closed before the first read and after one. Verified: with the two invocation.suspend() calls removed the file is 12 failed; with the util's close finalization moved outside _execution_context() the four close cases fail on the new assertion; on the branch 12 passed and the package suite is 215 passed, 3 skipped. Assisted-by: Doubao * Apply suggestion from @lmolkova * Update instrumentation/opentelemetry-instrumentation-genai-portkey/.changelog/890.fixed * style(portkey): apply ruff format to stream context tests Assisted-by: Doubao * build(portkey): drop the local util-genai editable install from oldest deps opentelemetry-util-genai 1.3b0.dev (the stream execution hook) is now on PyPI, so the temporary '-e util/opentelemetry-util-genai' pin in requirements.oldest.txt is no longer needed; the oldest factor resolves the declared floor from PyPI. Fixes the deps-check failure on this PR. * Fix(portkey-tests): wrap non-streaming mock response for portkey-ai 1.x GenericResponse portkey-ai 1.x parses prompts.completions.create responses into a GenericResponse pydantic model that requires success/data/warning fields. The oldest-env tox job (portkey-ai==1.1.2) failed with 24 pydantic ValidationErrors. Detect portkey-ai major version and wrap the mocked OpenAI-shaped body in the envelope only on 1.x; 2.x uses PromptCompletion/ChatCompletions directly. * fix(portkey): apply ruff format to test_token_usage.py * fix(portkey): wrap streaming SSE chunks in 1.x GenericResponse envelope * fix(portkey): repair SSE chunk string literal and apply ruff format * style(portkey): apply ruff format with line-length=79 (CI config) * test(anthropic): filter beta query parameter when matching VCR cassettes * style(anthropic): apply ruff format to vcr beta compat test --------- Co-authored-by: Liudmila Molkova <neskazu@gmail.com>

Description
Adds
_execution_context()to util-genai's stream wrappers, so the OpenAI chat completion and Responses API invocations are active during reads and cleanup and the caller's context is restored between chunks. The inference span is no longer current in the caller's context while a stream sits unconsumed. The first commit is the util and OpenAI part of @lmolkova's prototype for #677, with her authorship; the commits after it migrate the Responses API wrappers and expand the docs.Known gaps: spans a caller creates between chunks keep the caller's context instead of becoming children of the inference span. Needs the unreleased util-genai 1.3b0.dev floor, the way #632 did. The LangChain half follows in a separate PR on top of this. The other stream wrappers (anthropic, bedrock, google-genai, portkey, agno, qwen, smolagents) still hold the token across the return; I'll file a follow-up for them. Anthropic has no dedicated test against the new hook.
Refs #677.
Type of change
How has this been tested?
util/.../tests/test_stream.pycases cover the hook; OpenAItest_chat_stream_context.pyandtest_responses_stream_context.pycover cross-task reads, context restoration and the absence of ERROR logs fromopentelemetry.context. The Responses cases fail on the prototype alone.tox -e py312-test-util-genai-latest: 520 passed;py312-test-instrumentation-genai-openai-latest: 394 passed;-oldest: 191 passed; main's langchain suites pass unchanged on this branch (483 latest, 435 oldest); typecheck, precommit.Checklist