fixing streaming

tests/test_real_model_backend.py

@@ -4,6 +4,8 @@ import json
 import os
 from pathlib import Path
 import sys
+import threading
+import time
 import types
 import urllib.request
 
@@ -94,7 +96,7 @@ class _FakePipelineHeadBackend(_FakeBackend):
     tokenizer = _FakeChatTokenizer()
 
     def encode_prompt(self, prompt: str) -> TensorPayload:
-        assert prompt == "debug prompt"
+        assert prompt.startswith("debug prompt")
         return TensorPayload(
             body=b"\x00" * (1 * 6 * 8 * 2),
             shape=[1, 6, 8],
@@ -113,6 +115,19 @@ class _FakePipelineTailBackend(_FakeTailBackend):
         return " token"
 
 
+class _BlockingStreamingTailBackend(_FakeTailBackend):
+    def __init__(self, second_token_release: threading.Event) -> None:
+        self._release = second_token_release
+        self.calls = 0
+
+    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
+        self.calls += 1
+        if self.calls == 1:
+            return " first"
+        self._release.wait(timeout=3.0)
+        return " second"
+
+
 def test_quantization_flag_validation():
     assert validate_quantization("bfloat16") == "bfloat16"
     assert validate_quantization("int8") == "int8"
@@ -299,6 +314,56 @@ def test_pipeline_hop_logs_are_enabled_with_debug(capsys):
     assert "  [node] pipeline hop 0 returned text=' token'" in out
 
 
+def test_split_shard_chat_streams_each_generated_token_incrementally():
+    release_second = threading.Event()
+    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
+    tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second))
+    head_port = head.start()
+    tail_port = tail.start()
+    response = None
+    try:
+        payload = json.dumps({
+            "model": "fake-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "stream": True,
+            "max_tokens": 2,
+        }).encode()
+        req = urllib.request.Request(
+            f"http://127.0.0.1:{head_port}/v1/chat/completions",
+            data=payload,
+            headers={
+                "Content-Type": "application/json",
+                "X-Meshnet-Route": json.dumps([
+                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
+                ]),
+            },
+            method="POST",
+        )
+        response = urllib.request.urlopen(req, timeout=5)
+
+        first_token_line = ""
+        deadline = time.time() + 2.0
+        while time.time() < deadline:
+            line = response.readline().decode()
+            if '"content": " first"' in line:
+                first_token_line = line
+                break
+
+        assert first_token_line
+        assert not release_second.is_set()
+        release_second.set()
+        rest = response.read().decode()
+    finally:
+        release_second.set()
+        if response is not None:
+            response.close()
+        head.stop()
+        tail.stop()
+
+    assert '"content": " second"' in rest
+    assert "data: [DONE]" in rest
+
+
 def test_int_tensor_header_serializes_torch_tensors():
     torch = pytest.importorskip("torch")