fix(openai): guard fallback stream finalization

Author: hassiebp

langfuse/openai.py

@@ -1195,6 +1195,7 @@ def __init__(
         self.response = response
         self.generation = generation
         self.completion_start_time: Optional[datetime] = None
+        self._is_finalized = False
 
     def __iter__(self) -> Any:
         try:
@@ -1230,6 +1231,10 @@ def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         pass
 
     def _finalize(self) -> None:
+        if self._is_finalized:
+            return
+
+        self._is_finalized = True
         _finalize_stream_response(
             resource=self.resource,
             items=self.items,
@@ -1252,6 +1257,7 @@ def __init__(
         self.response = response
         self.generation = generation
         self.completion_start_time: Optional[datetime] = None
+        self._is_finalized = False
 
     async def __aiter__(self) -> Any:
         try:
@@ -1287,6 +1293,10 @@ async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None
         pass
 
     async def _finalize(self) -> None:
+        if self._is_finalized:
+            return
+
+        self._is_finalized = True
         _finalize_stream_response(
             resource=self.resource,
             items=self.items,

tests/unit/test_openai.py

@@ -3,6 +3,7 @@
 
 import pytest
 
+import langfuse.openai as lf_openai_module
 from langfuse._client.attributes import LangfuseOtelSpanAttributes
 from langfuse.openai import openai as lf_openai
 
@@ -39,6 +40,17 @@ async def _stream(self, items):
             yield item
 
 
+class DummyGeneration:
+    def __init__(self) -> None:
+        self.end_calls = 0
+
+    def update(self, **kwargs):
+        return self
+
+    def end(self) -> None:
+        self.end_calls += 1
+
+
 def _make_chat_stream_chunks():
     usage = SimpleNamespace(prompt_tokens=3, completion_tokens=1, total_tokens=4)
 
@@ -76,6 +88,24 @@ def _make_chat_stream_chunks():
     ]
 
 
+def _make_single_chunk_stream():
+    return SimpleNamespace(
+        model="gpt-4o-mini",
+        choices=[
+            SimpleNamespace(
+                delta=SimpleNamespace(
+                    role="assistant",
+                    content="2",
+                    function_call=None,
+                    tool_calls=None,
+                ),
+                finish_reason="stop",
+            )
+        ],
+        usage=None,
+    )
+
+
 def test_chat_completion_exports_generation_span(
     langfuse_memory_client, get_span, json_attr
 ):
@@ -439,6 +469,50 @@ async def test_openai_async_stream_supports_anext(
     }
 
 
+def test_fallback_sync_stream_finalizes_once():
+    resource = SimpleNamespace(object="Completions", type="chat")
+    generation = DummyGeneration()
+
+    def fallback_stream():
+        yield _make_single_chunk_stream()
+
+    wrapper = lf_openai_module.LangfuseResponseGeneratorSync(
+        resource=resource,
+        response=fallback_stream(),
+        generation=generation,
+    )
+
+    list(wrapper)
+
+    with pytest.raises(StopIteration):
+        next(wrapper)
+
+    assert generation.end_calls == 1
+
+
+@pytest.mark.asyncio
+async def test_fallback_async_stream_finalizes_once():
+    resource = SimpleNamespace(object="Completions", type="chat")
+    generation = DummyGeneration()
+
+    async def fallback_stream():
+        yield _make_single_chunk_stream()
+
+    wrapper = lf_openai_module.LangfuseResponseGeneratorAsync(
+        resource=resource,
+        response=fallback_stream(),
+        generation=generation,
+    )
+
+    async for _ in wrapper:
+        pass
+
+    with pytest.raises(StopAsyncIteration):
+        await wrapper.__anext__()
+
+    assert generation.end_calls == 1
+
+
 def test_embedding_exports_dimensions_and_count(
     langfuse_memory_client, get_span, json_attr
 ):