Merge pull request #1752 from gooddata/fix/retry-remote-protocol-error

Fix: retry httpx.RemoteProtocolError in the SSE chat client

Author: Tomkess

packages/gooddata-eval/src/gooddata_eval/core/chat/sse_client.py

@@ -89,6 +89,12 @@ def _is_retryable_exc(exc: Exception) -> bool:
         return True
     if isinstance(exc, httpx.HTTPStatusError):
         return exc.response.status_code in _RETRYABLE_STATUS_CODES
+    if isinstance(exc, httpx.RemoteProtocolError):
+        # Mid-stream disconnect ("peer closed connection without sending complete
+        # message body") -- pure network flake, not a real agent/content failure.
+        # Confirmed live: contaminated ~1-4% of visualization runs with a hard
+        # fail and zero retry attempts.
+        return True
     return False
 
 
@@ -223,7 +229,14 @@ def parse_sse_lines(lines: Iterable[str]) -> ChatResult:
             # Only a transport-level failure (e.g. connection drop mid-stream) is rescued
             # here -- a bug in the processing below must propagate uncaught, not get
             # mislabeled as a network error.
-            raise ChatError(f"SSE stream error: {exc}", partial_result=_build_chat_result(acc)) from exc
+            partial = _build_chat_result(acc)
+            if isinstance(exc, httpx.RemoteProtocolError):
+                # Same mid-stream disconnect _is_retryable_exc already retries when it happens
+                # at connect time -- here it surfaces from `next(it)` instead, so it must be
+                # raised as TransientChatError or the wrapping below would mask it as
+                # non-retryable and defeat the retry this class exists for.
+                raise TransientChatError(f"SSE stream error: {exc}", partial_result=partial) from exc
+            raise ChatError(f"SSE stream error: {exc}", partial_result=partial) from exc
         line = raw_line.decode("utf-8") if isinstance(raw_line, bytes) else raw_line
         if not line:
             current_event = "message"  # blank line ends one event block per the SSE spec

packages/gooddata-eval/tests/test_sse_client.py

@@ -85,6 +85,28 @@ def _lines():
     assert partial.tool_call_events[0].function_name == "create_key_driver_analysis"
 
 
+def test_parse_sse_lines_remote_protocol_error_mid_stream_is_retryable():
+    # httpx.RemoteProtocolError raised from `next(it)` (mid-stream, not at connect time) must
+    # come out as TransientChatError -- otherwise _is_retryable_exc never sees the raw
+    # RemoteProtocolError (only the ChatError parse_sse_lines wraps it in) and the retry this
+    # class exists for never fires. Same partial_result guarantee as any other transport error.
+    def _lines():
+        yield (
+            'data: {"item": {"role": "assistant", "content": '
+            + json.dumps({"type": "toolCall", "callId": "c1", "name": "create_key_driver_analysis"})
+            + "}}"
+        )
+        yield ""
+        raise httpx.RemoteProtocolError("peer closed connection without sending complete message body")
+
+    with pytest.raises(TransientChatError) as ei:
+        parse_sse_lines(_lines())
+    partial = ei.value.partial_result
+    assert partial is not None
+    assert len(partial.tool_call_events) == 1
+    assert partial.tool_call_events[0].function_name == "create_key_driver_analysis"
+
+
 def test_parse_sse_lines_a_real_parsing_bug_propagates_uncaught_not_as_a_chat_error():
     # A malformed payload (here: "item" is a string, not a dict) crashes the processing
     # code itself with a plain AttributeError -- must surface loudly as that bug, not get
@@ -380,6 +402,26 @@ def handler(request):
     assert sleeps == [5, 10]
 
 
+def test_create_conversation_retries_remote_protocol_error(monkeypatch):
+    # "peer closed connection without sending complete message body" -- a pure
+    # network flake, not a real agent/content failure. Previously not retried
+    # at all: hard-failed on the first occurrence with zero retry attempts.
+    sleeps = []
+    monkeypatch.setattr(sse_mod.time, "sleep", lambda s: sleeps.append(s))
+    calls = {"n": 0}
+
+    def handler(request):
+        calls["n"] += 1
+        if calls["n"] < 3:
+            raise httpx.RemoteProtocolError("peer closed connection without sending complete message body")
+        return httpx.Response(200, json={"conversationId": "abc"})
+
+    client = _client_with_handler(handler)
+    assert client.create_conversation() == "abc"
+    assert calls["n"] == 3
+    assert sleeps == [5, 10]
+
+
 def test_create_conversation_does_not_retry_4xx(monkeypatch):
     sleeps = []
     monkeypatch.setattr(sse_mod.time, "sleep", lambda s: sleeps.append(s))