logging, routing

tests/test_tracker_routing.py

@@ -502,6 +502,75 @@ def test_tracker_logs_stream_progress_before_request_completes():
         node_thread.join(timeout=1.0)
 
 
+def test_tracker_stream_survives_idle_gap_between_sse_chunks():
+    first_chunk_sent = threading.Event()
+
+    class IdleStreamingChatHandler(http.server.BaseHTTPRequestHandler):
+        def log_message(self, format, *args):
+            pass
+
+        def do_POST(self):
+            if self.path != "/v1/chat/completions":
+                self.send_response(404)
+                self.end_headers()
+                return
+            self.rfile.read(int(self.headers.get("Content-Length", 0)))
+            self.send_response(200)
+            self.send_header("Content-Type", "text/event-stream; charset=utf-8")
+            self.end_headers()
+            first = json.dumps({
+                "choices": [{"delta": {"content": "hello"}}],
+            }).encode()
+            second = json.dumps({
+                "choices": [{"delta": {"content": " world"}}],
+            }).encode()
+            self.wfile.write(b"data: " + first + b"\n\n")
+            self.wfile.flush()
+            first_chunk_sent.set()
+            time.sleep(1.0)
+            self.wfile.write(b"data: " + second + b"\n\n")
+            self.wfile.write(b"data: [DONE]\n\n")
+            self.wfile.flush()
+
+    node = http.server.HTTPServer(("127.0.0.1", 0), IdleStreamingChatHandler)
+    node_thread = threading.Thread(target=node.serve_forever, daemon=True)
+    node_thread.start()
+    tracker = TrackerServer(heartbeat_timeout=60.0)
+    tracker_port = tracker.start()
+    response = None
+    try:
+        _post_json(
+            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+            {"endpoint": f"http://127.0.0.1:{node.server_address[1]}",
+             "model": "idle-stream-model", "num_layers": 1,
+             "shard_start": 0, "shard_end": 0,
+             "hardware_profile": {}, "score": 1.0},
+        )
+        req = urllib.request.Request(
+            f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
+            data=json.dumps({
+                "model": "idle-stream-model",
+                "stream": True,
+                "messages": [{"role": "user", "content": "hi"}],
+            }).encode(),
+            headers={"Content-Type": "application/json"},
+            method="POST",
+        )
+        response = urllib.request.urlopen(req, timeout=3.0)
+        assert response.readline().startswith(b"data:")
+        assert first_chunk_sent.wait(timeout=1.0)
+
+        remaining = response.read().splitlines()
+        assert b"data: [DONE]" in remaining
+    finally:
+        if response is not None:
+            response.close()
+        tracker.stop()
+        node.shutdown()
+        node.server_close()
+        node_thread.join(timeout=1.0)
+
+
 def test_tracker_dashboard_can_cancel_inflight_proxy():
     chunk_sent = threading.Event()
     release = threading.Event()