From 10b2b11efae6c1c558ad4d808d8ef6232158746f Mon Sep 17 00:00:00 2001
From: kshitij <82637225+kshitijk4poor@users.noreply.github.com>
Date: Wed, 12 Aug 2026 12:34:59 +0530
Subject: [PATCH] fix: widen APIConnectionError handling to finalization drain
 loop

The PR added APIConnectionError handling to the main request and
iteration try blocks but missed the finalization drain loop (line ~1492).
That site catches httpx transport errors to preserve an already-completed,
already-billed response when the drain iterator fails. Without the
APIConnectionError handler, an SDK-wrapped transport error during drain
would propagate uncaught and discard the completed response.

Also strengthens the test's no-payload-leak assertion to check the full
request body and URL are absent from the log message, not just the
literal string 'payload'.
---
 agent/codex_runtime.py                        | 20 ++++++++++++++++++-
 ...est_codex_request_transport_diagnostics.py |  2 ++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/agent/codex_runtime.py b/agent/codex_runtime.py
index 65cc3cd1a0..e098751d98 100644
--- a/agent/codex_runtime.py
+++ b/agent/codex_runtime.py
@@ -1489,7 +1489,25 @@ def run_codex_stream(agent, api_kwargs: dict, client: Any = None, on_first_delta
                 try:
                     for _ignored in event_stream:
                         pass
-                except (_httpx.RemoteProtocolError, _httpx.ReadTimeout, _httpx.ConnectError, ConnectionError) as exc:
+                except (
+                    _httpx.RemoteProtocolError,
+                    _httpx.ReadTimeout,
+                    _httpx.ConnectError,
+                    ConnectionError,
+                ) as exc:
+                    logger.warning(
+                        "Codex Responses stream transport finalization failed "
+                        "after a terminal response was already received; "
+                        "returning the completed response instead of "
+                        "retrying. %s error=%s",
+                        agent._client_log_context(), exc,
+                    )
+                except _APIConnectionError as exc:
+                    _log_codex_request_failure(
+                        agent,
+                        exc,
+                        stream_opened=writer_token["value"] is not None,
+                    )
                     logger.warning(
                         "Codex Responses stream transport finalization failed "
                         "after a terminal response was already received; "
diff --git a/tests/agent/test_codex_request_transport_diagnostics.py b/tests/agent/test_codex_request_transport_diagnostics.py
index 0cb4567d2a..0a65155baf 100644
--- a/tests/agent/test_codex_request_transport_diagnostics.py
+++ b/tests/agent/test_codex_request_transport_diagnostics.py
@@ -59,3 +59,5 @@ def test_transport_failure_logs_exact_request_bytes_and_class_chain(caplog):
     assert "stream_opened=false" in message
     assert "exception_chain=APIConnectionError <- RemoteProtocolError" in message
     assert "payload" not in message
+    assert request_content.decode() not in message
+    assert "example.invalid" not in message