From 716a58e6205ec2e80a2aa4ca99e5f8025bbc93e1 Mon Sep 17 00:00:00 2001 From: fei <204683769+feiiiiii5@users.noreply.github.com> Date: Sat, 7 Nov 2026 00:10:49 +0800 Subject: [PATCH 1/2] fix(openai): end observations cancelled before a response --- langfuse/openai.py | 4 ++ tests/unit/test_openai.py | 94 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 98 insertions(+) diff --git a/langfuse/openai.py b/langfuse/openai.py index 32637569b..0bff6d6d6 100644 --- a/langfuse/openai.py +++ b/langfuse/openai.py @@ -23,6 +23,7 @@ See docs for more details: https://langfuse.com/docs/integrations/openai """ +import asyncio import json import os import types @@ -1408,6 +1409,9 @@ async def _wrap_async( ).end() return openai_response + except asyncio.CancelledError: + generation.end() + raise except Exception as ex: logger.warning(ex) model = kwargs.get("model", None) or None diff --git a/tests/unit/test_openai.py b/tests/unit/test_openai.py index a2a4f7f4f..234218835 100644 --- a/tests/unit/test_openai.py +++ b/tests/unit/test_openai.py @@ -1,5 +1,7 @@ import asyncio import base64 +import json +from contextlib import suppress from types import SimpleNamespace from unittest.mock import patch @@ -602,6 +604,98 @@ def test_chat_completion_error_marks_generation_error(langfuse_memory_client, ge assert LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT not in span.attributes +@pytest.mark.asyncio +@pytest.mark.parametrize( + "resource,stream", + [ + pytest.param("chat", False, id="chat"), + pytest.param("chat", True, id="chat-stream"), + pytest.param("responses", False, id="responses"), + pytest.param("responses", True, id="responses-stream"), + pytest.param("embeddings", None, id="embeddings"), + ], +) +async def test_async_request_cancelled_before_response_ends_observation( + resource, stream, langfuse_memory_client, find_spans +): + import httpx + + request_started = asyncio.Event() + release_response = asyncio.Event() + cancellations = [] + name = f"unit-openai-cancel-before-response-{resource}-{stream}" + model = "text-embedding-3-small" if resource == "embeddings" else "gpt-4o-mini" + + async def handler(request: httpx.Request) -> httpx.Response: + expected_path = { + "chat": "/v1/chat/completions", + "responses": "/v1/responses", + "embeddings": "/v1/embeddings", + }[resource] + assert request.url.path == expected_path + body = json.loads(request.content) + if resource == "embeddings": + assert "stream" not in body + else: + assert body["stream"] is stream + request_started.set() + try: + await release_response.wait() + except asyncio.CancelledError as exc: + cancellations.append(exc) + raise + raise AssertionError("Request resumed without cancellation") + + async with lf_openai.AsyncOpenAI( + api_key="test", + max_retries=0, + http_client=httpx.AsyncClient( + transport=httpx.MockTransport(handler), trust_env=False + ), + ) as client: + if resource == "chat": + call = client.chat.completions.create( + name=name, + model=model, + messages=[{"role": "user", "content": "hello"}], + stream=stream, + ) + elif resource == "responses": + call = client.responses.create( + name=name, model=model, input="hello", stream=stream + ) + else: + call = client.embeddings.create(name=name, model=model, input="hello") + + task = asyncio.create_task(call) + try: + await asyncio.wait_for(request_started.wait(), timeout=5) + task.cancel("caller cancelled") + with pytest.raises(asyncio.CancelledError, match="caller cancelled") as exc: + await task + assert task.cancelled() + assert len(cancellations) == 1 + assert exc.value is cancellations[0] + finally: + task.cancel() + with suppress(asyncio.CancelledError): + await task + + langfuse_memory_client.flush() + spans = find_spans(name) + assert len(spans) == 1 + span = spans[0] + assert span.end_time is not None + assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_MODEL] == model + assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_TYPE] == ( + "embedding" if resource == "embeddings" else "generation" + ) + assert span.attributes.get(LangfuseOtelSpanAttributes.OBSERVATION_LEVEL) != "ERROR" + assert LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT not in span.attributes + assert LangfuseOtelSpanAttributes.OBSERVATION_USAGE_DETAILS not in span.attributes + assert LangfuseOtelSpanAttributes.OBSERVATION_COST_DETAILS not in span.attributes + + def test_openai_stream_preserves_original_stream_contract( langfuse_memory_client, get_span, json_attr ): From 0685844a747832d2347cd79c37d6b4d409619d81 Mon Sep 17 00:00:00 2001 From: fei <204683769+feiiiiii5@users.noreply.github.com> Date: Sat, 7 Nov 2026 00:26:13 +0800 Subject: [PATCH 2/2] test(openai): move cancellation transport import to module scope --- tests/unit/test_openai.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/test_openai.py b/tests/unit/test_openai.py index 234218835..60c1c1860 100644 --- a/tests/unit/test_openai.py +++ b/tests/unit/test_openai.py @@ -5,6 +5,7 @@ from types import SimpleNamespace from unittest.mock import patch +import httpx import pytest from openai.types.chat import ChatCompletion from openai.types.responses import ParsedResponseOutputMessage, ParsedResponseOutputText @@ -618,8 +619,6 @@ def test_chat_completion_error_marks_generation_error(langfuse_memory_client, ge async def test_async_request_cancelled_before_response_ends_observation( resource, stream, langfuse_memory_client, find_spans ): - import httpx - request_started = asyncio.Event() release_response = asyncio.Event() cancellations = []