From a8a4b73592357749e442d3175587d36e60667378 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 5 Jun 2026 05:10:24 +0000 Subject: [PATCH] test(rmcp): cover live refresh token rotation --- .../tests/streamable_http_oauth_startup.rs | 198 ++++++++++++++++++ 1 file changed, 198 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 2486952c48..38a6b090c7 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -37,10 +37,13 @@ const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; const REFRESH_TOKEN: &str = "valid-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; const SECOND_REFRESHED_ACCESS_TOKEN: &str = "second-refreshed-access-token"; +const THIRD_REFRESHED_ACCESS_TOKEN: &str = "third-refreshed-access-token"; const REPLACEMENT_REFRESH_TOKEN: &str = "replacement-refresh-token"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; const OMITTED_REFRESH_TOKEN_CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_OMITTED_REFRESH_TOKEN_SERVER_URL"; +const LIVE_REFRESH_TOKEN_CHILD_SERVER_URL_ENV: &str = + "MCP_TEST_OAUTH_LIVE_REFRESH_TOKEN_SERVER_URL"; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> { @@ -383,3 +386,198 @@ async fn oauth_preserved_refresh_token_child() -> anyhow::Result<()> { ); Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn omitted_refresh_token_preserves_running_manager_state() -> anyhow::Result<()> { + fn respond_to_mcp(request: &Request) -> ResponseTemplate { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + 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": {}, + "serverInfo": { + "name": "oauth-live-refresh-token-test", + "version": "0.0.0-test", + }, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "tools": [], + }, + })), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + } + + 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; + + let refresh_count = Arc::new(AtomicUsize::new(0)); + let refresh_count_for_responder = Arc::clone(&refresh_count); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .respond_with(move |request: &Request| { + let body = String::from_utf8_lossy(&request.body); + match refresh_count_for_responder.fetch_add(1, Ordering::SeqCst) { + 0 if body.contains(&format!("refresh_token={REFRESH_TOKEN}")) => { + ResponseTemplate::new(200).set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 0, + })) + } + 1 if body.contains(&format!("refresh_token={REFRESH_TOKEN}")) => { + ResponseTemplate::new(200).set_body_json(json!({ + "access_token": SECOND_REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 0, + "refresh_token": REPLACEMENT_REFRESH_TOKEN, + })) + } + 2 if body.contains(&format!("refresh_token={REPLACEMENT_REFRESH_TOKEN}")) => { + ResponseTemplate::new(200).set_body_json(json!({ + "access_token": THIRD_REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + })) + } + request_count => ResponseTemplate::new(400).set_body_string(format!( + "unexpected refresh request {request_count}: {body}" + )), + } + }) + .expect(3) + .mount(&server) + .await; + + for access_token in [ + REFRESHED_ACCESS_TOKEN, + SECOND_REFRESHED_ACCESS_TOKEN, + THIRD_REFRESHED_ACCESS_TOKEN, + ] { + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {access_token}"))) + .respond_with(respond_to_mcp) + .expect(1) + .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_live_refresh_token_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(LIVE_REFRESH_TOKEN_CHILD_SERVER_URL_ENV, server_url) + .status() + .await?; + assert!( + status.success(), + "OAuth live refresh token child failed: {status}" + ); + assert_eq!(refresh_count.load(Ordering::SeqCst), 3); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by omitted_refresh_token_preserves_running_manager_state"] +async fn oauth_live_refresh_token_child() -> anyhow::Result<()> { + let server_url = std::env::var(LIVE_REFRESH_TOKEN_CHILD_SERVER_URL_ENV)?; + let codex_home = std::env::var("CODEX_HOME")?; + + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_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 tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.clone(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }; + save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?; + + 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 credentials_path = std::path::Path::new(&codex_home).join(".credentials.json"); + let persisted_after_replacement: Value = + serde_json::from_str(&std::fs::read_to_string(&credentials_path)?)?; + let persisted_entry = persisted_after_replacement + .as_object() + .and_then(|entries| entries.values().next()) + .expect("one persisted OAuth credential entry"); + assert_eq!( + persisted_entry.get("access_token"), + Some(&json!(SECOND_REFRESHED_ACCESS_TOKEN)) + ); + assert_eq!( + persisted_entry.get("refresh_token"), + Some(&json!(REPLACEMENT_REFRESH_TOKEN)) + ); + + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await?; + assert!(tools.tools.is_empty()); + + let persisted_after_third_refresh: Value = + serde_json::from_str(&std::fs::read_to_string(credentials_path)?)?; + let persisted_entry = persisted_after_third_refresh + .as_object() + .and_then(|entries| entries.values().next()) + .expect("one persisted OAuth credential entry"); + assert_eq!( + persisted_entry.get("access_token"), + Some(&json!(THIRD_REFRESHED_ACCESS_TOKEN)) + ); + assert_eq!( + persisted_entry.get("refresh_token"), + Some(&json!(REPLACEMENT_REFRESH_TOKEN)) + ); + + client.shutdown().await; + Ok(()) +}