Skip to content

fix(openai): preserve native v1 stream contract #1627

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hassiebp wants to merge 5 commits into
base: main
Choose a base branch
from
+516 −42
Open

fix(openai): preserve native v1 stream contract #1627

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b818885
fix(openai): preserve native v1 stream contract
hassiebp Apr 15, 2026
2dcc02b
fix(openai): finalize sync stream on early break
hassiebp Apr 16, 2026
514171b
Merge branch 'main' into hassieb/lfe-8788-bug-langfuseresponsegenerat…
hassiebp Apr 16, 2026
b5a7071
fix(openai): guard fallback stream finalization
hassiebp Apr 16, 2026
95c933d
fix(openai): satisfy mypy for stream hooks
hassiebp Apr 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
251 changes: 209 additions & 42 deletions langfuse/openai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -830,6 +830,178 @@
)


_openai_stream_iter_hook_installed = False


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()

setattr(openai.Stream, "__iter__", traced_iter)
_openai_stream_iter_hook_installed = True
Comment on lines +836 to +854
Copy link
Copy Markdown

@claude claude bot Apr 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 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

  1. User calls await openai.chat.completions.create(..., stream=True) → returns native openai.AsyncStream s with s._iterator = traced_iterator() (TI).
  2. User writes: async for chunk in s: break
  3. Python calls s.__aiter__() → outer async generator G. G executes async for item in self._iterator: yield item, reading s._iterator = TI.
  4. G yields the first item. break causes Python to call await G.aclose().
  5. GeneratorExit is thrown into G at the yield item point. G has no try/finally, so it terminates immediately.
  6. TI is still alive: s._iterator still holds a strong reference. TI.aclose() is NOT called. TI's finally does NOT run.
  7. finalize_once() is never called; generation.end() is never called.
  8. The generation leaks until s is GC-collected (at which point s._iterator is released and TI eventually finalized — non-deterministic).
  9. The test suite has test_openai_stream_break_still_finalizes_generation for the sync case but no equivalent async test, confirming the async path was overlooked.



def _finalize_stream_response(
*,
resource: OpenAiDefinition,
items: list[Any],
generation: LangfuseGeneration,
completion_start_time: Optional[datetime],
) -> None:
try:
model, completion, usage, metadata = (
_extract_streamed_response_api_response(items)
if resource.object == "Responses" or resource.object == "AsyncResponses"
else _extract_streamed_openai_response(resource, items)
)

_create_langfuse_update(
completion,
generation,
completion_start_time,
model=model,
usage=usage,
metadata=metadata,
)
except Exception:
pass
finally:
generation.end()


def _instrument_openai_stream(
*,
resource: OpenAiDefinition,
response: Any,
generation: LangfuseGeneration,
) -> Any:
if not hasattr(response, "_iterator"):
return LangfuseResponseGeneratorSync(
resource=resource,
response=response,
generation=generation,
)

items: list[Any] = []
raw_iterator = response._iterator
completion_start_time: Optional[datetime] = None
is_finalized = False
close = response.close

def finalize_once() -> None:
nonlocal is_finalized
if is_finalized:
return

is_finalized = True
_finalize_stream_response(
resource=resource,
items=items,
generation=generation,
completion_start_time=completion_start_time,
)

response._langfuse_finalize_once = finalize_once # type: ignore[attr-defined]

def traced_iterator() -> Any:
nonlocal completion_start_time
try:
for item in raw_iterator:
items.append(item)

if completion_start_time is None:
completion_start_time = _get_timestamp()

yield item
finally:
finalize_once()

def traced_close() -> Any:
try:
return close()
finally:
finalize_once()

response._iterator = traced_iterator()
response.close = traced_close
Comment thread
claude[bot] marked this conversation as resolved.

return response


def _instrument_openai_async_stream(
*,
resource: OpenAiDefinition,
response: Any,
generation: LangfuseGeneration,
) -> Any:
if not hasattr(response, "_iterator"):
return LangfuseResponseGeneratorAsync(
resource=resource,
response=response,
generation=generation,
)

items: list[Any] = []
raw_iterator = response._iterator
completion_start_time: Optional[datetime] = None
is_finalized = False
close = response.close

async def finalize_once() -> None:
nonlocal is_finalized
if is_finalized:
return

is_finalized = True
_finalize_stream_response(
resource=resource,
items=items,
generation=generation,
completion_start_time=completion_start_time,
)

async def traced_iterator() -> Any:
nonlocal completion_start_time
try:
async for item in raw_iterator:
items.append(item)

if completion_start_time is None:
completion_start_time = _get_timestamp()

yield item
finally:
await finalize_once()

async def traced_close() -> Any:
try:
return await close()
finally:
await finalize_once()

async def traced_aclose() -> Any:
return await traced_close()

response._iterator = traced_iterator()
response.close = traced_close
response.aclose = traced_aclose
Comment on lines +998 to +1000
Copy link
Copy Markdown

@chatgpt-codex-connector chatgpt-codex-connector bot Apr 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Finalize async native streams on early loop exit

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 👍 / 👎.


return response


@_langfuse_wrapper
def _wrap(
open_ai_resource: OpenAiDefinition, wrapped: Any, args: Any, kwargs: Any
Expand Down Expand Up @@ -863,7 +1035,13 @@
try:
openai_response = wrapped(**arg_extractor.get_openai_args())

if _is_streaming_response(openai_response):
if _is_openai_v1() and isinstance(openai_response, openai.Stream):
return _instrument_openai_stream(
resource=open_ai_resource,
response=openai_response,
generation=generation,
)
elif _is_streaming_response(openai_response):
return LangfuseResponseGeneratorSync(
resource=open_ai_resource,
response=openai_response,
Expand Down Expand Up @@ -934,7 +1112,13 @@
try:
openai_response = await wrapped(**arg_extractor.get_openai_args())

if _is_streaming_response(openai_response):
if _is_openai_v1() and isinstance(openai_response, openai.AsyncStream):
return _instrument_openai_async_stream(
resource=open_ai_resource,
response=openai_response,
generation=generation,
)
elif _is_streaming_response(openai_response):
return LangfuseResponseGeneratorAsync(
resource=open_ai_resource,
response=openai_response,
Expand Down Expand Up @@ -994,6 +1178,7 @@


register_tracing()
_install_openai_stream_iteration_hooks()


class LangfuseResponseGeneratorSync:
Expand All @@ -1010,6 +1195,7 @@
self.response = response
self.generation = generation
self.completion_start_time: Optional[datetime] = None
self._is_finalized = False

def __iter__(self) -> Any:
try:
Expand Down Expand Up @@ -1045,26 +1231,16 @@
pass

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,
)


class LangfuseResponseGeneratorAsync:
Comment thread
claude[bot] marked this conversation as resolved.
Expand All @@ -1081,6 +1257,7 @@
self.response = response
self.generation = generation
self.completion_start_time: Optional[datetime] = None
self._is_finalized = False

async def __aiter__(self) -> Any:
try:
Expand Down Expand Up @@ -1113,32 +1290,22 @@
return self.__aiter__()

async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
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.

Check notice on line 1308 in langfuse/openai.py

View check run for this annotation

Claude / Claude Code Review

Fallback async wrapper close/exit paths skip _finalize 

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 t
Comment on lines 1293 to 1308
Copy link
Copy Markdown

@claude claude bot Apr 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟣 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

  1. User's code receives a LangfuseResponseGeneratorAsync wrapper W (fallback path: async generator response without _iterator).
  2. User writes await W.aclose() without iterating.
  3. aclose() executes await self.response.aclose() and returns — no call to self._finalize().
  4. generation.end() is never called.
  5. The Langfuse generation span leaks until GC or explicit event loop shutdown.
  6. Alternatively, user writes async with W as gen: pass — __aenter__ returns self.__aiter__(), a fresh async generator that has not been iterated; __aexit__ runs pass; the async generator object goes out of scope. Without asyncio asyncgen hooks, its finally block is not guaranteed to run synchronously.
  7. langfuse_client.flush() immediately after the with block may not see the generation end event.


Automatically called if the response body is read to completion.
"""
Expand Down
Loading
Loading