From 0a7e4b6fafe939acc4bb7e4052581c9bf696f43d Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 5 Jun 2026 04:09:20 +0000 Subject: [PATCH] test(rmcp-client): reproduce external OAuth overwrite --- .../tests/streamable_http_oauth_startup.rs | 226 ++++++++++++++++++ 1 file changed, 226 insertions(+) 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 1c18f2c98b..6309f2a8cb 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -1,6 +1,11 @@ mod streamable_http_test_support; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; @@ -11,6 +16,7 @@ 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 serde_json::Value; @@ -33,6 +39,14 @@ const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; const REFRESH_TOKEN: &str = "valid-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; +const EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_EXTERNAL_UPDATE_SERVER_URL"; +const OLD_ACCESS_TOKEN: &str = "old-access-token"; +const OLD_REFRESH_TOKEN: &str = "old-refresh-token"; +const EXTERNAL_ACCESS_TOKEN: &str = "external-access-token"; +const EXTERNAL_REFRESH_TOKEN: &str = "external-refresh-token"; +const STALE_REFRESHED_ACCESS_TOKEN: &str = "stale-refreshed-access-token"; +const STALE_ROTATED_REFRESH_TOKEN: &str = "stale-rotated-refresh-token"; +const REFRESH_SKEW: Duration = Duration::from_secs(30); #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> { @@ -153,3 +167,215 @@ async fn oauth_startup_child() -> anyhow::Result<()> { initialize_client(&client).await?; Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn running_client_keeps_and_can_persist_over_external_oauth_update() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": [""], + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={OLD_REFRESH_TOKEN}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": STALE_REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": STALE_ROTATED_REFRESH_TOKEN, + }))) + .expect(1) + .mount(&server) + .await; + + let tool_call_count = Arc::new(AtomicUsize::new(0)); + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with({ + let tool_call_count = Arc::clone(&tool_call_count); + move |request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + let method = body.get("method").and_then(Value::as_str); + let authorization = request + .headers + .get("authorization") + .and_then(|value| value.to_str().ok()); + let expected_access_token = match method { + Some("tools/call") if tool_call_count.fetch_add(1, Ordering::SeqCst) == 1 => { + STALE_REFRESHED_ACCESS_TOKEN + } + _ => OLD_ACCESS_TOKEN, + }; + if authorization != Some(format!("Bearer {expected_access_token}").as_str()) { + return ResponseTemplate::new(401); + } + + match method { + Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": { + "tools": {}, + }, + "serverInfo": { + "name": "oauth-external-update-test", + "version": "0.0.0-test", + }, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/call") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "content": [], + "isError": false, + }, + })), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + } + }) + .expect(4) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let server_url = format!("{}/mcp", server.uri()); + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_external_update_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV, server_url) + .status() + .await?; + assert!( + status.success(), + "OAuth external update child failed: {status}" + ); + assert_eq!(tool_call_count.load(Ordering::SeqCst), 2); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by running_client_keeps_and_can_persist_over_external_oauth_update"] +async fn oauth_external_update_child() -> anyhow::Result<()> { + let server_url = std::env::var(EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV)?; + let old_expires_at = SystemTime::now() + .duration_since(UNIX_EPOCH)? + .checked_add(REFRESH_SKEW + Duration::from_secs(5)) + .expect("test expiry should fit") + .as_millis() as u64; + save_test_tokens( + &server_url, + OLD_ACCESS_TOKEN, + OLD_REFRESH_TOKEN, + old_expires_at, + )?; + + 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?; + initialize_client(&client).await?; + + let external_expires_at = SystemTime::now() + .duration_since(UNIX_EPOCH)? + .checked_add(Duration::from_secs(7200)) + .expect("test expiry should fit") + .as_millis() as u64; + save_test_tokens( + &server_url, + EXTERNAL_ACCESS_TOKEN, + EXTERNAL_REFRESH_TOKEN, + external_expires_at, + )?; + + client + .call_tool( + "first-call-after-external-update".to_string(), + /*arguments*/ None, + /*meta*/ None, + Some(Duration::from_secs(5)), + ) + .await?; + + let refresh_at = UNIX_EPOCH + + Duration::from_millis(old_expires_at) + .saturating_sub(REFRESH_SKEW) + .saturating_add(Duration::from_millis(250)); + if let Ok(wait) = refresh_at.duration_since(SystemTime::now()) { + tokio::time::sleep(wait).await; + } + + client + .call_tool( + "second-call-after-stale-refresh".to_string(), + /*arguments*/ None, + /*meta*/ None, + Some(Duration::from_secs(5)), + ) + .await?; + + let credentials = std::fs::read_to_string( + std::path::Path::new(&std::env::var("CODEX_HOME")?).join(".credentials.json"), + )?; + assert!(credentials.contains(STALE_REFRESHED_ACCESS_TOKEN)); + assert!(credentials.contains(STALE_ROTATED_REFRESH_TOKEN)); + assert!(!credentials.contains(EXTERNAL_ACCESS_TOKEN)); + assert!(!credentials.contains(EXTERNAL_REFRESH_TOKEN)); + Ok(()) +} + +fn save_test_tokens( + server_url: &str, + access_token: &str, + refresh_token: &str, + expires_at: u64, +) -> anyhow::Result<()> { + 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))); + save_oauth_tokens( + SERVER_NAME, + &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: Some(expires_at), + }, + OAuthCredentialsStoreMode::File, + ) +}