From 6a9b77ee3795b57dff250b67a3fceab3e7affc1c Mon Sep 17 00:00:00 2001 From: Apoze <158856608+Apoze@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:09:08 +0200 Subject: [PATCH] Studio: harden OpenAI-compatible GGUF streaming (#6950) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 16 +- studio/backend/routes/inference.py | 956 +++++++-- .../test_inference_dispatcher_resilience.py | 2 +- .../tests/test_llama_cpp_stream_cancel.py | 81 + .../tests/test_llama_route_timeouts.py | 84 + .../tests/test_openai_tool_passthrough.py | 1875 ++++++++++++++++- .../backend/tests/test_passthrough_healing.py | 1 + 7 files changed, 2814 insertions(+), 201 deletions(-) create mode 100644 studio/backend/tests/test_llama_cpp_stream_cancel.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7d782fbbd..c2e3adf81 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -100,6 +100,10 @@ class LlamaServerNotFoundError(RuntimeError): Subclasses RuntimeError so existing handlers still catch it.""" +class _LlamaStreamCancelled(Exception): + """Internal signal for an expected client/request cancellation.""" + + # Shared so the from_identifier preflight and the load-time raise stay in sync. LLAMA_SERVER_NOT_FOUND_DETAIL = ( "This is a GGUF model, but the llama.cpp runtime (llama-server) is not " @@ -8616,7 +8620,7 @@ class LlamaCppBackend: ): """Open one streaming POST and let cancel interrupt prefill or reads.""" if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled _cancel_closed = threading.Event() _response_ref: list = [None] @@ -8661,13 +8665,13 @@ class LlamaCppBackend: ) as response: _response_ref[0] = response if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled yield response return except (httpx.RequestError, RuntimeError): # Response was closed by the cancel watcher if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled raise finally: _cancel_closed.set() @@ -8870,6 +8874,8 @@ class LlamaCppBackend: "finish_reason": _metadata_finish_reason, } + except _LlamaStreamCancelled: + return except httpx.ConnectError as e: # Server already down. If this was an MTP+tensor crash, recover by # reloading without MTP (scheduled in the background) and fail this @@ -9994,6 +10000,8 @@ class LlamaCppBackend: break continue + except _LlamaStreamCancelled: + return except httpx.ConnectError: # Mark unresolved provisional cards as failed before raising. for _pid, _pname in provisional_started_tool_calls.items(): @@ -10176,6 +10184,8 @@ class LlamaCppBackend: if _meta is not None: yield _meta + except _LlamaStreamCancelled: + return except httpx.ConnectError: raise RuntimeError("Lost connection to llama-server") except Exception as e: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1e39c6d12..826669d92 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -13,7 +13,7 @@ from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse, JSONResponse, Response from starlette.requests import ClientDisconnect -from typing import Any, List, Optional, Union +from typing import Any, Callable, List, Optional, Union import json import httpx from loggers import get_logger @@ -268,10 +268,87 @@ def _effective_max_tokens(payload): ) +_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT" + + +def _positive_float_env(env_name: str, default): + """Parse a positive float from an env var. A parseable non-positive value + returns ``None`` (0 disables the guarded feature); only unparseable or unset + values fall back to ``default``.""" + raw_value = os.environ.get(env_name) + if raw_value is None or not raw_value.strip(): + return default + try: + value = float(raw_value.strip()) + except ValueError: + return default + return value if value > 0 else None + + +def _effective_openai_max_tokens_from_values(max_tokens, max_completion_tokens = None): + """Resolve the OpenAI-compatible generation cap from raw request values. + + Prefers ``max_completion_tokens`` over the deprecated ``max_tokens``, and + returns ``None`` when both are omitted so callers keep their context-window + default (OpenAI treats an omitted cap as bounded only by the context + window). Explicit client caps pass through unchanged. + """ + + def _validate_explicit(value, param: str): + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + f"'{param}' must be an integer.", + status = 400, + code = "invalid_type", + param = param, + ), + ) + # The legacy completions spec declares ``minimum: 0`` for max_tokens, + # so 0 is a valid (if degenerate) cap and only negatives are rejected. + # The chat fields never reach here with 0 (pydantic enforces ge=1). + if value < 0: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + f"'{param}' must be at least 0.", + status = 400, + code = "invalid_value", + param = param, + ), + ) + return value + + max_tokens = _validate_explicit(max_tokens, "max_tokens") + max_completion_tokens = _validate_explicit(max_completion_tokens, "max_completion_tokens") + return max_completion_tokens if max_completion_tokens is not None else max_tokens + + +def _effective_openai_max_tokens(payload): + return _effective_openai_max_tokens_from_values( + getattr(payload, "max_tokens", None), + getattr(payload, "max_completion_tokens", None), + ) + + def _wants_multiple_choices(payload) -> bool: return (payload.n or 1) > 1 +def _has_openai_tool_history(messages) -> bool: + for message in messages or []: + if isinstance(message, dict): + if message.get("role") == "tool" or message.get("tool_calls"): + return True + continue + if getattr(message, "role", None) == "tool" or getattr(message, "tool_calls", None): + return True + return False + + def _raise_unsupported_openai_parameter(param: str, message: str) -> None: raise HTTPException( status_code = 400, @@ -331,6 +408,14 @@ def _openai_stream_error_chunk(exc) -> dict: return openai_error_body(_friendly_error(exc), status = 500) +def _openai_stream_error_sse(error: dict) -> str: + return f"data: {json.dumps(error)}\n\ndata: [DONE]\n\n" + + +def _openai_stream_error_sse_bytes(error: dict) -> bytes: + return _openai_stream_error_sse(error).encode("utf-8") + + def _openai_passthrough_error(status_code, text) -> "HTTPException": """HTTPException for a non-200 upstream response on the OpenAI passthrough (tools / response_format). An over-context upstream error is mapped to a 400 @@ -545,20 +630,62 @@ def _drop_parallel_tool_call_deltas(chunk) -> bool: return changed -def _cap_parallel_tool_calls_sse_line(raw_line: str) -> str: - """Drop tool_call deltas whose index >= 1 from one streamed OpenAI SSE - ``data:`` line so only the first tool call survives (parallel_tool_calls=false, - best-effort). Non-tool / unparseable payloads are returned byte-for-byte.""" - payload = raw_line[len("data: ") :] +def _add_empty_content_to_reasoning_deltas(chunk: dict) -> bool: + """Make reasoning-only deltas palatable to strict OpenAI adapters. + + Some clients built on OpenAI-compatible streams ignore or reject chunks whose + delta only contains non-standard ``reasoning_content``. Preserve that field, + but add an empty standard ``content`` member so the chunk is still a valid + text-delta shape and downstream parsers keep the stream alive. + """ + changed = False + choices = chunk.get("choices") + if not isinstance(choices, list): + return False + for choice in choices: + if not isinstance(choice, dict): + continue + delta = choice.get("delta") + if not isinstance(delta, dict): + continue + if "reasoning_content" in delta and "content" not in delta: + delta["content"] = "" + changed = True + return changed + + +def _normalize_openai_passthrough_sse_line( + raw_line: str, *, cap_parallel_tool_calls: bool = False +) -> str: + """Normalize one passthrough OpenAI SSE ``data:`` line before relaying. + + The function is intentionally narrow: it leaves comments, blank events, + ``[DONE]``, and unparseable upstream bytes untouched; parsed chunks are + re-serialized only when a compatibility mutation is actually required. + """ + if not raw_line.startswith("data:"): + return raw_line + # Both mutations key off JSON object keys, so a line without either quoted + # key can never change; skip the parse on the per-token common case. + if '"reasoning_content"' not in raw_line and not ( + cap_parallel_tool_calls and '"tool_calls"' in raw_line + ): + return raw_line + payload = raw_line[len("data:") :].lstrip() if payload.strip() in ("", "[DONE]"): return raw_line try: obj = json.loads(payload) except Exception: return raw_line - if not _drop_parallel_tool_call_deltas(obj): + if not isinstance(obj, dict): return raw_line - return "data: " + json.dumps(obj, separators = (",", ":")) + changed = _add_empty_content_to_reasoning_deltas(obj) + if cap_parallel_tool_calls and _drop_parallel_tool_call_deltas(obj): + changed = True + if not changed: + return raw_line + return "data: " + json.dumps(obj, separators = (",", ":"), ensure_ascii = False) def _prompt_tokens_details(upstream): @@ -575,6 +702,49 @@ def _wants_stream_usage(payload) -> bool: return bool((payload.stream_options or {}).get("include_usage")) +_OPENAI_PASSTHROUGH_TERMINAL_GRACE_S = 2.0 +_SSE_DONE_LINE = "data: [DONE]" + + +def _openai_passthrough_sse_line_terminal_state(raw_line: str) -> Optional[str]: + """Classify OpenAI-compatible chat stream terminal markers. + + Some llama-server builds can emit the logical final chunk (``finish_reason``) + and optional usage chunk, then keep the HTTP stream open without sending the + OpenAI ``data: [DONE]`` sentinel. Classifying those chunks lets Studio close + the client stream promptly while preserving an optional trailing usage chunk. + """ + if not raw_line.startswith("data:"): + return None + data_str = raw_line[5:].lstrip() + if data_str == "[DONE]": + return "done" + try: + data = json.loads(data_str) + except json.JSONDecodeError: + return None + return _openai_passthrough_terminal_state_from_data(data) + + +def _openai_passthrough_terminal_state_from_data(data) -> Optional[str]: + """Dict-level core of ``_openai_passthrough_sse_line_terminal_state`` for + callers that already parsed the chunk (avoids a re-parse per relayed line).""" + if not isinstance(data, dict): + return None + if _monitor_openai_error_message(data): + return "error" + choices = data.get("choices") + if isinstance(choices, list): + if not choices and isinstance(data.get("usage"), dict): + return "usage" + for choice in choices: + if isinstance(choice, dict) and choice.get("finish_reason") is not None: + return "finish" + elif isinstance(data.get("usage"), dict): + return "usage" + return None + + def _openai_stream_usage_chunk( payload, completion_id, created, model_name, stream_usage, stream_timings ): @@ -641,12 +811,16 @@ def _chat_content_chunk(completion_id, created, model_name, text) -> str: def _chat_reasoning_chunk(completion_id, created, model_name, text) -> str: - """Like ``_chat_content_chunk`` but on ``reasoning_content`` (renders the UI thinking block).""" + """Like ``_chat_content_chunk`` but on ``reasoning_content`` (renders the UI thinking block). + + Carries ``content: ""`` alongside, like the GGUF and passthrough paths, so + strict OpenAI adapters don't drop the reasoning-only delta. + """ return _chat_chunk_sse( completion_id, created, model_name, - delta = ChoiceDelta(reasoning_content = text), + delta = ChoiceDelta(content = "", reasoning_content = text), finish_reason = None, ) @@ -881,8 +1055,11 @@ def _llama_streaming_generation_timeout() -> httpx.Timeout: def _set_stream_response_read_timeout( - response: httpx.Response, read_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S + response: httpx.Response, read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S ) -> None: + # ``read_timeout_s = None`` clears httpx's read timeout (wait indefinitely), + # used when the stall guard is disabled so a stale first-token deadline + # can't keep timing out post-first-chunk gaps. try: timeout_ext = response.request.extensions.get("timeout") if isinstance(timeout_ext, dict): @@ -893,6 +1070,31 @@ def _set_stream_response_read_timeout( _STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1 +_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S = 5.0 +_OPENAI_PASSTHROUGH_SSE_KEEPALIVE = ": keep-alive\n\n" + + +def _openai_compat_stream_stall_timeout(): + """Max silent gap after an OpenAI passthrough stream has produced data. + + If the socket goes silent after valid SSE data, this bounds how long the + client is kept open. Defaults to the backend-wide stall timeout so this + path stalls out like every sibling stream; set the env var to tighten it + for local serving, or to 0 to disable the guard. + """ + return _positive_float_env( + _OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, + _DEFAULT_STREAM_STALL_TIMEOUT_S, + ) + + +def _openai_passthrough_upstream_headers(*, llama_backend = None) -> dict: + headers = {} + auth_headers = getattr(llama_backend, "_auth_headers", None) + if isinstance(auth_headers, dict): + headers.update(auth_headers) + headers["Connection"] = "close" + return headers class _CompatSameTaskTimeout: @@ -1020,7 +1222,8 @@ async def _aclose_stream_resources( """Tear down an httpx streaming generator's resources in the required order: cancel + await each watcher task, then aclose() the byte/line iterator, the response, and the client. Each step swallows its own exceptions so teardown - always completes. See _anthropic_passthrough_stream for the ordering rationale.""" + always completes; a close-time CancelledError is re-raised only after every + step has run. See _anthropic_passthrough_stream for the ordering rationale.""" for watcher in watchers: if watcher is not None: watcher.cancel() @@ -1028,21 +1231,30 @@ async def _aclose_stream_resources( await watcher except (asyncio.CancelledError, Exception): pass + close_cancelled = False if iterator is not None: try: await iterator.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass if resp is not None: try: await resp.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass if client is not None: try: await client.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass + if close_cancelled: + raise asyncio.CancelledError() async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: @@ -1065,6 +1277,7 @@ async def _send_stream_with_preheader_cancel( req: httpx.Request, cancel_event = None, request: Optional[Request] = None, + mark_cancel_on_cancel: bool = True, ) -> Optional[httpx.Response]: if cancel_event is None and request is None: return await client.send(req, stream = True) @@ -1096,7 +1309,7 @@ async def _send_stream_with_preheader_cancel( await _stop_send_task() return None except asyncio.CancelledError: - if cancel_event is not None: + if mark_cancel_on_cancel and cancel_event is not None: cancel_event.set() await _stop_send_task() raise @@ -1115,11 +1328,19 @@ async def _aiter_llama_stream_items( request: Optional[Request] = None, first_token_deadline: Optional[float] = None, response: Optional[httpx.Response] = None, - post_first_item_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, + post_first_item_read_timeout_s: Optional[ + Union[float, Callable[[], Optional[float]]] + ] = _DEFAULT_STREAM_STALL_TIMEOUT_S, ): if first_token_deadline is None: first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S last_item_at: Optional[float] = None + + def _post_first_timeout_s() -> Optional[float]: + if callable(post_first_item_read_timeout_s): + return post_first_item_read_timeout_s() + return post_first_item_read_timeout_s + while True: if cancel_event is not None and cancel_event.is_set(): return @@ -1140,15 +1361,14 @@ async def _aiter_llama_stream_items( async with _same_task_timeout(remaining_s): item = await async_iter.__anext__() else: + timeout_s = _post_first_timeout_s() if ( request is not None and response is not None - and post_first_item_read_timeout_s is not None + and timeout_s is not None and last_item_at is not None ): - stall_remaining_s = post_first_item_read_timeout_s - ( - time.monotonic() - last_item_at - ) + stall_remaining_s = timeout_s - (time.monotonic() - last_item_at) if stall_remaining_s <= 0: raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") _set_stream_response_read_timeout(response, stall_remaining_s) @@ -1165,19 +1385,16 @@ async def _aiter_llama_stream_items( if now >= first_token_deadline: raise continue - if ( - request is not None - and post_first_item_read_timeout_s is not None - and now - last_item_at < post_first_item_read_timeout_s - ): + timeout_s = _post_first_timeout_s() + if request is not None and timeout_s is not None and now - last_item_at < timeout_s: continue raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") - if ( - last_item_at is None - and response is not None - and post_first_item_read_timeout_s is not None - ): - _set_stream_response_read_timeout(response, post_first_item_read_timeout_s) + if last_item_at is None and response is not None: + # The first-token read deadline no longer applies once a chunk has + # arrived: switch to the stall timeout, or clear the read timeout + # entirely when the stall guard is disabled (callable returns None) + # so a long gap can't trip the stale first-token deadline. + _set_stream_response_read_timeout(response, _post_first_timeout_s()) last_item_at = time.monotonic() yield item @@ -1556,6 +1773,20 @@ def _effective_enable_tools(payload) -> Optional[bool]: return policy if policy is not None else payload.enable_tools +def _explicit_studio_tool_loop_requested(payload) -> bool: + """True when the request itself asks Studio to execute local tools. + + Process-wide CLI policy can default Studio's tool loop on for ordinary chat, + but it must not steal OpenAI-compatible client tools or response_format + requests from the llama-server passthrough path. A policy of ``False`` + (--disable-tools) vetoes even an explicit ``enable_tools: true`` ask. + """ + from state.tool_policy import get_tool_policy + + policy = get_tool_policy() + return policy is not False and (payload.enable_tools is True or bool(payload.mcp_enabled)) + + # Cancel registry. Proxies (e.g. Colab) can swallow client fetch aborts so # is_disconnected() never fires. POST /inference/cancel looks up in-flight # cancel_events here by cancel_id (per-run) or session_id / completion_id @@ -1696,6 +1927,39 @@ async def _await_disconnect_then_cancel(request, cancel_event) -> None: return +def _cancelable_nonstreaming_client() -> httpx.AsyncClient: + return httpx.AsyncClient( + limits = httpx.Limits(max_connections = 1, max_keepalive_connections = 0), + trust_env = False, + ) + + +async def _await_cancel_or_disconnect_then_close_client( + *, cancel_event, request: Optional[Request], client: httpx.AsyncClient +) -> None: + """Close a dedicated non-streaming upstream client on cancel/disconnect. + + The shared ``nonstreaming_client()`` is pooled, so cancelable generation calls + use a per-request client. Closing it interrupts a blocked llama-server + request without affecting unrelated pooled non-streaming calls. + """ + try: + while True: + if cancel_event is not None and cancel_event.is_set(): + break + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + break + await asyncio.sleep(0.1) + try: + await client.aclose() + except Exception: + pass + except asyncio.CancelledError: + return + + async def _stop_local_disconnect_cancel_watcher(watcher) -> None: watcher.cancel() try: @@ -6090,11 +6354,11 @@ async def openai_chat_completions( # unaware of `role="tool"` messages and assistant messages that only # carry `tool_calls` (content=None) — both of which are valid in # multi-turn client-side tool loops. - effective_max_tokens = _effective_max_tokens(payload) + effective_max_tokens = _effective_openai_max_tokens(payload) normalized_stop = _normalize_stop_sequences(payload.stop) - _has_tool_messages = any(m.role == "tool" or m.tool_calls for m in payload.messages) + _has_tool_messages = _has_openai_tool_history(payload.messages) # Route guided-decoding requests through the verbatim passthrough so # ``response_format`` (JSON schema) reaches llama-server and the model's # GBNF-constrained output comes back unmodified. The non-passthrough GGUF @@ -6103,13 +6367,43 @@ async def openai_chat_completions( # free-form sampling. Guided decoding does not require ``supports_tools`` -- # the grammar machinery is independent of tool-call parsing. _has_response_format = _extract_response_format(payload) is not None - _tools_passthrough = getattr( + _has_tool_catalog = bool(payload.tools and len(payload.tools) > 0) + _has_active_tool_catalog = _has_tool_catalog and payload.tool_choice != "none" + _has_client_tool_contract = _has_active_tool_catalog or _has_tool_messages + # The Studio tool loop needs a tool-capable backend, so a request that asks + # for it on a backend that can't run it (DiffusionGemma forces supports_tools + # off) must not steal client tools from the passthrough (#6851). + _studio_tool_loop_requested = ( + _explicit_studio_tool_loop_requested(payload) and llama_backend.supports_tools + ) + _client_disabled_tool_calls = payload.tool_choice == "none" and not _studio_tool_loop_requested + _supports_tool_passthrough = getattr( llama_backend, "supports_tool_passthrough", llama_backend.supports_tools - ) and ((payload.tools and len(payload.tools) > 0) or _has_tool_messages) - # DiffusionGemma keeps supports_tools off, so the server-side tool loop can't - # claim the request; fall through to client passthrough, matching /v1/messages. - _server_tool_loop = _effective_enable_tools(payload) and llama_backend.supports_tools - if using_gguf and not _server_tool_loop and (_tools_passthrough or _has_response_format): + ) + _tools_passthrough = _supports_tool_passthrough and _has_client_tool_contract + if ( + using_gguf + and not _studio_tool_loop_requested + and _has_client_tool_contract + and not _supports_tool_passthrough + ): + raise _reject( + 400, + openai_error_body( + ( + "Client-supplied tools or tool-call history require a GGUF chat template " + "with tool-call support; the current model/template does not advertise tools." + ), + status = 400, + code = "unsupported_parameter", + param = "tools" if payload.tools else "messages", + ), + ) + if ( + using_gguf + and not _studio_tool_loop_requested + and (_tools_passthrough or _has_response_format) + ): if _wants_multiple_choices(payload): raise _reject_unsupported_n("GGUF tool or response_format passthrough") if payload.audio_base64: @@ -6153,12 +6447,20 @@ async def openai_chat_completions( completion_id, monitor_id = monitor_id, ) - return await _openai_passthrough_non_streaming( - llama_backend, - payload, - model_name, - monitor_id = monitor_id, - ) + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + try: + return await _openai_passthrough_non_streaming( + llama_backend, + payload, + model_name, + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + finally: + _tracker.__exit__(None, None, None) # ── Parse messages (handles multimodal content parts) ───── # Reuse the pre-hook parse when auto-switch did it, else parse now. @@ -6218,6 +6520,8 @@ async def openai_chat_completions( ) def _gguf_chat_delta_line(delta: ChoiceDelta, finish_reason = None) -> str: + if delta.reasoning_content is not None and delta.content is None: + delta = delta.model_copy(update = {"content": ""}) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -6241,8 +6545,12 @@ async def openai_chat_completions( from state.tool_policy import get_tool_policy as _get_tool_policy_g _cli_policy = _get_tool_policy_g() - _tools_on = _effective_enable_tools(payload) - _mcp_allowed = bool(payload.mcp_enabled) and _cli_policy is not False + _tools_on = False if _client_disabled_tool_calls else _effective_enable_tools(payload) + _mcp_allowed = ( + not _client_disabled_tool_calls + and bool(payload.mcp_enabled) + and _cli_policy is not False + ) use_tools = (_tools_on or _mcp_allowed) and llama_backend.supports_tools if use_tools: @@ -6497,7 +6805,7 @@ async def openai_chat_completions( # Recover if an MTP+tensor crash killed the server mid-stream. get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: await _stop_local_disconnect_cancel_watcher(disconnect_watcher) if gen is not None: @@ -6664,6 +6972,9 @@ async def openai_chat_completions( disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) + gen = None + next_task = None + stream_completed = False try: yield _chat_role_chunk(completion_id, created, model_name) @@ -6682,7 +6993,14 @@ async def openai_chat_completions( cancel_event.set() api_monitor.finish(monitor_id, "cancelled") return - cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) + next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _gguf_sentinel) + ) + try: + cumulative = await asyncio.shield(next_task) + finally: + if next_task.done(): + next_task = None if cumulative is _gguf_sentinel: break # Capture server metadata for the final usage chunk @@ -6751,6 +7069,7 @@ async def openai_chat_completions( api_monitor.finish( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) + stream_completed = True yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -6761,10 +7080,39 @@ async def openai_chat_completions( logger.error(f"Error during GGUF streaming: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: - await _stop_local_disconnect_cancel_watcher(disconnect_watcher) - _tracker.__exit__(None, None, None) + try: + if not stream_completed: + cancel_event.set() + task_to_drain = next_task + next_task = None + while task_to_drain is not None and not task_to_drain.done(): + try: + await asyncio.shield(task_to_drain) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if task_to_drain is not None and task_to_drain.done(): + try: + task_to_drain.exception() + except (asyncio.CancelledError, Exception): + pass + if gen is not None and not stream_completed: + try: + await asyncio.to_thread(gen.close) + except (RuntimeError, ValueError): + pass + except Exception: + logger.debug( + "Error closing GGUF stream generator during cleanup", + exc_info = True, + ) + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) + finally: + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( gguf_stream_chunks(), @@ -6780,54 +7128,91 @@ async def openai_chat_completions( try: # ``n`` requests several independent completions; the single # decode slot yields one at a time, so loop sequentially. - _n = payload.n or 1 + drain_task = None - _choices = [] - _monitor_replies = [] - _prompt_tokens = 0 - _sum_completion = 0 - _prompt_details = None - for _idx in range(_n): - # Stop spawning the remaining choices once cancelled. - if cancel_event.is_set(): - break - full_text = "" - completion_usage = None - completion_finish = None - for token in gguf_generate(_idx): - if isinstance(token, dict): - if token.get("type") == "metadata": - completion_usage = token.get("usage") - completion_finish = token.get("finish_reason") + async def _drain_cancelled_gguf_task(): + if drain_task is None: + return + while not drain_task.done(): + try: + await asyncio.shield(drain_task) + except asyncio.CancelledError: + cancel_event.set() continue - full_text = token + except Exception: + break + if drain_task.done(): + try: + drain_task.exception() + except (asyncio.CancelledError, Exception): + pass - reasoning_text, visible_text = _extract_responses_reasoning( - full_text, - parse_think_markers = _responses_should_parse_think_markers( - payload, - llama_backend, - ), - ) - message_kwargs = {"content": visible_text} - if reasoning_text: - message_kwargs["reasoning_content"] = reasoning_text - _choices.append( - CompletionChoice( - index = _idx, - message = CompletionMessage(**message_kwargs), - finish_reason = _clamp_finish_reason(completion_finish), + def _drain_gguf_choices(): + _n = payload.n or 1 + _choices = [] + _monitor_replies = [] + _prompt_tokens = 0 + _sum_completion = 0 + _prompt_details = None + for _idx in range(_n): + # Stop spawning the remaining choices once cancelled. + if cancel_event.is_set(): + break + full_text = "" + completion_usage = None + completion_finish = None + for token in gguf_generate(_idx): + if isinstance(token, dict): + if token.get("type") == "metadata": + completion_usage = token.get("usage") + completion_finish = token.get("finish_reason") + continue + full_text = token + + reasoning_text, visible_text = _extract_responses_reasoning( + full_text, + parse_think_markers = _responses_should_parse_think_markers( + payload, + llama_backend, + ), ) + message_kwargs = {"content": visible_text} + if reasoning_text: + message_kwargs["reasoning_content"] = reasoning_text + _choices.append( + CompletionChoice( + index = _idx, + message = CompletionMessage(**message_kwargs), + finish_reason = _clamp_finish_reason(completion_finish), + ) + ) + _monitor_replies.append(visible_text) + if completion_usage: + # The prompt is shared across all n choices, so count its + # tokens ONCE (OpenAI bills only generated tokens for each + # extra choice). Only completion_tokens accumulates. + _prompt_tokens = completion_usage.get("prompt_tokens") or _prompt_tokens + _sum_completion += completion_usage.get("completion_tokens") or 0 + if _prompt_details is None: + _prompt_details = completion_usage.get("prompt_tokens_details") + return ( + _n, + _choices, + _monitor_replies, + _prompt_tokens, + _sum_completion, + _prompt_details, ) - _monitor_replies.append(visible_text) - if completion_usage: - # The prompt is shared across all n choices, so count its - # tokens ONCE (OpenAI bills only generated tokens for each - # extra choice). Only completion_tokens accumulates. - _prompt_tokens = completion_usage.get("prompt_tokens") or _prompt_tokens - _sum_completion += completion_usage.get("completion_tokens") or 0 - if _prompt_details is None: - _prompt_details = completion_usage.get("prompt_tokens_details") + + drain_task = asyncio.create_task(asyncio.to_thread(_drain_gguf_choices)) + ( + _n, + _choices, + _monitor_replies, + _prompt_tokens, + _sum_completion, + _prompt_details, + ) = await asyncio.shield(drain_task) response = ChatCompletion( id = completion_id, @@ -6859,6 +7244,11 @@ async def openai_chat_completions( api_monitor.finish(monitor_id) return _model_json_response(response) + except asyncio.CancelledError: + cancel_event.set() + await _drain_cancelled_gguf_task() + api_monitor.finish(monitor_id, "cancelled") + raise except Exception as e: logger.error(f"Error during GGUF completion: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) @@ -6878,7 +7268,6 @@ async def openai_chat_completions( ), ) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) - # ── Standard Unsloth path ───────────────────────────────── # Decode image (from content parts OR legacy field) @@ -7198,7 +7587,7 @@ async def openai_chat_completions( "type": "server_error", }, } - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: await _stop_local_disconnect_cancel_watcher(disconnect_watcher) if gen is not None: @@ -7543,7 +7932,7 @@ async def openai_chat_completions( "type": "server_error", }, } - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: await _stop_local_disconnect_cancel_watcher(disconnect_watcher) _tracker.__exit__(None, None, None) @@ -8041,8 +8430,12 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge if not isinstance(body, dict): raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") - if body.get("max_tokens") is None: - body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR + _resolved_max_tokens = _effective_openai_max_tokens_from_values(body.get("max_tokens")) + body["max_tokens"] = ( + _resolved_max_tokens + if _resolved_max_tokens is not None + else (llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR) + ) target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) prompt_text = _flatten_monitor_prompt(body.get("prompt", "")) @@ -8135,7 +8528,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge logger.error("openai_completions stream error: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") + yield _openai_stream_error_sse_bytes(error_chunk) return api_monitor.finish(monitor_id, "cancelled") return @@ -8150,7 +8543,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge logger.error("openai_completions stream error: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") + yield _openai_stream_error_sse_bytes(error_chunk) return finally: await _aclose_stream_resources( @@ -10897,13 +11290,15 @@ def _build_passthrough_payload( ): body = { "messages": openai_messages, - "tools": openai_tools, - "tool_choice": tool_choice, "temperature": temperature, "top_p": top_p, "top_k": top_k, "stream": stream, } + if openai_tools: + body["tools"] = openai_tools + if tool_choice is not None: + body["tool_choice"] = tool_choice if seed is not None: body["seed"] = seed if stream and stream_options is not None: @@ -11110,13 +11505,17 @@ async def _anthropic_passthrough_stream( yield event return finally: - await _aclose_stream_resources( - watchers = (cancel_watcher, disconnect_watcher), - iterator = lines_iter, - resp = resp, - client = client, - ) - _tracker.__exit__(None, None, None) + # Same shape as the OpenAI passthrough: a close-time CancelledError + # re-raised by _aclose_stream_resources must not skip the tracker exit. + try: + await _aclose_stream_resources( + watchers = (cancel_watcher, disconnect_watcher), + iterator = lines_iter, + resp = resp, + client = client, + ) + finally: + _tracker.__exit__(None, None, None) for line in emitter.finish(): yield line @@ -11613,6 +12012,9 @@ def _build_openai_passthrough_body( system_prompt, _, _ = _extract_content_parts(payload.messages) messages = _set_or_prepend_system_message(messages, system_prompt) tool_choice = payload.tool_choice if payload.tool_choice is not None else "auto" + tools = payload.tools + if payload.tool_choice == "none" and not _has_openai_tool_history(payload.messages): + tools = None # Forward per-request reasoning fields (enable_thinking / reasoning_effort / # preserve_thinking) via chat_template_kwargs so the Jinja template renders # in the caller's mode, gated on the active template's capabilities exactly @@ -11628,12 +12030,12 @@ def _build_openai_passthrough_body( ) return _build_passthrough_payload( messages, - payload.tools, + tools, payload.temperature, payload.top_p, payload.top_k, # Honor max_completion_tokens on the tools/response_format passthrough too. - _effective_max_tokens(payload), + _effective_openai_max_tokens(payload), payload.stream, stop = payload.stop, min_p = payload.min_p, @@ -11660,30 +12062,23 @@ async def _openai_passthrough_stream( """Streaming client-side pass-through for /v1/chat/completions. Forwards the client's OpenAI function-calling request to llama-server and - relays the SSE stream back verbatim, preserving llama-server's native - response ``id``, ``finish_reason`` (including ``"tool_calls"``), - ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so - the client sees a standard OpenAI response. + relays the SSE stream back with minimal normalization (reasoning-only + deltas gain ``content: ""``; errors and missing terminal markers get a + closing ``[DONE]``), preserving llama-server's native response ``id``, + ``finish_reason`` (including ``"tool_calls"``), ``delta.tool_calls``, and + any client-requested trailing ``usage`` chunk so the client sees a + standard OpenAI response. Reasoning/tool-call splitting is delegated to llama-server (``--jinja --reasoning-format auto``), so ``delta.content`` carries no raw markup and is deliberately not re-parsed locally, unlike the ``/completion`` paths. """ - target_url = f"{llama_backend.base_url}/v1/chat/completions" - body = _build_openai_passthrough_body( - payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend - ) - # Text-form tool calls from small models get promoted to structured calls on - # the way back (declared client tools only); requests without tools or with - # auto_heal_tool_calls=false keep the verbatim relay. tool_choice constrains - # the allowlist ("none" disables, a forced function narrows to it). - _allowed_tools = heal_gate( - payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") - ) - _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() + target_url = f"{llama_backend.base_url}/v1/chat/completions" + upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend) + client = None resp = None send_task: Optional[asyncio.Task[Optional[httpx.Response]]] = None @@ -11705,6 +12100,17 @@ async def _openai_passthrough_stream( # Keep tracker cleanup paired if pre-header dispatch is cancelled. try: + body = _build_openai_passthrough_body( + payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend + ) + # Text-form tool calls from small models get promoted to structured calls on + # the way back (declared client tools only); requests without tools or with + # auto_heal_tool_calls=false keep the unhealed relay. tool_choice constrains + # the allowlist ("none" disables, a forced function narrows to it). + _allowed_tools = heal_gate( + payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") + ) + # Keep the pre-header window short so accepted SSE clients receive # immediate headers in the common timeout-reduced stall. client = httpx.AsyncClient( @@ -11718,12 +12124,16 @@ async def _openai_passthrough_stream( while True: try: - req = client.build_request( - "POST", target_url, json = body, headers = {"Connection": "close"} - ) + req = client.build_request("POST", target_url, json = body, headers = upstream_headers) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S send_task = asyncio.create_task( - _send_stream_with_preheader_cancel(client, req, cancel_event, request = request) + _send_stream_with_preheader_cancel( + client, + req, + cancel_event, + request = request, + mark_cancel_on_cancel = False, + ) ) done, _ = await asyncio.wait( {send_task}, @@ -11749,13 +12159,14 @@ async def _openai_passthrough_stream( if resp is None and send_task is not None and not send_task.done(): break if resp is None: + if cancel_event is not None: + cancel_event.set() api_monitor.finish(monitor_id, "cancelled") - await _aclose_send_task(send_task) try: - await client.aclose() - except Exception: - pass - _tracker.__exit__(None, None, None) + await _aclose_send_task(send_task) + await _aclose_stream_resources(client = client) + finally: + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( iter(()), media_type = "text/event-stream", @@ -11780,6 +12191,7 @@ async def _openai_passthrough_stream( await resp.aclose() except Exception: pass + resp = None # Opt-in overflow policy: shrink and retry instead of a fatal 400. if ( _truncate_budget > 0 @@ -11808,11 +12220,14 @@ async def _openai_passthrough_stream( disconnect_watcher = None nonlocal resp, send_task, first_token_deadline, _truncate_budget + nonlocal client monitor_done = False saw_finish_reason = False saw_done = False saw_stream_error = False + saw_stream_item = False saw_tool_call_delta = False + terminal_seen = False last_chunk_id = completion_id last_chunk_model = model_name last_chunk_created = int(time.time()) @@ -11873,6 +12288,13 @@ async def _openai_passthrough_stream( lines.append("data: " + json.dumps(chunk, ensure_ascii = False)) return lines + stall_timeout_s = _openai_compat_stream_stall_timeout() + + def _terminal_read_timeout_s() -> Optional[float]: + if terminal_seen: + return _OPENAI_PASSTHROUGH_TERMINAL_GRACE_S + return stall_timeout_s + def _heal_transform(chunk_data: dict, raw_line: str) -> list: """SSE lines to emit in place of one upstream line (healing on).""" choices = chunk_data.get("choices") @@ -11941,24 +12363,47 @@ async def _openai_passthrough_stream( try: while True: - if send_task is not None and not send_task.done(): - try: - resp = await send_task - except httpx.RequestError as e: - logger.error("openai passthrough stream: upstream unreachable: %s", e) - api_monitor.fail(monitor_id, _friendly_error(e)) - yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" - return - send_task = None - elif send_task is not None: - try: - resp = send_task.result() - except httpx.RequestError as e: - logger.error("openai passthrough stream: upstream unreachable: %s", e) - api_monitor.fail(monitor_id, _friendly_error(e)) - yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" - return - send_task = None + if send_task is not None: + last_keepalive_at = time.monotonic() + while not send_task.done(): + # Wake often enough that _preheader_cancelled keeps + # cancel/disconnect latency sub-second during prefill; + # keepalives still pace off last_keepalive_at. + wait_timeout = min( + _STREAM_DISCONNECT_POLL_TIMEOUT_S, + _OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S, + ) + done, _ = await asyncio.wait( + {send_task}, + timeout = wait_timeout, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task in done: + break + if await _preheader_cancelled(cancel_event, request): + api_monitor.finish(monitor_id, "cancelled") + return + # The downstream SSE response is already committed; + # keep strict clients and proxies from treating a long + # llama-server prefill/header wait as a dead stream. + now = time.monotonic() + if ( + now - last_keepalive_at + >= _OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S + ): + last_keepalive_at = now + yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE + if resp is None: + try: + resp = send_task.result() + except httpx.RequestError as e: + logger.error( + "openai passthrough stream: upstream unreachable: %s", e + ) + api_monitor.fail(monitor_id, _friendly_error(e)) + yield _openai_stream_error_sse(_openai_stream_error_chunk(e)) + return + send_task = None if resp is None: api_monitor.finish(monitor_id, "cancelled") @@ -11986,12 +12431,16 @@ async def _openai_passthrough_stream( ): _truncate_budget -= 1 req = client.build_request( - "POST", target_url, json = body, headers = {"Connection": "close"} + "POST", target_url, json = body, headers = upstream_headers ) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S send_task = asyncio.create_task( _send_stream_with_preheader_cancel( - client, req, cancel_event, request = request + client, + req, + cancel_event, + request = request, + mark_cancel_on_cancel = False, ) ) continue @@ -12006,7 +12455,7 @@ async def _openai_passthrough_stream( ) ) api_monitor.fail(monitor_id, err_text[:500]) - yield f"data: {json.dumps(error_payload)}\n\n" + yield _openai_stream_error_sse(error_payload) return cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) @@ -12020,12 +12469,14 @@ async def _openai_passthrough_stream( request = request, first_token_deadline = first_token_deadline, response = resp, + post_first_item_read_timeout_s = _terminal_read_timeout_s, ): if not raw_line: continue - if not raw_line.startswith("data: "): + if not raw_line.startswith("data:"): continue - data_text = raw_line[6:].strip() + saw_stream_item = True + data_text = raw_line[5:].strip() if data_text == "[DONE]": saw_done = True # Upstream ended without a finish chunk: heal the residue @@ -12057,13 +12508,11 @@ async def _openai_passthrough_stream( 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() + raw_line = _normalize_openai_passthrough_sse_line( + raw_line, + cap_parallel_tool_calls = payload.parallel_tool_calls is False, + ) + data_text = raw_line[5:].strip() try: chunk_data = json.loads(data_text) except json.JSONDecodeError: @@ -12090,8 +12539,9 @@ async def _openai_passthrough_stream( if _monitor_openai_error_message(chunk_data): saw_stream_error = True # With healing active, a content-bearing line may be replaced by - # held/promoted chunks; otherwise the single upstream line - # relays verbatim (monitored exactly as emitted either way). + # held/promoted chunks; otherwise the single (already + # normalized) line relays unchanged (monitored exactly as + # emitted either way). if ( healer is not None and not healer.dormant @@ -12141,6 +12591,27 @@ async def _openai_passthrough_stream( yield out_line + "\n\n" if monitor_event == "done": monitor_done = True + break + terminal_state = ( + _openai_passthrough_terminal_state_from_data(chunk_data) + if out_line is raw_line + else _openai_passthrough_sse_line_terminal_state(out_line) + ) + if terminal_state == "usage" or ( + terminal_state == "finish" and not _wants_stream_usage(payload) + ): + done_line = _SSE_DONE_LINE + _monitor_openai_sse_line( + monitor_id, + done_line, + llama_backend.context_length, + ) + yield done_line + "\n\n" + saw_done = True + monitor_done = True + break + if terminal_state == "finish": + terminal_seen = True if monitor_done: break if not saw_done and not saw_stream_error and not cancel_event.is_set(): @@ -12162,7 +12633,7 @@ async def _openai_passthrough_stream( llama_backend.context_length, ) yield finish_line + "\n\n" - done_line = "data: [DONE]" + done_line = _SSE_DONE_LINE _monitor_openai_sse_line( monitor_id, done_line, @@ -12178,6 +12649,29 @@ async def _openai_passthrough_stream( except asyncio.CancelledError: api_monitor.finish(monitor_id, "cancelled") raise + except httpx.ReadTimeout as e: + if terminal_seen and not saw_stream_error and not cancel_event.is_set(): + done_line = _SSE_DONE_LINE + _monitor_openai_sse_line( + monitor_id, + done_line, + llama_backend.context_length, + ) + yield done_line + "\n\n" + api_monitor.finish(monitor_id) + return + if cancel_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") + return + logger.error( + "openai passthrough stream %s: %s", + "stalled mid-response" if saw_stream_item else "timeout", + e, + ) + api_monitor.fail(monitor_id, _friendly_error(e)) + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) + err = _openai_stream_error_chunk(e) + yield _openai_stream_error_sse(err) except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: # Watcher closed resp on cancel. Emit nothing extra; the client # initiated the cancel or already disconnected. @@ -12186,6 +12680,16 @@ async def _openai_passthrough_stream( get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) raise api_monitor.finish(monitor_id, "cancelled") + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error_payload = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error_payload) except Exception as e: if cancel_event.is_set(): api_monitor.finish(monitor_id, "cancelled") @@ -12195,25 +12699,31 @@ async def _openai_passthrough_stream( api_monitor.fail(monitor_id, _friendly_error(e)) get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) err = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(err)}\n\n" + yield _openai_stream_error_sse(err) finally: - await _aclose_send_task(send_task) - await _aclose_stream_resources( - watchers = (cancel_watcher, disconnect_watcher), - iterator = lines_iter, - resp = resp, - client = client, - ) - _tracker.__exit__(None, None, None) + # _aclose_stream_resources re-raises a close-time CancelledError + # only after finishing teardown, and the tracker exits either way. + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources( + watchers = (cancel_watcher, disconnect_watcher), + iterator = lines_iter, + resp = resp, + client = client, + ) + finally: + _tracker.__exit__(None, None, None) async def _unstarted_cleanup() -> None: # Client disconnected before the body stream started, so _stream()'s # finally never ran. Release the eagerly-opened upstream resp/client # and the cancel-registry entry here; the watchers and line iterator # are created inside _stream(), so there is nothing else to close. - await _aclose_send_task(send_task) - await _aclose_stream_resources(resp = resp, client = client) - _tracker.__exit__(None, None, None) + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) + finally: + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( _stream(), @@ -12225,10 +12735,19 @@ async def _openai_passthrough_stream( }, unstarted_cleanup = _unstarted_cleanup, ) - except BaseException: - await _aclose_send_task(send_task) - await _aclose_stream_resources(resp = resp, client = client) - _tracker.__exit__(None, None, None) + except BaseException as exc: + if isinstance(exc, asyncio.CancelledError): + if cancel_event is not None: + cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") + else: + detail = exc.detail if isinstance(exc, HTTPException) else _friendly_error(exc) + api_monitor.fail(monitor_id, str(detail)) + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) + finally: + _tracker.__exit__(None, None, None) raise @@ -12237,6 +12756,9 @@ async def _openai_passthrough_non_streaming( payload, model_name, monitor_id: Optional[str] = None, + *, + request: Optional[Request] = None, + cancel_event = None, ): """Non-streaming client-side pass-through for /v1/chat/completions. @@ -12245,20 +12767,67 @@ async def _openai_passthrough_non_streaming( ``tool_calls``, and accurate ``usage`` token counts. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" + upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend) body = _build_openai_passthrough_body( payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) + body["stream"] = False + body.pop("stream_options", None) _truncate_budget = ( _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 ) - while True: - try: - resp = await nonstreaming_client().post( + + async def _post(body_to_send): + if cancel_event is None and request is None: + return await nonstreaming_client().post( target_url, - json = body, + json = body_to_send, + headers = upstream_headers, timeout = _llama_non_streaming_generation_timeout(), ) + + if cancel_event is None: + cancel = threading.Event() + else: + cancel = cancel_event + client = _cancelable_nonstreaming_client() + watcher = asyncio.create_task( + _await_cancel_or_disconnect_then_close_client( + cancel_event = cancel, + request = request, + client = client, + ) + ) + try: + try: + response = await client.post( + target_url, + json = body_to_send, + headers = upstream_headers, + timeout = _llama_non_streaming_generation_timeout(), + ) + except httpx.RequestError: + if cancel.is_set(): + raise asyncio.CancelledError() + raise + if cancel.is_set(): + raise asyncio.CancelledError() + return response + finally: + watcher.cancel() + try: + await watcher + except (asyncio.CancelledError, Exception): + pass + try: + await client.aclose() + except Exception: + pass + + while True: + try: + resp = await _post(body) except asyncio.CancelledError: api_monitor.finish(monitor_id, "cancelled") raise @@ -12326,15 +12895,14 @@ async def _openai_passthrough_non_streaming( "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], } try: - retry_resp = await nonstreaming_client().post( - target_url, - json = retry_body, - timeout = _llama_non_streaming_generation_timeout(), - ) + retry_resp = await _post(retry_body) if retry_resp.status_code == 200: retry_data = retry_resp.json() if response_has_promotable_calls(retry_data, _allowed_tools, body.get("tools")): resp, data = retry_resp, retry_data + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise except (httpx.RequestError, ValueError) as exc: logger.warning("tool-call nudge retry failed; keeping original: %s", exc) diff --git a/studio/backend/tests/test_inference_dispatcher_resilience.py b/studio/backend/tests/test_inference_dispatcher_resilience.py index 60d750348..6184496d7 100644 --- a/studio/backend/tests/test_inference_dispatcher_resilience.py +++ b/studio/backend/tests/test_inference_dispatcher_resilience.py @@ -112,7 +112,7 @@ def test_route_llama_streaming_async_clients_disable_proxy_env(): continue calls.append(node) - assert len(calls) == 4 + assert len(calls) == 5 for call in calls: assert any( kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False diff --git a/studio/backend/tests/test_llama_cpp_stream_cancel.py b/studio/backend/tests/test_llama_cpp_stream_cancel.py new file mode 100644 index 000000000..1c73f4d17 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_stream_cancel.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import contextlib +import os +import sys +import threading + +import httpx +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference.llama_cpp import LlamaCppBackend, _LlamaStreamCancelled + + +def _backend_stub() -> LlamaCppBackend: + backend = LlamaCppBackend.__new__(LlamaCppBackend) + backend._process = object() + backend._healthy = True + backend._port = 48848 + backend._effective_context_length = 4096 + backend._supports_reasoning = False + backend._reasoning_always_on = False + backend._reasoning_style = "enable_thinking" + backend._supports_preserve_thinking = False + return backend + + +def test_stream_cancel_uses_internal_exception_not_generator_exit(): + class FakeResponse: + status_code = 200 + + def close(self): + pass + + class FakeStream: + def __enter__(self): + return FakeResponse() + + def __exit__(self, *_args): + return False + + class FakeClient: + def stream(self, *_args, **_kwargs): + return FakeStream() + + cancel_event = threading.Event() + + with pytest.raises(Exception) as exc_info: + with LlamaCppBackend._stream_with_retry( + FakeClient(), + "http://llama.test/v1/chat/completions", + {}, + cancel_event, + ): + cancel_event.set() + raise httpx.ReadError("client closed") + + assert exc_info.type is _LlamaStreamCancelled + assert not issubclass(exc_info.type, GeneratorExit) + + +def test_generate_chat_completion_swallows_internal_stream_cancel(monkeypatch): + backend = _backend_stub() + + @contextlib.contextmanager + def fake_open_stream(*_args, **_kwargs): + raise _LlamaStreamCancelled + + monkeypatch.setattr(backend, "_open_stream", fake_open_stream) + + chunks = list( + backend.generate_chat_completion( + [{"role": "user", "content": "hi"}], + cancel_event = threading.Event(), + ) + ) + + assert chunks == [] diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py index 24866bd03..b4666e3d1 100644 --- a/studio/backend/tests/test_llama_route_timeouts.py +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -203,3 +203,87 @@ def test_preheader_send_cleanup_on_disconnect_and_cancel(): asyncio.run(_run(False)) asyncio.run(_run(True)) + + +def test_stream_stall_timeout_callable_re_resolved_each_read(): + # The OpenAI passthrough passes a callable so the stall bound can switch to + # the short post-terminal grace mid-stream; it must be re-resolved per read, + # not captured once at generator start. + async def _run(): + response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) + values = iter([100.0, 2.0]) + seen = [] + + class _Request: + async def is_disconnected(self): + return False + + class _Items: + def __init__(self): + self.count = 0 + + async def __anext__(self): + self.count += 1 + if self.count > 3: + raise StopAsyncIteration + return "data: {}" + + async for _ in inf_mod._aiter_llama_stream_items( + _Items(), + cancel_event = threading.Event(), + request = _Request(), + response = response, + first_token_deadline = time.monotonic() + 1, + post_first_item_read_timeout_s = lambda: next(values, 5.0), + ): + seen.append(response.request.extensions["timeout"].get("read")) + + assert len(seen) == 3 + # The callable is resolved right after the first item (arming the + # post-first window) and again before each later read, consuming + # successive values. + assert seen[0] == 100.0 + assert 1.0 <= seen[1] <= 2.0 + assert 4.0 <= seen[2] <= 5.0 + + asyncio.run(_run()) + + +def test_stream_stall_timeout_disabled_clears_read_timeout(): + # UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT=0 disables the stall guard, so + # the callable returns None. Once a chunk has arrived the leftover + # first-token read timeout must be cleared, else a long post-first-chunk gap + # trips a stale deadline the operator asked to turn off. + async def _run(): + response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) + seen = [] + + class _Request: + async def is_disconnected(self): + return False + + class _Items: + def __init__(self): + self.count = 0 + + async def __anext__(self): + self.count += 1 + if self.count > 2: + raise StopAsyncIteration + return "data: {}" + + async for _ in inf_mod._aiter_llama_stream_items( + _Items(), + cancel_event = threading.Event(), + request = _Request(), + response = response, + first_token_deadline = time.monotonic() + 5, + post_first_item_read_timeout_s = lambda: None, + ): + seen.append(response.request.extensions["timeout"].get("read")) + + # The first-token path armed a finite read timeout; after the first chunk + # with the guard disabled, it is cleared to None on every subsequent read. + assert seen == [None, None], seen + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 05e017ba7..10cbd2aa0 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -8,6 +8,7 @@ import sys import asyncio import json import threading +import time from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") @@ -30,6 +31,7 @@ from core.inference.anthropic_compat import ( ) from core.inference.api_monitor import ApiMonitor from routes.inference import ( + _aclose_stream_resources, _build_chat_request, _build_openai_passthrough_body, _build_passthrough_payload, @@ -38,24 +40,58 @@ from routes.inference import ( _coalesce_consecutive_user_turns, _drop_empty_assistant_sentinels, _effective_max_tokens, + _effective_openai_max_tokens, + _effective_openai_max_tokens_from_values, _extract_content_parts, _friendly_error, _friendly_upstream_error, _merge_user_content, _monitor_openai_chunk, _monitor_openai_sse_event, + _normalize_openai_passthrough_sse_line, + _openai_compat_stream_stall_timeout, _openai_messages_for_gguf_chat, + _openai_passthrough_sse_line_terminal_state, + _openai_passthrough_upstream_headers, _openai_passthrough_non_streaming, _openai_passthrough_stream, + _openai_stream_error_sse, _openai_stream_usage_chunk, _proxy_to_external_provider, _SameTaskStreamingResponse, + _OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, _set_or_prepend_system_message, openai_completions, openai_embeddings, openai_chat_completions, ) -from state.tool_policy import reset_tool_policy +from state.tool_policy import reset_tool_policy, set_tool_policy + + +def test_aclose_stream_resources_attempts_remaining_closes_after_cancel(): + class Closeable: + def __init__(self, *, cancel = False): + self.cancel = cancel + self.closed = False + + async def aclose(self): + self.closed = True + if self.cancel: + raise asyncio.CancelledError() + + async def _run(): + iterator = Closeable(cancel = True) + resp = Closeable() + client = Closeable() + + with pytest.raises(asyncio.CancelledError): + await _aclose_stream_resources(iterator = iterator, resp = resp, client = client) + + assert iterator.closed + assert resp.closed + assert client.closed + + asyncio.run(_run()) class TestFriendlyUpstreamError: @@ -548,6 +584,263 @@ class TestChatCompletionRequestToolFields: assert "n > 1 is not supported" in entry["error"] assert monitor.active_count() == 0 + def test_client_tools_rejected_when_gguf_template_has_no_tool_support(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must not fall through to the standard GGUF path") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + self._assert_unsupported_param(resp, "tools") + assert "does not advertise tools" in resp.json()["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "does not advertise tools" in entry["error"] + assert monitor.active_count() == 0 + + def test_client_tools_use_passthrough_capability_when_tool_loop_is_disabled(self, monkeypatch): + import routes.inference as inference_route + + captured = {} + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + supports_tool_passthrough = True + is_vision = False + _is_audio = False + context_length = 4096 + base_url = "http://llama.passthrough-capability.test" + _request_reasoning_kwargs = lambda *_args, **_kwargs: None + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must use passthrough") + + def generate_chat_completion_with_tools(self, **_kwargs): + raise AssertionError("Studio tool loop must stay disabled") + + async def fake_passthrough(llama_backend, payload, model_name, **kwargs): + captured["body"] = inference_route._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + inference_route.api_monitor.finish(kwargs.get("monitor_id")) + return inference_route.JSONResponse({"ok": True, "model": model_name}) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + monkeypatch.setattr( + inference_route, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "use client tool"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + assert resp.status_code == 200 + assert resp.json()["ok"] is True + assert captured["body"]["tools"][0]["function"]["name"] == "lookup" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + def test_enable_tools_on_non_tool_backend_keeps_client_tools_on_passthrough(self, monkeypatch): + # DiffusionGemma forces supports_tools off while passthrough stays + # available (#6851): enable_tools=True must not steal client tools + # from the passthrough into a Studio tool loop that cannot run. + import routes.inference as inference_route + + captured = {} + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + supports_tool_passthrough = True + is_vision = False + _is_audio = False + context_length = 4096 + base_url = "http://llama.passthrough-capability.test" + _request_reasoning_kwargs = lambda *_args, **_kwargs: None + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must use passthrough") + + def generate_chat_completion_with_tools(self, **_kwargs): + raise AssertionError("Studio tool loop cannot run on a non-tool backend") + + async def fake_passthrough(llama_backend, payload, model_name, **kwargs): + captured["body"] = inference_route._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + inference_route.api_monitor.finish(kwargs.get("monitor_id")) + return inference_route.JSONResponse({"ok": True, "model": model_name}) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + monkeypatch.setattr( + inference_route, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "use client tool"}], + "enable_tools": True, + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + assert resp.status_code == 200 + assert resp.json()["ok"] is True + assert captured["body"]["tools"][0]["function"]["name"] == "lookup" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + def test_tool_choice_none_allows_tool_catalog_without_tool_template(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **kwargs): + assert kwargs["max_tokens"] is None + yield "plain response" + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + "tool_choice": "none", + }, + ) + + assert resp.status_code == 200 + assert resp.json()["choices"][0]["message"]["content"] == "plain response" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "plain response" + assert monitor.active_count() == 0 + + def test_tool_call_history_rejected_when_gguf_template_has_no_tool_support(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **_kwargs): + raise AssertionError( + "tool-call history must not fall through to the standard GGUF path" + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [ + {"role": "user", "content": "use a tool"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "{}"}, + ], + }, + ) + + self._assert_unsupported_param(resp, "messages") + assert "does not advertise tools" in resp.json()["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "does not advertise tools" in entry["error"] + assert monitor.active_count() == 0 + def test_n_rejected_for_non_gguf_path(self, monkeypatch): class _NoGGUFBackend: is_loaded = False @@ -748,11 +1041,32 @@ class TestBuildPassthroughPayloadToolChoice: ) assert body.get("stream_options") == {"include_usage": False} + def test_response_format_without_tools_omits_tool_fields(self): + args = self._args() + args["openai_tools"] = None + + body = _build_passthrough_payload( + **args, + response_format = {"type": "json_object"}, + ) + + assert body["response_format"] == {"type": "json_object"} + assert "tools" not in body + assert "tool_choice" not in body + def test_repetition_penalty_renamed(self): body = _build_passthrough_payload(**self._args(), repetition_penalty = 1.1) assert body.get("repeat_penalty") == 1.1 assert "repetition_penalty" not in body + def test_omitted_passthrough_max_tokens_uses_backend_context(self): + args = self._args() + args["max_tokens"] = None + + body = _build_passthrough_payload(**args, backend_ctx = 4096) + + assert body["max_tokens"] == 4096 + def test_passthrough_body_merges_system_and_developer_messages(self): payload = ChatCompletionRequest( model = "default", @@ -772,6 +1086,74 @@ class TestBuildPassthroughPayloadToolChoice: ] +class TestOpenAIPassthroughSSETerminalState: + def test_done_sentinel(self): + assert _openai_passthrough_sse_line_terminal_state("data: [DONE]") == "done" + + def test_finish_reason_with_space(self): + line = 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + assert _openai_passthrough_sse_line_terminal_state(line) == "finish" + + def test_finish_reason_without_space(self): + line = 'data:{"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}' + assert _openai_passthrough_sse_line_terminal_state(line) == "finish" + + def test_usage_chunk(self): + line = 'data: {"choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2}}' + assert _openai_passthrough_sse_line_terminal_state(line) == "usage" + + def test_error_chunk(self): + line = 'data: {"error":{"message":"boom"}}' + assert _openai_passthrough_sse_line_terminal_state(line) == "error" + + def test_cap_parallel_tool_calls_accepts_no_space_after_data_colon(self): + line = ( + 'data:{"choices":[{"delta":{"tool_calls":[' + '{"index":0,"function":{"name":"a"}},' + '{"index":1,"function":{"name":"b"}}]}}]}' + ) + + capped = _normalize_openai_passthrough_sse_line(line, cap_parallel_tool_calls = True) + + data = json.loads(capped[len("data:") :].lstrip()) + assert data["choices"][0]["delta"]["tool_calls"] == [ + {"index": 0, "function": {"name": "a"}} + ] + + def test_plain_content_line_is_returned_identically(self): + # The relay dispatches terminal classification on `out_line is raw_line`, + # so the no-mutation path must return the identical string object. + line = 'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":null}]}' + assert _normalize_openai_passthrough_sse_line(line) is line + assert _normalize_openai_passthrough_sse_line(line, cap_parallel_tool_calls = True) is line + + def test_reasoning_key_inside_content_text_keeps_line_identical(self): + # Fast-path substring gate fires, but the parse finds nothing to change: + # the original object must come back so the relay stays byte-identical. + line = ( + 'data: {"choices":[{"index":0,"delta":{"content":' + '"mentions \\"reasoning_content\\" in text"},"finish_reason":null}]}' + ) + assert _normalize_openai_passthrough_sse_line(line) is line + + def test_reasoning_only_delta_gets_empty_content(self): + line = ( + 'data: {"choices":[{"index":0,' + '"delta":{"reasoning_content":"thinking"},' + '"finish_reason":null}]}' + ) + + normalized = _normalize_openai_passthrough_sse_line(line) + + data = json.loads(normalized[len("data:") :].lstrip()) + delta = data["choices"][0]["delta"] + assert delta["reasoning_content"] == "thinking" + assert delta["content"] == "" + + def test_reasoning_normalization_preserves_done_sentinel(self): + assert _normalize_openai_passthrough_sse_line("data: [DONE]") == "data: [DONE]" + + # ===================================================================== # Passthrough reasoning kwargs — enable_thinking / reasoning_effort / # preserve_thinking must reach llama-server via chat_template_kwargs, @@ -892,6 +1274,98 @@ class TestOpenAICompatibilityHelpers: payload = SimpleNamespace(max_tokens = 128, max_completion_tokens = 64) assert _effective_max_tokens(payload) == 64 + def test_openai_compat_max_tokens_returns_none_when_omitted(self): + payload = SimpleNamespace(max_tokens = None, max_completion_tokens = None) + assert _effective_openai_max_tokens(payload) is None + + @pytest.mark.parametrize( + ("payload", "expected"), + [ + (SimpleNamespace(max_tokens = 8192, max_completion_tokens = None), 8192), + (SimpleNamespace(max_tokens = 8192, max_completion_tokens = 256), 256), + ], + ) + def test_openai_compat_explicit_values_pass_through(self, payload, expected): + assert _effective_openai_max_tokens(payload) == expected + + @pytest.mark.parametrize( + ("payload", "param"), + [ + (SimpleNamespace(max_tokens = "128", max_completion_tokens = None), "max_tokens"), + (SimpleNamespace(max_tokens = True, max_completion_tokens = None), "max_tokens"), + (SimpleNamespace(max_tokens = 12.5, max_completion_tokens = None), "max_tokens"), + ( + SimpleNamespace(max_tokens = None, max_completion_tokens = "128"), + "max_completion_tokens", + ), + ], + ) + def test_openai_compat_max_tokens_rejects_non_integer_explicit_values(self, payload, param): + with pytest.raises(HTTPException) as exc: + _effective_openai_max_tokens(payload) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["param"] == param + assert exc.value.detail["error"]["code"] == "invalid_type" + + def test_openai_compat_max_tokens_zero_is_valid_and_negative_rejected(self): + # Legacy completions spec: max_tokens has minimum 0, so 0 must pass + # through; only negatives are invalid_value. + assert _effective_openai_max_tokens_from_values(0) == 0 + + with pytest.raises(HTTPException) as exc: + _effective_openai_max_tokens_from_values(-1) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["code"] == "invalid_value" + assert exc.value.detail["error"]["param"] == "max_tokens" + + def test_chat_reasoning_chunk_carries_empty_content(self): + from routes.inference import _chat_reasoning_chunk + + line = _chat_reasoning_chunk("chatcmpl-test", 123, "gguf", "thinking...") + chunk = json.loads(line[len("data: ") :]) + delta = chunk["choices"][0]["delta"] + + assert delta["reasoning_content"] == "thinking..." + assert delta["content"] == "" + + def test_passthrough_upstream_headers_include_backend_auth(self): + headers = _openai_passthrough_upstream_headers( + llama_backend = SimpleNamespace(_auth_headers = {"Authorization": "Bearer secret"}), + ) + + assert headers["Authorization"] == "Bearer secret" + assert headers["Connection"] == "close" + + def test_openai_compat_stream_stall_timeout_uses_default(self, monkeypatch): + monkeypatch.delenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raising = False) + assert _openai_compat_stream_stall_timeout() == 120.0 + + def test_openai_compat_stream_stall_timeout_uses_env_override(self, monkeypatch): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, "4.5") + assert _openai_compat_stream_stall_timeout() == 4.5 + + @pytest.mark.parametrize("raw_value", ["", "not-a-float"]) + def test_openai_compat_stream_stall_timeout_invalid_env_uses_default( + self, monkeypatch, raw_value + ): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raw_value) + assert _openai_compat_stream_stall_timeout() == 120.0 + + @pytest.mark.parametrize("raw_value", ["0", "-1"]) + def test_openai_compat_stream_stall_timeout_non_positive_env_disables( + self, monkeypatch, raw_value + ): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raw_value) + assert _openai_compat_stream_stall_timeout() is None + + def test_openai_stream_error_sse_closes_with_done(self): + error = {"error": {"message": "boom"}} + assert _openai_stream_error_sse(error) == ( + 'data: {"error": {"message": "boom"}}\n\n' "data: [DONE]\n\n" + ) + @pytest.mark.parametrize( "finish_reason", ["stop", "length", "tool_calls", "content_filter", "function_call"], @@ -1515,9 +1989,330 @@ class TestGgufVisionToolRouting: assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" assert "".join(d.get("content", "") for d in deltas) == "visible" assert all("" not in d.get("content", "") for d in deltas) + assert all("content" in d for d in deltas if "reasoning_content" in d) [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" + def test_global_enable_tools_does_not_preempt_response_format_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + captured = {} + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("Studio tool loop should not steal response_format") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "json"}], + response_format = {"type": "json_object"}, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["response_format"] == {"type": "json_object"} + assert "tools" not in captured["body"] + assert "tool_choice" not in captured["body"] + finally: + reset_tool_policy() + + def test_global_enable_tools_does_not_replace_client_tools_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + captured = {} + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("Studio tool loop should not replace client tools") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "use client tool"}], + tools = client_tools, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["tools"] == client_tools + assert captured["body"]["tool_choice"] == "auto" + finally: + reset_tool_policy() + + def test_global_enable_tools_honors_client_tool_choice_none(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**kwargs): + assert kwargs["max_tokens"] is None + yield "plain response" + + def _tools(**_kwargs): + raise AssertionError("tool_choice='none' must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "do not use tools"}], + tools = client_tools, + tool_choice = "none", + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["choices"][0]["message"]["content"] == "plain response" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "plain response" + assert monitor.active_count() == 0 + finally: + reset_tool_policy() + + def test_enabled_tools_without_enable_tools_keeps_response_format_passthrough( + self, monkeypatch + ): + import routes.inference as inf_mod + + reset_tool_policy() + captured = {} + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("enabled_tools alone must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.enabled-tools.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_openai_passthrough_non_streaming", fake_passthrough) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "json"}], + enabled_tools = ["web_search"], + response_format = {"type": "json_object"}, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["response_format"] == {"type": "json_object"} + + def test_enabled_tools_without_enable_tools_keeps_client_tools_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + captured = {} + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("enabled_tools alone must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.enabled-tools.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_openai_passthrough_non_streaming", fake_passthrough) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "use client tool"}], + enabled_tools = ["web_search"], + tools = client_tools, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["tools"] == client_tools + assert captured["body"]["tool_choice"] == "auto" + def test_reasoning_capable_gguf_stream_splits_reasoning_by_default(self, monkeypatch): def _generate(**_kwargs): yield "planvisible" @@ -1691,6 +2486,58 @@ class TestGgufVisionToolRouting: assert entry["completion_tokens"] == 3 assert monitor.active_count() == 0 + def test_non_streaming_gguf_cancel_drains_worker(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + started = threading.Event() + released = threading.Event() + + def _generate(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert await asyncio.to_thread(started.wait, 1.0) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch): import routes.inference as inf_mod @@ -1801,7 +2648,12 @@ class TestApiMonitorProviderAndCompletionStreams: async def is_disconnected(self): return False - async def _run_passthrough_stream(self, monkeypatch, lines): + async def _run_passthrough_stream( + self, + monkeypatch, + lines, + stream_options = None, + ): import routes.inference as inf_mod class Request: @@ -1829,6 +2681,7 @@ class TestApiMonitorProviderAndCompletionStreams: model = "default", messages = [ChatMessage(role = "user", content = "hi")], stream = True, + stream_options = stream_options, tools = [ { "type": "function", @@ -1913,6 +2766,134 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_forwards_backend_auth_headers(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured_headers = {} + + async def fake_send(_client, req, *_args, **_kwargs): + captured_headers.update(dict(req.headers)) + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + 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, + _auth_headers = {"Authorization": "Bearer secret"}, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + + assert "data: [DONE]\n\n" in "".join(chunks) + assert captured_headers["authorization"] == "Bearer secret" + assert captured_headers["connection"] == "close" + + asyncio.run(_run()) + + def test_passthrough_stream_keepalive_while_upstream_headers_are_pending(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, + "_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S", + 0.01, + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + + first = await asyncio.wait_for(response.body_iterator.__anext__(), timeout = 0.2) + assert first == ": keep-alive\n\n" + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data: [DONE]\n\n" in body + + asyncio.run(_run()) + def test_passthrough_stream_preheader_non_200_in_window(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -2055,6 +3036,7 @@ class TestApiMonitorProviderAndCompletionStreams: body = "".join(chunks) assert "data:" in body assert '"error"' in body + assert "data: [DONE]" in body [entry] = monitor.snapshot() assert entry["status"] == "error" assert "bad" in entry["error"] @@ -2117,7 +3099,13 @@ class TestApiMonitorProviderAndCompletionStreams: async for chunk in response.body_iterator ] body = "".join(chunks) - payload = json.loads(body.removeprefix("data: ").strip()) + events = [ + line.removeprefix("data: ") + for line in body.splitlines() + if line.startswith("data: ") + ] + assert events[-1] == "[DONE]" + payload = json.loads(events[0]) assert payload["error"]["code"] == "context_length_exceeded" assert payload["error"]["param"] == "messages" assert isinstance(payload["error"], dict) @@ -2207,6 +3195,105 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_preheader_immediate_context_retry_adopts_delayed_response( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + calls = [] + err_body = json.dumps( + { + "error": { + "message": "request (10000 tokens) exceeds the available context size (2048 tokens)", + "n_prompt_tokens": 10000, + "n_ctx": 2048, + } + } + ).encode("utf-8") + ok_lines = [ + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,' + '"model":"gguf","choices":[{"index":0,"delta":{"content":"OK"},' + '"finish_reason":null}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,' + '"model":"gguf","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + + async def fake_send(_client, req, *_args, **_kwargs): + calls.append(json.loads(req.content.decode("utf-8"))) + if len(calls) == 1: + return httpx.Response(400, content = err_body) + await gate.wait() + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + for line in ok_lines: + yield line + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + 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) + + messages = [ + ChatMessage(role = "system", content = "system"), + *[ + ChatMessage(role = "user", content = f"turn {idx} " + ("x" * 1000)) + for idx in range(8) + ], + ] + payload = ChatCompletionRequest( + model = "default", + messages = messages, + stream = True, + context_overflow = "truncate_middle", + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + + assert "OK" in body + assert "context_length_exceeded" not in body + assert len(calls) == 2 + assert len(calls[1]["messages"]) < len(calls[0]["messages"]) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + def test_passthrough_stream_preheader_delayed_request_error_cleans_up(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -2647,6 +3734,150 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_completions_omitted_max_tokens_falls_back_to_context(self, monkeypatch): + # With no env knobs set, an omitted max_tokens must forward the + # backend's context length, exactly as on main. + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False} + + captured = [] + + class CapturingClient: + async def post(self, _url, *, json, **_kwargs): + captured.append(dict(json)) + return httpx.Response( + 200, + json = { + "id": "cmpl-test", + "choices": [{"text": "ok"}], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + await openai_completions(Request(), current_subject = "test") + + assert captured[0]["max_tokens"] == 4096 + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_forwards_spec_valid_zero_max_tokens(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False, "max_tokens": 0} + + captured = [] + + class CapturingClient: + async def post(self, _url, *, json, **_kwargs): + captured.append(dict(json)) + return httpx.Response( + 200, + json = { + "id": "cmpl-test", + "choices": [{"text": "", "finish_reason": "length"}], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 0, + "total_tokens": 1, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + await openai_completions(Request(), current_subject = "test") + + assert captured[0]["max_tokens"] == 0 + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_rejects_non_integer_max_tokens_before_forwarding(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False, "max_tokens": "128"} + + class UnusedClient: + async def post(self, *_args, **_kwargs): + raise AssertionError("invalid max_tokens must not reach llama-server") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: UnusedClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + with pytest.raises(HTTPException) as exc: + await openai_completions(Request(), current_subject = "test") + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["param"] == "max_tokens" + assert exc.value.detail["error"]["code"] == "invalid_type" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_monitor_openai_chunk_records_all_choice_replies(self, monkeypatch): import routes.inference as inf_mod @@ -2793,6 +4024,8 @@ class TestApiMonitorProviderAndCompletionStreams: yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' await asyncio.sleep(3600) + cancel_id = "passthrough-stream-delete-cancel" + monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) @@ -2807,6 +4040,7 @@ class TestApiMonitorProviderAndCompletionStreams: model = "default", messages = [ChatMessage(role = "user", content = "hi")], stream = True, + cancel_id = cancel_id, tools = [ { "type": "function", @@ -2824,6 +4058,7 @@ class TestApiMonitorProviderAndCompletionStreams: SimpleNamespace( base_url = "http://llama.test", context_length = 4096, + _auth_headers = {"Authorization": "Bearer secret"}, _request_reasoning_kwargs = lambda *_args, **_kwargs: None, ), payload, @@ -2835,6 +4070,7 @@ class TestApiMonitorProviderAndCompletionStreams: iterator = response.body_iterator first = await anext(iterator) assert "hello" in first + assert cancel_id in inf_mod._CANCEL_REGISTRY pending = asyncio.create_task(anext(iterator)) await asyncio.sleep(0) @@ -2846,6 +4082,7 @@ class TestApiMonitorProviderAndCompletionStreams: assert entry["status"] == "cancelled" assert entry["reply"] == "hello" assert monitor.active_count() == 0 + assert cancel_id not in inf_mod._CANCEL_REGISTRY asyncio.run(_run()) @@ -2931,6 +4168,27 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_usage_done_are_separate_sse_events(self, monkeypatch): + async def _run(): + result = await self._run_passthrough_stream( + monkeypatch, + [ + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}', + ], + stream_options = {"include_usage": True}, + ) + + assert ( + '"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2' in result.body + ) + assert "data: [DONE]" in result.body + assert "}\n\ndata: [DONE]\n\n" in result.body + assert "}\ndata: [DONE]\n\n" not in result.body + + asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -2990,6 +4248,407 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_closes_blocked_upstream_post(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + async def is_disconnected(self): + return False + + client = HangingCancelableClient() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_cancelable_nonstreaming_client", + lambda: client, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + cancel_event = threading.Event() + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + request = Request(), + cancel_event = cancel_event, + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + cancel_event.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_route_registers_cancel_id(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + + async def is_disconnected(self): + return False + + cancel_id = "passthrough-nonstream-cancel-id" + client = HangingCancelableClient() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client) + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + ), + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + cancel_id = cancel_id, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + assert cancel_id in inf_mod._CANCEL_REGISTRY + assert inf_mod._cancel_by_cancel_id_or_stash(cancel_id) == 1 + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + assert cancel_id not in inf_mod._CANCEL_REGISTRY + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_disconnect_closes_blocked_upstream_post(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + def __init__(self): + self.disconnected = False + + async def is_disconnected(self): + return self.disconnected + + client = HangingCancelableClient() + request = Request() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_cancelable_nonstreaming_client", + lambda: client, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + cancel_event = threading.Event() + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + request.disconnected = True + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + assert cancel_event.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_forwards_backend_auth_headers(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured = {} + + class FakeNonStreamingClient: + async def post(self, *_args, **kwargs): + captured["headers"] = kwargs.get("headers") + return httpx.Response( + 200, + json = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 123, + "model": "gguf", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "nonstreaming_client", + lambda: FakeNonStreamingClient(), + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _auth_headers = {"Authorization": "Bearer secret"}, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + ) + + assert json.loads(response.body)["choices"][0]["message"]["content"] == "OK" + assert captured["headers"]["Authorization"] == "Bearer secret" + assert captured["headers"]["Connection"] == "close" + + asyncio.run(_run()) + + def test_passthrough_non_streaming_forces_upstream_stream_false(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured = {} + + class FakeNonStreamingClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def post(self, *_args, **kwargs): + captured["json"] = kwargs.get("json") + return httpx.Response( + 200, + json = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 123, + "model": "gguf", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "nonstreaming_client", + lambda: FakeNonStreamingClient(), + ) + 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, + stream_options = {"include_usage": True}, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + ) + + assert captured["json"]["stream"] is False + assert "stream_options" not in captured["json"] + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + def test_passthrough_clean_eof_finalizes_monitor(self, monkeypatch): async def _run(): result = await self._run_passthrough_stream( @@ -3009,6 +4668,216 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_finish_without_done_closes_stream_early(self, monkeypatch): + # Some llama-server builds emit the finish chunk and then hold the HTTP + # stream open without sending [DONE]; the terminal classifier must end + # the client stream promptly instead of hanging on the open socket. + 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: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}' + yield 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + await asyncio.Event().wait() # upstream never closes + + 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, + ) + + async def _consume(): + return [chunk async for chunk in response.body_iterator] + + chunks = await asyncio.wait_for(_consume(), timeout = 2) + body = "".join(chunks) + + assert '"finish_reason":"stop"' in body.replace(" ", "") + assert body.endswith("data: [DONE]\n\n") + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stall_after_finish_closes_cleanly(self, monkeypatch): + # include_usage keeps the stream open past the finish chunk waiting for + # the usage chunk; if that never arrives, the post-terminal grace path + # must close with a clean [DONE], not an in-band error. + 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: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}' + yield 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + raise httpx.ReadTimeout("usage chunk never arrived") + + 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, + stream_options = {"include_usage": 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 '"type":"api_error"' not in body.replace(" ", "") + assert body.endswith("data: [DONE]\n\n") + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_stall_after_data_emits_error(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: {"choices":[{"delta":{"content":"hello"}}]}' + raise httpx.ReadTimeout("upstream went silent") + + 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 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' in body + assert '"finish_reason"' not in body.replace(" ", "") + assert '"type":"api_error"' in body.replace(" ", "") + assert "still processing the prompt" in body + assert body.endswith("data: [DONE]\n\n") + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "still processing the prompt" in entry["error"] + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + class TestApiMonitorSafetensorsUsage: class _Request: diff --git a/studio/backend/tests/test_passthrough_healing.py b/studio/backend/tests/test_passthrough_healing.py index 83bcc5864..9b22ab8b0 100644 --- a/studio/backend/tests/test_passthrough_healing.py +++ b/studio/backend/tests/test_passthrough_healing.py @@ -515,6 +515,7 @@ class ScriptedClient: _url, json = None, timeout = None, + headers = None, ): self.posts.append(json) return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)])