Quiet expected client disconnects

Author: gaoguobin

src/codex_fast_proxy/proxy.py

@@ -35,6 +35,7 @@
 BODY_METHODS = {"POST", "PUT", "PATCH"}
 RESPONSES_PATH = "/v1/responses"
 HEALTH_PATH = "/__codex_fast_proxy/health"
+CLIENT_DISCONNECT_ERRORS = (BrokenPipeError, ConnectionAbortedError, ConnectionResetError)
 
 
 def source_fingerprint(paths: Iterable[Path]) -> str:
@@ -274,7 +275,7 @@ def proxy(self) -> None:
                 self.forward_response(response)
             finally:
                 connection.close()
-        except BrokenPipeError:
+        except CLIENT_DISCONNECT_ERRORS:
             error_type = "client_disconnected"
         except Exception as exc:
             error_type = type(exc).__name__

tests/test_proxy.py

@@ -3,11 +3,13 @@
 import json
 import shutil
 import sys
+import threading
 import unittest
 import uuid
 from io import BytesIO
 from pathlib import Path
 from types import SimpleNamespace
+from typing import Any
 
 
 ROOT = Path(__file__).resolve().parents[1]
@@ -20,6 +22,7 @@
     safe_url_display,
 )
 from codex_fast_proxy.proxy import (  # noqa: E402
+    FastProxyHandler,
     copy_request_headers,
     dashboard_requested,
     parse_args,
@@ -50,6 +53,26 @@ def read(self, _size: int) -> bytes:
         raise AssertionError("SSE should use line-buffered reads")
 
 
+class FakeConnection:
+    def __init__(self) -> None:
+        self.closed = False
+
+    def request(self, _method: str, _path: str, body: bytes, headers: dict[str, str]) -> None:
+        self.body = body
+        self.headers = headers
+
+    def getresponse(self) -> Any:
+        return SimpleNamespace(
+            status=200,
+            getheader=lambda name, default=None: "text/event-stream"
+            if name.lower() == "content-type"
+            else default,
+        )
+
+    def close(self) -> None:
+        self.closed = True
+
+
 class ProxyPatchTests(unittest.TestCase):
     def test_injects_missing_service_tier_without_changing_payload_fields(self) -> None:
         payload = {
@@ -139,6 +162,55 @@ def test_sse_payload_bytes_are_forwarded_without_event_rewrite(self) -> None:
 
         self.assertEqual(writer.getvalue(), b"event: response.output_text.delta\ndata: {\"x\":1}\n\n")
 
+    def test_client_disconnect_during_stream_is_logged_without_bad_gateway(self) -> None:
+        temp_root = ROOT / ".test_tmp"
+        temp_root.mkdir(exist_ok=True)
+        temp_dir = temp_root / f"client-disconnect-{uuid.uuid4().hex}"
+        temp_dir.mkdir()
+        try:
+            log_path = temp_dir / "fast_proxy.jsonl"
+            body = json.dumps({"model": "gpt-test", "stream": True}).encode("utf-8")
+            connection = FakeConnection()
+            handler = FastProxyHandler.__new__(FastProxyHandler)
+            handler.command = "POST"
+            handler.path = "/v1/responses"
+            handler.headers = {
+                "Content-Type": "application/json",
+                "Content-Length": str(len(body)),
+            }
+            handler.rfile = BytesIO(body)
+            handler.server = SimpleNamespace(
+                proxy_base="/v1",
+                upstream_base_path="/v1",
+                upstream_netloc="api.example.test",
+                service_tier="priority",
+                log_path=log_path,
+                log_lock=threading.Lock(),
+                verbose=False,
+                open_connection=lambda: connection,
+            )
+
+            def raise_client_disconnect(_response: Any) -> None:
+                raise ConnectionAbortedError("client closed SSE")
+
+            def fail_bad_gateway() -> None:
+                raise AssertionError("client disconnect should not be converted to a 502 response")
+
+            handler.forward_response = raise_client_disconnect
+            handler.respond_bad_gateway = fail_bad_gateway
+
+            FastProxyHandler.proxy(handler)
+            event = json.loads(log_path.read_text(encoding="utf-8"))
+        finally:
+            shutil.rmtree(temp_dir)
+
+        self.assertTrue(connection.closed)
+        self.assertEqual(event["status"], 200)
+        self.assertEqual(event["error_type"], "client_disconnected")
+        self.assertEqual(event["service_tier_before"], "<absent>")
+        self.assertEqual(event["service_tier_after"], "priority")
+        self.assertTrue(event["service_tier_injected"])
+
     def test_foreground_proxy_default_log_dir_uses_runtime_state_dir(self) -> None:
         args = parse_args(["--upstream-base", "https://api.example.test/v1"])
         expected_suffixes = (