Merge remote-tracking branch 'origin/master' into ralph/proxy-stream-cancellation

packages/tracker/meshnet_tracker/server.py

@@ -2814,6 +2814,42 @@ def _set_upstream_read_timeout(upstream: Any, timeout: float | None) -> None:
         sock.settimeout(timeout)
 
 
+def _upstream_wait_readable(upstream: Any, sock: Any, timeout: float) -> bool:
+    """Return True when the next ``upstream.readline()`` can make progress.
+
+    Cancellation of a direct SSE proxy is polled with a bounded wait, but a
+    bounded socket timeout cannot be used for the read itself: ``socket``'s
+    makefile reader poisons itself after a single timeout, which would abort any
+    idle gap between frames.  So the read stays blocking (timeout ``None``) and
+    this helper does the bounded waiting instead.
+
+    ``select()`` alone is insufficient: ``urlopen``'s header parsing can read
+    the first SSE frame ahead into the response ``BufferedReader``, leaving the
+    raw socket with no new bytes -- ``select()`` would then withhold that frame
+    until the upstream produced more output.  A non-blocking ``peek()`` surfaces
+    such already-buffered bytes first; only when nothing is buffered do we wait
+    on the raw socket.
+    """
+    fp = getattr(upstream, "fp", None)
+    if fp is not None and sock is not None:
+        sock.setblocking(False)
+        try:
+            if fp.peek(1):
+                return True
+        except (BlockingIOError, InterruptedError):
+            pass
+        except (OSError, ValueError):
+            # Let the blocking readline surface the real error/EOF.
+            return True
+        finally:
+            # Restore blocking mode (timeout None) so readline never poisons.
+            sock.setblocking(True)
+    if sock is None:
+        return True
+    readable, _, _ = select.select([sock], [], [], timeout)
+    return bool(readable)
+
+
 def _clear_proxy_progress_log_state(server: "_TrackerHTTPServer", request_id: str) -> None:
     state = getattr(server, "_proxy_progress_log_state", None)
     if state is not None:
@@ -2980,6 +3016,35 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         except BrokenPipeError:
             pass
 
+    def _start_sse_response(self, status: int, content_type: str) -> tuple[Any, Any]:
+        """Start a chunked SSE response and return its write/finish callbacks.
+
+        A streaming proxy cannot use EOF as its framing boundary: HTTP clients
+        and intermediaries are then free to retain an SSE frame until the
+        upstream closes.  Chunked framing makes each flushed frame observable
+        immediately while still giving cancellation a clean response end.
+        """
+        # The tracker otherwise serves finite HTTP/1.0-style responses.  Limit
+        # HTTP/1.1 to this framed response so those existing response contracts
+        # do not become persistent connections accidentally.
+        self.protocol_version = "HTTP/1.1"
+        self.send_response(status)
+        self.send_header("Content-Type", content_type)
+        self.send_header("Cache-Control", "no-cache")
+        self.send_header("X-Accel-Buffering", "no")
+        self.send_header("Transfer-Encoding", "chunked")
+        self.end_headers()
+
+        def write(chunk: bytes) -> None:
+            self.wfile.write(f"{len(chunk):X}\r\n".encode() + chunk + b"\r\n")
+            self.wfile.flush()
+
+        def finish() -> None:
+            self.wfile.write(b"0\r\n\r\n")
+            self.wfile.flush()
+
+        return write, finish
+
     # ---- unified auth boundary (ADR-0017) ----
 
     def _resolve_identity(self) -> tuple[str | None, dict | None]:
@@ -4035,11 +4100,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             upstream = upstream_result[0]
             with proxy_ctx.upstream_lock:
                 proxy_ctx.upstream = upstream
+            # Keep the read blocking (timeout None); cancellation is polled by
+            # _upstream_wait_readable() rather than by a bounded socket timeout,
+            # which would poison the makefile reader across idle SSE gaps.
+            _set_upstream_read_timeout(upstream, None)
             upstream_sock = _upstream_socket(upstream)
