diff --git a/c/openai_server.py b/c/openai_server.py index 84a26574..496de3c0 100644 --- a/c/openai_server.py +++ b/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") diff --git a/c/tests/test_openai_server.py b/c/tests/test_openai_server.py index 6cdf3cb9..fb8235a3 100644 --- a/c/tests/test_openai_server.py +++ b/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"") # 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("", 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("", 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("", 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("", 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("", 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()