From c3da68e6cf904ba7b9487bb36c64f51ef91cf399 Mon Sep 17 00:00:00 2001 From: Katia Bazzi Date: Wed, 10 Jun 2026 21:27:17 -0700 Subject: [PATCH] Preserve Guardian stream-disconnect classification --- codex-rs/core/src/guardian/review.rs | 18 +++++++ codex-rs/core/src/guardian/review_session.rs | 54 ++++++++++++++++++++ codex-rs/protocol/src/error.rs | 3 ++ codex-rs/protocol/src/error_tests.rs | 21 ++++++++ 4 files changed, 96 insertions(+) diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index 99f9c5d999..0dcfdab6fa 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -910,6 +910,7 @@ fn should_retry_guardian_review(outcome: &GuardianReviewOutcome) -> bool { #[cfg(test)] mod review_tests { use super::*; + use codex_protocol::error::CodexErr; use std::time::Duration; #[test] @@ -1014,6 +1015,23 @@ mod review_tests { } } + #[test] + fn guardian_review_retries_codex_stream_disconnect_errors() { + let error_event = CodexErr::Stream( + "WebSocket protocol error: Connection reset without closing handshake".to_string(), + None, + ) + .to_error_event(Some("Error running remote compact task".to_string())); + let outcome = GuardianReviewOutcome::Error(GuardianReviewError::session_with_error_info( + anyhow::anyhow!(error_event.message), + error_event + .codex_error_info + .expect("stream disconnect should preserve structured error info"), + )); + + assert!(should_retry_guardian_review(&outcome)); + } + #[tokio::test] async fn guardian_review_retry_wait_honors_cancellation() { let cancel_token = CancellationToken::new(); diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 63b778838e..b3cb020be9 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -1066,6 +1066,7 @@ async fn interrupt_and_drain_turn(codex: &Codex, expected_turn_id: &str) -> anyh #[cfg(test)] mod tests { use super::*; + use codex_protocol::error::CodexErr; use codex_protocol::protocol::AgentStatus; use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::Submission; @@ -1652,6 +1653,59 @@ mod tests { assert!(capture_token_usage); } + #[tokio::test] + async fn wait_for_guardian_review_preserves_remote_compact_stream_disconnect() { + let (review_session, tx_event, _rx_sub) = test_review_session().await; + let stream_error = CodexErr::Stream( + "WebSocket protocol error: Connection reset without closing handshake".to_string(), + None, + ); + tx_event + .send(Event { + id: "current-turn".to_string(), + msg: EventMsg::Error( + stream_error + .to_error_event(Some("Error running remote compact task".to_string())), + ), + }) + .await + .expect("queue guardian error"); + tx_event + .send(turn_complete_event( + "current-turn", + /*last_agent_message*/ None, + Some(42), + )) + .await + .expect("queue current turn completion"); + + let mut analytics_result = GuardianReviewAnalyticsResult::without_session(); + let (outcome, keep_review_session, capture_token_usage) = wait_for_guardian_review( + &review_session, + "current-turn", + tokio::time::Instant::now() + Duration::from_secs(1), + /*external_cancel*/ None, + &mut analytics_result, + ) + .await; + + let GuardianReviewSessionOutcome::SessionFailed { error, error_info } = outcome else { + panic!("expected structured session failure"); + }; + assert_eq!( + error.to_string(), + "Error running remote compact task: stream disconnected before completion: WebSocket protocol error: Connection reset without closing handshake" + ); + assert_eq!( + error_info, + Some(CodexErrorInfo::ResponseStreamDisconnected { + http_status_code: None, + }) + ); + assert!(keep_review_session); + assert!(capture_token_usage); + } + #[tokio::test] async fn wait_for_guardian_review_ignores_prior_turn_aborts() { let (review_session, tx_event, _rx_sub) = test_review_session().await; diff --git a/codex-rs/protocol/src/error.rs b/codex-rs/protocol/src/error.rs index d7e953af0b..2dbe2a532f 100644 --- a/codex-rs/protocol/src/error.rs +++ b/codex-rs/protocol/src/error.rs @@ -228,6 +228,9 @@ impl CodexErr { CodexErr::RetryLimit(_) => CodexErrorInfo::ResponseTooManyFailedAttempts { http_status_code: self.http_status_code_value(), }, + CodexErr::Stream(..) => CodexErrorInfo::ResponseStreamDisconnected { + http_status_code: self.http_status_code_value(), + }, CodexErr::ConnectionFailed(_) => CodexErrorInfo::HttpConnectionFailed { http_status_code: self.http_status_code_value(), }, diff --git a/codex-rs/protocol/src/error_tests.rs b/codex-rs/protocol/src/error_tests.rs index bd7a5f0058..f234d6336f 100644 --- a/codex-rs/protocol/src/error_tests.rs +++ b/codex-rs/protocol/src/error_tests.rs @@ -190,6 +190,27 @@ fn to_error_event_handles_response_stream_failed() { ); } +#[test] +fn to_error_event_handles_response_stream_disconnected() { + let err = CodexErr::Stream( + "WebSocket protocol error: Connection reset without closing handshake".to_string(), + None, + ); + + let event = err.to_error_event(Some("Error running remote compact task".to_string())); + + assert_eq!( + event.message, + "Error running remote compact task: stream disconnected before completion: WebSocket protocol error: Connection reset without closing handshake" + ); + assert_eq!( + event.codex_error_info, + Some(CodexErrorInfo::ResponseStreamDisconnected { + http_status_code: None, + }) + ); +} + #[test] fn sandbox_denied_reports_exit_code_when_no_output_available() { let output = ExecToolCallOutput {