-            if upstream_sock is not None:
-                _set_upstream_read_timeout(upstream, None)
-            else:
-                _set_upstream_read_timeout(upstream, 0.5)
             _tracker_log(server, "info", "proxy connected", request_id=request_id, target_url=target_url)
         except urllib.error.HTTPError as exc:
             # Relay error status + body from node
@@ -4088,11 +4153,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             content_type = upstream.headers.get("Content-Type", "application/json")
             if is_stream or "text/event-stream" in content_type:
                 # Relay SSE stream chunk-by-chunk
-                self.send_response(200)
-                self.send_header("Content-Type", "text/event-stream; charset=utf-8")
-                self.send_header("Cache-Control", "no-cache")
-                self.send_header("X-Accel-Buffering", "no")
-                self.end_headers()
+                write_sse, finish_sse = self._start_sse_response(
+                    200, "text/event-stream; charset=utf-8"
+                )
                 stream_usage: dict | None = None
                 observed_stream_tokens = 0
                 client_gone = False
@@ -4100,22 +4163,14 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                     while True:
                         if proxy_ctx.cancel_event.is_set():
                             break
-                        if upstream_sock is not None:
-                            readable, _, _ = select.select([upstream_sock], [], [], 0.5)
-                            if not readable:
-                                continue
-                        try:
-                            line = upstream.readline()
-                        except TimeoutError:
+                        if not _upstream_wait_readable(upstream, upstream_sock, 0.5):
                             continue
+                        line = upstream.readline()
                         if not line:
-                            if proxy_ctx.cancel_event.is_set():
-                                break
                             break
                         if not client_gone:
                             try:
-                                self.wfile.write(line)
-                                self.wfile.flush()
+                                write_sse(line)
                             except (BrokenPipeError, ConnectionResetError):
                                 # Keep draining upstream so completed node work is still billed.
                                 client_gone = True
@@ -4135,6 +4190,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                             stream_usage = usage
                 except (BrokenPipeError, ConnectionResetError):
                     client_gone = True
+                except (OSError, ValueError):
+                    # Cancellation closes the upstream from another handler
+                    # thread to break the blocking readline safely.
+                    if not proxy_ctx.cancel_event.is_set():
+                        raise
                 if self._finalize_proxy_cancel(
                     proxy_ctx=proxy_ctx,
                     server=server,
@@ -4150,6 +4210,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                     observed_stream_tokens=observed_stream_tokens,
                     stream_usage=stream_usage,
                 ):
+                    if not client_gone:
+                        try:
+                            finish_sse()
+                        except (BrokenPipeError, ConnectionResetError):
+                            pass
                     return
                 elapsed = time.monotonic() - started
                 # Bill even on client disconnect — the nodes did the work.
@@ -4182,6 +4247,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                     api_key, model, tokens, node_work,
                     input_tokens=in_tokens, output_tokens=out_tokens,
                 )
+                if not client_gone:
+                    try:
+                        finish_sse()
+                    except (BrokenPipeError, ConnectionResetError):
+                        pass
             else:
                 # Non-streaming: buffer and relay
                 resp_body = upstream.read()
@@ -4371,11 +4441,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         """Forward a streamed relay response (US-036) to the client as SSE,
         billing with the same accounting as the direct stream path."""
         headers = first.get("headers") if isinstance(first.get("headers"), dict) else {}
-        self.send_response(int(first.get("status", 200)))
-        self.send_header("Content-Type", headers.get("Content-Type", "text/event-stream; charset=utf-8"))
-        self.send_header("Cache-Control", "no-cache")
-        self.send_header("X-Accel-Buffering", "no")
-        self.end_headers()
+        write_sse, finish_sse = self._start_sse_response(
+            int(first.get("status", 200)),
+            headers.get("Content-Type", "text/event-stream; charset=utf-8"),
+        )
         server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
         stream_usage: dict | None = None
         observed_stream_tokens = 0
@@ -4389,8 +4458,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             data = chunk.encode()
             if not client_gone:
                 try:
-                    self.wfile.write(data)
-                    self.wfile.flush()
+                    write_sse(data)
                 except BrokenPipeError:
                     # Keep draining frames — the nodes did the work; bill it.
                     client_gone = True
@@ -4429,6 +4497,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                 stream_usage=stream_usage,
             )
         ):
