openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-03 19:53:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refresh MCP HTTP helper headers after authorization failures (#41400)

Browse Source 
## 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
This commit is contained in:
xl-openai
2026-08-28 19:29:51 +00:00
committed by copyberry
parent bb998aeb8e
commit d9511fb788
9 changed files with 697 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/tests/suite/rmcp_client.rs
+53 -6
View File
@@ -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;
codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs
+15 -1
View File
@@ -118,7 +118,21 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
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(());
}
_ => {}
codex-rs/rmcp-client/src/http_client_adapter.rs
+1 -3
View File
@@ -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";
codex-rs/rmcp-client/src/http_headers.rs
+219 -56
View File
@@ -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<BoxFuture<'static, std::result::Result<Arc<HeaderMap
struct HttpHeadersProvider {
server_origin: Origin,
cached: CachedHeaders,
command: String,
cwd: PathBuf,
cache: Mutex<HeadersCache>,
}
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<CachedHeaders>,
}
struct RequestHeaders {
refresh_epoch: u64,
values: Arc<HeaderMap>,
}
struct HttpHeadersClient {
@@ -118,29 +136,89 @@ impl Drop for HelperProcess {
impl HttpHeadersProvider {
fn new(server_url: &str, command: &str, cwd: PathBuf) -> Result<Self> {
let command = command.to_string();
let cached = async move {
run_helper(&command, &cwd)
.await
.map(Arc::new)
.map_err(|error| Arc::<str>::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<Arc<HeaderMap>, 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::<str>::from(error.to_string()))
});
tasks
.join_next()
.await
.ok_or_else(|| Arc::<str>::from("MCP HTTP headers helper task was unavailable"))?
.map_err(|_| Arc::<str>::from("MCP HTTP headers helper task failed"))?
}
.boxed()
.shared()
}
async fn headers(&self) -> Result<RequestHeaders, ExecServerError> {
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<Arc<HeaderMap>, 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<dyn HttpClient>,
server_url: &str,
@@ -154,16 +232,14 @@ pub fn with_http_headers_helper(
impl HttpHeadersClient {
async fn prepare_request(
&self,
mut params: HttpRequestParams,
) -> Result<HttpRequestParams, ExecServerError> {
params: HttpRequestParams,
) -> Result<(HttpRequestParams, Option<RequestHeaders>, Option<Instant>), ExecServerError> {
let Ok(url) = Url::parse(&params.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<Instant>,
) -> Result<HttpRequestParams, ExecServerError> {
// 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<Instant>,
response: &HttpRequestResponse,
) -> Option<HttpRequestParams> {
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<T, ExecServerError>>,
response_headers: impl Fn(&T) -> &HttpRequestResponse,
) -> Result<T, ExecServerError> {
// 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(&params.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(&params);
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<HttpRequestResponse, ExecServerError>> {
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(&params.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(&params.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<u8>) -> Result<HeaderMap> {
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"
codex-rs/rmcp-client/src/http_headers_tests.rs
+338 -5
View File
@@ -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<Arc<Barrier>>,
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<dyn HttpClient> = 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<dyn HttpClient> = 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);
}
codex-rs/rmcp-client/src/lib.rs
+1
View File
@@ -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;
codex-rs/rmcp-client/src/http_client_adapter/www_authenticate.rs → codex-rs/rmcp-client/src/www_authenticate.rs
+4 -4
View File
@@ -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<String>,
pub(crate) struct InsufficientScopeChallenge {
pub(crate) www_authenticate_header: String,
pub(crate) required_scope: Option<String>,
}
#[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<InsufficientScopeChallenge> {
headers
codex-rs/rmcp-client/src/http_client_adapter/www_authenticate_tests.rs → codex-rs/rmcp-client/src/www_authenticate_tests.rs
View File
codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs
+66 -16
View File
@@ -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(())
}