fix: retry and classify remote protocol disconnects

Retry one-shot OpenAI-compatible requests when httpx reports transient "Server disconnected" remote protocol errors so ZAI coding-plan traffic can recover without unnecessary failover. Map unrecovered RemoteProtocolError cases to BackendError(502) and lock behavior with targeted retry/error-mapping tests.

Made-with: Cursor

Author: Mateusz

src/connectors/openai.py

@@ -305,6 +305,20 @@ def _raise_for_httpx_request_error(
             status_code=502,
         ) from exc
 
+    if isinstance(exc, httpx.RemoteProtocolError):
+        if logger.isEnabledFor(logging.WARNING):
+            logger.warning(
+                "Upstream protocol error (remote disconnect): %s: %s",
+                url,
+                exc,
+                extra=log_extra,
+            )
+        raise BackendError(
+            message=f"Upstream protocol error: remote server disconnected ({exc!s})",
+            details={"url": url, "reason": "remote_protocol_error"},
+            status_code=502,
+        ) from exc
+
     logger.error(
         "Request failed to %s: %s",
         url,
@@ -322,7 +336,10 @@ def _is_retryable_http2_stream_termination(exc: httpx.RequestError) -> bool:
     if not isinstance(exc, httpx.RemoteProtocolError):
         return False
     message = str(exc)
-    return "ConnectionTerminated" in message and "ErrorCodes.NO_ERROR" in message
+    lowered = message.lower()
+    if "ConnectionTerminated" in message and "ErrorCodes.NO_ERROR" in message:
+        return True
+    return "server disconnected" in lowered
 
 
 class OpenAIConnector(LLMBackend):

tests/unit/connectors/test_openai_canonical.py

@@ -600,11 +600,59 @@ async def test_stream_completion_retries_once_on_http2_no_error_termination(
                 chunk async for chunk in openai_connector.stream_completion(request)
             ]
 
-        assert chunks == [
-            'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n',
-            "data: [DONE]\n\n",
-        ]
-        assert openai_connector._capture_http_client.send.await_count == 2
+        assert chunks == [
+            'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n',
+            "data: [DONE]\n\n",
+        ]
+        assert openai_connector._capture_http_client.send.await_count == 2
+
+    @pytest.mark.asyncio
+    async def test_stream_completion_retries_once_on_server_disconnected(
+        self, openai_connector
+    ):
+        request = CanonicalChatRequest(
+            model="gpt-4",
+            messages=[ChatMessage(role="user", content="Hello")],
+            max_tokens=100,
+            stream=True,
+        )
+        http_request = httpx.Request(
+            "POST",
+            "https://api.openai.com/v1/chat/completions",
+            json={"model": "gpt-4", "stream": True},
+        )
+        streamed_response = MagicMock()
+        streamed_response.status_code = 200
+        streamed_response.headers = {"content-type": "text/event-stream"}
+        streamed_response.aiter_bytes = lambda: _aiter_bytes(
+            b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n',
+            b"data: [DONE]\n\n",
+        )
+        streamed_response.aclose = AsyncMock()
+
+        openai_connector.client.build_request.return_value = http_request
+        with patch.object(
+            openai_connector,
+            "_prepare_payload",
+            new_callable=AsyncMock,
+            return_value={"model": "gpt-4", "messages": []},
+        ):
+            openai_connector._capture_http_client.send = AsyncMock(
+                side_effect=[
+                    httpx.RemoteProtocolError("Server disconnected"),
+                    streamed_response,
+                ]
+            )
+
+            chunks = [
+                chunk async for chunk in openai_connector.stream_completion(request)
+            ]
+
+        assert chunks == [
+            'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n',
+            "data: [DONE]\n\n",
+        ]
+        assert openai_connector._capture_http_client.send.await_count == 2
 
 
 class TestOpenAIPayloadCleaning:

tests/unit/openai_connector_tests/test_read_error_handling.py

@@ -113,6 +113,24 @@ def test_unknown_error_maps_to_service_unavailable_503(self) -> None:
         assert ctx.value.status_code == 503
         assert "Could not connect to backend" in ctx.value.message
 
+    def test_remote_protocol_error_maps_to_backend_error_502(self) -> None:
+        """httpx.RemoteProtocolError should map to BackendError(502)."""
+        exc = httpx.RemoteProtocolError(
+            "Server disconnected",
+            request=httpx.Request("POST", "https://example.com/v1/chat/completions"),
+        )
+
+        with pytest.raises(BackendError) as ctx:
+            _raise_for_httpx_request_error(
+                exc,
+                url="https://example.com/v1/chat/completions",
+                log_extra=None,
+            )
+
+        assert ctx.value.status_code == 502
+        assert ctx.value.details.get("reason") == "remote_protocol_error"
+        assert "remote server disconnected" in ctx.value.message.lower()
+
 
 class TestRetryableHttp2StreamTermination:
     """Test HTTP/2 stream termination retry detection."""
@@ -130,15 +148,14 @@ def test_read_error_is_not_retryable(self) -> None:
 
         assert not _is_retryable_http2_stream_termination(exc)
 
-    def test_remote_protocol_error_with_no_error_is_retryable(self) -> None:
-        """httpx.RemoteProtocolError with NO_ERROR termination is retryable."""
+    def test_remote_protocol_error_with_server_disconnected_is_retryable(self) -> None:
+        """httpx.RemoteProtocolError with server disconnected message is retryable."""
         exc = httpx.RemoteProtocolError(
-            "Server disconnected without sending a response.",
+            "Server disconnected",
             request=httpx.Request("POST", "https://example.com/v1/chat/completions"),
         )
 
-        # This is NOT retryable because it doesn't have ConnectionTerminated/NO_ERROR
-        assert not _is_retryable_http2_stream_termination(exc)
+        assert _is_retryable_http2_stream_termination(exc)
 
     def test_remote_protocol_error_with_graceful_termination_is_retryable(self) -> None:
         """httpx.RemoteProtocolError with graceful HTTP/2 termination is retryable."""
@@ -147,5 +164,4 @@ def test_remote_protocol_error_with_graceful_termination_is_retryable(self) -> N
             "Server disconnected without sending a response. ConnectionTerminated",
             request=httpx.Request("POST", "https://example.com/v1/chat/completions"),
         )
-        # Note: This doesn't match the full pattern because NO_ERROR is missing
-        assert not _is_retryable_http2_stream_termination(exc)
+        assert _is_retryable_http2_stream_termination(exc)