+            if not client_gone:
+                try:
+                    finish_sse()
+                except BrokenPipeError:
+                    pass
             return
         elapsed = time.monotonic() - started
         in_tokens, out_tokens = _stream_billable_split(
@@ -4459,6 +4532,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             api_key, model, tokens, node_work,
             input_tokens=in_tokens, output_tokens=out_tokens,
         )
+        if not client_gone:
+            try:
+                finish_sse()
+            except BrokenPipeError:
+                pass
 
     def _send_relayed_response(self, response: dict) -> None:
         status = int(response.get("status", 503))

tests/test_tracker_routing.py

@@ -635,6 +635,7 @@ def test_tracker_dashboard_can_cancel_inflight_proxy():
     "Tracker dashboard can cancel inflight proxy\n\nTags: http, routing, tracker"
     chunk_sent = threading.Event()
     release = threading.Event()
+    upstream_finished = threading.Event()
 
     class StreamingChatHandler(http.server.BaseHTTPRequestHandler):
         def log_message(self, fmt, *args):
@@ -655,9 +656,10 @@ def test_tracker_dashboard_can_cancel_inflight_proxy():
             self.wfile.write(b"data: " + payload + b"\n\n")
             self.wfile.flush()
             chunk_sent.set()
-            release.wait(timeout=3.0)
+            release.wait()
             self.wfile.write(b"data: [DONE]\n\n")
             self.wfile.flush()
+            upstream_finished.set()
 
     node = http.server.HTTPServer(("127.0.0.1", 0), StreamingChatHandler)
     node_thread = threading.Thread(target=node.serve_forever, daemon=True)
@@ -677,6 +679,7 @@ def test_tracker_dashboard_can_cancel_inflight_proxy():
         req = urllib.request.Request(
             f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
             data=json.dumps({
+                "id": "direct-sse-cancel",
                 "model": "cancel-proxy-model",
                 "stream": True,
                 "messages": [{"role": "user", "content": "hi"}],
@@ -688,14 +691,8 @@ def test_tracker_dashboard_can_cancel_inflight_proxy():
         first_line = response.readline()
         assert first_line.startswith(b"data:")
         assert chunk_sent.wait(timeout=1.0)
-
-        console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
-        selected = [
-            event for event in console["events"]
-            if event["message"] == "proxy route selected"
-        ]
-        assert selected
-        request_id = selected[-1]["fields"]["request_id"]
+        assert not upstream_finished.is_set()
+        request_id = "direct-sse-cancel"
 
         cancel = _post_json(
             f"http://127.0.0.1:{tracker_port}/v1/proxy/requests/{urllib.parse.quote(request_id, safe='')}/cancel",
@@ -703,19 +700,27 @@ def test_tracker_dashboard_can_cancel_inflight_proxy():
         )
         assert cancel["status"] == "canceled"
 
-        deadline = time.time() + 5.0
-        canceled_events = []
-        while time.time() < deadline:
-            console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
-            canceled_events = [
-                event for event in console["events"]
-                if event["message"] == "proxy canceled"
-                and event["fields"].get("request_id") == request_id
-            ]
-            if canceled_events:
-                break
-            time.sleep(0.05)
+        # Chunked SSE termination makes this return as soon as the tracker has
+        # observed cancellation, rather than blocking on the held upstream. The
+        # only bytes that may still arrive are the blank-line terminator of the
+        # already-relayed first frame; the upstream's held "[DONE]" must never
+        # reach the client, and the upstream must still be blocked mid-stream.
+        remainder = response.read()
+        assert b"[DONE]" not in remainder
+        assert remainder.strip() == b""
+        assert not upstream_finished.is_set()
+        console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
+        canceled_events = [
+            event for event in console["events"]
+            if event["message"] == "proxy canceled"
+            and event["fields"].get("request_id") == request_id
+        ]
         assert canceled_events
+        assert len(canceled_events) == 1
+        assert tracker._server is not None
+        assert request_id not in tracker._server.active_proxies
+        registered = next(iter(tracker._server.registry.values()))
+        assert registered.proxy_inflight == 0
     finally:
         release.set()
         if response is not None: