From 490340ffcf593518bcd843ee903d1ea2160d5363 Mon Sep 17 00:00:00 2001 From: ssetty-oai Date: Tue, 9 Jun 2026 15:49:48 -0700 Subject: [PATCH] [codex] Retry streamable HTTP initialize failures (#25147) ## Summary - Retry transient streamable HTTP failures during RMCP startup when the failure happens while sending the initialize request. - Retry transient streamable HTTP failures for tools/list, which is read-only and safe to replay. - Cover both retryable HTTP statuses and request-layer failures where no HTTP status is returned. - Surface retryable HTTP statuses from the streamable HTTP adapter as typed client errors. - Add integration coverage for initialize retry, tools/list retry, no-status request failure retry, and non-retryable initialize status. ## Root cause The observed codex_apps failures can happen before normal tool execution: RMCP startup fails while sending initialize, or the first read-only tools/list fails after startup. Retrying hosted_apps_bridge tools/call would not cover initialize and would risk replaying side-effecting tool calls. This change retries the streamable HTTP handshake itself, recreates the transport between initialize attempts, and retries only tools/list among post-initialize service operations. ## Validation - cargo fmt --package codex-rmcp-client - cargo test -p codex-rmcp-client --test streamable_http_recovery --- .../src/bin/test_streamable_http_server.rs | 119 ++++++-- .../rmcp-client/src/http_client_adapter.rs | 72 ++++- codex-rs/rmcp-client/src/rmcp_client.rs | 154 +++++++++-- .../rmcp-client/src/streamable_http_retry.rs | 239 +++++++++++++++++ .../src/streamable_http_retry_tests.rs | 74 +++++ .../tests/streamable_http_recovery.rs | 253 +++++++++++++++++- .../tests/streamable_http_test_support.rs | 115 +++++++- 7 files changed, 986 insertions(+), 40 deletions(-) create mode 100644 codex-rs/rmcp-client/src/streamable_http_retry.rs create mode 100644 codex-rs/rmcp-client/src/streamable_http_retry_tests.rs diff --git a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs index 2384d3947..5232c04dd 100644 --- a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs @@ -8,6 +8,7 @@ use std::time::Duration; use axum::Router; use axum::body::Body; +use axum::body::to_bytes; use axum::extract::Json; use axum::extract::State; use axum::http::HeaderMap; @@ -48,6 +49,7 @@ use rmcp::transport::StreamableHttpServerConfig; use rmcp::transport::StreamableHttpService; use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; use serde::Deserialize; +use serde_json::Value; use serde_json::json; use tokio::sync::Mutex; use tokio::task; @@ -64,18 +66,32 @@ const MEMO_URI: &str = "memo://codex/example-note"; const MEMO_CONTENT: &str = "This is a sample MCP resource served by the rmcp test server."; const MCP_SESSION_ID_HEADER: &str = "mcp-session-id"; const SESSION_POST_FAILURE_CONTROL_PATH: &str = "/test/control/session-post-failure"; +const INITIALIZE_POST_FAILURE_CONTROL_PATH: &str = "/test/control/initialize-post-failure"; +const INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH: &str = + "/test/control/initialized-notification-post-failure"; +const MAX_MCP_POST_BODY_BYTES: usize = 1024 * 1024; #[derive(Clone, Default)] -struct SessionFailureState { +struct PostFailureState { armed_failure: Arc>>, } +#[derive(Clone, Copy, Debug)] +enum ArmedFailureTarget { + Initialize, + InitializedNotification, + Session, +} + #[derive(Clone, Debug)] struct ArmedFailure { + target: ArmedFailureTarget, status: StatusCode, remaining: usize, /// Raw `WWW-Authenticate` challenge header field values returned with the failure. www_authenticate_headers: Vec, + content_type: Option, + body: Option, } #[derive(Debug, Deserialize)] @@ -85,6 +101,8 @@ struct ArmSessionPostFailureRequest { /// Raw `WWW-Authenticate` challenge header field values to add to the failure. #[serde(default)] www_authenticate_headers: Vec, + content_type: Option, + body: Option, } #[derive(Deserialize)] @@ -97,7 +115,7 @@ struct EchoArgs { #[tokio::main] async fn main() -> Result<(), Box> { let bind_addr = parse_bind_addr()?; - let session_failure_state = SessionFailureState::default(); + let post_failure_state = PostFailureState::default(); const MAX_BIND_RETRIES: u32 = 20; const BIND_RETRY_DELAY: Duration = Duration::from_millis(50); @@ -129,6 +147,14 @@ async fn main() -> Result<(), Box> { SESSION_POST_FAILURE_CONTROL_PATH, post(arm_session_post_failure), ) + .route( + INITIALIZE_POST_FAILURE_CONTROL_PATH, + post(arm_initialize_post_failure), + ) + .route( + INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH, + post(arm_initialized_notification_post_failure), + ) .route( "/.well-known/oauth-authorization-server/mcp", get({ @@ -162,10 +188,10 @@ async fn main() -> Result<(), Box> { ), ) .layer(middleware::from_fn_with_state( - session_failure_state.clone(), - fail_session_post_when_armed, + post_failure_state.clone(), + fail_mcp_post_when_armed, )) - .with_state(session_failure_state); + .with_state(post_failure_state); let router = if let Ok(token) = std::env::var("MCP_EXPECT_BEARER") { let expected = Arc::new(format!("Bearer {token}")); @@ -404,8 +430,30 @@ async fn require_bearer( } async fn arm_session_post_failure( - State(state): State, + State(state): State, Json(request): Json, +) -> Result { + arm_post_failure(state, request, ArmedFailureTarget::Session).await +} + +async fn arm_initialize_post_failure( + State(state): State, + Json(request): Json, +) -> Result { + arm_post_failure(state, request, ArmedFailureTarget::Initialize).await +} + +async fn arm_initialized_notification_post_failure( + State(state): State, + Json(request): Json, +) -> Result { + arm_post_failure(state, request, ArmedFailureTarget::InitializedNotification).await +} + +async fn arm_post_failure( + state: PostFailureState, + request: ArmSessionPostFailureRequest, + target: ArmedFailureTarget, ) -> Result { let status = StatusCode::from_u16(request.status).map_err(|_| StatusCode::BAD_REQUEST)?; let www_authenticate_headers = request @@ -413,46 +461,76 @@ async fn arm_session_post_failure( .into_iter() .map(|value| HeaderValue::from_str(&value).map_err(|_| StatusCode::BAD_REQUEST)) .collect::, _>>()?; + let content_type = request + .content_type + .map(|value| HeaderValue::from_str(&value).map_err(|_| StatusCode::BAD_REQUEST)) + .transpose()?; let armed_failure = if request.remaining == 0 { None } else { Some(ArmedFailure { + target, status, remaining: request.remaining, www_authenticate_headers, + content_type, + body: request.body, }) }; *state.armed_failure.lock().await = armed_failure; Ok(StatusCode::NO_CONTENT) } -async fn fail_session_post_when_armed( - State(state): State, +async fn fail_mcp_post_when_armed( + State(state): State, request: Request, next: Next, ) -> Response { - if request.uri().path() != "/mcp" - || request.method() != Method::POST - || !request.headers().contains_key(MCP_SESSION_ID_HEADER) - { + if request.uri().path() != "/mcp" || request.method() != Method::POST { return next.run(request).await; } + let (parts, body) = request.into_parts(); + let body_bytes = match to_bytes(body, MAX_MCP_POST_BODY_BYTES).await { + Ok(body_bytes) => body_bytes, + Err(_) => { + let mut response = Response::new(Body::from("failed to read request body")); + *response.status_mut() = StatusCode::BAD_REQUEST; + return response; + } + }; + let has_session_id = parts.headers.contains_key(MCP_SESSION_ID_HEADER); + let mcp_method = request_mcp_method(&body_bytes); { let mut armed_failure = state.armed_failure.lock().await; if let Some(failure) = armed_failure.as_mut() && failure.remaining > 0 + && match failure.target { + ArmedFailureTarget::Initialize => !has_session_id, + ArmedFailureTarget::InitializedNotification => { + has_session_id && mcp_method.as_deref() == Some("notifications/initialized") + } + ArmedFailureTarget::Session => { + has_session_id && mcp_method.as_deref() != Some("notifications/initialized") + } + } { failure.remaining -= 1; let status = failure.status; let www_authenticate_headers = failure.www_authenticate_headers.clone(); + let content_type = failure.content_type.clone(); + let body = failure + .body + .clone() + .unwrap_or_else(|| format!("forced session failure with status {status}")); if failure.remaining == 0 { *armed_failure = None; } - let mut response = Response::new(Body::from(format!( - "forced session failure with status {status}" - ))); + let mut response = Response::new(Body::from(body)); *response.status_mut() = status; + if let Some(content_type) = content_type { + response.headers_mut().insert(CONTENT_TYPE, content_type); + } for www_authenticate_header in www_authenticate_headers { response .headers_mut() @@ -462,5 +540,14 @@ async fn fail_session_post_when_armed( } } - next.run(request).await + next.run(Request::from_parts(parts, Body::from(body_bytes))) + .await +} + +fn request_mcp_method(body: &[u8]) -> Option { + serde_json::from_slice::(body) + .ok()? + .get("method")? + .as_str() + .map(ToString::to_string) } diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 6f98f7892..a9cdc7b88 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -28,6 +28,8 @@ use reqwest::header::CONTENT_TYPE; use reqwest::header::HeaderMap; use reqwest::header::HeaderName; use rmcp::model::ClientJsonRpcMessage; +use rmcp::model::ClientNotification; +use rmcp::model::ConstString; use rmcp::model::JsonRpcMessage; use rmcp::model::ServerJsonRpcMessage; use rmcp::transport::streamable_http_client::AuthRequiredError; @@ -185,6 +187,25 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { let content_type = response_header(&response.headers, CONTENT_TYPE); let session_id = response_header(&response.headers, HEADER_SESSION_ID); + if !status_is_success(response.status) { + let body = collect_body(&mut body_stream).await?; + if !retryable_post_response_status(mcp_method.as_deref(), response.status) + && content_type + .as_deref() + .is_some_and(|content_type| content_type.starts_with(JSON_MIME_TYPE)) + && let Some(message) = parse_json_rpc_error(&body) + { + return Ok(StreamableHttpPostResponse::Json(message, session_id)); + } + return Err(StreamableHttpError::UnexpectedServerResponse( + format!( + "HTTP {}: {}", + response.status, + body_preview(String::from_utf8_lossy(&body).to_string()) + ) + .into(), + )); + } match content_type.as_deref() { Some(content_type) if content_type.starts_with(EVENT_STREAM_MIME_TYPE) => { let event_stream = sse_stream_from_body(body_stream); @@ -380,7 +401,26 @@ fn client_jsonrpc_message_fields( Some(request.id.to_string()), ), JsonRpcMessage::Response(response) => (None, Some(response.id.to_string())), - JsonRpcMessage::Notification(_) => (None, None), + JsonRpcMessage::Notification(notification) => { + let method = match ¬ification.notification { + ClientNotification::CancelledNotification(notification) => { + notification.method.as_str() + } + ClientNotification::ProgressNotification(notification) => { + notification.method.as_str() + } + ClientNotification::InitializedNotification(notification) => { + notification.method.as_str() + } + ClientNotification::RootsListChangedNotification(notification) => { + notification.method.as_str() + } + ClientNotification::CustomNotification(notification) => { + notification.method.as_str() + } + }; + (Some(method.to_string()), None) + } JsonRpcMessage::Error(error) => (None, error.id.as_ref().map(ToString::to_string)), } } @@ -463,6 +503,36 @@ fn status_is_success(status: u16) -> bool { StatusCode::from_u16(status).is_ok_and(|status| status.is_success()) } +fn retryable_post_response_status(mcp_method: Option<&str>, status: u16) -> bool { + let Ok(status) = StatusCode::from_u16(status) else { + return false; + }; + is_retryable_http_status(status) + && matches!( + mcp_method, + Some("initialize" | "notifications/initialized" | "tools/list") + ) +} + +fn is_retryable_http_status(status: StatusCode) -> bool { + matches!( + status, + StatusCode::REQUEST_TIMEOUT + | StatusCode::TOO_MANY_REQUESTS + | StatusCode::INTERNAL_SERVER_ERROR + | StatusCode::BAD_GATEWAY + | StatusCode::SERVICE_UNAVAILABLE + | StatusCode::GATEWAY_TIMEOUT + ) +} + +fn parse_json_rpc_error(body: &[u8]) -> Option { + match serde_json::from_slice::(body) { + Ok(message @ JsonRpcMessage::Error(_)) => Some(message), + _ => None, + } +} + async fn collect_body( body_stream: &mut HttpResponseBodyStream, ) -> std::result::Result, StreamableHttpError> { diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 90b09d724..a28403e67 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -74,6 +74,13 @@ use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use codex_config::types::OAuthCredentialsStoreMode; +#[path = "streamable_http_retry.rs"] +mod streamable_http_retry; + +use self::streamable_http_retry::HandshakeError; +use self::streamable_http_retry::STREAMABLE_HTTP_RETRY_DELAYS_MS; +use self::streamable_http_retry::sleep_with_retry_deadline; + enum PendingTransport { InProcess { transport: tokio::io::DuplexStream, @@ -223,6 +230,25 @@ enum ClientOperationError { Timeout { label: String, duration: Duration }, } +fn remaining_operation_timeout( + label: &str, + timeout: Option, + deadline: Option, +) -> std::result::Result, ClientOperationError> { + let Some(deadline) = deadline else { + return Ok(None); + }; + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + Err(ClientOperationError::Timeout { + label: label.to_string(), + duration: timeout.unwrap_or(remaining), + }) + } else { + Ok(Some(remaining)) + } +} + pub type Elicitation = CreateElicitationRequestParams; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -396,9 +422,13 @@ impl RmcpClient { } }; - let (service, oauth_persistor) = - Self::connect_pending_transport(pending_transport, client_service.clone(), timeout) - .await?; + let (service, oauth_persistor) = self + .connect_pending_transport_with_initialize_retries( + pending_transport, + client_service.clone(), + timeout, + ) + .await?; let initialize_result_rmcp = service .peer() @@ -845,14 +875,31 @@ impl RmcpClient { ), }; - let service = match timeout { - Some(duration) => time::timeout(duration, transport) - .await - .map_err(|_| anyhow!("timed out handshaking with MCP server after {duration:?}"))? - .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?, + let service_result = match timeout { + Some(duration) => match time::timeout(duration, transport).await { + Ok(result) => { + result.map_err(|source| anyhow::Error::from(HandshakeError { source })) + } + Err(_elapsed) => Err(anyhow!( + "timed out handshaking with MCP server after {duration:?}" + )), + }, None => transport .await - .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?, + .map_err(|source| anyhow::Error::from(HandshakeError { source })), + }; + let service = match service_result { + Ok(service) => service, + Err(error) => { + if let Some(runtime) = oauth_persistor.as_ref() + && let Err(persist_error) = runtime.persist_if_needed().await + { + warn!( + "failed to persist OAuth tokens after failed initialize: {persist_error}" + ); + } + return Err(error); + } }; Ok((Arc::new(service), oauth_persistor)) @@ -869,7 +916,7 @@ impl RmcpClient { Fut: std::future::Future>, { let service = self.service().await?; - match Self::run_service_operation_once( + match Self::run_service_operation_with_transient_retries( Arc::clone(&service), label, timeout, @@ -882,7 +929,7 @@ impl RmcpClient { Err(error) if Self::is_session_expired_404(&error) => { self.reinitialize_after_session_expiry(&service).await?; let recovered_service = self.service().await?; - Self::run_service_operation_once( + Self::run_service_operation_with_transient_retries( recovered_service, label, timeout, @@ -896,6 +943,62 @@ impl RmcpClient { } } + async fn run_service_operation_with_transient_retries( + service: Arc>, + label: &str, + timeout: Option, + pause_state: ElicitationPauseState, + operation: &F, + ) -> std::result::Result + where + F: Fn(Arc>) -> Fut, + Fut: std::future::Future>, + { + let retry_deadline = timeout.map(|duration| Instant::now() + duration); + for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS + .iter() + .copied() + .map(Some) + .chain(std::iter::once(None)) + .enumerate() + { + let attempt_timeout = remaining_operation_timeout(label, timeout, retry_deadline)?; + match Self::run_service_operation_once( + Arc::clone(&service), + label, + attempt_timeout, + pause_state.clone(), + operation, + ) + .await + { + Ok(result) => return Ok(result), + Err(error) if Self::is_retryable_tools_list_error(label, &error) => { + let Some(retry_delay_ms) = retry_delay_ms else { + return Err(error); + }; + let delay = Duration::from_millis(retry_delay_ms); + warn!( + attempt = attempt + 1, + max_attempts = STREAMABLE_HTTP_RETRY_DELAYS_MS.len() + 1, + delay_ms = delay.as_millis(), + error = %error, + "streamable HTTP MCP tools/list failed with a retryable error; retrying" + ); + if !sleep_with_retry_deadline(delay, retry_deadline).await { + return Err(ClientOperationError::Timeout { + label: label.to_string(), + duration: timeout.unwrap_or(delay), + }); + } + } + Err(error) => return Err(error), + } + } + + unreachable!("service operation retry loop should return on success or final error") + } + async fn run_service_operation_once( service: Arc>, label: &str, @@ -921,6 +1024,22 @@ impl RmcpClient { } } + fn is_retryable_tools_list_error(label: &str, error: &ClientOperationError) -> bool { + if label != "tools/list" { + return false; + } + let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) = + error + else { + return false; + }; + + error + .error + .downcast_ref::>() + .is_some_and(Self::is_retryable_streamable_http_error) + } + fn is_session_expired_404(error: &ClientOperationError) -> bool { let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) = error @@ -974,12 +1093,13 @@ impl RmcpClient { .clone() .ok_or_else(|| anyhow!("MCP client cannot recover before initialize succeeds"))?; let pending_transport = Self::create_pending_transport(&self.transport_recipe).await?; - let (service, oauth_persistor) = Self::connect_pending_transport( - pending_transport, - initialize_context.client_service, - initialize_context.timeout, - ) - .await?; + let (service, oauth_persistor) = self + .connect_pending_transport_with_initialize_retries( + pending_transport, + initialize_context.client_service, + initialize_context.timeout, + ) + .await?; { let mut guard = self.state.lock().await; diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs new file mode 100644 index 000000000..73da95de5 --- /dev/null +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -0,0 +1,239 @@ +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Result; +use anyhow::anyhow; +use codex_exec_server::ExecServerError; +use reqwest::StatusCode; +use rmcp::service::RoleClient; +use rmcp::service::RunningService; +use rmcp::transport::streamable_http_client::StreamableHttpError; +use tokio::time; +use tracing::warn; + +use crate::elicitation_client_service::ElicitationClientService; +use crate::http_client_adapter::StreamableHttpClientAdapterError; +use crate::oauth::OAuthPersistor; + +use super::PendingTransport; +use super::RmcpClient; + +const JSON_RPC_INTERNAL_ERROR_CODE: i64 = -32603; +pub(super) const STREAMABLE_HTTP_RETRY_DELAYS_MS: [u64; 2] = [250, 1_000]; + +impl RmcpClient { + pub(super) async fn connect_pending_transport_with_initialize_retries( + &self, + initial_transport: PendingTransport, + client_service: ElicitationClientService, + timeout: Option, + ) -> Result<( + Arc>, + Option, + )> { + let should_retry = match &initial_transport { + PendingTransport::InProcess { .. } | PendingTransport::Stdio { .. } => false, + PendingTransport::StreamableHttp { .. } + | PendingTransport::StreamableHttpWithOAuth { .. } => true, + }; + let retry_deadline = timeout.map(|duration| Instant::now() + duration); + let mut pending_transport = Some(initial_transport); + + for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS + .iter() + .copied() + .map(Some) + .chain(std::iter::once(None)) + .enumerate() + { + let transport = match pending_transport.take() { + Some(transport) => transport, + None => { + let remaining = remaining_initialize_timeout(timeout, retry_deadline)?; + match remaining { + Some(remaining) => time::timeout( + remaining, + Self::create_pending_transport(&self.transport_recipe), + ) + .await + .map_err(|_| initialize_timeout_error(timeout, remaining))??, + None => Self::create_pending_transport(&self.transport_recipe).await?, + } + } + }; + let attempt_timeout = remaining_initialize_timeout(timeout, retry_deadline)?; + + match Self::connect_pending_transport( + transport, + client_service.clone(), + attempt_timeout, + ) + .await + { + Ok(result) => return Ok(result), + Err(error) if should_retry && Self::is_retryable_initialize_error(&error) => { + let Some(retry_delay_ms) = retry_delay_ms else { + return Err(error); + }; + let delay = Duration::from_millis(retry_delay_ms); + warn!( + attempt = attempt + 1, + max_attempts = STREAMABLE_HTTP_RETRY_DELAYS_MS.len() + 1, + delay_ms = delay.as_millis(), + error = %error, + "streamable HTTP MCP initialize failed with a retryable error; retrying" + ); + if !sleep_with_retry_deadline(delay, retry_deadline).await { + let duration = timeout.unwrap_or(delay); + return Err(anyhow!( + "timed out handshaking with MCP server after {duration:?}" + )); + } + } + Err(error) => return Err(error), + } + } + + unreachable!("initialize retry loop should return on success or final error") + } + + fn is_retryable_initialize_error(error: &anyhow::Error) -> bool { + error.chain().any(|source| { + source + .downcast_ref::() + .is_some_and(|error| Self::is_retryable_client_initialize_error(&error.source)) + || source + .downcast_ref::() + .is_some_and(Self::is_retryable_client_initialize_error) + }) + } + + fn is_retryable_client_initialize_error(error: &rmcp::service::ClientInitializeError) -> bool { + match error { + rmcp::service::ClientInitializeError::TransportError { error, context } + if context.as_ref() == "send initialize request" => + { + error + .error + .downcast_ref::>() + .is_some_and(Self::is_retryable_streamable_http_error) + } + rmcp::service::ClientInitializeError::TransportError { error, context } + if context.as_ref() == "send initialized notification" => + { + error + .error + .downcast_ref::>() + .is_some_and(|error| { + matches!(error, StreamableHttpError::TransportChannelClosed) + || Self::is_retryable_streamable_http_error(error) + }) + } + _ => false, + } + } + + pub(super) fn is_retryable_streamable_http_error( + error: &StreamableHttpError, + ) -> bool { + match error { + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::HttpRequest(_), + )) => true, + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::Server { code, message }, + )) => { + *code == JSON_RPC_INTERNAL_ERROR_CODE && message.starts_with("http/request failed:") + } + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::Protocol(message), + )) => message.starts_with("http response stream `") && message.contains("` failed:"), + StreamableHttpError::UnexpectedServerResponse(message) => { + is_retryable_unexpected_server_response(message.as_ref()) + } + StreamableHttpError::AuthRequired(_) + | StreamableHttpError::InsufficientScope(_) + | StreamableHttpError::SessionExpired + | StreamableHttpError::UnexpectedContentType(_) + | StreamableHttpError::ServerDoesNotSupportSse + | StreamableHttpError::Deserialize(_) + | StreamableHttpError::Client(StreamableHttpClientAdapterError::SessionExpired404) + | StreamableHttpError::Client(StreamableHttpClientAdapterError::Header(_)) => false, + _ => false, + } + } +} + +fn is_retryable_unexpected_server_response(message: &str) -> bool { + let Some(message) = message.strip_prefix("HTTP ") else { + return false; + }; + let status_code = message + .chars() + .take_while(char::is_ascii_digit) + .collect::(); + let Ok(status) = status_code.parse::() else { + return false; + }; + let Ok(status) = StatusCode::from_u16(status) else { + return false; + }; + is_retryable_http_status(status) +} + +fn is_retryable_http_status(status: StatusCode) -> bool { + matches!( + status, + StatusCode::REQUEST_TIMEOUT + | StatusCode::TOO_MANY_REQUESTS + | StatusCode::INTERNAL_SERVER_ERROR + | StatusCode::BAD_GATEWAY + | StatusCode::SERVICE_UNAVAILABLE + | StatusCode::GATEWAY_TIMEOUT + ) +} + +fn remaining_initialize_timeout( + timeout: Option, + deadline: Option, +) -> Result> { + let Some(deadline) = deadline else { + return Ok(None); + }; + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + Err(initialize_timeout_error(timeout, remaining)) + } else { + Ok(Some(remaining)) + } +} + +fn initialize_timeout_error(timeout: Option, fallback: Duration) -> anyhow::Error { + let duration = timeout.unwrap_or(fallback); + anyhow!("timed out handshaking with MCP server after {duration:?}") +} + +pub(super) async fn sleep_with_retry_deadline(delay: Duration, deadline: Option) -> bool { + if let Some(deadline) = deadline { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return false; + } + time::timeout(remaining, time::sleep(delay)).await.is_ok() + } else { + time::sleep(delay).await; + true + } +} + +#[derive(Debug, thiserror::Error)] +#[error("handshaking with MCP server failed: {source}")] +pub(super) struct HandshakeError { + #[source] + pub(super) source: rmcp::service::ClientInitializeError, +} + +#[cfg(test)] +#[path = "streamable_http_retry_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs b/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs new file mode 100644 index 000000000..ed0d8030e --- /dev/null +++ b/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs @@ -0,0 +1,74 @@ +use std::any::TypeId; + +use codex_exec_server::ExecServerError; +use pretty_assertions::assert_eq; +use rmcp::transport::DynamicTransportError; +use rmcp::transport::streamable_http_client::StreamableHttpError; + +use crate::http_client_adapter::StreamableHttpClientAdapterError; + +use super::*; + +#[test] +fn retryable_initialize_error_includes_initialized_notification_context() { + let contexts = [ + "send initialize request", + "send initialized notification", + "receive initialize response", + ]; + + assert_eq!( + contexts.map(|context| { + RmcpClient::is_retryable_client_initialize_error(&retryable_initialize_error(context)) + }), + [true, true, false], + ); +} + +#[test] +fn retryable_streamable_http_error_includes_remote_body_stream_failure() { + let errors = [ + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::HttpRequest("error sending request for url".to_string()), + )), + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::Server { + code: JSON_RPC_INTERNAL_ERROR_CODE, + message: "http/request failed: error sending request for url".to_string(), + }, + )), + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::Protocol( + "http response stream `http-1` failed: exec-server transport disconnected" + .to_string(), + ), + )), + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::Protocol( + "http response stream `http-1` received seq 2, expected 1".to_string(), + ), + )), + StreamableHttpError::UnexpectedServerResponse("HTTP 502: upstream failure".into()), + StreamableHttpError::UnexpectedServerResponse("HTTP 400: bad request".into()), + ]; + + assert_eq!( + errors.map(|error| RmcpClient::is_retryable_streamable_http_error(&error)), + [true, true, true, false, true, false], + ); +} + +fn retryable_initialize_error(context: &'static str) -> rmcp::service::ClientInitializeError { + rmcp::service::ClientInitializeError::TransportError { + error: DynamicTransportError::from_parts( + "streamable_http", + TypeId::of::<()>(), + Box::new(StreamableHttpError::Client( + StreamableHttpClientAdapterError::HttpRequest(ExecServerError::HttpRequest( + "error sending request for url".to_string(), + )), + )), + ), + context: context.into(), + } +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs index 087d3d00d..29d2404ed 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs @@ -1,13 +1,229 @@ mod streamable_http_test_support; -use pretty_assertions::assert_eq; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; +use codex_exec_server::Environment; +use codex_exec_server::ExecServerError; +use codex_exec_server::HttpClient; +use codex_exec_server::HttpRequestParams; +use codex_exec_server::HttpRequestResponse; +use codex_exec_server::HttpResponseBodyStream; +use futures::FutureExt as _; +use futures::future::BoxFuture; +use pretty_assertions::assert_eq; +use serde_json::Value; + +use streamable_http_test_support::arm_initialize_post_failure; +use streamable_http_test_support::arm_initialize_post_json_rpc_failure; +use streamable_http_test_support::arm_initialized_notification_post_json_rpc_failure; use streamable_http_test_support::arm_session_post_failure; +use streamable_http_test_support::arm_session_post_json_rpc_failure; use streamable_http_test_support::call_echo_tool; use streamable_http_test_support::create_client; +use streamable_http_test_support::create_client_with_http_client; use streamable_http_test_support::expected_echo_result; use streamable_http_test_support::spawn_streamable_http_server; +const JSON_RPC_INTERNAL_ERROR_CODE: i64 = -32603; +const SIMULATED_NO_RESPONSE_MESSAGE: &str = + "http/request failed: error sending request for url (simulated no response)"; + +#[derive(Clone)] +struct FailFirstInitializeHttpClient { + inner: Arc, + failures_remaining: Arc, + initialize_attempts: Arc, +} + +impl FailFirstInitializeHttpClient { + fn new(inner: Arc, failures_remaining: usize) -> Self { + Self { + inner, + failures_remaining: Arc::new(AtomicUsize::new(failures_remaining)), + initialize_attempts: Arc::new(AtomicUsize::new(0)), + } + } + + fn initialize_attempts(&self) -> usize { + self.initialize_attempts.load(Ordering::SeqCst) + } + + fn fail_next_initialize(&self) { + self.failures_remaining.store(1, Ordering::SeqCst); + } +} + +impl HttpClient for FailFirstInitializeHttpClient { + fn http_request( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + self.inner.http_request(params) + } + + fn http_request_stream( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { + let inner = Arc::clone(&self.inner); + let failures_remaining = Arc::clone(&self.failures_remaining); + let initialize_attempts = Arc::clone(&self.initialize_attempts); + + async move { + if is_initialize_post(¶ms) { + initialize_attempts.fetch_add(1, Ordering::SeqCst); + if failures_remaining.swap(0, Ordering::SeqCst) > 0 { + return Err(ExecServerError::Server { + code: JSON_RPC_INTERNAL_ERROR_CODE, + message: SIMULATED_NO_RESPONSE_MESSAGE.to_string(), + }); + } + } + + inner.http_request_stream(params).await + } + .boxed() + } +} + +fn is_initialize_post(params: &HttpRequestParams) -> bool { + params.method.eq_ignore_ascii_case("POST") + && params + .body + .as_ref() + .and_then(|body| serde_json::from_slice::(&body.0).ok()) + .and_then(|body| { + body.get("method") + .and_then(Value::as_str) + .map(|method| method == "initialize") + }) + .unwrap_or(false) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_initialize_retries_remote_no_response_error() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + let http_client = FailFirstInitializeHttpClient::new( + Environment::default_for_tests().get_http_client(), + /*failures_remaining*/ 1, + ); + + let client = create_client_with_http_client(&base_url, Arc::new(http_client.clone())).await?; + let result = call_echo_tool(&client, "after-init-retry").await?; + + assert_eq!(http_client.initialize_attempts(), 2); + assert_eq!(result, expected_echo_result("after-init-retry")); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_initialize_retries_transient_http_status() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + + arm_initialize_post_failure(&base_url, /*status*/ 502, /*remaining*/ 1).await?; + + let client = create_client(&base_url).await?; + let result = call_echo_tool(&client, "after-status-retry").await?; + + assert_eq!(result, expected_echo_result("after-status-retry")); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_initialize_retries_json_rpc_transient_status() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + + arm_initialize_post_json_rpc_failure(&base_url, /*status*/ 502, /*remaining*/ 1).await?; + + let client = create_client(&base_url).await?; + let result = call_echo_tool(&client, "after-json-status-retry").await?; + + assert_eq!(result, expected_echo_result("after-json-status-retry")); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_retries_initialized_notification_status() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + + arm_initialized_notification_post_json_rpc_failure( + &base_url, /*status*/ 502, /*remaining*/ 1, + ) + .await?; + + let client = create_client(&base_url).await?; + let result = call_echo_tool(&client, "after-notification-status-retry").await?; + + assert_eq!( + result, + expected_echo_result("after-notification-status-retry") + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_tools_list_retries_transient_http_status() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + let client = create_client(&base_url).await?; + + let expected = client + .list_tools( + /*params*/ None, + /*timeout*/ Some(Duration::from_secs(5)), + ) + .await?; + arm_session_post_failure( + &base_url, + /*status*/ 502, + /*remaining*/ 1, + /*www_authenticate_headers*/ &[], + ) + .await?; + + let result = client + .list_tools( + /*params*/ None, + /*timeout*/ Some(Duration::from_secs(5)), + ) + .await?; + + assert_eq!(result, expected); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_tools_list_retries_json_rpc_transient_status() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + let client = create_client(&base_url).await?; + + let expected = client + .list_tools( + /*params*/ None, + /*timeout*/ Some(Duration::from_secs(5)), + ) + .await?; + arm_session_post_json_rpc_failure(&base_url, /*status*/ 502, /*remaining*/ 1).await?; + + let result = client + .list_tools( + /*params*/ None, + /*timeout*/ Some(Duration::from_secs(5)), + ) + .await?; + + assert_eq!(result, expected); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn streamable_http_404_session_expiry_recovers_and_retries_once() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; @@ -30,6 +246,34 @@ async fn streamable_http_404_session_expiry_recovers_and_retries_once() -> anyho Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_session_recovery_retries_initialize_failure() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + let http_client = FailFirstInitializeHttpClient::new( + Environment::default_for_tests().get_http_client(), + /*failures_remaining*/ 0, + ); + let client = create_client_with_http_client(&base_url, Arc::new(http_client.clone())).await?; + + let warmup = call_echo_tool(&client, "warmup").await?; + assert_eq!(warmup, expected_echo_result("warmup")); + + arm_session_post_failure( + &base_url, + /*status*/ 404, + /*remaining*/ 1, + /*www_authenticate_headers*/ &[], + ) + .await?; + http_client.fail_next_initialize(); + + let recovered = call_echo_tool(&client, "recovered-after-retry").await?; + assert_eq!(http_client.initialize_attempts(), 3); + assert_eq!(recovered, expected_echo_result("recovered-after-retry")); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn streamable_http_401_does_not_trigger_recovery() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; @@ -129,11 +373,10 @@ async fn streamable_http_404_recovery_only_retries_once() -> anyhow::Result<()> .await?; let error = call_echo_tool(&client, "double-404").await.unwrap_err(); + let error_message = error.to_string(); assert!( - error - .to_string() - .contains("handshaking with MCP server failed") - || error.to_string().contains("Transport channel closed") + error_message.contains("404") || error_message.contains("session expired"), + "expected session-expiry error, got: {error:#}" ); let recovered = call_echo_tool(&client, "after-double-404").await?; diff --git a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs index 03ee79392..6b5df6e29 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -20,6 +20,7 @@ use anyhow::Context as _; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; use codex_exec_server::ExecServerClient; +use codex_exec_server::HttpClient; use codex_exec_server::RemoteExecServerConnectArgs; use codex_rmcp_client::ElicitationAction; use codex_rmcp_client::ElicitationResponse; @@ -44,6 +45,9 @@ use tokio::process::Command; use tokio::time::sleep; const SESSION_POST_FAILURE_CONTROL_PATH: &str = "/test/control/session-post-failure"; +const INITIALIZE_POST_FAILURE_CONTROL_PATH: &str = "/test/control/initialize-post-failure"; +const INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH: &str = + "/test/control/initialized-notification-post-failure"; fn streamable_http_server_bin() -> Result { codex_utils_cargo_bin::cargo_bin("test_streamable_http_server") @@ -74,6 +78,14 @@ pub(crate) fn expected_echo_result(message: &str) -> CallToolResult { } pub(crate) async fn create_client(base_url: &str) -> anyhow::Result { + create_client_with_http_client(base_url, Environment::default_for_tests().get_http_client()) + .await +} + +pub(crate) async fn create_client_with_http_client( + base_url: &str, + http_client: Arc, +) -> anyhow::Result { let client = RmcpClient::new_streamable_http_client( "test-streamable-http", &format!("{base_url}/mcp"), @@ -81,7 +93,7 @@ pub(crate) async fn create_client(base_url: &str) -> anyhow::Result /*http_headers*/ None, /*env_http_headers*/ None, OAuthCredentialsStoreMode::File, - Environment::default_for_tests().get_http_client(), + http_client, /*auth_provider*/ None, ) .await?; @@ -183,6 +195,107 @@ pub(crate) async fn arm_session_post_failure( Ok(()) } +pub(crate) async fn arm_session_post_json_rpc_failure( + base_url: &str, + status: u16, + remaining: usize, +) -> anyhow::Result<()> { + let response = reqwest::Client::new() + .post(format!("{base_url}{SESSION_POST_FAILURE_CONTROL_PATH}")) + .json(&json!({ + "status": status, + "remaining": remaining, + "content_type": "application/json", + "body": json!({ + "jsonrpc": "2.0", + "id": 1, + "error": { + "code": -32000, + "message": "transient session failure", + }, + }).to_string(), + })) + .send() + .await?; + + assert_eq!(response.status(), reqwest::StatusCode::NO_CONTENT); + Ok(()) +} + +pub(crate) async fn arm_initialized_notification_post_json_rpc_failure( + base_url: &str, + status: u16, + remaining: usize, +) -> anyhow::Result<()> { + let response = reqwest::Client::new() + .post(format!( + "{base_url}{INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH}" + )) + .json(&json!({ + "status": status, + "remaining": remaining, + "content_type": "application/json", + "body": json!({ + "jsonrpc": "2.0", + "id": 1, + "error": { + "code": -32000, + "message": "transient session failure", + }, + }).to_string(), + })) + .send() + .await?; + + assert_eq!(response.status(), reqwest::StatusCode::NO_CONTENT); + Ok(()) +} + +pub(crate) async fn arm_initialize_post_failure( + base_url: &str, + status: u16, + remaining: usize, +) -> anyhow::Result<()> { + let response = reqwest::Client::new() + .post(format!("{base_url}{INITIALIZE_POST_FAILURE_CONTROL_PATH}")) + .json(&json!({ + "status": status, + "remaining": remaining, + })) + .send() + .await?; + + assert_eq!(response.status(), reqwest::StatusCode::NO_CONTENT); + Ok(()) +} + +pub(crate) async fn arm_initialize_post_json_rpc_failure( + base_url: &str, + status: u16, + remaining: usize, +) -> anyhow::Result<()> { + let response = reqwest::Client::new() + .post(format!("{base_url}{INITIALIZE_POST_FAILURE_CONTROL_PATH}")) + .json(&json!({ + "status": status, + "remaining": remaining, + "content_type": "application/json", + "body": json!({ + "jsonrpc": "2.0", + "id": 1, + "error": { + "code": -32000, + "message": "transient initialize failure", + }, + }).to_string(), + })) + .send() + .await?; + + assert_eq!(response.status(), reqwest::StatusCode::NO_CONTENT); + Ok(()) +} + pub(crate) async fn spawn_streamable_http_server() -> anyhow::Result<(Child, String)> { let listener = TcpListener::bind("127.0.0.1:0")?; let port = listener.local_addr()?.port();