Add explicit cancellation logic to OpenAI streaming responses in litert-lm serve.

LiteRT-LM-PiperOrigin-RevId: 908385328

Author: advaitjain

python/litert_lm_cli/serve.py

@@ -386,7 +386,6 @@ def do_POST(self) -> None:  # pylint: disable=invalid-name
           )
           return
 
-        # TODO: b/507147993 - Handle client early disconnects robustly.
         # Handle streaming response using Server-Sent Events (SSE).
         # We send response.created, response.output_text.delta, and
         # response.completed events.
@@ -433,6 +432,7 @@ def do_POST(self) -> None:  # pylint: disable=invalid-name
           self.wfile.flush()
         except Exception as e:
           click.echo(click.style(f"Error during streaming: {e!r}", fg="red"))
+          conv.cancel_process()
           try:
             self.wfile.write(
                 "event: response.error\ndata:"

python/litert_lm_cli/serve_openai_streaming_test.py

@@ -1,7 +1,9 @@
 import collections.abc
+import http.client
 import http.server
 import json
 import pathlib
+import socket
 import threading
 from unittest import mock
 import urllib.request
@@ -55,9 +57,6 @@ def tearDown(self):
     super().tearDown()
 
   def test_openai_responses_streaming(self):
-    self.assertTrue(
-        self.model_path.exists(), f"Model not found at {self.model_path}"
-    )
 
     mock_from_id = self.enter_context(
         mock.patch.object(model.Model, "from_model_id", autospec=True)
@@ -114,6 +113,62 @@ def test_openai_responses_streaming(self):
       with self.subTest(name="Verify DONE message"):
         self.assertIn("data: [DONE]", lines)
 
+  def test_openai_responses_streaming_client_disconnect(self):
+
+    mock_from_id = self.enter_context(
+        mock.patch.object(model.Model, "from_model_id", autospec=True)
+    )
+    mock_from_id.return_value = model.Model(
+        model_id="gemma3", model_path=str(self.model_path)
+    )
+
+    data = json.dumps(
+        {"model": "gemma3", "input": "Count to 50", "stream": True}
+    ).encode("utf-8")
+
+    req = urllib.request.Request(
+        f"http://localhost:{self.port}/v1/responses",
+        data=data,
+        headers={"Content-Type": "application/json"},
+    )
+
+    response = urllib.request.urlopen(req, timeout=60)
+    self.assertEqual(response.getcode(), 200)
+
+    for line in response:
+      line_str = line.decode("utf-8")
+      if line_str.startswith("event: response.output_text.delta"):
+        data_line = next(response).decode("utf-8")
+        self.assertStartsWith(data_line, "data: ")
+        break
+    else:
+      self.fail("Stream ended early without delta event")
+
+    # This tests a scenario where a client makes a request and exits before the
+    # response is completed. Note: this assumes prefill is already complete.
+    # TODO: b/508348544 - There are other scenarios where a client can cause the
+    # server to hang.
+    response.close()
+
+    conn = http.client.HTTPConnection("localhost", self.port, timeout=15)
+    try:
+      conn.request(
+          "POST",
+          "/v1/responses",
+          body=json.dumps({"model": "gemma3", "input": "Hi"}).encode("utf-8"),
+          headers={"Content-Type": "application/json"},
+      )
+      try:
+        response2 = conn.getresponse()
+      except Exception as e:
+        self.fail(f"Second request failed (timed out as expected?): {e!r}")
+
+      self.assertEqual(response2.status, 200)
+      res_body2 = json.loads(response2.read().decode("utf-8"))
+      self.assertIn("id", res_body2)
+    finally:
+      conn.close()
+
 
 if __name__ == "__main__":
   absltest.main()