From d9511fb7888d98f89526d4ae019dd9be2f14199e Mon Sep 17 00:00:00 2001 From: xl-openai Date: Fri, 28 Aug 2026 18:32:44 +0000 Subject: [PATCH] Refresh MCP HTTP helper headers after authorization failures (#41400) ## What changed - Re-run the HTTP headers helper after a same-origin `POST` receives a `401` or `403`, and retry the request once when the effective headers changed. - Share refreshes across concurrent rejected requests, retain the previous headers when refresh fails, and keep the retry within the original timeout. - Allow helpers to supply `Authorization` while preserving explicit OAuth or bearer credentials, OAuth insufficient-scope challenges, and redirect safety checks. ## Testing - Cover rotating helper headers for MCP tool calls and same-origin OAuth token refreshes. - Cover concurrent rejections, failed and unchanged refreshes, retry limits, timeouts, and redirects. GitOrigin-RevId: c662edefdb1cd2876d9d02545a72d5bd6739864f --- codex-rs/core/tests/suite/rmcp_client.rs | 59 ++- .../src/bin/test_streamable_http_server.rs | 16 +- .../rmcp-client/src/http_client_adapter.rs | 4 +- codex-rs/rmcp-client/src/http_headers.rs | 275 +++++++++++--- .../rmcp-client/src/http_headers_tests.rs | 343 +++++++++++++++++- codex-rs/rmcp-client/src/lib.rs | 1 + .../www_authenticate.rs | 8 +- .../www_authenticate_tests.rs | 0 .../tests/streamable_http_oauth_startup.rs | 82 ++++- 9 files changed, 697 insertions(+), 91 deletions(-) rename codex-rs/rmcp-client/src/{http_client_adapter => }/www_authenticate.rs (97%) rename codex-rs/rmcp-client/src/{http_client_adapter => }/www_authenticate_tests.rs (100%) diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index b2187ef3e6..74d0d6e641 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -3254,13 +3254,28 @@ impl StreamableHttpTestServer { } } +enum HeadersHelperMode { + None, + Static, + Rotating, + RotatingAuthorization, +} + /// What this tests: Codex can discover and call a Streamable HTTP MCP tool in /// both local and remote-aware placements, and the tool observes the expected /// environment value from the server process that actually handled the request. -#[test_case(false; "plain")] -#[test_case(true; "headers helper")] +#[test_case(HeadersHelperMode::None; "plain")] +#[test_case(HeadersHelperMode::Static; "headers helper")] +#[test_case(HeadersHelperMode::Rotating; "headers helper refreshes rejected tool call")] +#[test_case(HeadersHelperMode::RotatingAuthorization; "Authorization helper refreshes rejected tool call")] #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn streamable_http_tool_call_round_trip(with_headers_helper: bool) -> anyhow::Result<()> { +async fn streamable_http_tool_call_round_trip(mode: HeadersHelperMode) -> anyhow::Result<()> { + let with_headers_helper = !matches!(mode, HeadersHelperMode::None); + let helper_authorization = matches!(mode, HeadersHelperMode::RotatingAuthorization); + let refresh_rejected_call = matches!( + mode, + HeadersHelperMode::Rotating | HeadersHelperMode::RotatingAuthorization + ); skip_if_no_network!(Ok(())); if with_headers_helper && is_remote_test_environment() { return Ok(()); @@ -3306,16 +3321,31 @@ async fn streamable_http_tool_call_round_trip(with_headers_helper: bool) -> anyh let expected_env_value = "propagated-env-http"; let Some(http_server) = start_streamable_http_test_server( expected_env_value, - /*expected_token*/ None, - with_headers_helper.then_some("gateway-token"), + helper_authorization.then_some("gateway-token"), + (with_headers_helper && !helper_authorization).then_some("gateway-token"), ) .await? else { return Ok(()); }; let server_url = http_server.url().to_string(); + let helper_directory = tempdir()?; + let helper_invocations = helper_directory.path().join("helper-invocations"); let http_headers_helper = with_headers_helper.then(|| { - if cfg!(windows) { + if refresh_rejected_call { + let authorization_arg = if helper_authorization { + " --authorization" + } else { + "" + }; + format!( + "\"{}\" --http-headers-helper \"{}\"{authorization_arg}", + cargo_bin("test_streamable_http_server") + .expect("streamable HTTP helper binary") + .display(), + helper_invocations.display(), + ) + } else if cfg!(windows) { r#"echo {"Proxy-Authorization":"Bearer gateway-token"}"#.to_string() } else { r#"printf '{"Proxy-Authorization":"Bearer gateway-token"}'"#.to_string() @@ -3347,6 +3377,20 @@ async fn streamable_http_tool_call_round_trip(with_headers_helper: bool) -> anyh .await?; wait_for_mcp_server(&fixture.codex, server_name).await?; + if refresh_rejected_call { + let control_url = http_server + .url() + .replace("/mcp", "/test/control/session-post-failure"); + let response = HttpClientBuilder::new() + .build_direct()? + .post(control_url) + .bearer_auth("gateway-token") + .json(&json!({ "status": 401, "remaining": 1 })) + .send() + .await?; + assert_eq!(response.status(), StatusCode::NO_CONTENT); + } + // Phase 4: submit the user turn that should trigger the MCP tool call. fixture .codex @@ -3405,6 +3449,9 @@ async fn streamable_http_tool_call_round_trip(with_headers_helper: bool) -> anyh .and_then(Value::as_str) .expect("env snapshot inserted"); assert_eq!(env_value, expected_env_value); + if refresh_rejected_call { + assert_eq!(fs::read_to_string(helper_invocations)?, "xx"); + } // Phase 7: verify the scripted model calls were consumed and clean up the // placement-aware MCP server. wait_for_event(&fixture.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; 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 a0fc7b4f7c..9702794bd6 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 @@ -118,7 +118,21 @@ async fn main() -> Result<(), Box> { if std::env::var_os("MCP_TEST_AMBIENT_SECRET").is_some() { return Err("helper inherited ambient secret".into()); } - println!(r#"{{"Proxy-Authorization":"Bearer gateway-token"}}"#); + if let Some(invocations) = args.next() { + let count = fs::read_to_string(&invocations).unwrap_or_default().len() + 1; + fs::write(invocations, "x".repeat(count))?; + let header = + if args.next().as_deref() == Some(std::ffi::OsStr::new("--authorization")) { + "Authorization" + } else { + "Proxy-Authorization" + }; + println!( + r#"{{"{header}":"Bearer gateway-token","X-Helper-Generation":"{count}"}}"# + ); + } else { + println!(r#"{{"Proxy-Authorization":"Bearer gateway-token"}}"#); + } return Ok(()); } _ => {} diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 0d31522cc6..b887a682b4 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -59,9 +59,7 @@ use crate::event_notification_transport::MAX_EVENT_NOTIFICATION_BYTES; use crate::http_client_redirect::SameOriginRedirectHttpClient; use crate::local_stdio_transport::MAX_MCP_STDIO_LINE_BYTES; -mod www_authenticate; - -use self::www_authenticate::insufficient_scope_challenge; +use crate::www_authenticate::insufficient_scope_challenge; const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream"; const JSON_MIME_TYPE: &str = "application/json"; diff --git a/codex-rs/rmcp-client/src/http_headers.rs b/codex-rs/rmcp-client/src/http_headers.rs index 12118e50f6..d0b2d28c98 100644 --- a/codex-rs/rmcp-client/src/http_headers.rs +++ b/codex-rs/rmcp-client/src/http_headers.rs @@ -6,6 +6,8 @@ use std::path::Path; use std::path::PathBuf; use std::process::Stdio; use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; use std::time::Duration; use anyhow::Result; @@ -38,6 +40,7 @@ use url::Origin; use url::Url; use crate::utils::create_env_for_mcp_server; +use crate::www_authenticate::insufficient_scope_challenge; const HELPER_TIMEOUT: Duration = Duration::from_secs(10); const MAX_HELPER_OUTPUT_BYTES: usize = 64 * 1024; @@ -45,7 +48,22 @@ type CachedHeaders = Shared, +} + +struct HeadersCache { + // Identifies a cohort of rejected requests, not a credential version. Advance even after + // failed or unchanged refreshes so concurrent rejections share one helper invocation. + refresh_epoch: u64, + current: CachedHeaders, + refresh: Option, +} + +struct RequestHeaders { + refresh_epoch: u64, + values: Arc, } struct HttpHeadersClient { @@ -118,29 +136,89 @@ impl Drop for HelperProcess { impl HttpHeadersProvider { fn new(server_url: &str, command: &str, cwd: PathBuf) -> Result { let command = command.to_string(); - let cached = async move { - run_helper(&command, &cwd) - .await - .map(Arc::new) - .map_err(|error| Arc::::from(error.to_string())) - } - .boxed() - .shared(); + let cached = Self::helper_attempt(command.clone(), cwd.clone()); Ok(Self { server_origin: Url::parse(server_url)?.origin(), - cached, + command, + cwd, + cache: Mutex::new(HeadersCache { + refresh_epoch: 0, + current: cached, + refresh: None, + }), }) } - async fn headers(&self) -> Result, ExecServerError> { - self.cached - .clone() + fn helper_attempt(command: String, cwd: PathBuf) -> CachedHeaders { + async move { + // Keep helper cleanup running after caller cancellation; dropping the set aborts it. + let mut tasks = tokio::task::JoinSet::new(); + tasks.spawn(async move { + run_helper(&command, &cwd) + .await + .map(Arc::new) + .map_err(|error| Arc::::from(error.to_string())) + }); + tasks + .join_next() + .await + .ok_or_else(|| Arc::::from("MCP HTTP headers helper task was unavailable"))? + .map_err(|_| Arc::::from("MCP HTTP headers helper task failed"))? + } + .boxed() + .shared() + } + + async fn headers(&self) -> Result { + let (refresh_epoch, current) = { + let cache = self.cache.lock().unwrap_or_else(PoisonError::into_inner); + (cache.refresh_epoch, cache.current.clone()) + }; + current .await + .map(|values| RequestHeaders { + refresh_epoch, + values, + }) .map_err(|error| ExecServerError::HttpRequest(error.to_string())) } + + async fn refresh(&self, rejected_epoch: u64) -> Result, ExecServerError> { + let attempt = { + let mut cache = self.cache.lock().unwrap_or_else(PoisonError::into_inner); + if cache.refresh_epoch != rejected_epoch { + cache.current.clone() + } else { + cache + .refresh + .get_or_insert_with(|| { + Self::helper_attempt(self.command.clone(), self.cwd.clone()) + }) + .clone() + } + }; + let result = attempt.clone().await; + let mut cache = self.cache.lock().unwrap_or_else(PoisonError::into_inner); + if cache.refresh_epoch == rejected_epoch + && cache + .refresh + .as_ref() + .is_some_and(|refresh| refresh.ptr_eq(&attempt)) + { + cache.refresh_epoch = rejected_epoch.saturating_add(/*rhs*/ 1); + if result.is_ok() { + cache.current = attempt; + } + cache.refresh = None; + } + result.map_err(|error| ExecServerError::HttpRequest(error.to_string())) + } } -/// No rejection-driven refresh: 401/403 may be OAuth challenges, and reconnecting loses sessions. +/// Refreshes helper headers once after a same-origin POST returns 401/403, retrying only +/// when helper-provided values change. OAuth challenges survive failed or unchanged refreshes. +/// Helper Authorization is used only when the request has no explicit bearer/OAuth credential. +/// The helper is a short-lived process; independently managed daemons remain caller-owned. pub fn with_http_headers_helper( inner: Arc, server_url: &str, @@ -154,16 +232,14 @@ pub fn with_http_headers_helper( impl HttpHeadersClient { async fn prepare_request( &self, - mut params: HttpRequestParams, - ) -> Result { + params: HttpRequestParams, + ) -> Result<(HttpRequestParams, Option, Option), ExecServerError> { let Ok(url) = Url::parse(¶ms.url) else { - return Ok(params); + return Ok((params, None, None)); }; if self.provider.server_origin != url.origin() { - return Ok(params); + return Ok((params, None, None)); } - // TODO: Follow same-origin redirects once later hops cannot leak helper headers. - params.redirect_policy = HttpRedirectPolicy::Stop; let deadline = params .timeout_ms @@ -176,7 +252,28 @@ impl HttpHeadersClient { })??, None => self.provider.headers().await?, }; - for (name, value) in headers.iter() { + let params = self.apply_headers(params, &headers.values, deadline)?; + Ok((params, Some(headers), deadline)) + } + + fn apply_headers( + &self, + mut params: HttpRequestParams, + headers: &HeaderMap, + deadline: Option, + ) -> Result { + // TODO: Follow same-origin redirects once later hops cannot leak helper headers. + params.redirect_policy = HttpRedirectPolicy::Stop; + for (name, value) in headers { + // Explicit bearer tokens, MCP OAuth and token-endpoint client authentication win. + if name == http::header::AUTHORIZATION + && params + .headers + .iter() + .any(|header| header.name.eq_ignore_ascii_case("authorization")) + { + continue; + } params .headers .retain(|header| !header.name.eq_ignore_ascii_case(name.as_str())); @@ -215,6 +312,96 @@ impl HttpHeadersClient { } Ok(()) } + + async fn retry_request( + &self, + params: HttpRequestParams, + rejected: RequestHeaders, + deadline: Option, + response: &HttpRequestResponse, + ) -> Option { + if !matches!(response.status, 401 | 403) + || response.status == 403 && insufficient_scope_challenge(&response.headers).is_some() + { + return None; + } + // A rejection may be an OAuth challenge, so preserve it if the helper cannot improve it. + let refreshed = match deadline { + Some(deadline) => { + tokio::time::timeout_at(deadline, self.provider.refresh(rejected.refresh_epoch)) + .await + .ok()? + .ok()? + } + None => self.provider.refresh(rejected.refresh_epoch).await.ok()?, + }; + // Header names are user-defined. Compare effective helper values, excluding an + // Authorization that would be overridden, without treating JSON key order as a change. + let explicit_authorization = params + .headers + .iter() + .any(|header| header.name.eq_ignore_ascii_case("authorization")); + let differs = |left: &HeaderMap, right: &HeaderMap| { + left.iter().any(|(name, value)| { + (!explicit_authorization || name != http::header::AUTHORIZATION) + && right.get(name) != Some(value) + }) + }; + if !differs(&rejected.values, &refreshed) && !differs(&refreshed, &rejected.values) { + return None; + } + self.apply_headers(params, &refreshed, deadline).ok() + } + + async fn request<'a, T>( + &'a self, + params: HttpRequestParams, + send: impl Fn(HttpRequestParams) -> BoxFuture<'a, Result>, + response_headers: impl Fn(&T) -> &HttpRequestResponse, + ) -> Result { + // OAuth validates redirect responses itself; only intercept MCP replay. + let mcp_redirect_was_stopped = params.redirect_policy == HttpRedirectPolicy::Stop + && !params.request_id.starts_with("oauth-request-"); + let original_headers = params + .method + .eq_ignore_ascii_case("POST") + .then(|| params.headers.clone()); + let (params, headers, deadline) = self.prepare_request(params).await?; + let original = original_headers + .filter(|_| headers.is_some()) + .map(|headers| { + let mut original = params.clone(); + original.headers = headers; + original + }); + let needs_redirect_check = |params: &HttpRequestParams| { + mcp_redirect_was_stopped + && Url::parse(¶ms.url).is_ok_and(|url| url.scheme() == "http") + && params + .headers + .iter() + .any(|header| header.name.eq_ignore_ascii_case("proxy-authorization")) + }; + let prevent_proxy_authorization_redirect = needs_redirect_check(¶ms); + let response = send(params).await?; + if prevent_proxy_authorization_redirect { + Self::reject_proxy_authorization_redirect(response_headers(&response))?; + } + if let (Some(original), Some(headers)) = (original, headers) + && let Some(retry) = self + .retry_request(original, headers, deadline, response_headers(&response)) + .await + { + drop(response); + let prevent_proxy_authorization_redirect = needs_redirect_check(&retry); + let response = send(retry).await?; + if prevent_proxy_authorization_redirect { + Self::reject_proxy_authorization_redirect(response_headers(&response))?; + } + return Ok(response); + } + Ok(response) + } } impl HttpClient for HttpHeadersClient { @@ -222,23 +409,11 @@ impl HttpClient for HttpHeadersClient { &self, params: HttpRequestParams, ) -> BoxFuture<'_, Result> { - async move { - // OAuth validates redirect responses itself; only intercept MCP replay. - let mcp_redirect_was_stopped = params.redirect_policy == HttpRedirectPolicy::Stop - && !params.request_id.starts_with("oauth-request-"); - let params = self.prepare_request(params).await?; - let prevent_proxy_authorization_redirect = mcp_redirect_was_stopped - && Url::parse(¶ms.url).is_ok_and(|url| url.scheme() == "http") - && params - .headers - .iter() - .any(|header| header.name.eq_ignore_ascii_case("proxy-authorization")); - let response = self.inner.http_request(params).await?; - if prevent_proxy_authorization_redirect { - Self::reject_proxy_authorization_redirect(&response)?; - } - Ok(response) - } + self.request( + params, + |params| self.inner.http_request(params), + |response| response, + ) .boxed() } @@ -246,22 +421,11 @@ impl HttpClient for HttpHeadersClient { &self, params: HttpRequestParams, ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { - async move { - let mcp_redirect_was_stopped = params.redirect_policy == HttpRedirectPolicy::Stop - && !params.request_id.starts_with("oauth-request-"); - let params = self.prepare_request(params).await?; - let prevent_proxy_authorization_redirect = mcp_redirect_was_stopped - && Url::parse(¶ms.url).is_ok_and(|url| url.scheme() == "http") - && params - .headers - .iter() - .any(|header| header.name.eq_ignore_ascii_case("proxy-authorization")); - let (response, stream) = self.inner.http_request_stream(params).await?; - if prevent_proxy_authorization_redirect { - Self::reject_proxy_authorization_redirect(&response)?; - } - Ok((response, stream)) - } + self.request( + params, + |params| self.inner.http_request_stream(params), + |response| &response.0, + ) .boxed() } } @@ -364,13 +528,12 @@ fn parse_helper_output(stdout: Vec) -> Result { for (name, value) in headers.entries { let name = HeaderName::from_bytes(name.as_bytes()) .map_err(|_| anyhow!("MCP HTTP headers helper returned an invalid header name"))?; - // Helper values replace same-name configured headers; bearer/OAuth owns Authorization. + // Helper values replace same-name configured headers, except explicit Authorization. // Google IAP uses Proxy-Authorization alongside application Authorization. For HTTPS MCP // URLs it is sent through the forward-proxy tunnel to IAP, not used as CONNECT auth. if matches!( name.as_str(), "accept" - | "authorization" | "connection" | "content-encoding" | "content-length" diff --git a/codex-rs/rmcp-client/src/http_headers_tests.rs b/codex-rs/rmcp-client/src/http_headers_tests.rs index 9b9fa6098b..c27c5d1002 100644 --- a/codex-rs/rmcp-client/src/http_headers_tests.rs +++ b/codex-rs/rmcp-client/src/http_headers_tests.rs @@ -4,7 +4,7 @@ use pretty_assertions::assert_eq; #[test] fn helper_output_errors_do_not_echo_secrets() { for output in [ - br#"{"Authorization":"secret"}"#.as_slice(), + br#"{"Host":"secret"}"#.as_slice(), br#"{"secret":"secret","secret":"secret"}"#.as_slice(), br#"{"secret":"secret","Secret":"secret"}"#.as_slice(), ] { @@ -24,13 +24,16 @@ async fn helper_attempt_is_shared_after_cancellation() { .canonicalize() .expect("canonical helper directory"); let cancelled_invocations = cwd.join("cancelled-invocations"); + let cancelled_finished = cwd.join("cancelled-helper-finished"); let cancelled = HttpHeadersProvider::new( "https://example.com", &format!( "test \"$(pwd)\" = '{0}'; test -n \"$HOME\"; test -n \"$PATH\"; \ - printf x >> '{1}'; sleep 0.2; printf '{{\"X-Gateway\":\"token\"}}'", + printf x >> '{1}'; sleep 0.2; printf x > '{2}'; \ + printf '{{\"X-Gateway\":\"token\"}}'", cwd.display(), cancelled_invocations.display(), + cancelled_finished.display(), ), cwd.clone(), ) @@ -40,11 +43,33 @@ async fn helper_attempt_is_shared_after_cancellation() { .await .is_err() ); + tokio::time::sleep(Duration::from_millis(500)).await; + assert!(cancelled_finished.exists()); assert!(cancelled.headers().await.is_ok()); assert_eq!( - std::fs::read_to_string(cancelled_invocations).expect("cancelled invocation count"), + std::fs::read_to_string(&cancelled_invocations).expect("cancelled invocation count"), "x" ); + std::fs::remove_file(&cancelled_finished).expect("reset helper completion marker"); + let headers = cancelled.headers().await.expect("cached headers"); + assert!( + tokio::time::timeout( + Duration::from_millis(/*millis*/ 20), + cancelled.refresh(headers.refresh_epoch) + ) + .await + .is_err() + ); + tokio::time::sleep(Duration::from_millis(/*millis*/ 500)).await; + assert!(cancelled_finished.exists()); + assert_eq!( + cancelled.refresh(headers.refresh_epoch).await.unwrap(), + headers.values + ); + assert_eq!( + std::fs::read_to_string(cancelled_invocations).expect("refresh invocation count"), + "xx" + ); let dropped_started = cwd.join("dropped-helper-started"); let dropped_finished = cwd.join("dropped-helper-finished"); @@ -90,13 +115,40 @@ async fn nonzero_helper_exit_is_cached() { temp.path().to_path_buf(), ) .expect("failed provider"); - let first = failed.headers().await.expect_err("failed helper"); - let second = failed.headers().await.expect_err("cached failure"); + let first = failed.headers().await.err().expect("failed helper"); + let second = failed.headers().await.err().expect("cached failure"); assert_eq!(first.to_string(), second.to_string()); let invocations = std::fs::read_to_string(failed_invocations).expect("failed invocation count"); assert_eq!(invocations.lines().count(), 1); } +#[cfg(unix)] +#[tokio::test] +async fn failed_refresh_is_shared_without_replacing_credentials() { + let temp = tempfile::tempdir().expect("temporary helper directory"); + let invocations = temp.path().join("invocations"); + let command = format!( + "printf x >> '{0}'; count=$(wc -c < '{0}'); \ + if [ \"$count\" -gt 1 ]; then exit 23; fi; \ + printf '{{\"X-Gateway\":\"token\"}}'", + invocations.display(), + ); + let provider = HttpHeadersProvider::new( + "https://example.com/mcp", + &command, + temp.path().to_path_buf(), + ) + .expect("headers provider"); + let original = provider.headers().await.expect("initial helper headers"); + assert!(provider.refresh(original.refresh_epoch).await.is_err()); + let stale = provider.refresh(original.refresh_epoch).await.unwrap(); + assert_eq!(stale, original.values); + let next = provider.headers().await.expect("preserved credentials"); + assert_eq!(next.values, original.values); + assert!(provider.refresh(next.refresh_epoch).await.is_err()); + assert_eq!(std::fs::read_to_string(invocations).unwrap(), "xxx"); +} + #[cfg(unix)] #[tokio::test] async fn connection_headers_are_cached_and_origin_bound() { @@ -198,3 +250,284 @@ async fn connection_headers_are_cached_and_origin_bound() { "x" ); } + +#[cfg(unix)] +#[tokio::test] +async fn concurrent_rejected_posts_share_one_headers_refresh() { + use axum::Router; + use axum::extract::State; + use axum::http::StatusCode; + use axum::routing::post; + use codex_exec_server::RouteAwareHttpClient; + use codex_http_client::HttpClientFactory; + use codex_http_client::OutboundProxyPolicy; + use tempfile::tempdir; + use tokio::net::TcpListener; + use tokio::sync::Barrier; + + for (rejection, credential_header) in [ + (StatusCode::UNAUTHORIZED, "Proxy-Authorization"), + (StatusCode::FORBIDDEN, "x-litellm-api-key"), + (StatusCode::UNAUTHORIZED, "Authorization"), + ] { + let temp = tempdir().expect("temporary helper directory"); + let invocations = temp.path().join("invocations"); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/mcp", listener.local_addr().unwrap()); + let barrier = Arc::new(Barrier::new(/*n*/ 2)); + let app = + Router::new() + .route( + "/mcp", + post( + move |State(barrier): State>, + headers: HeaderMap, + body: String| async move { + assert_eq!(body, "original request body"); + assert_eq!( + headers.get("mcp-session-id"), + Some(&HeaderValue::from_static("existing-session")) + ); + match headers[credential_header].to_str().unwrap() { + "Bearer token-1" => { + barrier.wait().await; + rejection + } + "Bearer token-2" => { + assert_eq!(headers["x-helper-only"], "configured"); + StatusCode::NO_CONTENT + } + unexpected => panic!("unexpected test credential: {unexpected}"), + } + }, + ), + ) + .with_state(barrier); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let command = format!( + "printf x >> '{0}'; count=$(wc -c < '{0}'); count=$((count)); \ + if [ \"$count\" -eq 1 ]; then \ + printf '{{\"{credential_header}\":\"Bearer token-1\",\"X-Helper-Only\":\"stale\"}}'; \ + else printf '{{\"{credential_header}\":\"Bearer token-%s\"}}' \"$count\"; fi", + invocations.display(), + ); + let inner: Arc = Arc::new(RouteAwareHttpClient::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + )); + let client = with_http_headers_helper(inner, &url, &command, temp.path().to_path_buf()) + .expect("headers helper client"); + let request = |request_id: &str| HttpRequestParams { + method: "POST".to_string(), + url: url.clone(), + headers: vec![ + HttpHeader { + name: "X-Helper-Only".to_string(), + value: "configured".to_string(), + value_env_var: None, + }, + HttpHeader { + name: "mcp-session-id".to_string(), + value: "existing-session".to_string(), + value_env_var: None, + }, + ], + body: Some(b"original request body".to_vec().into()), + timeout_ms: Some(5_000), + redirect_policy: HttpRedirectPolicy::Follow, + request_id: request_id.to_string(), + stream_response: true, + }; + + let (left, right) = tokio::join!( + client.http_request(request("left")), + client.http_request_stream(request("right")), + ); + assert_eq!(left.expect("left request").status, 204); + assert_eq!(right.expect("right request").0.status, 204); + assert_eq!( + std::fs::read_to_string(invocations).expect("helper invocation count"), + "xx" + ); + } +} + +#[cfg(unix)] +#[tokio::test] +async fn helper_refresh_preserves_oauth_challenges_and_retries_at_most_once() { + use codex_exec_server::RouteAwareHttpClient; + use codex_http_client::HttpClientFactory; + use codex_http_client::OutboundProxyPolicy; + use tempfile::tempdir; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::header; + use wiremock::matchers::method; + + let unchanged = "printf '{\"Proxy-Authorization\":\"Bearer token\"}'"; + let changed = "printf '{\"Proxy-Authorization\":\"Bearer token-%s\"}' \"$count\""; + let failed = "if [ \"$count\" -gt 1 ]; then exit 23; fi; printf '{\"Proxy-Authorization\":\"Bearer token\"}'"; + let ignored_authorization = "if [ \"$count\" -eq 1 ]; then \ + printf '{\"Authorization\":\"old\",\"X-A\":\"a\",\"X-B\":\"b\"}'; \ + else printf '{\"X-B\":\"b\",\"X-A\":\"a\",\"Authorization\":\"new\"}'; fi"; + for streamed in [false, true] { + for (helper_output, status, expected_mcp_requests) in [ + (unchanged, 401, 1), + (failed, 401, 1), + (changed, 401, 2), + (changed, 403, 1), + (ignored_authorization, 401, 1), + ] { + let (challenge, expected_helper_invocations) = if status == 403 { + ( + r#"Bearer error="insufficient_scope", scope="tools:write""#, + "x", + ) + } else { + (r#"Bearer realm="oauth""#, "xx") + }; + let temp = tempdir().expect("temporary helper directory"); + let invocations = temp.path().join("invocations"); + let server = MockServer::start().await; + let url = format!("{}/mcp", server.uri()); + Mock::given(method("POST")) + .and(header("authorization", "Bearer oauth-token")) + .respond_with( + ResponseTemplate::new(status) + .insert_header("www-authenticate", challenge) + .set_body_string("original OAuth challenge"), + ) + .expect(expected_mcp_requests) + .mount(&server) + .await; + let command = format!( + "printf x >> '{0}'; count=$(wc -c < '{0}'); {helper_output}", + invocations.display(), + ); + let inner: Arc = Arc::new(RouteAwareHttpClient::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + )); + let client = with_http_headers_helper(inner, &url, &command, temp.path().to_path_buf()) + .expect("headers helper client"); + let params = HttpRequestParams { + method: "POST".to_string(), + url: url.clone(), + headers: vec![HttpHeader { + name: "aUtHoRiZaTiOn".to_string(), + value: "Bearer oauth-token".to_string(), + value_env_var: None, + }], + body: None, + timeout_ms: Some(5_000), + redirect_policy: HttpRedirectPolicy::Follow, + request_id: "oauth-challenge".to_string(), + stream_response: streamed, + }; + let (response, body) = if streamed { + let (response, mut body) = client + .http_request_stream(params) + .await + .expect("original OAuth response"); + let mut bytes = Vec::new(); + while let Some(chunk) = body.recv().await.expect("response body chunk") { + bytes.extend(chunk); + } + (response, bytes) + } else { + let response = client + .http_request(params) + .await + .expect("original OAuth response"); + let bytes = response.body.0.clone(); + (response, bytes) + }; + + assert_eq!(response.status, status); + assert!(response.headers.iter().any(|header| { + header.name.eq_ignore_ascii_case("www-authenticate") && header.value == challenge + })); + assert_eq!(body, b"original OAuth challenge"); + assert_eq!( + std::fs::read_to_string(invocations).expect("helper invocation count"), + expected_helper_invocations + ); + } + } +} + +#[cfg(unix)] +#[tokio::test] +async fn refresh_retry_rechecks_deadline_and_redirects() { + use codex_exec_server::RouteAwareHttpClient; + use codex_http_client::HttpClientFactory; + use codex_http_client::OutboundProxyPolicy; + + let temp = tempfile::tempdir().expect("helper directory"); + let command = "if [ -e invoked ]; then printf '{\"Proxy-Authorization\":\"Bearer fresh\"}'; else touch invoked; printf '{}'; fi"; + let client = HttpHeadersClient { + inner: Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + ))), + provider: HttpHeadersProvider::new( + "http://example.com/mcp", + command, + temp.path().to_path_buf(), + ) + .expect("provider"), + }; + let requests = Mutex::new(Vec::new()); + let error = client + .request( + HttpRequestParams { + method: "POST".to_string(), + url: "http://example.com/mcp".to_string(), + headers: Vec::new(), + body: Some(b"original request".to_vec().into()), + timeout_ms: Some(5_000), + redirect_policy: HttpRedirectPolicy::Stop, + request_id: "mcp-request".to_string(), + stream_response: false, + }, + |params| { + let mut requests = requests.lock().expect("requests lock"); + let status = if requests.is_empty() { 401 } else { 307 }; + requests.push(params); + async move { + tokio::time::sleep(Duration::from_millis(/*millis*/ 20)).await; + Ok(HttpRequestResponse { + status, + headers: vec![HttpHeader { + name: "Location".to_string(), + value: "http://example.com/redirected".to_string(), + value_env_var: None, + }], + body: Vec::new().into(), + }) + } + .boxed() + }, + |response| response, + ) + .await + .expect_err("refreshed proxy redirect must fail"); + assert!( + error + .to_string() + .contains("cannot safely replay Proxy-Authorization") + ); + let requests = requests.into_inner().expect("recorded requests"); + let [original, retry] = requests.as_slice() else { + panic!("expected exactly two requests"); + }; + assert!(retry.timeout_ms < original.timeout_ms); + let mut expected_retry = original.clone(); + expected_retry.timeout_ms = retry.timeout_ms; + expected_retry.headers = vec![HttpHeader { + name: "proxy-authorization".to_string(), + value: "Bearer fresh".to_string(), + value_env_var: None, + }]; + assert_eq!(retry, &expected_retry); +} diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 9b3a2ae6ce..3ca80072e6 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -24,6 +24,7 @@ mod service_error; mod startup_error; mod stdio_server_launcher; mod utils; +mod www_authenticate; pub use auth_status::McpAuthState; pub use auth_status::McpLoginRequirement; diff --git a/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs b/codex-rs/rmcp-client/src/www_authenticate.rs similarity index 97% rename from codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs rename to codex-rs/rmcp-client/src/www_authenticate.rs index 374d11d067..c4382c8b1b 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs +++ b/codex-rs/rmcp-client/src/www_authenticate.rs @@ -2,9 +2,9 @@ use codex_exec_server::HttpHeader; use http::header::WWW_AUTHENTICATE; #[derive(Debug, PartialEq, Eq)] -pub(super) struct InsufficientScopeChallenge { - pub(super) www_authenticate_header: String, - pub(super) required_scope: Option, +pub(crate) struct InsufficientScopeChallenge { + pub(crate) www_authenticate_header: String, + pub(crate) required_scope: Option, } #[derive(Debug, PartialEq, Eq)] @@ -64,7 +64,7 @@ impl BearerChallenge { /// Finds a Bearer insufficient-scope challenge among all `WWW-Authenticate` /// response header field values. -pub(super) fn insufficient_scope_challenge( +pub(crate) fn insufficient_scope_challenge( headers: &[HttpHeader], ) -> Option { headers diff --git a/codex-rs/rmcp-client/src/http_client_adapter/www_authenticate_tests.rs b/codex-rs/rmcp-client/src/www_authenticate_tests.rs similarity index 100% rename from codex-rs/rmcp-client/src/http_client_adapter/www_authenticate_tests.rs rename to codex-rs/rmcp-client/src/www_authenticate_tests.rs diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs index e6b98304d2..b1855e8a1e 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -64,6 +64,7 @@ const REFRESHABLE_SERVER_URL: &str = "https://refreshable.example/mcp"; enum OAuthStartupScenario { DirectAuthorizationMetadata, GatewayHeadersHelper, + SameOriginGatewayHeadersHelper, ProtectedResourceMetadata, } @@ -83,6 +84,11 @@ async fn refreshes_oauth_with_gateway_headers_helper() -> anyhow::Result<()> { assert_expired_token_refresh(OAuthStartupScenario::GatewayHeadersHelper).await } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn refreshes_oauth_after_gateway_rejects_token_request() -> anyhow::Result<()> { + assert_expired_token_refresh(OAuthStartupScenario::SameOriginGatewayHeadersHelper).await +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refresh_uses_discovered_protected_resource_audience() -> anyhow::Result<()> { assert_expired_token_refresh(OAuthStartupScenario::ProtectedResourceMetadata).await @@ -91,10 +97,20 @@ async fn refresh_uses_discovered_protected_resource_audience() -> anyhow::Result async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow::Result<()> { let server = MockServer::start().await; let authorization_server = MockServer::start().await; + let same_origin_gateway = matches!( + scenario, + OAuthStartupScenario::SameOriginGatewayHeadersHelper + ); + let token_server = if same_origin_gateway { + &server + } else { + &authorization_server + }; let resource_url = format!("{}/mcp", server.uri()); let (server_url, mcp_path, authorization_metadata_path) = match scenario { OAuthStartupScenario::DirectAuthorizationMetadata - | OAuthStartupScenario::GatewayHeadersHelper => ( + | OAuthStartupScenario::GatewayHeadersHelper + | OAuthStartupScenario::SameOriginGatewayHeadersHelper => ( resource_url.clone(), "/mcp", "/.well-known/oauth-authorization-server/mcp", @@ -130,8 +146,8 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow: let mut authorization_metadata = json!({ "issuer": server_url.clone(), - "authorization_endpoint": format!("{}/oauth/authorize", authorization_server.uri()), - "token_endpoint": format!("{}/oauth/token", authorization_server.uri()), + "authorization_endpoint": format!("{}/oauth/authorize", token_server.uri()), + "token_endpoint": format!("{}/oauth/token", token_server.uri()), "scopes_supported": [""], }); if matches!(scenario, OAuthStartupScenario::ProtectedResourceMetadata) { @@ -142,6 +158,8 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow: .ok_or_else(|| anyhow::anyhow!("authorization metadata should include issuer"))? .to_string(); + let helper_directory = TempDir::new()?; + let helper_invocations = helper_directory.path().join("helper-invocations"); Mock::given(method("GET")) .and(path(authorization_metadata_path)) .and(header("user-agent", RESOURCE_USER_AGENT)) @@ -152,7 +170,14 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow: .await; Mock::given(method("POST")) .and(path("/oauth/token")) - .and(header("user-agent", MCP_USER_AGENT)) + .and(header( + "user-agent", + if same_origin_gateway { + RESOURCE_USER_AGENT + } else { + MCP_USER_AGENT + }, + )) .and(body_string_contains("grant_type=refresh_token")) .and(body_string_contains(format!( "refresh_token={REFRESH_TOKEN}" @@ -164,14 +189,25 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow: .any(|(name, value)| name == "resource" && value == expected_resource) } }) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "access_token": REFRESHED_ACCESS_TOKEN, - "token_type": "Bearer", - "expires_in": 7200, - "refresh_token": REFRESH_TOKEN, - }))) - .expect(1) - .mount(&authorization_server) + .respond_with(move |request: &Request| { + if same_origin_gateway + && request + .headers + .get("x-helper-generation") + .is_some_and(|generation| generation == "1") + { + ResponseTemplate::new(/*s*/ 401) + } else { + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": REFRESH_TOKEN, + })) + } + }) + .expect(if same_origin_gateway { 2 } else { 1 }) + .mount(token_server) .await; Mock::given(method("POST")) .and(path(mcp_path)) @@ -214,7 +250,11 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow: .await; let codex_home = TempDir::new()?; - let with_headers_helper = matches!(scenario, OAuthStartupScenario::GatewayHeadersHelper); + let with_headers_helper = matches!( + scenario, + OAuthStartupScenario::GatewayHeadersHelper + | OAuthStartupScenario::SameOriginGatewayHeadersHelper + ); // Credential storage resolves CODEX_HOME from the process environment. // Run the client half of the test in an ignored helper test so it can use @@ -231,14 +271,19 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow: command.env( CHILD_HELPER_COMMAND_ENV, format!( - "\"{}\" --http-headers-helper", + "\"{}\" --http-headers-helper \"{}\"", cargo_bin("test_streamable_http_server")?.display(), + helper_invocations.display(), ), ); } let status = command.status().await?; assert!(status.success(), "OAuth startup child failed: {status}"); if with_headers_helper { + assert_eq!( + std::fs::read_to_string(helper_invocations)?, + if same_origin_gateway { "xx" } else { "x" } + ); let requests = server.received_requests().await.unwrap_or_default(); assert!(requests.iter().all(|request| { request @@ -251,6 +296,12 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow: .received_requests() .await .ok_or_else(|| anyhow::anyhow!("authorization server should record requests"))?; + server.verify().await; + authorization_server.verify().await; + if same_origin_gateway { + assert!(authorization_requests.is_empty()); + return Ok(()); + } assert_eq!(authorization_requests.len(), 1); assert_eq!(authorization_requests[0].headers.get("x-api-key"), None); assert_eq!( @@ -264,8 +315,7 @@ async fn assert_expired_token_refresh(scenario: OAuthStartupScenario) -> anyhow: .map(http::HeaderValue::as_bytes), Some(MCP_USER_AGENT.as_bytes()) ); - server.verify().await; - authorization_server.verify().await; + Ok(()) }