fix: flush direct SSE proxy frames before completion

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

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:
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: