[rmcp-client] Refresh OAuth tokens before initialize

codex-rs/cli/src/mcp_cmd.rs

@@ -26,7 +26,7 @@ use codex_mcp::oauth_login_support;
 use codex_mcp::resolve_oauth_scopes;
 use codex_mcp::should_retry_without_scopes;
 use codex_protocol::protocol::McpAuthStatus;
-use codex_rmcp_client::delete_oauth_tokens;
+use codex_rmcp_client::delete_oauth_tokens_async;
 use codex_rmcp_client::perform_oauth_login;
 use codex_utils_cli::CliConfigOverrides;
 use codex_utils_cli::format_env_display;
@@ -514,7 +514,7 @@ async fn run_logout(config_overrides: &CliConfigOverrides, logout_args: LogoutAr
         _ => bail!("OAuth logout is only supported for streamable_http transports."),
     };
 
-    match delete_oauth_tokens(&name, &url, config.mcp_oauth_credentials_store_mode) {
+    match delete_oauth_tokens_async(&name, &url, config.mcp_oauth_credentials_store_mode).await {
         Ok(true) => println!("Removed OAuth credentials for '{name}'."),
         Ok(false) => println!("No OAuth credentials stored for '{name}'."),
         Err(err) => return Err(anyhow!("failed to delete OAuth credentials: {err}")),

codex-rs/rmcp-client/Cargo.toml

@@ -21,6 +21,7 @@ codex-config = { workspace = true }
 codex-exec-server = { workspace = true }
 codex-keyring-store = { workspace = true }
 codex-protocol = { workspace = true }
+codex-utils-path = { workspace = true }
 codex-utils-pty = { workspace = true }
 codex-utils-home-dir = { workspace = true }
 bytes = { workspace = true }

codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs

@@ -4,10 +4,13 @@ use std::fs;
 use std::io::ErrorKind;
 use std::net::SocketAddr;
 use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
 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;
@@ -24,15 +27,33 @@ use axum::middleware::Next;
 use axum::response::Response;
 use axum::routing::get;
 use axum::routing::post;
+use codex_config::types::OAuthCredentialsStoreMode;
+use codex_exec_server::Environment;
+use codex_rmcp_client::ElicitationAction;
+use codex_rmcp_client::ElicitationResponse;
+use codex_rmcp_client::RmcpClient;
+use codex_rmcp_client::StoredOAuthTokens;
+use codex_rmcp_client::WrappedOAuthTokenResponse;
+use codex_rmcp_client::save_oauth_tokens_async;
+use futures::FutureExt as _;
+use oauth2::AccessToken;
+use oauth2::RefreshToken;
+use oauth2::basic::BasicTokenType;
 use rmcp::ErrorData as McpError;
 use rmcp::handler::server::ServerHandler;
 use rmcp::model::CallToolRequestParams;
 use rmcp::model::CallToolResult;
+use rmcp::model::ClientCapabilities;
+use rmcp::model::ElicitationCapability;
+use rmcp::model::FormElicitationCapability;
+use rmcp::model::Implementation;
+use rmcp::model::InitializeRequestParams;
 use rmcp::model::JsonObject;
 use rmcp::model::ListResourceTemplatesResult;
 use rmcp::model::ListResourcesResult;
 use rmcp::model::ListToolsResult;
 use rmcp::model::PaginatedRequestParams;
+use rmcp::model::ProtocolVersion;
 use rmcp::model::RawResource;
 use rmcp::model::RawResourceTemplate;
 use rmcp::model::ReadResourceRequestParams;
@@ -46,12 +67,17 @@ use rmcp::model::Tool;
 use rmcp::model::ToolAnnotations;
 use rmcp::transport::StreamableHttpServerConfig;
 use rmcp::transport::StreamableHttpService;
+use rmcp::transport::auth::OAuthTokenResponse;
+use rmcp::transport::auth::VendorExtraTokenFields;
 use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
 use serde::Deserialize;
 use serde_json::json;
 use tokio::sync::Mutex;
 use tokio::task;
 use tokio::time::sleep;
+use urlencoding::decode;
+
+static REFRESH_TOKEN_USES: AtomicUsize = AtomicUsize::new(0);
 
 #[derive(Clone)]
 struct TestToolServer {
@@ -96,6 +122,24 @@ struct EchoArgs {
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    if let Ok(server_url) = std::env::var("MCP_TEST_OAUTH_CLIENT_URL") {
+        let server_name = std::env::var("MCP_TEST_OAUTH_SERVER_NAME")?;
+        return run_oauth_test_client(&server_name, &server_url).await;
+    }
+    if let Ok(server_name) = std::env::var("MCP_TEST_OAUTH_WRITE_SERVER_NAME") {
+        let server_url = std::env::var("MCP_TEST_OAUTH_WRITE_SERVER_URL")?;
+        let access_token = std::env::var("MCP_TEST_OAUTH_WRITE_ACCESS_TOKEN")?;
+        let refresh_token = std::env::var("MCP_TEST_OAUTH_WRITE_REFRESH_TOKEN")?;
+        if let Ok(barrier) = std::env::var("MCP_TEST_OAUTH_WRITE_BARRIER") {
+            while !std::path::Path::new(&barrier).exists() {
+                std::thread::sleep(Duration::from_millis(10));
+            }
+        }
+        write_oauth_test_credentials(&server_name, &server_url, &access_token, &refresh_token)
+            .await?;
+        return Ok(());
+    }
+
     let bind_addr = parse_bind_addr()?;
     let session_failure_state = SessionFailureState::default();
     const MAX_BIND_RETRIES: u32 = 20;
@@ -129,6 +173,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({
@@ -165,6 +210,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             session_failure_state.clone(),
             fail_session_post_when_armed,
         ))
+        .layer(middleware::from_fn(delay_mcp_initialize_when_configured))
         .with_state(session_failure_state);
 
     let router = if let Ok(token) = std::env::var("MCP_EXPECT_BEARER") {
@@ -179,6 +225,76 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     Ok(())
 }
 
+async fn run_oauth_test_client(
+    server_name: &str,
+    server_url: &str,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let client = RmcpClient::new_streamable_http_client(
+        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 mut capabilities = ClientCapabilities::default();
+    capabilities.elicitation = Some(ElicitationCapability {
+        form: Some(FormElicitationCapability {
+            schema_validation: None,
+        }),
+        url: None,
+    });
+    let params = InitializeRequestParams::new(
+        capabilities,
+        Implementation::new("codex-test", "0.0.0-test").with_title("Codex rmcp OAuth process test"),
+    )
+    .with_protocol_version(ProtocolVersion::V_2025_06_18);
+    client
+        .initialize(
+            params,
+            Some(Duration::from_secs(15)),
+            Box::new(|_, _| {
+                async {
+                    Ok(ElicitationResponse {
+                        action: ElicitationAction::Accept,
+                        content: Some(json!({})),
+                        meta: None,
+                    })
+                }
+                .boxed()
+            }),
+        )
+        .await?;
+    Ok(())
+}
+
+async fn write_oauth_test_credentials(
+    server_name: &str,
+    server_url: &str,
+    access_token: &str,
+    refresh_token: &str,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let mut response = OAuthTokenResponse::new(
+        AccessToken::new(access_token.to_string()),
+        BasicTokenType::Bearer,
+        VendorExtraTokenFields::default(),
+    );
+    response.set_refresh_token(Some(RefreshToken::new(refresh_token.to_string())));
+    response.set_expires_in(Some(&Duration::from_secs(7200)));
+    let stored = StoredOAuthTokens {
+        server_name: server_name.to_string(),
+        url: server_url.to_string(),
+        client_id: "test-client-id".to_string(),
+        token_response: WrappedOAuthTokenResponse(response),
+        expires_at: None,
+    };
+    save_oauth_tokens_async(server_name, &stored, OAuthCredentialsStoreMode::File).await?;
+    Ok(())
+}
+
 impl ServerHandler for TestToolServer {
     fn get_info(&self) -> ServerInfo {
         ServerInfo::new(
@@ -389,7 +505,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 +520,110 @@ 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)?;
+
+    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)?;
+
+    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 oauth_error_response(StatusCode::BAD_REQUEST, "invalid_grant");
+    }
+
+    if let Ok(max_uses) = std::env::var("MCP_REFRESH_TOKEN_MAX_USES")
+        && let Ok(max_uses) = max_uses.parse::<usize>()
+    {
+        let use_count = REFRESH_TOKEN_USES.fetch_add(1, Ordering::SeqCst) + 1;
+        if use_count > max_uses {
+            return oauth_error_response(StatusCode::BAD_REQUEST, "invalid_grant");
+        }
+    }
+
+    if let Ok(error) = std::env::var("MCP_REFRESH_ERROR") {
+        let status = if error == "server_error" {
+            StatusCode::INTERNAL_SERVER_ERROR
+        } else {
+            StatusCode::BAD_REQUEST
+        };
+        return oauth_error_response(status, &error);
+    }
+
+    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 mut token_response = json!({
+        "access_token": access_token,
+        "token_type": "Bearer",
+        "expires_in": 7200,
+    });
+    if std::env::var("MCP_OMIT_ROTATED_REFRESH_TOKEN").is_err() {
+        let refresh_token =
+            std::env::var("MCP_ROTATED_REFRESH_TOKEN").map_err(|_| StatusCode::BAD_REQUEST)?;
+        token_response["refresh_token"] = json!(refresh_token);
+    }
+
+    #[expect(clippy::expect_used)]
+    Ok(Response::builder()
+        .status(StatusCode::OK)
+        .header(CONTENT_TYPE, "application/json")
+        .body(Body::from(
+            serde_json::to_vec(&token_response).expect("failed to serialize token response"),
+        ))
+        .expect("valid token response"))
+}
+
+fn oauth_error_response(status: StatusCode, error: &str) -> Result<Response, StatusCode> {
+    #[expect(clippy::expect_used)]
+    Ok(Response::builder()
+        .status(status)
+        .header(CONTENT_TYPE, "application/json")
+        .body(Body::from(
+            serde_json::to_vec(&json!({ "error": error }))
+                .expect("failed to serialize OAuth error response"),
+        ))
+        .expect("valid OAuth error response"))
+}
+
+async fn delay_mcp_initialize_when_configured(request: Request<Body>, next: Next) -> Response {
+    if request.uri().path() == "/mcp"
+        && request.method() == Method::POST
+        && !request.headers().contains_key(MCP_SESSION_ID_HEADER)
+        && let Ok(delay_ms) = std::env::var("MCP_INITIALIZE_DELAY_MS")
+        && let Ok(delay_ms) = delay_ms.parse::<u64>()
+    {
+        sleep(Duration::from_millis(delay_ms)).await;
+    }
+
+    next.run(request).await
+}
+
+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>,

codex-rs/rmcp-client/src/lib.rs

@@ -5,6 +5,7 @@ mod http_client_adapter;
 mod in_process_transport;
 mod logging_client_handler;
 mod oauth;
+mod oauth_lock;
 mod perform_oauth_login;
 mod program_resolver;
 mod rmcp_client;
@@ -20,8 +21,10 @@ pub use in_process_transport::InProcessTransportFactory;
 pub use oauth::StoredOAuthTokens;
 pub use oauth::WrappedOAuthTokenResponse;
 pub use oauth::delete_oauth_tokens;
+pub use oauth::delete_oauth_tokens_async;
 pub(crate) use oauth::load_oauth_tokens;
 pub use oauth::save_oauth_tokens;
+pub use oauth::save_oauth_tokens_async;
 pub use perform_oauth_login::OAuthProviderError;
 pub use perform_oauth_login::OauthLoginHandle;
 pub use perform_oauth_login::perform_oauth_login;