agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

[codex] Retry streamable HTTP initialize failures (#25147)

Browse Source 
## 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
This commit is contained in:
ssetty-oai
2026-06-09 15:49:48 -07:00
committed by GitHub
Unverified
parent 9316acf9b2
commit 490340ffcf
7 changed files with 986 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs
+103 -16
View File
@@ -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<Mutex<Option<ArmedFailure>>>,
}
#[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<HeaderValue>,
content_type: Option<HeaderValue>,
body: Option<String>,
}
#[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<String>,
content_type: Option<String>,
body: Option<String>,
}
#[derive(Deserialize)]
@@ -97,7 +115,7 @@ struct EchoArgs {
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
),
)
.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<SessionFailureState>,
State(state): State<PostFailureState>,
Json(request): Json<ArmSessionPostFailureRequest>,
) -> Result<StatusCode, StatusCode> {
arm_post_failure(state, request, ArmedFailureTarget::Session).await
}
async fn arm_initialize_post_failure(
State(state): State<PostFailureState>,
Json(request): Json<ArmSessionPostFailureRequest>,
) -> Result<StatusCode, StatusCode> {
arm_post_failure(state, request, ArmedFailureTarget::Initialize).await
}
async fn arm_initialized_notification_post_failure(
State(state): State<PostFailureState>,
Json(request): Json<ArmSessionPostFailureRequest>,
) -> Result<StatusCode, StatusCode> {
arm_post_failure(state, request, ArmedFailureTarget::InitializedNotification).await
}
async fn arm_post_failure(
state: PostFailureState,
request: ArmSessionPostFailureRequest,
target: ArmedFailureTarget,
) -> Result<StatusCode, StatusCode> {
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::<Result<Vec<_>, _>>()?;
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<SessionFailureState>,
async fn fail_mcp_post_when_armed(
State(state): State<PostFailureState>,
request: Request<Body>,
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<String> {
serde_json::from_slice::<Value>(body)
.ok()?
.get("method")?
.as_str()
.map(ToString::to_string)
}
codex-rs/rmcp-client/src/http_client_adapter.rs
+71 -1
View File
@@ -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 &notification.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<ServerJsonRpcMessage> {
match serde_json::from_slice::<ServerJsonRpcMessage>(body) {
Ok(message @ JsonRpcMessage::Error(_)) => Some(message),
_ => None,
}
}
async fn collect_body(
body_stream: &mut HttpResponseBodyStream,
) -> std::result::Result<Vec<u8>, StreamableHttpError<StreamableHttpClientAdapterError>> {
codex-rs/rmcp-client/src/rmcp_client.rs
+137 -17
View File
@@ -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<Duration>,
deadline: Option<Instant>,
) -> std::result::Result<Option<Duration>, 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<Output = std::result::Result<T, rmcp::service::ServiceError>>,
{
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<T, F, Fut>(
service: Arc<RunningService<RoleClient, ElicitationClientService>>,
label: &str,
timeout: Option<Duration>,
pause_state: ElicitationPauseState,
operation: &F,
) -> std::result::Result<T, ClientOperationError>
where
F: Fn(Arc<RunningService<RoleClient, ElicitationClientService>>) -> Fut,
Fut: std::future::Future<Output = std::result::Result<T, rmcp::service::ServiceError>>,
{
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<T, F, Fut>(
service: Arc<RunningService<RoleClient, ElicitationClientService>>,
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::<StreamableHttpError<StreamableHttpClientAdapterError>>()
.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;
codex-rs/rmcp-client/src/streamable_http_retry.rs
+239
View File
@@ -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<Duration>,
) -> Result<(
Arc<RunningService<RoleClient, ElicitationClientService>>,
Option<OAuthPersistor>,
)> {
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::<HandshakeError>()
.is_some_and(|error| Self::is_retryable_client_initialize_error(&error.source))
|| source
.downcast_ref::<rmcp::service::ClientInitializeError>()
.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::<StreamableHttpError<StreamableHttpClientAdapterError>>()
.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::<StreamableHttpError<StreamableHttpClientAdapterError>>()
.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<StreamableHttpClientAdapterError>,
) -> 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::<String>();
let Ok(status) = status_code.parse::<u16>() 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<Duration>,
deadline: Option<Instant>,
) -> Result<Option<Duration>> {
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<Duration>, 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<Instant>) -> 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;
codex-rs/rmcp-client/src/streamable_http_retry_tests.rs
+74
View File
@@ -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(),
}
}
codex-rs/rmcp-client/tests/streamable_http_recovery.rs
+248 -5
View File
@@ -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<dyn HttpClient>,
failures_remaining: Arc<AtomicUsize>,
initialize_attempts: Arc<AtomicUsize>,
}
impl FailFirstInitializeHttpClient {
fn new(inner: Arc<dyn HttpClient>, 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<HttpRequestResponse, ExecServerError>> {
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(&params) {
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::<Value>(&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?;
codex-rs/rmcp-client/tests/streamable_http_test_support.rs
+114 -1
View File
@@ -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<PathBuf, CargoBinError> {
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<RmcpClient> {
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<dyn HttpClient>,
) -> anyhow::Result<RmcpClient> {
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<RmcpClient>
/*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();