diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index f58f3aeaba..ee2aeaebf0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4748,7 +4748,7 @@ async def openai_chat_completions( finally: _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( audio_input_stream(), media_type = "text/event-stream", headers = { @@ -5239,7 +5239,7 @@ async def openai_chat_completions( pass _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( gguf_tool_stream(), media_type = "text/event-stream", headers = { @@ -5393,7 +5393,7 @@ async def openai_chat_completions( finally: _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( gguf_stream_chunks(), media_type = "text/event-stream", headers = { @@ -5842,7 +5842,7 @@ async def openai_chat_completions( _sf_tracker.__exit__(None, None, None) if payload.stream: - return StreamingResponse( + return _SameTaskStreamingResponse( sf_tool_stream(), media_type = "text/event-stream", headers = { @@ -6062,7 +6062,7 @@ async def openai_chat_completions( finally: _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( stream_chunks(), media_type = "text/event-stream", headers = { @@ -9547,7 +9547,7 @@ async def _openai_passthrough_stream( except Exception: pass _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( iter(()), media_type = "text/event-stream", headers = { @@ -9598,6 +9598,26 @@ async def _openai_passthrough_stream( _await_disconnect_then_close(request, resp, cancel_event) ) monitor_done = False + saw_finish_reason = False + saw_done = False + last_chunk_id = completion_id + last_chunk_model = model_name + last_chunk_created = int(time.time()) + + def _synthetic_finish_line() -> str: + chunk = ChatCompletionChunk( + id = last_chunk_id, + created = last_chunk_created, + model = last_chunk_model, + choices = [ + ChunkChoice( + delta = ChoiceDelta(), + finish_reason = "stop", + ) + ], + ) + return f"data: {chunk.model_dump_json(exclude_none = True)}" + try: lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( @@ -9611,6 +9631,37 @@ async def _openai_passthrough_stream( continue if not raw_line.startswith("data: "): continue + data_text = raw_line[6:].strip() + if data_text == "[DONE]": + saw_done = True + if not saw_finish_reason and not cancel_event.is_set(): + finish_line = _synthetic_finish_line() + _monitor_openai_sse_line( + monitor_id, + finish_line, + llama_backend.context_length, + ) + yield finish_line + "\n\n" + saw_finish_reason = True + yield raw_line + "\n\n" + monitor_done = True + break + try: + chunk_data = json.loads(data_text) + except json.JSONDecodeError: + chunk_data = None + if isinstance(chunk_data, dict): + if isinstance(chunk_data.get("id"), str): + last_chunk_id = chunk_data["id"] + if isinstance(chunk_data.get("model"), str): + last_chunk_model = chunk_data["model"] + if isinstance(chunk_data.get("created"), int): + last_chunk_created = chunk_data["created"] + 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 @@ -9625,9 +9676,19 @@ async def _openai_passthrough_stream( # Relay verbatim to preserve llama-server's native id, # finish_reason, delta.tool_calls, and usage chunks. yield raw_line + "\n\n" - if monitor_event == "done" or raw_line[6:].strip() == "[DONE]": + if monitor_event == "done": monitor_done = True break + if not saw_done and not saw_finish_reason and not cancel_event.is_set(): + finish_line = _synthetic_finish_line() + _monitor_openai_sse_line( + monitor_id, + finish_line, + llama_backend.context_length, + ) + yield finish_line + "\n\n" + yield "data: [DONE]\n\n" + monitor_done = True if not monitor_done: api_monitor.finish( monitor_id, @@ -9680,7 +9741,7 @@ async def _openai_passthrough_stream( pass _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( _stream(), media_type = "text/event-stream", headers = { diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 6d2341b716..03f353a063 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -48,6 +48,7 @@ from routes.inference import ( _openai_passthrough_stream, _openai_stream_usage_chunk, _proxy_to_external_provider, + _SameTaskStreamingResponse, _set_or_prepend_system_message, openai_completions, openai_embeddings, @@ -2239,6 +2240,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ) + assert isinstance(response, _SameTaskStreamingResponse) iterator = response.body_iterator first = await anext(iterator) assert "hello" in first @@ -2256,6 +2258,67 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_synthesizes_missing_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":{"content":"hello"}}]}' + 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, + ) + chunks = [chunk async for chunk in response.body_iterator] + body = "".join(chunks) + + assert '"finish_reason":"stop"' in body.replace(" ", "") + assert "data: [DONE]" in body + + asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 2e6b5f14da..35a7258a65 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -146,24 +146,42 @@ def test_async_generators_cleanup_tracker_in_finally(): ) -def test_streaming_responses_have_no_background_task(): - top = None - for n in ast.walk(_TREE): - if isinstance(n, ast.AsyncFunctionDef) and n.name == "openai_chat_completions": - top = n - break - assert top is not None +def test_chat_completions_streams_avoid_starlette_task_group(): + top = _async_function("openai_chat_completions") + legacy_calls = [] + same_task_calls = 0 for sub in ast.walk(top): if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)): continue - if sub.func.id != "StreamingResponse": + if sub.func.id == "StreamingResponse": + legacy_calls.append(sub.lineno) + if sub.func.id == "_SameTaskStreamingResponse": + same_task_calls += 1 + assert not legacy_calls, ( + "Streaming /v1/chat/completions must use _SameTaskStreamingResponse, " + "not Starlette's legacy task-group StreamingResponse. Lines: " + f"{legacy_calls}" + ) + assert same_task_calls >= 5 + + +def test_openai_passthrough_stream_avoids_starlette_task_group(): + top = _async_function("_openai_passthrough_stream") + legacy_calls = [] + same_task_calls = 0 + for sub in ast.walk(top): + if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)): continue - kwargs = {kw.arg for kw in sub.keywords if kw.arg} - assert "background" not in kwargs, ( - "StreamingResponse in openai_chat_completions must not pass " - "`background=` -- cleanup now lives in the generator's finally " - "block; a BackgroundTask would be skipped on abrupt disconnect" - ) + if sub.func.id == "StreamingResponse": + legacy_calls.append(sub.lineno) + if sub.func.id == "_SameTaskStreamingResponse": + same_task_calls += 1 + assert not legacy_calls, ( + "OpenAI passthrough streams must use _SameTaskStreamingResponse, " + "not Starlette's legacy task-group StreamingResponse. Lines: " + f"{legacy_calls}" + ) + assert same_task_calls >= 2 def test_direct_llama_server_streams_install_disconnect_watcher():