Repository navigation
fix(util-genai,openai): scope stream invocation context - #817
sangkyoonnam wants to merge 11 commits into
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 statusWaiting on reviewers · refreshed 2026-10-07 05:04 UTC Review the latest changes. 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.
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