openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(rmcp-client): refresh oauth credentials before initialize

Browse Source
This commit is contained in:
Casey Chow
2026-06-03 16:13:37 -04:00
parent f6e529656f
commit bfe92ebd0e
4 changed files with 276 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs
View File

@@ -8,6 +8,7 @@ use std::time::Duration;
use axum::Router;
use axum::body::Body;
use axum::body::to_bytes;
use axum::extract::Json;
use axum::extract::State;
use axum::http::HeaderMap;
@@ -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<dyn std::error::Error>> {
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<Body>,
next: Next,
) -> Result<Response, StatusCode> {
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<Body>) -> Result<Response, StatusCode> {
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::<u64>()
{
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<HashMap<String, String>, 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<SessionFailureState>,
Json(request): Json<ArmSessionPostFailureRequest>,

34
codex-rs/rmcp-client/src/rmcp_client.rs
View File

@@ -823,6 +823,8 @@ impl RmcpClient {
Arc<RunningService<RoleClient, ElicitationClientService>>,
Option<OAuthPersistor>,
)> {
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}"))?,

161
codex-rs/rmcp-client/tests/streamable_http_recovery.rs
View File

@@ -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<OsString>,
dir: TempDir,
}
impl TempCodexHome {
fn new() -> anyhow::Result<Self> {
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,
)
}

47
codex-rs/rmcp-client/tests/streamable_http_test_support.rs
View File

@@ -86,10 +86,23 @@ pub(crate) async fn create_client(base_url: &str) -> anyhow::Result<RmcpClient>
)
.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<Duration>,
) -> 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<RmcpClient>
}),
)
.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))