serve: fix HTTP/1.1 keep-alive desync and SSE framing (#597 item 3)

c/openai_server.py

@@ -1359,10 +1359,59 @@ class APIHandler(BaseHTTPRequestHandler):
     timeout = 30   # per-request socket timeout: a slowloris client that dribbles its
                    # request line/body can't pin a worker thread (and a slot) forever
     server_version = "colibri"
+    _committed = False    # status line already on the wire; reset per request below
+    _body_read = False    # request body fully consumed, so nothing is left to drain
 
     def log_message(self, fmt, *args):
         sys.stderr.write("[api] %s - %s\n" % (self.address_string(), fmt % args))
 
+    def handle_one_request(self):
+        """Per-request bookkeeping for HTTP/1.1 persistence (#597 item 3).
+
+        One handler instance serves every request on a keep-alive connection, so both flags
+        reset here rather than in do_POST. The drain afterwards is the whole fix for the
+        reported `Bad request syntax ('{...json...}POST /v1/...')`: any early rejection --
+        403 Host, 401 auth, a bad or oversized Content-Length -- returns before read_json(),
+        leaving the body in the socket, where the next readline() eats it as a request line.
+        Draining once at the request boundary covers every such path, present and future,
+        instead of asking each early return to remember."""
+        self._committed = False
+        self._body_read = False
+        super().handle_one_request()
+        if not self.close_connection:
+            self._drain_request_body()
+
+    def send_response(self, code, message=None):
+        """Single choke point for "the status line is out". Overriding here rather than
+        tracking it at each call site means no responder can forget (#597 item 3)."""
+        self._committed = True
+        super().send_response(code, message)
+
+    def _drain_request_body(self):
+        """Consume any unread request body so the next request line is at the head of the
+        stream. Where the body can't be swallowed safely, close instead: an unreusable
+        connection is correct, a desynchronised one is not."""
+        if self._body_read:
+            return
+        self._body_read = True
+        if self.headers.get("Transfer-Encoding"):
+            self.close_connection = True   # not framed by Content-Length; we don't de-chunk
+            return
+        try:
+            remaining = int(self.headers.get("Content-Length", "0"))
+        except ValueError:
+            self.close_connection = True   # unparseable framing: the body length is unknown
+            return
+        if remaining < 0 or remaining > MAX_BODY:
+            self.close_connection = True   # don't burn bandwidth just to keep a socket warm
+            return
+        while remaining > 0:
+            chunk = self.rfile.read(min(remaining, 65536))
+            if not chunk:
+                self.close_connection = True
+                return
+            remaining -= len(chunk)
+
     def send_json(self, status, body, request_id=None, headers=None):
         data = json.dumps(body, ensure_ascii=False, separators=(",", ":")).encode()
         self.send_response(status)
@@ -1436,8 +1485,12 @@ def read_json(self):
             raise APIError(400, "Invalid Content-Length header.")
         if length < 1 or length > MAX_BODY:
             raise APIError(400, f"Request body must be between 1 and {MAX_BODY} bytes.")
+        raw = self.rfile.read(length)
+        # Only a full read leaves nothing to drain; a short read means the peer went away
+        # mid-body, and the drain will notice the EOF and close (#597 item 3).
+        self._body_read = len(raw) == length
         try:
-            body = json.loads(self.rfile.read(length))
+            body = json.loads(raw)
         except (json.JSONDecodeError, UnicodeDecodeError):
             raise APIError(400, "Request body must be valid JSON.")
         if not isinstance(body, dict):
@@ -1559,20 +1612,29 @@ def do_POST(self):
             else:
                 raise APIError(404, "Not found.", None, "not_found")
         except APIError as error:
-            self.send_json(error.status, self.error_body(error), request_id, error.headers)
+            self._fail(error, request_id)
         except ClientCancelled:
             pass
         except (BrokenPipeError, ConnectionResetError):
             pass
         except Exception as error:
             self.log_error("request failed: %s", error)
-            api_error = APIError(500, "The colibri engine failed to process the request.",
-                                 None, "engine_error", "server_error")
             try:
-                self.send_json(500, self.error_body(api_error), request_id)
+                self._fail(APIError(500, "The colibri engine failed to process the request.",
+                                    None, "engine_error", "server_error"), request_id)
             except OSError:
                 pass
 
+    def _fail(self, error, request_id):
+        """Report an error, unless the response is already on the wire. Once a streaming 200
+        is committed, a second status line would be framed as SSE body -- clients saw a whole
+        `HTTP/1.1 500` spliced into the event stream. All we can still do is stop talking; the
+        stream ends at the close, which the 200 already announced (#597 item 3)."""
+        if self._committed:
+            self.close_connection = True
+            return
+        self.send_json(error.status, self.error_body(error), request_id, error.headers)
+
     def error_body(self, error):
         """Anthropic clients parse a different error envelope; the OpenAI one is unchanged."""
         if urlsplit(self.path).path != "/v1/messages":
@@ -1743,6 +1805,13 @@ def start_stream(_accept_info=None):
                 self.send_header("Content-Type", "text/event-stream")
                 self.send_header("Cache-Control", "no-cache")
                 self.send_header("X-Accel-Buffering", "no")
+                # An SSE body has neither Content-Length nor chunked framing, so end-of-message
+                # IS the close -- HTTP/1.1 requires us to say so, or the client waits for a
+                # length that never comes and then tries to reuse a socket we are about to drop.
+                # Set close_connection HERE, not after the last event: if generation raises once
+                # the 200 is out, the connection must still not be offered for reuse (#597 item 3).
+                self.send_header("Connection", "close")
+                self.close_connection = True
                 self.send_header("x-request-id", request_id)
                 for name, value in queue_headers.items(): self.send_header(name, value)
                 self.send_cors_headers()
@@ -1842,7 +1911,7 @@ def emit_plain(chunk):
                         self.wfile.flush()
                     except OSError:
                         pass
-            self.close_connection = True
+            # close_connection was already set when the 200 was committed (#597 item 3).
 
     def client_disconnected(self):
         try:
@@ -1976,6 +2045,8 @@ def blocks_and_stop(text, stats):
             self.send_header("Content-Type", "text/event-stream")
             self.send_header("Cache-Control", "no-cache")
             self.send_header("X-Accel-Buffering", "no")
+            self.send_header("Connection", "close")   # see the OpenAI path: SSE is close-framed
+            self.close_connection = True
             self.send_header("x-request-id", request_id)
             for name, value in queue_headers.items():
                 self.send_header(name, value)
@@ -2107,7 +2178,7 @@ def on_text(chunk):
                 "delta": {"stop_reason": stop_reason, "stop_sequence": None},
                 "usage": {"output_tokens": stats["completion_tokens"]}})
             send_event("message_stop", {"type": "message_stop"})
-            self.close_connection = True
+            # close_connection was already set when the 200 was committed (#597 item 3).
 
     def completion(self, body, request_id):
         prompt = body.get("prompt")

c/tests/test_openai_server.py

@@ -1247,5 +1247,183 @@ def test_streaming_context_exceeded_is_clean_400(self):
         self.assertEqual(body["error"]["param"], "messages")
 
 
+class _ExplodingEngine(FakeEngine):
+    """ACCEPTs the prompt (committing the streaming 200), then dies mid-generation."""
+    def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0,
+                 cancelled=None, grammar=None, stopped=None, on_accept=None):
+        if on_accept is not None:
+            on_accept({"prompt_tokens": 7})
+        on_text("partial")
+        raise RuntimeError("engine died mid-stream")
+
+
+class KeepAliveFramingTest(unittest.TestCase):
+    """#597 item 3: HTTP/1.1 persistence must not desynchronise.
+
+    The report was `Bad request syntax ('{...json body...}POST /v1/chat/completions HTTP/1.1')`
+    -- a previous body being parsed as the next request line. Two independent causes: an early
+    rejection returning before the body is read, and a streaming 200 that neither announced
+    close-framing nor stopped offering the socket for reuse when generation failed."""
+
+    CHAT = {"model": "test-model", "messages": [{"role": "user", "content": "x"}]}
+
+    def _server(self, engine=None, **kw):
+        server = APIServer(("127.0.0.1", 0), engine or FakeEngine(), "test-model", **kw)
+        thread = threading.Thread(target=server.serve_forever, daemon=True)
+        thread.start()
+        self.addCleanup(thread.join, 2)
+        self.addCleanup(server.server_close)
+        self.addCleanup(server.shutdown)
+        self.addCleanup(server.scheduler.close)
+        return server
+
+    def _conn(self, server):
+        conn = http.client.HTTPConnection("127.0.0.1", server.server_port, timeout=3)
+        self.addCleanup(conn.close)
+        return conn
+
+    def _post(self, conn, body=None, headers=None, path="/v1/chat/completions"):
+        payload = json.dumps(self.CHAT if body is None else body)
+        head = {"Content-Type": "application/json"}
+        head.update(headers or {})
+        conn.request("POST", path, body=payload, headers=head)
+        response = conn.getresponse()
+        return response.status, response.read()
+
+    def _raw(self, server, request_bytes, read=4096):
+        """Byte-level exchange, for assertions about framing that a client library hides."""
+        sock = socket.create_connection(("127.0.0.1", server.server_port), timeout=3)
+        self.addCleanup(sock.close)
+        sock.sendall(request_bytes)
+        chunks = []
+        try:
+            while True:
+                chunk = sock.recv(read)
+                if not chunk:
+                    break                      # server closed: the SSE message boundary
+                chunks.append(chunk)
+        except socket.timeout:
+            chunks.append(b"<STILL-OPEN>")     # no EOF: the connection was left reusable
+        return b"".join(chunks).decode("utf-8", "replace")
+
+    def _request_bytes(self, body, host="127.0.0.1"):
+        payload = json.dumps(body)
+        return (f"POST /v1/chat/completions HTTP/1.1\r\nHost: {host}\r\n"
+                f"Content-Type: application/json\r\nContent-Length: {len(payload)}\r\n"
+                f"\r\n{payload}").encode()
+
+    # --- an early rejection must not leave its body in the socket -------------------------
+
+    def test_rejected_host_does_not_desync_the_next_request(self):
+        server = self._server()
+        conn = self._conn(server)
+        status, _ = self._post(conn, headers={"Host": "evil.example.com"})
+        self.assertEqual(status, 403)
+        status, payload = self._post(conn)                  # same connection, valid request
+        self.assertEqual(status, 200, "the 403's unread body desynchronised the connection")
+        self.assertEqual(json.loads(payload)["object"], "chat.completion")
+
+    def test_rejected_auth_does_not_desync_the_next_request(self):
+        server = self._server(api_key="secret")
+        conn = self._conn(server)
+        status, _ = self._post(conn)                        # no Authorization header
+        self.assertEqual(status, 401)
+        status, payload = self._post(conn, headers={"Authorization": "Bearer secret"})
+        self.assertEqual(status, 200, "the 401's unread body desynchronised the connection")
+        self.assertEqual(json.loads(payload)["object"], "chat.completion")
+
+    def test_unknown_model_does_not_desync_the_next_request(self):
+        server = self._server()
+        conn = self._conn(server)
+        status, _ = self._post(conn, body=dict(self.CHAT, model="nope"))
+        self.assertEqual(status, 404)
+        status, _ = self._post(conn)
+        self.assertEqual(status, 200)
+
+    def test_each_body_is_consumed_exactly_once_across_reused_requests(self):
+        """The engine sees one prompt per request, with no body bytes bleeding between them."""
+        engine = FakeEngine()
+        server = self._server(engine)
+        conn = self._conn(server)
+        for index in range(4):
+            body = {"model": "test-model",
+                    "messages": [{"role": "user", "content": f"question-{index}"}]}
+            status, _ = self._post(conn, body=body)
+            self.assertEqual(status, 200)
+        self.assertEqual(len(engine.calls), 4)
+        for index, call in enumerate(engine.calls):
+            self.assertIn(f"question-{index}", call[0])
+            self.assertNotIn("question-", call[0].split(f"question-{index}")[1],
+                             "a later body leaked into an earlier prompt")
+
+    def test_interleaved_rejections_and_successes_stay_in_sync(self):
+        server = self._server(api_key="secret")
+        conn = self._conn(server)
+        good = {"Authorization": "Bearer secret"}
+        for _ in range(3):
+            self.assertEqual(self._post(conn)[0], 401)
+            self.assertEqual(self._post(conn, headers={"Host": "evil.example.com", **good})[0], 403)
+            self.assertEqual(self._post(conn, headers=good)[0], 200)
+
+    # --- bodies we refuse to swallow must close rather than desynchronise -----------------
+
+    def test_oversized_content_length_closes_instead_of_desyncing(self):
+        server = self._server()
+        raw = self._raw(server, (b"POST /v1/chat/completions HTTP/1.1\r\nHost: 127.0.0.1\r\n"
+                                 b"Content-Type: application/json\r\n"
+                                 b"Content-Length: 99999999\r\n\r\n{}"))
+        self.assertIn(" 400 ", raw.splitlines()[0])
+        self.assertNotIn("<STILL-OPEN>", raw,
+                         "an over-limit body must close the connection, not keep it alive")
+
+    def test_unparseable_content_length_closes_instead_of_desyncing(self):
+        server = self._server()
+        raw = self._raw(server, (b"POST /v1/chat/completions HTTP/1.1\r\nHost: 127.0.0.1\r\n"
+                                 b"Content-Type: application/json\r\n"
+                                 b"Content-Length: abc\r\n\r\n{}"))
+        self.assertIn("400", raw.splitlines()[0])
+        self.assertNotIn("<STILL-OPEN>", raw)
+
+    # --- a streaming 200 is close-framed, and says so ------------------------------------
+
+    def test_streaming_response_announces_close_framing(self):
+        server = self._server()
+        raw = self._raw(server, self._request_bytes(dict(self.CHAT, stream=True)))
+        headers = raw.split("\r\n\r\n", 1)[0].lower()
+        self.assertIn("content-type: text/event-stream", headers)
+        self.assertIn("connection: close", headers,
+                      "SSE has no Content-Length, so the close IS the boundary and must be declared")
+        self.assertIn("data: [DONE]", raw)
+        self.assertNotIn("<STILL-OPEN>", raw)
+
+    def test_anthropic_stream_announces_close_framing(self):
+        server = self._server()
+        payload = json.dumps({"model": "test-model", "stream": True, "max_tokens": 16,
+                              "messages": [{"role": "user", "content": "x"}]})
+        raw = self._raw(server, (f"POST /v1/messages HTTP/1.1\r\nHost: 127.0.0.1\r\n"
+                                 f"Content-Type: application/json\r\n"
+                                 f"Content-Length: {len(payload)}\r\n\r\n{payload}").encode())
+        self.assertIn("connection: close", raw.split("\r\n\r\n", 1)[0].lower())
+        self.assertIn("event: message_stop", raw)
+        self.assertNotIn("<STILL-OPEN>", raw)
+
+    def test_engine_failure_after_commit_does_not_splice_a_second_response(self):
+        """Once the 200 is out, a 500 status line would land inside the event stream."""
+        server = self._server(_ExplodingEngine())
+        raw = self._raw(server, self._request_bytes(dict(self.CHAT, stream=True)))
+        self.assertEqual(raw.count("HTTP/1."), 1,
+                         "a second HTTP response was spliced into the committed SSE stream")
+        self.assertIn("partial", raw)            # the events sent before the failure survive
+        self.assertNotIn("<STILL-OPEN>", raw)
+
+    def test_non_streaming_response_still_reuses_the_connection(self):
+        """The fix must not turn every response into a close: plain JSON stays persistent."""
+        server = self._server()
+        conn = self._conn(server)
+        self.assertEqual(self._post(conn)[0], 200)
+        self.assertEqual(self._post(conn)[0], 200)
+        self.assertIsNotNone(conn.sock, "the JSON path should keep the connection open")
+
+
 if __name__ == "__main__":
     unittest.main()