diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index d7e3173f37..093a92e38d 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8606,7 +8606,8 @@ class LlamaCppBackend: @staticmethod def _shutdown_active_httpx_sockets(client: "httpx.Client") -> None: - """Best-effort interrupt for a sync httpx request blocked before headers.""" + """Best-effort interrupt for a sync httpx read blocked in recv(), whether + parked before headers (prefill) or mid-stream.""" try: pool = getattr(getattr(client, "_transport", None), "_pool", None) connections = list(getattr(pool, "_connections", []) or []) @@ -8654,10 +8655,11 @@ class LlamaCppBackend: while not _cancel_closed.is_set(): r = _response_ref[0] try: + # response.close() can't wake a read already blocked in + # recv(); only a socket shutdown does, so shut down first. + LlamaCppBackend._shutdown_active_httpx_sockets(client) if r is not None: r.close() - else: - LlamaCppBackend._shutdown_active_httpx_sockets(client) return except Exception as e: logger.debug(f"Error closing request in cancel watcher: {e}") diff --git a/studio/backend/tests/test_llama_cpp_stream_cancel.py b/studio/backend/tests/test_llama_cpp_stream_cancel.py index 1c73f4d17c..f87ce86450 100644 --- a/studio/backend/tests/test_llama_cpp_stream_cancel.py +++ b/studio/backend/tests/test_llama_cpp_stream_cancel.py @@ -3,8 +3,10 @@ import contextlib import os +import socket import sys import threading +import time import httpx import pytest @@ -79,3 +81,112 @@ def test_generate_chat_completion_swallows_internal_stream_cancel(monkeypatch): ) assert chunks == [] + + +class _StallUpstream: + """Raw HTTP/1.1 server that streams one chunked SSE chunk, then holds the + socket open and silent so the client's next read blocks in recv() until its + side is torn down. Reproduces a mid-stream stall (llama-server goes quiet).""" + + def __init__(self): + self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self._sock.bind(("127.0.0.1", 0)) + self._sock.listen(1) + self.port = self._sock.getsockname()[1] + self._stop = threading.Event() + self._thread = threading.Thread(target = self._serve, daemon = True) + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.port}/v1/chat/completions" + + def __enter__(self): + self._thread.start() + return self + + def __exit__(self, *_exc): + self._stop.set() + try: + self._sock.close() + except OSError: + pass + self._thread.join(timeout = 5) + + def _serve(self) -> None: + try: + conn, _ = self._sock.accept() + except OSError: + return + with conn: + conn.settimeout(5) + try: + buf = b"" + while b"\r\n\r\n" not in buf: + data = conn.recv(4096) + if not data: + return + buf += data + head, _, body = buf.partition(b"\r\n\r\n") + content_length = 0 + for line in head.split(b"\r\n"): + if line.lower().startswith(b"content-length:"): + content_length = int(line.split(b":", 1)[1].strip()) + break + while len(body) < content_length: + data = conn.recv(4096) + if not data: + break + body += data + except OSError: + return + conn.sendall( + b"HTTP/1.1 200 OK\r\n" + b"Content-Type: text/event-stream\r\n" + b"Transfer-Encoding: chunked\r\n" + b"\r\n" + ) + chunk = b"data: hello\n\n" + conn.sendall(b"%x\r\n%s\r\n" % (len(chunk), chunk)) + # Stall: stay open and silent until the client shuts its side down. + while not self._stop.wait(timeout = 0.05): + try: + conn.settimeout(0.05) + if conn.recv(1) == b"": + return + except socket.timeout: + continue + except OSError: + return + + +def test_cancel_interrupts_a_read_blocked_on_a_mid_stream_stall(): + # Mid-stream stall: the reader is parked in recv() on a long bound read timeout, + # so response.close() alone can't wake it; the watcher must shut the socket down. + # Assert cancel lands in seconds, not at the far-off deadline (pre-fix: hung ~30s). + with _StallUpstream() as server: + cancel_event = threading.Event() + + def _cancel_soon(): + time.sleep(0.3) + cancel_event.set() + + threading.Thread(target = _cancel_soon, daemon = True).start() + + started = time.monotonic() + with httpx.Client( + limits = httpx.Limits(max_keepalive_connections = 0), trust_env = False + ) as client: + with pytest.raises(_LlamaStreamCancelled): + with LlamaCppBackend._stream_with_retry( + client, + server.url, + {}, + cancel_event, + first_token_deadline = started + 30, + ) as response: + for _chunk in response.iter_text(): + pass # first chunk arrives, then the read blocks silently + elapsed = time.monotonic() - started + + assert elapsed < 10, f"cancel took {elapsed:.1f}s; the blocked read was not interrupted"