fix(openai): preserve native v1 stream contract#1627
Conversation
|
@claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| def _install_openai_stream_iteration_hooks() -> None: | ||
| global _openai_stream_iter_hook_installed | ||
|
|
||
| if not _is_openai_v1(): | ||
| return | ||
|
|
||
| if not _openai_stream_iter_hook_installed: | ||
| original_iter = openai.Stream.__iter__ | ||
|
|
||
| def traced_iter(self: Any) -> Any: | ||
| try: | ||
| yield from original_iter(self) | ||
| finally: | ||
| finalize_once = getattr(self, "_langfuse_finalize_once", None) | ||
| if finalize_once is not None: | ||
| finalize_once() | ||
|
|
||
| openai.Stream.__iter__ = traced_iter | ||
| _openai_stream_iter_hook_installed = True |
There was a problem hiding this comment.
🔴 The sync early-break fix (commit 2dcc02b) has no async equivalent: _install_openai_stream_iteration_hooks only patches openai.Stream.__iter__ and never installs a hook on openai.AsyncStream.__aiter__, and _instrument_openai_async_stream never sets response._langfuse_finalize_once. When a user does async for chunk in stream: break, Python closes the outer __aiter__ generator but does not propagate aclose() to self._iterator (an attribute reference, not a frame-local), so traced_iterator()'s finally block never fires and generation.end() is never called. Fix by installing an analogous async hook on openai.AsyncStream.__aiter__ and storing finalize_once as response._langfuse_finalize_once in _instrument_openai_async_stream.
Extended reasoning...
What the bug is and how it manifests
The PR fixes sync stream early-break finalization via two coordinated changes: (1) _install_openai_stream_iteration_hooks patches openai.Stream.__iter__ with a traced_iter that wraps the original with yield from original_iter(self) and calls self._langfuse_finalize_once() in a finally block, and (2) _instrument_openai_stream stores response._langfuse_finalize_once = finalize_once so the hook can reach the per-instance callback. Neither of these two steps was applied to the async path. _install_openai_stream_iteration_hooks never patches openai.AsyncStream.__aiter__, and _instrument_openai_async_stream never sets response._langfuse_finalize_once.
The specific code path that triggers it
openai.AsyncStream.__aiter__ (confirmed from OpenAI SDK _streaming.py) is an async generator: async for item in self._iterator: yield item. When a user does async for chunk in stream: break, Python creates an outer async generator G from __aiter__. G iterates self._iterator — which is traced_iterator() (call it TI) installed by _instrument_openai_async_stream. After yielding the first chunk, break causes Python to call await G.aclose(), throwing GeneratorExit into G at the yield item point.
Why existing code does not prevent it
G terminates on GeneratorExit, but self._iterator (TI) is stored on the response object as an attribute — it is not a local variable in G's frame. Python only automatically propagates aclose() to generators stored as frame-local variables (via yield from / async for on a local). Here, TI's refcount does not drop to zero synchronously, so TI.aclose() is never called. TI's finally: await finalize_once() never executes. traced_aclose / traced_close are only invoked by explicit stream.aclose() / stream.close() calls — they are not triggered by early break.
What the impact would be
generation.end() is never called immediately. The Langfuse generation leaks until either the user explicitly calls await stream.aclose() or the stream object is garbage-collected (non-deterministic). Users who read only the first N chunks from an async stream without remembering to call aclose() will silently lose Langfuse traces. This is an observability correctness regression: the previous LangfuseResponseGeneratorAsync.__aiter__ had a try/finally directly in the generator the caller iterates, guaranteeing deterministic finalization on any exit path including break.
How to fix it
Two symmetric changes are needed: (1) in _install_openai_stream_iteration_hooks, after patching openai.Stream.__iter__, also install an analogous async hook on openai.AsyncStream.__aiter__:
original_aiter = openai.AsyncStream.__aiter__
async def traced_aiter(self):
try:
async for x in original_aiter(self):
yield x
finally:
cb = getattr(self, "_langfuse_finalize_once", None)
if cb is not None:
await cb()
openai.AsyncStream.__aiter__ = traced_aiter(2) in _instrument_openai_async_stream, add response._langfuse_finalize_once = finalize_once immediately after defining finalize_once, mirroring the sync path (line 917 in the new code). A corresponding test test_openai_async_stream_break_still_finalizes_generation should also be added.
Step-by-step proof
- User calls
await openai.chat.completions.create(..., stream=True)→ returns nativeopenai.AsyncStream swiths._iterator = traced_iterator()(TI). - User writes:
async for chunk in s: break - Python calls
s.__aiter__()→ outer async generator G. G executesasync for item in self._iterator: yield item, readings._iterator = TI. - G yields the first item.
breakcauses Python to callawait G.aclose(). GeneratorExitis thrown into G at theyield itempoint. G has notry/finally, so it terminates immediately.- TI is still alive:
s._iteratorstill holds a strong reference.TI.aclose()is NOT called. TI'sfinallydoes NOT run. finalize_once()is never called;generation.end()is never called.- The generation leaks until
sis GC-collected (at which points._iteratoris released and TI eventually finalized — non-deterministic). - The test suite has
test_openai_stream_break_still_finalizes_generationfor the sync case but no equivalent async test, confirming the async path was overlooked.
…orasync-object-has-no-attribute
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b5a7071be7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| response._iterator = traced_iterator() | ||
| response.close = traced_close | ||
| response.aclose = traced_aclose |
There was a problem hiding this comment.
In _instrument_openai_async_stream, we only swap response._iterator and patch close/aclose, but unlike the sync path there is no AsyncStream iteration hook to finalize on async for ... break. Because the stream keeps a strong reference to response._iterator, breaking early leaves traced_iterator() unclosed, so its finally never runs and generation.end() is skipped until GC or explicit aclose(). In practice this can leave observations open and cause langfuse.flush() to miss partial async stream traces whenever callers stop consuming early.
Useful? React with 👍 / 👎.
| pass | ||
|
|
||
| async def _finalize(self) -> None: | ||
| try: | ||
| model, completion, usage, metadata = ( | ||
| _extract_streamed_response_api_response(self.items) | ||
| if self.resource.object == "Responses" | ||
| or self.resource.object == "AsyncResponses" | ||
| else _extract_streamed_openai_response(self.resource, self.items) | ||
| ) | ||
|
|
||
| _create_langfuse_update( | ||
| completion, | ||
| self.generation, | ||
| self.completion_start_time, | ||
| model=model, | ||
| usage=usage, | ||
| metadata=metadata, | ||
| ) | ||
| except Exception: | ||
| pass | ||
| finally: | ||
| self.generation.end() | ||
| if self._is_finalized: | ||
| return | ||
|
|
||
| self._is_finalized = True | ||
| _finalize_stream_response( | ||
| resource=self.resource, | ||
| items=self.items, | ||
| generation=self.generation, | ||
| completion_start_time=self.completion_start_time, | ||
| ) | ||
|
|
||
| async def close(self) -> None: | ||
| """Close the response and release the connection. |
There was a problem hiding this comment.
🟣 The fallback async wrapper's close(), aclose(), and __aexit__ methods never call self._finalize(), so if a user receives a LangfuseResponseGeneratorAsync wrapper (for non-OpenAI generator-style async streams) and closes or context-manages it without iterating, generation.end() is never called and the Langfuse span leaks. The same issue exists for LangfuseResponseGeneratorSync.__exit__. This is a pre-existing issue that predates this PR; the PR's _is_finalized guard improvements to _finalize() make the fix straightforward — add await self._finalize() to close()/aclose()/__aexit__ and self._finalize() to __exit__.
Extended reasoning...
What the bug is and how it manifests
LangfuseResponseGeneratorAsync exposes close() and aclose() methods (lines ~1307–1319 in the modified file) that only delegate to self.response.close()/self.response.aclose() without calling self._finalize(). The __aexit__ method (line ~1293) is implemented as pass. For the sync counterpart, LangfuseResponseGeneratorSync.__exit__ is similarly a no-op pass (line ~1231). If a user receives the fallback wrapper (returned when the response lacks a _iterator attribute, i.e. for non-native-OpenAI generator-style streams) and calls await stream.aclose() or uses it as an async context manager without iterating, generation.end() is never called.
The specific code path that triggers it
The only paths in LangfuseResponseGeneratorAsync that call self._finalize() are: (1) __aiter__'s finally block (line ~1270), and (2) __anext__'s except StopAsyncIteration handler (line ~1285). Both require actual iteration. close() executes await self.response.close() and returns; aclose() calls await self.response.aclose(); __aexit__ does pass. None of these reach _finalize(). For the async generator returned by __aiter__, Python's async generator finalization is non-deterministic without asyncio loop shutdown hooks (PEP 525), so GC-based finalization is not reliable.
Why existing code does not prevent it
The PR added an _is_finalized idempotency guard to _finalize() (clearly intending to improve finalization correctness), but did not extend that pass to wire up close(), aclose(), or __aexit__/__exit__ as callers of _finalize(). The contrast with _instrument_openai_async_stream (the new native-stream path in this same PR) is clear: its traced_close and traced_aclose both call finalize_once() in a finally block. The fallback wrapper was not updated with the same pattern.
What the impact would be
Any code using the fallback path (non-OpenAI async generators fed through the instrumented API) that calls await stream.aclose() or uses async with stream: without exhausting the stream will silently leak the Langfuse generation span. generation.end() is never called, the span remains open indefinitely, and calls to langfuse.flush() will race against GC-based finalization.
How to fix it
Add await self._finalize() at the end of close() and aclose(), and replace pass in __aexit__ with await self._finalize(). For the sync class, replace pass in __exit__ with self._finalize(). The _is_finalized guard (already in place from this PR) prevents double-finalization if the stream was also consumed via iteration.
Step-by-step proof
- User's code receives a
LangfuseResponseGeneratorAsyncwrapperW(fallback path: async generator response without_iterator). - User writes
await W.aclose()without iterating. aclose()executesawait self.response.aclose()and returns — no call toself._finalize().generation.end()is never called.- The Langfuse generation span leaks until GC or explicit event loop shutdown.
- Alternatively, user writes
async with W as gen: pass—__aenter__returnsself.__aiter__(), a fresh async generator that has not been iterated;__aexit__runspass; the async generator object goes out of scope. Without asyncio asyncgen hooks, itsfinallyblock is not guaranteed to run synchronously. langfuse_client.flush()immediately after thewithblock may not see the generation end event.
1 participant
Summary
Fix OpenAI v1 streaming instrumentation to preserve the original
openai.Streamandopenai.AsyncStreamobjects instead of replacing them with Langfuse wrapper types.This keeps the native stream contract intact for
_iterator,.response,close()/aclose(),async for, and manual__anext__()while still collecting streamed chunks for Langfuse generation updates.Changes
async forand__anext__()Verification
uv run --frozen ruff check langfuse/openai.py tests/unit/test_openai.pyuv run --frozen pytest tests/unit/test_openai.pyRefs LFE-8788.
Disclaimer: Experimental PR review
Greptile Summary
This PR fixes OpenAI v1 streaming instrumentation by patching
_iteratorandclose/acloseon the nativeopenai.Streamandopenai.AsyncStreamobjects in place, rather than replacing them with Langfuse wrapper types. The approach is sound:_iteratorreplacement correctly intercepts all iteration paths (__iter__,__next__,__anext__,async for) as confirmed by the SDK's implementation, and theis_finalizedflag reliably prevents double-finalization.Confidence Score: 5/5
Safe to merge — core patching logic is correct, double-finalization is guarded, and tests cover all three iteration patterns.
All remaining findings are P2 style suggestions. The
_finalize_stream_response_asyncwrapper is harmless. No P0/P1 logic bugs found; the_iteratorreplacement approach is confirmed sound by the OpenAI SDK's implementation.No files require special attention.
Important Files Changed
_instrument_openai_stream/_instrument_openai_async_streamto monkey-patch_iteratorandclose/acloseon nativeopenai.Streamobjects, preserving type identity while still collecting chunks for Langfuse; logic is sound but_finalize_stream_response_asyncis a no-op async wrapper.__anext__usage;DummyOpenAIStream/DummyOpenAIAsyncStreamcorrectly bypasssuper().__init__and set_iteratordirectly, which is appropriate for unit-level testing.Sequence Diagram
sequenceDiagram participant User participant openai.Stream participant traced_iterator (generator) participant raw_iterator participant finalize_once participant LangfuseGeneration Note over User,LangfuseGeneration: _instrument_openai_stream patches _iterator and close in place User->>openai.Stream: for chunk in stream (or next()) openai.Stream->>traced_iterator (generator): __next__() via self._iterator traced_iterator (generator)->>raw_iterator: next item raw_iterator-->>traced_iterator (generator): chunk traced_iterator (generator)->>traced_iterator (generator): items.append(chunk), set completion_start_time traced_iterator (generator)-->>openai.Stream: yield chunk openai.Stream-->>User: chunk Note over traced_iterator (generator): On exhaustion or close() traced_iterator (generator)->>finalize_once: finally block (is_finalized guard) finalize_once->>LangfuseGeneration: _finalize_stream_response → generation.end() alt User calls stream.close() explicitly User->>openai.Stream: close() [patched to traced_close] openai.Stream->>openai.Stream: await original_close() (closes HTTP) openai.Stream->>finalize_once: finally in traced_close finalize_once->>finalize_once: is_finalized=True → no-op if already finalized endReviews (1): Last reviewed commit: "fix(openai): preserve native v1 stream c..." | Re-trigger Greptile