try to fix streaming responses

tests/test_tracker_as_node.py

@@ -142,6 +142,21 @@ def _send_chat_request(gateway_url: str, prompt: str) -> dict:
         return json.loads(r.read())
 
 
+def _send_streaming_chat_request(gateway_url: str, prompt: str):
+    data = json.dumps({
+        "model": GPT2_MODEL,
+        "messages": [{"role": "user", "content": prompt}],
+        "stream": True,
+    }).encode()
+    req = urllib.request.Request(
+        f"{gateway_url}/v1/chat/completions",
+        data=data,
+        headers={"Content-Type": "application/json"},
+        method="POST",
+    )
+    return urllib.request.urlopen(req)
+
+
 def test_all_responses_valid_openai_format(tracker_node_setup):
     """Ten requests via gateway all return valid OpenAI chat completion format."""
     gateway_url, _, _ = tracker_node_setup
@@ -155,6 +170,30 @@ def test_all_responses_valid_openai_format(tracker_node_setup):
         assert isinstance(message.get("content"), str), f"request {i}: content must be a string"
 
 
+def test_streaming_head_worker_response_is_not_buffered_with_content_length(tracker_node_setup):
+    """Gateway must relay head-worker SSE as a live stream, not a buffered JSON-sized body."""
+    gateway_url, _, _ = tracker_node_setup
+
+    with _send_streaming_chat_request(gateway_url, "stream through head worker") as resp:
+        assert resp.status == 200
+        assert "text/event-stream" in resp.headers["Content-Type"]
+        assert "Content-Length" not in resp.headers
+        data_lines = []
+        while len(data_lines) < 4:
+            line = resp.readline().decode().strip()
+            if line.startswith("data: "):
+                data_lines.append(line)
+            if line == "data: [DONE]":
+                break
+
+    assert data_lines[-1] == "data: [DONE]"
+    content = "".join(
+        json.loads(line[6:])["choices"][0].get("delta", {}).get("content", "")
+        for line in data_lines[:-1]
+    )
+    assert "head-worker" in content
+
+
 def test_both_tracker_nodes_receive_load(tracker_node_setup):
     """Both head workers handle at least one request each out of ten."""
     gateway_url, tracker_node_a, tracker_node_b = tracker_node_setup