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 2384d39473..2b389a5258 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; @@ -52,6 +53,7 @@ use serde_json::json; use tokio::sync::Mutex; use tokio::task; use tokio::time::sleep; +use urlencoding::decode; #[derive(Clone)] struct TestToolServer { @@ -129,6 +131,7 @@ async fn main() -> Result<(), Box> { SESSION_POST_FAILURE_CONTROL_PATH, post(arm_session_post_failure), ) + .route("/oauth/token", post(exchange_refresh_token)) .route( "/.well-known/oauth-authorization-server/mcp", get({ @@ -389,7 +392,8 @@ async fn require_bearer( request: Request, next: Next, ) -> Result { - if request.uri().path().contains("/.well-known/") { + let path = request.uri().path(); + if path.contains("/.well-known/") || path.starts_with("/oauth/") { return Ok(next.run(request).await); } if request @@ -403,6 +407,64 @@ async fn require_bearer( } } +async fn exchange_refresh_token(request: Request) -> Result { + let body = to_bytes(request.into_body(), 16 * 1024) + .await + .map_err(|_| StatusCode::BAD_REQUEST)?; + let params = parse_form_body(&body)?; + + if let Ok(delay_ms) = std::env::var("MCP_REFRESH_TOKEN_DELAY_MS") + && let Ok(delay_ms) = delay_ms.parse::() + { + sleep(Duration::from_millis(delay_ms)).await; + } + + let expected_refresh_token = + std::env::var("MCP_EXPECT_REFRESH_TOKEN").map_err(|_| StatusCode::BAD_REQUEST)?; + let access_token = + std::env::var("MCP_REFRESH_ACCESS_TOKEN").map_err(|_| StatusCode::BAD_REQUEST)?; + let refresh_token = + std::env::var("MCP_ROTATED_REFRESH_TOKEN").map_err(|_| StatusCode::BAD_REQUEST)?; + + if params.get("grant_type").map(String::as_str) != Some("refresh_token") + || params.get("refresh_token").map(String::as_str) != Some(expected_refresh_token.as_str()) + { + return Err(StatusCode::UNAUTHORIZED); + } + + #[expect(clippy::expect_used)] + Ok(Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "application/json") + .body(Body::from( + serde_json::to_vec(&json!({ + "access_token": access_token, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": refresh_token, + })) + .expect("failed to serialize token response"), + )) + .expect("valid token response")) +} + +fn parse_form_body(body: &[u8]) -> Result, StatusCode> { + let body = std::str::from_utf8(body).map_err(|_| StatusCode::BAD_REQUEST)?; + body.split('&') + .filter(|part| !part.is_empty()) + .map(|part| { + let (name, value) = part.split_once('=').ok_or(StatusCode::BAD_REQUEST)?; + let name = decode(name) + .map_err(|_| StatusCode::BAD_REQUEST)? + .into_owned(); + let value = decode(value) + .map_err(|_| StatusCode::BAD_REQUEST)? + .into_owned(); + Ok((name, value)) + }) + .collect() +} + async fn arm_session_post_failure( State(state): State, Json(request): Json, diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 90b09d724c..2a6f78d6a6 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -823,6 +823,8 @@ impl RmcpClient { Arc>, Option, )> { + let connect_started_at = Instant::now(); + let mut remaining_timeout = timeout; let (transport, oauth_persistor) = match pending_transport { PendingTransport::InProcess { transport } => ( service::serve_client(client_service, transport).boxed(), @@ -839,17 +841,33 @@ impl RmcpClient { PendingTransport::StreamableHttpWithOAuth { transport, oauth_persistor, - } => ( - service::serve_client(client_service, transport).boxed(), - Some(oauth_persistor), - ), + } => { + match timeout { + Some(duration) => time::timeout(duration, oauth_persistor.refresh_if_needed()) + .await + .map_err(|_| { + anyhow!("timed out handshaking with MCP server after {duration:?}") + })??, + None => oauth_persistor.refresh_if_needed().await?, + } + remaining_timeout = + timeout.map(|duration| duration.saturating_sub(connect_started_at.elapsed())); + ( + service::serve_client(client_service, transport).boxed(), + Some(oauth_persistor), + ) + } }; 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}"))?, + Some(total_duration) => { + time::timeout(remaining_timeout.unwrap_or(total_duration), transport) + .await + .map_err(|_| { + anyhow!("timed out handshaking with MCP server after {total_duration:?}") + })? + .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))? + } None => transport .await .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?, diff --git a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs index 087d3d00df..c4a9e6036b 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs @@ -1,14 +1,40 @@ mod streamable_http_test_support; +use std::ffi::OsString; +use std::time::Duration; + +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::save_oauth_tokens; +use oauth2::AccessToken; +use oauth2::RefreshToken; +use oauth2::basic::BasicTokenType; use pretty_assertions::assert_eq; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::VendorExtraTokenFields; +use serial_test::serial; +use tempfile::TempDir; use streamable_http_test_support::arm_session_post_failure; use streamable_http_test_support::call_echo_tool; use streamable_http_test_support::create_client; use streamable_http_test_support::expected_echo_result; +use streamable_http_test_support::initialize_client; +use streamable_http_test_support::initialize_client_with_timeout; use streamable_http_test_support::spawn_streamable_http_server; +use streamable_http_test_support::spawn_streamable_http_server_with_env; + +const OAUTH_TEST_SERVER_NAME: &str = "test-streamable-http-oauth"; +const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; +const VALID_REFRESH_TOKEN: &str = "valid-refresh-token"; +const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; +const ROTATED_REFRESH_TOKEN: &str = "rotated-refresh-token"; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[serial(oauth_credentials_env)] async fn streamable_http_404_session_expiry_recovers_and_retries_once() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; let client = create_client(&base_url).await?; @@ -31,6 +57,7 @@ async fn streamable_http_404_session_expiry_recovers_and_retries_once() -> anyho } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[serial(oauth_credentials_env)] async fn streamable_http_401_does_not_trigger_recovery() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; let client = create_client(&base_url).await?; @@ -58,6 +85,7 @@ async fn streamable_http_401_does_not_trigger_recovery() -> anyhow::Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[serial(oauth_credentials_env)] async fn streamable_http_403_scope_challenge_returns_insufficient_scope() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; let client = create_client(&base_url).await?; @@ -84,6 +112,7 @@ async fn streamable_http_403_scope_challenge_returns_insufficient_scope() -> any } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[serial(oauth_credentials_env)] async fn streamable_http_403_finds_bearer_challenge_in_later_header_value() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; let client = create_client(&base_url).await?; @@ -113,6 +142,7 @@ async fn streamable_http_403_finds_bearer_challenge_in_later_header_value() -> a } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[serial(oauth_credentials_env)] async fn streamable_http_404_recovery_only_retries_once() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; let client = create_client(&base_url).await?; @@ -143,6 +173,86 @@ async fn streamable_http_404_recovery_only_retries_once() -> anyhow::Result<()> } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[serial(oauth_credentials_env)] +async fn streamable_http_oauth_refreshes_expired_token_before_initialize() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server_with_env(&[ + ("MCP_EXPECT_BEARER", REFRESHED_ACCESS_TOKEN), + ("MCP_EXPECT_REFRESH_TOKEN", VALID_REFRESH_TOKEN), + ("MCP_REFRESH_ACCESS_TOKEN", REFRESHED_ACCESS_TOKEN), + ("MCP_ROTATED_REFRESH_TOKEN", ROTATED_REFRESH_TOKEN), + ]) + .await?; + let codex_home = TempCodexHome::new()?; + let server_url = format!("{base_url}/mcp"); + save_expired_oauth_tokens(&server_url)?; + + let client = RmcpClient::new_streamable_http_client( + OAUTH_TEST_SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + initialize_client(&client).await?; + + let result = call_echo_tool(&client, "after-refresh").await?; + assert_eq!(result, expected_echo_result("after-refresh")); + + let credentials = std::fs::read_to_string(codex_home.dir.path().join(".credentials.json"))?; + assert!(credentials.contains(REFRESHED_ACCESS_TOKEN)); + assert!(credentials.contains(ROTATED_REFRESH_TOKEN)); + assert!(!credentials.contains(EXPIRED_ACCESS_TOKEN)); + assert!(!credentials.contains(VALID_REFRESH_TOKEN)); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[serial(oauth_credentials_env)] +async fn streamable_http_oauth_refresh_respects_initialize_timeout() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server_with_env(&[ + ("MCP_EXPECT_BEARER", REFRESHED_ACCESS_TOKEN), + ("MCP_EXPECT_REFRESH_TOKEN", VALID_REFRESH_TOKEN), + ("MCP_REFRESH_ACCESS_TOKEN", REFRESHED_ACCESS_TOKEN), + ("MCP_ROTATED_REFRESH_TOKEN", ROTATED_REFRESH_TOKEN), + ("MCP_REFRESH_TOKEN_DELAY_MS", "200"), + ]) + .await?; + let _codex_home = TempCodexHome::new()?; + let server_url = format!("{base_url}/mcp"); + save_expired_oauth_tokens(&server_url)?; + + let client = RmcpClient::new_streamable_http_client( + OAUTH_TEST_SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + + let error = initialize_client_with_timeout(&client, Some(Duration::from_millis(50))) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("timed out handshaking with MCP server after 50ms"), + "expected initialize timeout, got: {error:#}" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[serial(oauth_credentials_env)] async fn streamable_http_non_session_failure_does_not_trigger_recovery() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; let client = create_client(&base_url).await?; @@ -168,3 +278,54 @@ async fn streamable_http_non_session_failure_does_not_trigger_recovery() -> anyh Ok(()) } + +struct TempCodexHome { + original: Option, + dir: TempDir, +} + +impl TempCodexHome { + fn new() -> anyhow::Result { + let original = std::env::var_os("CODEX_HOME"); + let dir = TempDir::new()?; + unsafe { + std::env::set_var("CODEX_HOME", dir.path()); + } + Ok(Self { original, dir }) + } +} + +impl Drop for TempCodexHome { + fn drop(&mut self) { + unsafe { + if let Some(original) = &self.original { + std::env::set_var("CODEX_HOME", original); + } else { + std::env::remove_var("CODEX_HOME"); + } + } + } +} + +fn save_expired_oauth_tokens(server_url: &str) -> anyhow::Result<()> { + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(VALID_REFRESH_TOKEN.to_string()))); + response.set_expires_in(Some(&Duration::from_secs(7200))); + + let tokens = StoredOAuthTokens { + server_name: OAUTH_TEST_SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }; + save_oauth_tokens( + OAUTH_TEST_SERVER_NAME, + &tokens, + OAuthCredentialsStoreMode::File, + ) +} 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 822acef1a2..938acb4bac 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -86,10 +86,23 @@ pub(crate) async fn create_client(base_url: &str) -> anyhow::Result ) .await?; + initialize_client(&client).await?; + + Ok(client) +} + +pub(crate) async fn initialize_client(client: &RmcpClient) -> anyhow::Result<()> { + initialize_client_with_timeout(client, Some(Duration::from_secs(5))).await +} + +pub(crate) async fn initialize_client_with_timeout( + client: &RmcpClient, + timeout: Option, +) -> anyhow::Result<()> { client .initialize( init_params(), - Some(Duration::from_secs(5)), + timeout, Box::new(|_, _| { async { Ok(ElicitationResponse { @@ -102,8 +115,7 @@ pub(crate) async fn create_client(base_url: &str) -> anyhow::Result }), ) .await?; - - Ok(client) + Ok(()) } /// Creates a Streamable HTTP RMCP client that sends traffic through the remote @@ -124,22 +136,7 @@ pub(crate) async fn create_remote_client( ) .await?; - client - .initialize( - init_params(), - Some(Duration::from_secs(5)), - Box::new(|_, _| { - async { - Ok(ElicitationResponse { - action: ElicitationAction::Accept, - content: Some(json!({})), - meta: None, - }) - } - .boxed() - }), - ) - .await?; + initialize_client(&client).await?; Ok(client) } @@ -179,16 +176,24 @@ pub(crate) async fn arm_session_post_failure( } pub(crate) async fn spawn_streamable_http_server() -> anyhow::Result<(Child, String)> { + spawn_streamable_http_server_with_env(&[]).await +} + +pub(crate) async fn spawn_streamable_http_server_with_env( + env: &[(&str, &str)], +) -> anyhow::Result<(Child, String)> { let listener = TcpListener::bind("127.0.0.1:0")?; let port = listener.local_addr()?.port(); drop(listener); let bind_addr = format!("127.0.0.1:{port}"); let base_url = format!("http://{bind_addr}"); - let mut child = Command::new(streamable_http_server_bin()?) + let mut command = Command::new(streamable_http_server_bin()?); + command .kill_on_drop(true) .env("MCP_STREAMABLE_HTTP_BIND_ADDR", &bind_addr) - .spawn()?; + .envs(env.iter().copied()); + let mut child = command.spawn()?; wait_for_streamable_http_server(&mut child, &bind_addr, Duration::from_secs(5)).await?; Ok((child, base_url))