try to fix streaming responses

This commit is contained in:
Dobromir Popov
2026-07-07 22:19:22 +03:00
parent a0b37ad1b9
commit f0dc3bd93f
4 changed files with 140 additions and 26 deletions

View File

@@ -286,7 +286,7 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
self._send_json(200, completion) self._send_json(200, completion)
def _proxy_to_head_worker(self, url: str, body_bytes: bytes) -> None: def _proxy_to_head_worker(self, url: str, body_bytes: bytes) -> None:
"""Forward a raw request body to a head worker and stream the response back.""" """Forward a raw request body to a head worker and relay SSE without buffering."""
target_url = f"{url}/v1/chat/completions" target_url = f"{url}/v1/chat/completions"
req = urllib.request.Request( req = urllib.request.Request(
target_url, target_url,
@@ -297,6 +297,19 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
try: try:
with urllib.request.urlopen(req, timeout=30.0) as r: with urllib.request.urlopen(req, timeout=30.0) as r:
content_type = r.headers.get("Content-Type", "application/json") content_type = r.headers.get("Content-Type", "application/json")
if "text/event-stream" in content_type:
self.send_response(r.status)
self.send_header("Content-Type", content_type)
self.send_header("Cache-Control", "no-cache")
self.send_header("X-Accel-Buffering", "no")
self.end_headers()
while True:
line = r.readline()
if not line:
break
self.wfile.write(line)
self.wfile.flush()
return
resp_body = r.read() resp_body = r.read()
status = r.status status = r.status
except urllib.error.HTTPError as exc: except urllib.error.HTTPError as exc:

View File

@@ -1241,7 +1241,7 @@ async function sendChat() {
.map(msg => ({ role: msg.role, content: msg.content })), .map(msg => ({ role: msg.role, content: msg.content })),
{ role: "user", content: prompt }, { role: "user", content: prompt },
], ],
stream: false, stream: true,
max_tokens: 15120, max_tokens: 15120,
}; };
chatBusy = true; chatBusy = true;
@@ -1249,36 +1249,83 @@ async function sendChat() {
promptEl.value = ""; promptEl.value = "";
promptEl.style.height = "auto"; promptEl.style.height = "auto";
chatHistory.push({ role: "user", content: prompt, model: selectedChatModel }); chatHistory.push({ role: "user", content: prompt, model: selectedChatModel });
const assistantMessage = { role: "assistant", content: "", model: selectedChatModel };
chatHistory.push(assistantMessage);
renderChatHistory(); renderChatHistory();
persistActiveChatSession(); persistActiveChatSession();
renderChatStatus("sending request…"); renderChatStatus("streaming response…");
const r = await apiCall("/v1/chat/completions", "POST", body, bearerToken);
chatBusy = false; try {
$("chat-send").disabled = false; const headers = { "Content-Type": "application/json" };
if (!r.ok) { if (bearerToken) headers["Authorization"] = "Bearer " + bearerToken;
const error = r.data && r.data.error const response = await fetch("/v1/chat/completions", {
? (typeof r.data.error === "string" ? r.data.error : r.data.error.message || "request failed") method: "POST",
: "request failed"; headers,
chatHistory.push({ role: "error", content: error, model: selectedChatModel }); credentials: "same-origin",
body: JSON.stringify(body),
});
if (!response.ok) {
const data = await response.json().catch(() => ({}));
const error = data && data.error
? (typeof data.error === "string" ? data.error : data.error.message || "request failed")
: "request failed";
assistantMessage.role = "error";
assistantMessage.content = error;
renderChatHistory();
persistActiveChatSession();
renderChatStatus(error);
return;
}
if (!response.body) throw new Error("stream response body is missing");
const reader = response.body.getReader();
const decoder = new TextDecoder();
let buffer = "";
let receivedAny = false;
while (true) {
const { value, done } = await reader.read();
if (done) break;
buffer += decoder.decode(value, { stream: true });
const events = buffer.split("\n\n");
buffer = events.pop() || "";
for (const eventText of events) {
for (const line of eventText.split("\n")) {
if (!line.startsWith("data: ")) continue;
const data = line.slice(6).trim();
if (data === "[DONE]") {
buffer = "";
break;
}
try {
const chunk = JSON.parse(data);
const delta = chunk.choices && chunk.choices[0] && chunk.choices[0].delta;
if (delta && delta.content) {
assistantMessage.content += delta.content;
receivedAny = true;
renderChatHistory();
persistActiveChatSession();
}
} catch {
// Ignore malformed SSE keepalive/event lines.
}
}
}
}
if (!receivedAny) assistantMessage.content = "(empty response)";
renderChatHistory(); renderChatHistory();
persistActiveChatSession(); persistActiveChatSession();
renderChatStatus(error); renderChatStatus("done");
} catch (err) {
assistantMessage.role = "error";
assistantMessage.content = err && err.message ? err.message : "request failed";
renderChatHistory();
persistActiveChatSession();
renderChatStatus(assistantMessage.content);
} finally {
chatBusy = false;
$("chat-send").disabled = false;
promptEl.focus(); promptEl.focus();
return;
} }
const reply = (r.data && r.data.choices && r.data.choices[0] && r.data.choices[0].message && r.data.choices[0].message.content) || "";
const usage = r.data && r.data.usage;
chatHistory.push({
role: "assistant",
content: reply || "(empty response)",
model: selectedChatModel,
});
renderChatHistory();
persistActiveChatSession();
renderChatStatus(usage
? `done: ${usage.total_tokens ?? "?"} tokens`
: "done");
promptEl.focus();
} }
function bindChatPromptShortcuts() { function bindChatPromptShortcuts() {

View File

@@ -33,6 +33,21 @@ def test_dashboard_served_with_all_panels():
tracker.stop() tracker.stop()
def test_dashboard_chat_uses_streaming_fetch():
tracker = TrackerServer(billing=BillingLedger())
port = tracker.start()
try:
html = urllib.request.urlopen(
f"http://127.0.0.1:{port}/dashboard"
).read().decode()
finally:
tracker.stop()
assert "stream: true" in html
assert ".body.getReader()" in html
assert 'data === "[DONE]"' in html
def test_dashboard_served_by_follower(): def test_dashboard_served_by_follower():
"""A tracker that is not the leader (unreachable peers → never elected) """A tracker that is not the leader (unreachable peers → never elected)
still serves the dashboard from its own replicated state.""" still serves the dashboard from its own replicated state."""

View File

@@ -142,6 +142,21 @@ def _send_chat_request(gateway_url: str, prompt: str) -> dict:
return json.loads(r.read()) 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): def test_all_responses_valid_openai_format(tracker_node_setup):
"""Ten requests via gateway all return valid OpenAI chat completion format.""" """Ten requests via gateway all return valid OpenAI chat completion format."""
gateway_url, _, _ = tracker_node_setup 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" 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): def test_both_tracker_nodes_receive_load(tracker_node_setup):
"""Both head workers handle at least one request each out of ten.""" """Both head workers handle at least one request each out of ten."""
gateway_url, tracker_node_a, tracker_node_b = tracker_node_setup gateway_url, tracker_node_a, tracker_node_b = tracker_node_setup