From 6fd9d93e4bc9cd7c07be7cede59a9140461fb55e Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Thu, 16 Jul 2026 18:06:50 +0300 Subject: [PATCH] fix: flush direct SSE proxy frames before completion --- packages/tracker/meshnet_tracker/server.py | 132 ++++++++++++++++----- tests/test_tracker_routing.py | 47 ++++---- 2 files changed, 131 insertions(+), 48 deletions(-) diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 0a428ef..8159ce7 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -2744,6 +2744,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: @@ -2910,6 +2946,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]: @@ -3956,11 +4021,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 @@ -4009,11 +4074,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 @@ -4021,22 +4084,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 @@ -4056,6 +4111,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, @@ -4071,6 +4131,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. @@ -4103,6 +4168,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() @@ -4292,11 +4362,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 @@ -4310,8 +4379,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 @@ -4350,6 +4418,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( @@ -4380,6 +4453,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)) diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index 09cd63f..5f938b5 100644 --- a/tests/test_tracker_routing.py +++ b/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: