popov/neuron-tai
1
0
Fork 0
Code Issues 53 Pull Requests Actions Packages Projects Releases Wiki Activity

fix: flush direct SSE proxy frames before completion

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-16 18:06:50 +03:00
parent 5c9a2f6c97
commit 6fd9d93e4b
2 changed files with 131 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

130
packages/tracker/meshnet_tracker/server.py
View File

@@ -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
upstream_sock = _upstream_socket(upstream)
if upstream_sock is not None:
# 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)
else:
_set_upstream_read_timeout(upstream, 0.5)
upstream_sock = _upstream_socket(upstream)
_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:
if not _upstream_wait_readable(upstream, upstream_sock, 0.5):
continue
try:
line = upstream.readline()
except TimeoutError:
continue
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))

35
tests/test_tracker_routing.py
View File

@@ -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:
# 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
]
if canceled_events:
break
time.sleep(0.05)
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: