PLUS: SSE chunked fix

plus/sse_proxy.py

@@ -314,14 +314,17 @@ class Proxy(BaseHTTPRequestHandler):
         stream = content_type.startswith("text/event-stream")
         for key, value in answer.getheaders():
             low = key.lower()
-            if low in HOP_HEADERS or (stream and low == "content-length"):
+            if low in HOP_HEADERS or low in ("server", "date"):
+                continue
+            if stream and low in ("cache-control", "x-accel-buffering"):
                 continue
             self.send_header(key, value)
         if stream:
-            self.send_header("Cache-Control", "no-cache")
-            self.send_header("X-Accel-Buffering", "no")
-            self.end_headers()
-            self.stream_events(answer)
+            self.begin_stream()
+            try:
+                self.stream_events(answer)
+            finally:
+                self.end_stream()
         else:
             data = answer.read()
             content_type = (answer.getheader("Content-Type") or "").lower()
@@ -334,12 +337,41 @@ class Proxy(BaseHTTPRequestHandler):
                     and data[:4] == b"MThd"):
                 data, stats = apply_velocity(self.last_audio, data)
                 print("velocity: %s" % stats, flush=True)
-            self.send_header("Content-Length", str(len(data)))
+            if method == "HEAD":
+                length = answer.getheader("Content-Length")
+                self.send_header("Content-Length", length if length is not None else "0")
+            elif answer.status not in (204, 304) and answer.status >= 200:
+                self.send_header("Content-Length", str(len(data)))
             self.end_headers()
-            if method != "HEAD":
+            if method != "HEAD" and answer.status not in (204, 304):
                 self.wfile.write(data)
         upstream.close()
 
+    def begin_stream(self) -> None:
+        """Заголовки SSE: chunked + Connection: close, чтобы туннель (HTTP/2)
+        видел конец ответа и не рвал поток с ERR_HTTP2_PROTOCOL_ERROR."""
+        self.send_header("Cache-Control", "no-cache, no-transform")
+        self.send_header("X-Accel-Buffering", "no")
+        self.send_header("Transfer-Encoding", "chunked")
+        self.send_header("Connection", "close")
+        self.end_headers()
+        self.close_connection = True
+        self.emit(b": open\n\n")
+
+    def emit(self, data: bytes) -> None:
+        if not data:
+            return
+        self.wfile.write(b"%x\r\n" % len(data) + data + b"\r\n")
+        self.wfile.flush()
+
+    def end_stream(self) -> None:
+        self.close_connection = True
+        try:
+            self.wfile.write(b"0\r\n\r\n")
+            self.wfile.flush()
+        except Exception:
+            pass
+
     def refine(self, notes: list[dict], skip_filter: bool = False) -> list[dict]:
         """Пересчитать концы нот и артикуляцию по записи."""
         if not notes or not self.last_audio:
@@ -416,34 +448,26 @@ class Proxy(BaseHTTPRequestHandler):
         threading.Thread(target=worker, daemon=True).start()
         self.send_response(200)
         self.send_header("Content-Type", "text/event-stream")
-        self.send_header("Cache-Control", "no-cache")
-        self.send_header("X-Accel-Buffering", "no")
-        self.end_headers()
         try:
-            self.wfile.write(b": open\n\n")
-            self.wfile.flush()
+            self.begin_stream()
         except Exception:
             return
         payload = None
-        while True:
-            try:
-                payload = outcome.get(timeout=self.ping_seconds)
-                break
-            except queue.Empty:
-                try:
-                    self.wfile.write(b": ping\n\n")
-                    self.wfile.flush()
-                except Exception:
-                    return
-        if payload is None:
-            payload = fusion.to_sse([])
         try:
-            self.wfile.write(payload)
-            self.wfile.write(b": done\n\n")
-            self.wfile.flush()
+            while True:
+                try:
+                    payload = outcome.get(timeout=self.ping_seconds)
+                    break
+                except queue.Empty:
+                    self.emit(b": ping\n\n")
+            if payload is None:
+                payload = fusion.to_sse([])
+            self.emit(payload)
+            self.emit(b": done\n\n")
         except Exception:
             pass
-        self.close_connection = True
+        finally:
+            self.end_stream()
 
     def stream_events(self, answer):
         """Переливаем SSE, вставляя ": ping" во время молчания модели."""
@@ -462,23 +486,19 @@ class Proxy(BaseHTTPRequestHandler):
                 lines.put(None)
 
         threading.Thread(target=pump, daemon=True).start()
-        self.wfile.write(b": open\n\n")
-        self.wfile.flush()
         while True:
             try:
                 line = lines.get(timeout=PING_SECONDS)
             except queue.Empty:
                 try:
-                    self.wfile.write(b": ping\n\n")
-                    self.wfile.flush()
+                    self.emit(b": ping\n\n")
                 except Exception:
                     return
                 continue
             if line is None:
                 return
             try:
-                self.wfile.write(line)
-                self.wfile.flush()
+                self.emit(line)
             except Exception:
                 return