From b53015be858a671c3c43cd3be2a68c254aca3045 Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 19:40:29 +0200 Subject: [PATCH] Address OpenAI stream review issues --- studio/backend/routes/inference.py | 92 ++++++++++++++++--- .../tests/test_openai_tool_passthrough.py | 74 ++++++++++++++- .../tests/test_responses_tool_passthrough.py | 5 +- .../test_stream_cancel_registration_timing.py | 8 ++ 4 files changed, 163 insertions(+), 16 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ee2aeaebf0..f0cd4d3b3e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1303,6 +1303,16 @@ async def _await_disconnect_then_close(request, resp, cancel_event) -> None: return +async def _await_disconnect_then_cancel(request, cancel_event) -> None: + """Set ``cancel_event`` when a same-task local stream disconnects.""" + try: + while not await request.is_disconnected(): + await asyncio.sleep(0.1) + cancel_event.set() + except asyncio.CancelledError: + return + + # Centralized local/server tool nudge. Keep render_html guidance gated to turns # where the canvas tool is actually present in the tool schema; otherwise # small local models can hallucinate a missing tool call instead of following @@ -5095,6 +5105,9 @@ async def openai_chat_completions( async def gguf_tool_stream(): gen = None + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -5232,6 +5245,11 @@ async def openai_chat_completions( error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: + disconnect_watcher.cancel() + try: + await disconnect_watcher + except (asyncio.CancelledError, Exception): + pass if gen is not None: try: gen.close() @@ -5283,6 +5301,9 @@ async def openai_chat_completions( _tracker.__enter__() async def gguf_stream_chunks(): + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: # First chunk: role first_chunk = ChatCompletionChunk( @@ -5391,6 +5412,11 @@ async def openai_chat_completions( error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: + disconnect_watcher.cancel() + try: + await disconnect_watcher + except (asyncio.CancelledError, Exception): + pass _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( @@ -5710,6 +5736,9 @@ async def openai_chat_completions( async def sf_tool_stream(): gen = None + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -5834,6 +5863,11 @@ async def openai_chat_completions( } yield f"data: {json.dumps(error_chunk)}\n\n" finally: + disconnect_watcher.cancel() + try: + await disconnect_watcher + except (asyncio.CancelledError, Exception): + pass if gen is not None: try: gen.close() @@ -5952,6 +5986,9 @@ async def openai_chat_completions( _tracker.__enter__() async def stream_chunks(): + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -6060,6 +6097,11 @@ async def openai_chat_completions( } yield f"data: {json.dumps(error_chunk)}\n\n" finally: + disconnect_watcher.cancel() + try: + await disconnect_watcher + except (asyncio.CancelledError, Exception): + pass _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( @@ -7115,8 +7157,6 @@ async def _responses_non_streaming( # the model produced content, so clients expecting a pure tool-call turn # (finish_reason="tool_calls") don't see a spurious empty message item. output_items: list[dict] = [] - if reasoning_text and not text and not tool_calls: - text = reasoning_text if reasoning_text: output_items.append(_responses_reasoning_output_item(reasoning_text)) if text: @@ -7421,6 +7461,7 @@ async def _responses_stream( client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout()) resp = None lines_iter = None + disconnect_watcher = None disconnect_event = threading.Event() try: req = client.build_request( @@ -7482,6 +7523,9 @@ async def _responses_stream( return lines_iter = resp.aiter_lines() + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_close(request, resp, disconnect_event) + ) async for raw_line in _aiter_llama_stream_items( lines_iter, cancel_event = disconnect_event, @@ -7643,6 +7687,12 @@ async def _responses_stream( ) return finally: + if disconnect_watcher is not None: + disconnect_watcher.cancel() + try: + await disconnect_watcher + except (asyncio.CancelledError, Exception): + pass if lines_iter is not None: try: await lines_iter.aclose() @@ -9600,11 +9650,13 @@ async def _openai_passthrough_stream( monitor_done = False saw_finish_reason = False saw_done = False + saw_tool_call_delta = False last_chunk_id = completion_id last_chunk_model = model_name last_chunk_created = int(time.time()) def _synthetic_finish_line() -> str: + finish_reason = "tool_calls" if saw_tool_call_delta else "stop" chunk = ChatCompletionChunk( id = last_chunk_id, created = last_chunk_created, @@ -9612,7 +9664,7 @@ async def _openai_passthrough_stream( choices = [ ChunkChoice( delta = ChoiceDelta(), - finish_reason = "stop", + finish_reason = finish_reason, ) ], ) @@ -9643,9 +9695,21 @@ async def _openai_passthrough_stream( ) yield finish_line + "\n\n" saw_finish_reason = True + _monitor_openai_sse_line( + monitor_id, + raw_line, + llama_backend.context_length, + ) yield raw_line + "\n\n" monitor_done = True break + # Honor parallel_tool_calls=false (best-effort): drop tool_call + # deltas with index>=1 so only the first call streams. Only + # lines carrying tool_calls are reparsed; everything else is + # relayed byte-for-byte. + if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: + raw_line = _cap_parallel_tool_calls_sse_line(raw_line) + data_text = raw_line[6:].strip() try: chunk_data = json.loads(data_text) except json.JSONDecodeError: @@ -9660,14 +9724,12 @@ async def _openai_passthrough_stream( choices = chunk_data.get("choices") if isinstance(choices, list) and choices: choice = choices[0] - if isinstance(choice, dict) and choice.get("finish_reason"): - saw_finish_reason = True - # Honor parallel_tool_calls=false (best-effort): drop tool_call - # deltas with index>=1 so only the first call streams. Only - # lines carrying tool_calls are reparsed; everything else is - # relayed byte-for-byte. - if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: - raw_line = _cap_parallel_tool_calls_sse_line(raw_line) + if isinstance(choice, dict): + if choice.get("finish_reason"): + saw_finish_reason = True + delta = choice.get("delta") + if isinstance(delta, dict) and delta.get("tool_calls"): + saw_tool_call_delta = True monitor_event = _monitor_openai_sse_line( monitor_id, raw_line, @@ -9687,7 +9749,13 @@ async def _openai_passthrough_stream( llama_backend.context_length, ) yield finish_line + "\n\n" - yield "data: [DONE]\n\n" + done_line = "data: [DONE]" + _monitor_openai_sse_line( + monitor_id, + done_line, + llama_backend.context_length, + ) + yield done_line + "\n\n" monitor_done = True if not monitor_done: api_monitor.finish( diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 03f353a063..6d5d9badc1 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -2316,6 +2316,75 @@ class TestApiMonitorProviderAndCompletionStreams: assert '"finish_reason":"stop"' in body.replace(" ", "") assert "data: [DONE]" in body + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_synthesizes_tool_call_finish_reason(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield ( + 'data: {"id":"upstream","created":123,"model":"gguf",' + '"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,' + '"id":"call_1","type":"function","function":{"name":"lookup",' + '"arguments":"{}"}}]}}]}' + ) + yield "data: [DONE]" + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + body = "".join([chunk async for chunk in response.body_iterator]) + compact = body.replace(" ", "") + + assert '"finish_reason":"tool_calls"' in compact + assert '"finish_reason":"stop"' not in compact + assert "data: [DONE]" in body + assert monitor.active_count() == 0 asyncio.run(_run()) @@ -2434,7 +2503,10 @@ class TestApiMonitorProviderAndCompletionStreams: async for chunk in response.body_iterator: chunks.append(chunk) - assert chunks == ['data: {"choices":[{"delta":{"content":"hello"}}]}\n\n'] + assert chunks[0] == 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' + compact = "".join(chunks).replace(" ", "") + assert '"finish_reason":"stop"' in compact + assert chunks[-1] == "data: [DONE]\n\n" [entry] = monitor.snapshot() assert entry["status"] == "completed" assert entry["reply"] == "hello" diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 9a9cb024d6..4147746b54 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -992,7 +992,7 @@ class TestResponsesNonStreamingAdapter: assert [item["type"] for item in body["output"]] == ["message"] assert body["output"][0]["content"][0]["text"] == "33" - def test_reasoning_only_is_also_visible_message_text(self, monkeypatch): + def test_reasoning_only_stays_out_of_visible_message_text(self, monkeypatch): payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) body = self._run_with_message( monkeypatch, @@ -1000,9 +1000,8 @@ class TestResponsesNonStreamingAdapter: payload = payload, ) - assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert [item["type"] for item in body["output"]] == ["reasoning"] assert body["output"][0]["content"][0]["text"] == "plan" - assert body["output"][1]["content"][0]["text"] == "plan" # ===================================================================== diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 35a7258a65..9df54ba23c 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -184,6 +184,14 @@ def test_openai_passthrough_stream_avoids_starlette_task_group(): assert same_task_calls >= 2 +def test_local_chat_streams_install_same_task_disconnect_watcher(): + top = _async_function("openai_chat_completions") + assert _calls_name(top, "_await_disconnect_then_cancel"), ( + "Local same-task streams must watch request disconnects themselves; " + "do not restore Starlette's task-group StreamingResponse for this." + ) + + def test_direct_llama_server_streams_install_disconnect_watcher(): required = { "openai_completions",