Repository navigation
fix(portkey): suspend the invocation while a completion stream is unconsumed - #890
Conversation
…onsumed 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
The fragment records the behaviour change for the released package, as towncrier requires for a fix.
Pull request dashboard statusMerged · refreshed 2026-10-28 00:56 UTC Status above doesn't look right?
|
There was a problem hiding this comment.
🟢 Approval recommended
The implementation follows the established stream-wrapper pattern and adequately covers sync and async lifecycle paths.
0 open findings
What changed in this PR
Fixes Portkey stream context leakage by suspending inference invocations until stream activity occurs.
Changes:
- Suspend sync and async streaming invocations.
- Add context-restoration coverage across completion paths.
- Require the supporting util-genai version.
| File | Description |
|---|---|
wrappers.py |
Suspends streaming invocations. |
test_stream_context.py |
Tests context scoping and finalization. |
pyproject.toml |
Raises the util-genai minimum version. |
requirements.oldest.txt |
Uses the unreleased local utility. |
.changelog/890.fixed |
Documents the fix. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| if asynchronous: | ||
| await stream.aclose() | ||
| else: | ||
| stream.close() |
There was a problem hiding this comment.
Could we also assert that the inference span is current inside the underlying stream’s close() / aclose() ? The producer assertions cover activation during reads, but these close cases currently only check context restoration and that a span was exported. A small fake stream with an assertion in its close method would protect that part of the contract, including closing before the first read.
There was a problem hiding this comment.
Done in 5d91f8c: the fake stream now asserts the inference span is current inside close() / aclose(), for sync and async, both with and without a first read.
…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
sangkyoonnam
left a comment
sangkyoonnam
left a comment
There was a problem hiding this comment.
I checked this with real Portkey/AsyncPortkey clients and a mocked SSE response, and found one case worth knowing about: create the stream inside a caller span, let that span's with block exit, then drain the stream in the same context. On main the already-ended caller span becomes current again after the drain. This branch restores the context that was current right before the drain, for both sync and async streams.
LGTM.
Assisted-by: Doubao
Head branch was pushed to by a user without write access
…t 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.
|
Quick status on the latest run: the deps-check failure is fixed in f4b8b40 (opentelemetry-util-genai 1.3b0.dev is on PyPI now, so the temporary The three remaining failures are all in the anthropic package ( |
|
One more data point to help the merge decision: run 40171297541 (before the deps-check fix) failed 5 jobs, run 40190657824 (after) fails 4 -- the only delta is deps-check, which is now green. The anthropic-conformance/latest failures and the 3s |
….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.
github-project-automation
Bot
moved this to Done
in GenAI Semantic Conventions and Instrumentation libraries
Description
Suspends the Portkey invocation before the completion stream reaches the caller and re-activates it only while a chunk is read or the stream is cleaned up, by calling
invocation.suspend()inPortkeyStreamWrapper.__init__andAsyncPortkeyStreamWrapper.__init__, matching the OpenAI chat wrapper from #817. Portkey is one entry of #872's list; anthropic, bedrock, google-genai, agno, qwen-agent and smolagents are still open there. The package'sopentelemetry-util-genaifloor moves to>= 1.3b0.devand the oldest env installs the local util, since the hook is unreleased.Measured on the branch: with the two
suspend()calls removed,tests/test_stream_context.pyfails all 8 cases at line 78 (context.get_current() is before, checked right after the wrapper is built) and logsFailed to detach context ... was created in a different Context; with them, 8 passed.Part of #872.
Type of change
How has this been tested?
pytest instrumentation/opentelemetry-instrumentation-genai-portkey/tests/test_stream_context.py— 8 passed (sync and async x drain, close, close before any read, error); the same 8 cases fail wheninvocation.suspend()is removed.pytest instrumentation/opentelemetry-instrumentation-genai-portkey/tests/test_streaming.py— 11 passed; the whole portkeytests/— 211 passed, 3 skipped.ruff check,ruff format --check,codespell,uv lock --check— clean.Known gaps: spans a caller creates between chunks stay in the caller's context instead of becoming children of the inference span — the same limitation #817 documents. The conformance and oldest tox envs were not run locally.
Checklist