codex/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs

mod streamable_http_test_support;

use std::collections::HashMap;
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::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_exec_server::Environment;
use codex_rmcp_client::McpAuthState;
use codex_rmcp_client::McpLoginRequirement;
use codex_rmcp_client::McpOAuthRefreshMode;
use codex_rmcp_client::McpProtocolMode;
use codex_rmcp_client::OAuthDiscoveryTimeout;
use codex_rmcp_client::RmcpClient;
use codex_rmcp_client::StoredOAuthTokens;
use codex_rmcp_client::StreamableHttpRedirectMode;
use codex_rmcp_client::WrappedOAuthTokenResponse;
use codex_rmcp_client::determine_streamable_http_auth_status;
use codex_rmcp_client::is_authentication_required_error;
use codex_rmcp_client::save_oauth_tokens;
use codex_rmcp_client::with_http_headers_helper;
use codex_utils_cargo_bin::cargo_bin;
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;
use serde_json::json;
use tempfile::TempDir;
use tokio::process::Command;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::Request;
use wiremock::ResponseTemplate;
use wiremock::matchers::body_string_contains;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;

use streamable_http_test_support::initialize_client;
use streamable_http_test_support::initialize_client_with_timeout;

const SERVER_NAME: &str = "test-streamable-http-oauth-startup";
const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token";
const REFRESH_TOKEN: &str = "valid-refresh-token";
const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token";
const RESOURCE_API_KEY: &str = "resource-api-key-secret";
const RESOURCE_USER_AGENT: &str = "resource-only-user-agent";
const MCP_USER_AGENT: &str = concat!("codex-mcp-client/", env!("CARGO_PKG_VERSION"));
const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL";
const CHILD_REFRESH_MODE_ENV: &str = "MCP_TEST_OAUTH_REFRESH_MODE";
const CHILD_HELPER_COMMAND_ENV: &str = "MCP_TEST_OAUTH_STARTUP_HELPER_COMMAND";
const CHILD_RESOURCE_API_KEY_ENV: &str = "MCP_TEST_OAUTH_STARTUP_RESOURCE_API_KEY";
const CHILD_STORED_ISSUER_ENV: &str = "MCP_TEST_OAUTH_STARTUP_STORED_ISSUER";
const CHILD_ACCESS_TOKEN_EXPIRY_ENV: &str = "MCP_TEST_OAUTH_STARTUP_ACCESS_TOKEN_EXPIRY";
const CHILD_REFRESH_SUCCEEDS_ENV: &str = "MCP_TEST_REFRESH_SUCCEEDS";
const LEGACY_REFRESHABLE_SERVER_URL: &str = "https://legacy-refreshable.example/mcp";
const UNEXPIRED_SERVER_URL: &str = "https://unexpired.example/mcp";
const REFRESHABLE_SERVER_URL: &str = "https://refreshable.example/mcp";

#[derive(Clone, Copy)]
enum OAuthStartupScenario {
    DirectAuthorizationMetadata,
    GatewayHeadersHelper,
    SameOriginGatewayHeadersHelper,
    ProtectedResourceMetadata,
}

#[derive(Clone, Copy)]
enum IssuerMismatchAccessToken {
    Expired,
    Unexpired,
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn tool_call_preserves_challenge_only_after_silent_refresh_fails() -> anyhow::Result<()> {
    for refresh_succeeds in [true, false] {
        let server = MockServer::start().await;
        let server_url = format!("{}/mcp", server.uri());
        Mock::given(method("GET"))
            .and(path("/.well-known/oauth-authorization-server/mcp"))
            .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({
                "issuer": server_url,
                "authorization_endpoint": format!("{}/authorize", server.uri()),
                "token_endpoint": format!("{}/token", server.uri()),
            })))
            .mount(&server)
            .await;
        Mock::given(method("POST"))
            .and(path("/token"))
            .and(body_string_contains("grant_type=refresh_token"))
            .and(body_string_contains(format!(
                "refresh_token={REFRESH_TOKEN}"
            )))
            .respond_with(if refresh_succeeds {
                ResponseTemplate::new(/*s*/ 200).set_body_json(json!({
                    "access_token": REFRESHED_ACCESS_TOKEN,
                    "token_type": "Bearer",
                    "expires_in": 7200,
                    "refresh_token": REFRESH_TOKEN,
                }))
            } else {
                ResponseTemplate::new(/*s*/ 400).set_body_json(json!({"error": "invalid_grant"}))
            })
            .expect(1)
            .mount(&server)
            .await;
        Mock::given(method("POST"))
            .and(path("/mcp"))
            .respond_with(|request: &Request| {
                let body: Value = request.body_json().unwrap();
                let result = match body["method"].as_str() {
                    Some("initialize") => json!({
                        "protocolVersion": body["params"]["protocolVersion"],
                        "capabilities": {},
                        "serverInfo": {"name": "oauth-tool-call-test", "version": "1"},
                    }),
                    Some("notifications/initialized") => {
                        return ResponseTemplate::new(/*s*/ 202);
                    }
                    Some("tools/call") => {
                        if request.headers.get("authorization").unwrap()
                            != format!("Bearer {REFRESHED_ACCESS_TOKEN}").as_str()
                        {
                            return ResponseTemplate::new(/*s*/ 401).insert_header(
                                "www-authenticate",
                                r#"Bearer error="invalid_token""#,
                            );
                        }
                        json!({"content": [{"type": "text", "text": "refreshed"}]})
                    }
                    _ => return ResponseTemplate::new(/*s*/ 400),
                };
                ResponseTemplate::new(/*s*/ 200).set_body_json(json!({
                    "jsonrpc": "2.0", "id": body["id"], "result": result,
                }))
            })
            .mount(&server)
            .await;

        // The child owns its credential store; the parallel test runner's environment is unchanged.
        let codex_home = TempDir::new()?;
        let status = Command::new(std::env::current_exe()?)
            .args([
                "oauth_tool_call_child",
                "--exact",
                "--ignored",
                "--nocapture",
            ])
            .env("CODEX_HOME", codex_home.path())
            .env(CHILD_SERVER_URL_ENV, server_url)
            .env(CHILD_REFRESH_SUCCEEDS_ENV, refresh_succeeds.to_string())
            .status()
            .await?;
        assert!(status.success(), "OAuth tool call child failed: {status}");
        server.verify().await;
        let tool_calls = server
            .received_requests()
            .await
            .unwrap()
            .into_iter()
            .filter(|request| {
                request
                    .body_json::<Value>()
                    .ok()
                    .is_some_and(|body| body["method"] == "tools/call")
            })
            .count();
        assert_eq!(tool_calls, if refresh_succeeds { 2 } else { 1 });
    }
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[ignore = "spawned by tool_call_preserves_challenge_only_after_silent_refresh_fails"]
async fn oauth_tool_call_child() -> anyhow::Result<()> {
    let server_url = std::env::var(CHILD_SERVER_URL_ENV)?;
    let refresh_succeeds = std::env::var(CHILD_REFRESH_SUCCEEDS_ENV)? == "true";
    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(/*secs*/ 7200)));
    save_oauth_tokens(
        SERVER_NAME,
        &StoredOAuthTokens {
            server_name: SERVER_NAME.to_string(),
            url: server_url.clone(),
            issuer: Some(server_url.clone()),
            client_id: "test-client-id".to_string(),
            token_response: WrappedOAuthTokenResponse(response),
            expires_at: Some(
                (SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() + 7_200_000) as u64,
            ),
        },
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
    )
    .await?;
    let client = RmcpClient::new_streamable_http_client(
        SERVER_NAME,
        &server_url,
        /*bearer_token*/ None,
        /*http_headers*/ None,
        /*env_http_headers*/ None,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
        Environment::default_for_tests().get_http_client(),
        /*auth_provider*/ None,
    )
    .await?;
    initialize_client(&client).await?;
    let result = client
        .call_tool(
            "probe".to_string(),
            /*arguments*/ None,
            /*meta*/ None,
            Some(Duration::from_secs(/*secs*/ 5)),
        )
        .await?;
    if refresh_succeeds {
        assert_eq!(
            serde_json::to_value(result)?,
            json!({
                "content": [{"type": "text", "text": "refreshed"}],
            })
        );
    } else {
        assert_eq!(result.is_error, Some(true));
        assert_eq!(
            serde_json::to_value(result.meta)?,
            json!({
                "mcp/www_authenticate": [r#"Bearer error="invalid_token""#],
            })
        );
    }
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> {
    assert_expired_token_refresh(
        OAuthStartupScenario::DirectAuthorizationMetadata,
        McpOAuthRefreshMode::Coordinated,
    )
    .await
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn refreshes_oauth_with_gateway_headers_helper() -> anyhow::Result<()> {
    assert_expired_token_refresh(
        OAuthStartupScenario::GatewayHeadersHelper,
        McpOAuthRefreshMode::Legacy,
    )
    .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,
        McpOAuthRefreshMode::Legacy,
    )
    .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,
        McpOAuthRefreshMode::Legacy,
    )
    .await
}

async fn assert_expired_token_refresh(
    scenario: OAuthStartupScenario,
    refresh_mode: McpOAuthRefreshMode,
) -> 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::SameOriginGatewayHeadersHelper => (
            resource_url.clone(),
            "/mcp",
            "/.well-known/oauth-authorization-server/mcp",
        ),
        OAuthStartupScenario::ProtectedResourceMetadata => (
            format!("{resource_url}/?oauth=initialize"),
            "/mcp/",
            "/.well-known/oauth-authorization-server",
        ),
    };

    if matches!(scenario, OAuthStartupScenario::ProtectedResourceMetadata) {
        let resource_metadata_url = format!("{}/resource-metadata", server.uri());
        Mock::given(method("GET"))
            .and(path(mcp_path))
            .respond_with(ResponseTemplate::new(401).insert_header(
                "www-authenticate",
                format!("Bearer resource_metadata=\"{resource_metadata_url}\""),
            ))
            .expect(2)
            .mount(&server)
            .await;
        Mock::given(method("GET"))
            .and(path("/resource-metadata"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "resource": resource_url,
                "authorization_servers": [server.uri()],
            })))
            .expect(2)
            .mount(&server)
            .await;
    }

    let mut authorization_metadata = json!({
        "issuer": server_url.clone(),
        "authorization_endpoint": format!("{}/oauth/authorize", token_server.uri()),
        "token_endpoint": format!("{}/oauth/token", token_server.uri()),
        "scopes_supported": [""],
    });
    if matches!(scenario, OAuthStartupScenario::ProtectedResourceMetadata) {
        authorization_metadata["issuer"] = json!(server.uri());
    }
    let authorization_server_issuer = authorization_metadata["issuer"]
        .as_str()
        .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))
        .and(header("x-api-key", RESOURCE_API_KEY))
        .respond_with(ResponseTemplate::new(200).set_body_json(authorization_metadata))
        .expect(2)
        .mount(&server)
        .await;
    Mock::given(method("POST"))
        .and(path("/oauth/token"))
        .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}"
        )))
        .and({
            let expected_resource = resource_url.clone();
            move |request: &Request| {
                url::form_urlencoded::parse(&request.body)
                    .any(|(name, value)| name == "resource" && value == expected_resource)
            }
        })
        .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 {
                let response = ResponseTemplate::new(/*s*/ 200).set_body_json(json!({
                    "access_token": REFRESHED_ACCESS_TOKEN,
                    "token_type": "Bearer",
                    "expires_in": 7200,
                    "refresh_token": REFRESH_TOKEN,
                }));
                if refresh_mode == McpOAuthRefreshMode::Coordinated {
                    // Longer than the child's handshake timeout: refresh must finish first.
                    response.set_delay(Duration::from_secs(/*secs*/ 2))
                } else {
                    response
                }
            }
        })
        .expect(if same_origin_gateway { 2 } else { 1 })
        .mount(token_server)
        .await;
    Mock::given(method("POST"))
        .and(path(mcp_path))
        .and(header("user-agent", RESOURCE_USER_AGENT))
        .and(header("x-api-key", RESOURCE_API_KEY))
        .and(header(
            "authorization",
            format!("Bearer {REFRESHED_ACCESS_TOKEN}"),
        ))
        .respond_with(|request: &Request| {
            let body: Value = match request.body_json() {
                Ok(body) => body,
                Err(_) => {
                    return ResponseTemplate::new(400).set_body_string("invalid 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-startup-test",
                            "version": "0.0.0-test",
                        },
                    },
                })),
                Some("notifications/initialized") => ResponseTemplate::new(202),
                method => ResponseTemplate::new(400)
                    .set_body_string(format!("unexpected JSON-RPC method: {method:?}")),
            }
        })
        .expect(2)
        .mount(&server)
        .await;

    let codex_home = TempDir::new()?;
    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
    // an isolated home without mutating the parent test runner's environment.
    let mut command = Command::new(std::env::current_exe()?);
    command
        .args(["oauth_startup_child", "--exact", "--ignored", "--nocapture"])
        .env("CODEX_HOME", codex_home.path())
        .env(CHILD_SERVER_URL_ENV, server_url)
        .env(CHILD_STORED_ISSUER_ENV, authorization_server_issuer)
        .env(CHILD_RESOURCE_API_KEY_ENV, RESOURCE_API_KEY)
        .env(
            CHILD_REFRESH_MODE_ENV,
            if refresh_mode == McpOAuthRefreshMode::Coordinated {
                "coordinated"
            } else {
                "legacy"
            },
        )
        .env("MCP_TEST_AMBIENT_SECRET", "must-not-reach-helper");
    if with_headers_helper {
        command.env(
            CHILD_HELPER_COMMAND_ENV,
            format!(
                "\"{}\" --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
                .headers
                .get("proxy-authorization")
                .is_some_and(|value| value == "Bearer gateway-token")
        }));
    }
    let authorization_requests = authorization_server
        .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!(
        authorization_requests[0].headers.get("proxy-authorization"),
        None
    );
    assert_eq!(
        authorization_requests[0]
            .headers
            .get("user-agent")
            .map(http::HeaderValue::as_bytes),
        Some(MCP_USER_AGENT.as_bytes())
    );

    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn rejects_refresh_when_authorization_server_issuer_changes_before_startup()
-> anyhow::Result<()> {
    assert_issuer_mismatch_startup(IssuerMismatchAccessToken::Expired).await
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn does_not_refresh_unexpired_token_after_issuer_change_401() -> anyhow::Result<()> {
    assert_issuer_mismatch_startup(IssuerMismatchAccessToken::Unexpired).await
}

async fn assert_issuer_mismatch_startup(
    access_token: IssuerMismatchAccessToken,
) -> anyhow::Result<()> {
    let issuer_a = MockServer::start().await;
    let issuer_b = MockServer::start().await;
    let mcp_server = MockServer::start().await;
    let server_url = format!("{}/mcp", mcp_server.uri());
    let resource_metadata_url = format!("{}/resource-metadata", mcp_server.uri());

    Mock::given(method("GET"))
        .and(path("/mcp"))
        .respond_with(ResponseTemplate::new(401).insert_header(
            "www-authenticate",
            format!("Bearer resource_metadata=\"{resource_metadata_url}\""),
        ))
        .expect(1)
        .mount(&mcp_server)
        .await;
    Mock::given(method("GET"))
        .and(path("/resource-metadata"))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "resource": server_url,
            "authorization_servers": [issuer_b.uri()],
        })))
        .expect(1)
        .mount(&mcp_server)
        .await;
    Mock::given(method("GET"))
        .and(path("/.well-known/oauth-authorization-server"))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "issuer": issuer_b.uri(),
            "authorization_endpoint": format!("{}/oauth/authorize", issuer_b.uri()),
            "token_endpoint": format!("{}/oauth/token", issuer_b.uri()),
            "scopes_supported": [""],
        })))
        .expect(1)
        .mount(&issuer_b)
        .await;
    Mock::given(method("POST"))
        .and(path("/mcp"))
        .and(header(
            "authorization",
            format!("Bearer {EXPIRED_ACCESS_TOKEN}"),
        ))
        .respond_with(ResponseTemplate::new(401).insert_header(
            "www-authenticate",
            format!("Bearer resource_metadata=\"{resource_metadata_url}\""),
        ))
        .expect(match access_token {
            IssuerMismatchAccessToken::Expired => 0,
            IssuerMismatchAccessToken::Unexpired => 1,
        })
        .mount(&mcp_server)
        .await;

    let stored_issuer = issuer_a.uri();
    let status = run_issuer_startup_child(
        &server_url,
        &stored_issuer,
        match access_token {
            IssuerMismatchAccessToken::Expired => "expired",
            IssuerMismatchAccessToken::Unexpired => "unexpired",
        },
    )
    .await?;

    assert!(
        status.success(),
        "issuer mismatch startup child failed: {status}"
    );
    let issuer_b_requests = issuer_b.received_requests().await.unwrap_or_default();
    assert!(
        issuer_b_requests
            .iter()
            .all(|request| request.url.path() != "/oauth/token"),
        "stored refresh token must not be posted to the replacement issuer"
    );
    mcp_server.verify().await;
    issuer_b.verify().await;
    Ok(())
}

async fn run_issuer_startup_child(
    server_url: &str,
    stored_issuer: &str,
    access_token_expiry: &str,
) -> anyhow::Result<std::process::ExitStatus> {
    let codex_home = TempDir::new()?;
    Ok(Command::new(std::env::current_exe()?)
        .args([
            "issuer_startup_child",
            "--exact",
            "--ignored",
            "--nocapture",
        ])
        .env("CODEX_HOME", codex_home.path())
        .env(CHILD_SERVER_URL_ENV, server_url)
        .env(CHILD_STORED_ISSUER_ENV, stored_issuer)
        .env(CHILD_ACCESS_TOKEN_EXPIRY_ENV, access_token_expiry)
        .status()
        .await?)
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn does_not_refresh_against_metadata_from_second_discovery() -> anyhow::Result<()> {
    let issuer_a = MockServer::start().await;
    let issuer_b = MockServer::start().await;
    let mcp_server = MockServer::start().await;
    let server_url = format!("{}/mcp", mcp_server.uri());
    let resource_metadata_url = format!("{}/resource-metadata", mcp_server.uri());
    let discovery_count = Arc::new(AtomicUsize::new(0));
    let issuer_a_url = issuer_a.uri();
    let issuer_b_url = issuer_b.uri();

    Mock::given(method("GET"))
        .and(path("/mcp"))
        .respond_with(ResponseTemplate::new(401).insert_header(
            "www-authenticate",
            format!("Bearer resource_metadata=\"{resource_metadata_url}\""),
        ))
        .expect(1)
        .mount(&mcp_server)
        .await;
    Mock::given(method("GET"))
        .and(path("/resource-metadata"))
        .respond_with({
            let discovery_count = Arc::clone(&discovery_count);
            let issuer_a_url = issuer_a_url.clone();
            let issuer_b_url = issuer_b_url.clone();
            let metadata_resource_url = server_url.clone();
            move |_request: &Request| {
                let issuer = if discovery_count.fetch_add(1, Ordering::SeqCst) == 0 {
                    &issuer_a_url
                } else {
                    &issuer_b_url
                };
                ResponseTemplate::new(200).set_body_json(json!({
                    "resource": metadata_resource_url,
                    "authorization_servers": [issuer],
                }))
            }
        })
        .expect(1)
        .mount(&mcp_server)
        .await;
    Mock::given(method("GET"))
        .and(path("/.well-known/oauth-authorization-server"))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "issuer": issuer_a.uri(),
            "authorization_endpoint": format!("{}/oauth/authorize", issuer_a.uri()),
            "token_endpoint": format!("{}/oauth/token", issuer_a.uri()),
            "scopes_supported": [""],
        })))
        .expect(1)
        .mount(&issuer_a)
        .await;
    Mock::given(method("GET"))
        .and(path("/.well-known/oauth-authorization-server"))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "issuer": issuer_b.uri(),
            "authorization_endpoint": format!("{}/oauth/authorize", issuer_b.uri()),
            "token_endpoint": format!("{}/oauth/token", issuer_b.uri()),
            "scopes_supported": [""],
        })))
        .mount(&issuer_b)
        .await;
    Mock::given(method("POST"))
        .and(path("/mcp"))
        .and(header(
            "authorization",
            format!("Bearer {EXPIRED_ACCESS_TOKEN}"),
        ))
        .respond_with(ResponseTemplate::new(401).insert_header(
            "www-authenticate",
            format!("Bearer resource_metadata=\"{resource_metadata_url}\""),
        ))
        .expect(1)
        .mount(&mcp_server)
        .await;
    Mock::given(method("POST"))
        .and(path("/oauth/token"))
        .and(body_string_contains(format!(
            "refresh_token={REFRESH_TOKEN}"
        )))
        .respond_with(ResponseTemplate::new(400).set_body_json(json!({
            "error": "invalid_grant",
        })))
        .expect(1)
        .mount(&issuer_a)
        .await;
    Mock::given(method("POST"))
        .and(path("/oauth/token"))
        .respond_with(ResponseTemplate::new(500))
        .expect(0)
        .mount(&issuer_b)
        .await;

    let status = run_issuer_startup_child(&server_url, &issuer_a_url, "unexpired").await?;

    assert!(
        status.success(),
        "metadata swap startup child failed: {status}"
    );
    mcp_server.verify().await;
    issuer_a.verify().await;
    issuer_b.verify().await;
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn reports_auth_status_for_persisted_credentials() -> anyhow::Result<()> {
    let codex_home = TempDir::new()?;

    let status = Command::new(std::env::current_exe()?)
        .args([
            "persisted_credentials_auth_status_child",
            "--exact",
            "--ignored",
            "--nocapture",
        ])
        .env("CODEX_HOME", codex_home.path())
        .status()
        .await?;

    assert!(
        status.success(),
        "persisted credentials auth status child failed: {status}"
    );
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn identifies_expired_unrefreshable_token_startup_error() -> 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()),
        })))
        .expect(1)
        .mount(&server)
        .await;

    let codex_home = TempDir::new()?;
    let status = Command::new(std::env::current_exe()?)
        .args([
            "expired_unrefreshable_startup_child",
            "--exact",
            "--ignored",
            "--nocapture",
        ])
        .env("CODEX_HOME", codex_home.path())
        .env(CHILD_SERVER_URL_ENV, format!("{}/mcp", server.uri()))
        .status()
        .await?;

    assert!(
        status.success(),
        "expired OAuth startup child failed: {status}"
    );
    server.verify().await;
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[ignore = "spawned by reports_auth_status_for_persisted_credentials"]
async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> {
    let first_login_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", first_login_server.uri()),
            "token_endpoint": format!("{}/oauth/token", first_login_server.uri()),
        })))
        .expect(1)
        .mount(&first_login_server)
        .await;

    let status = auth_status(&format!("{}/mcp", first_login_server.uri())).await?;
    assert_eq!(status, McpAuthState::LoggedOut(McpLoginRequirement::Login));
    first_login_server.verify().await;

    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())));
    let tokens = StoredOAuthTokens {
        server_name: SERVER_NAME.to_string(),
        url: LEGACY_REFRESHABLE_SERVER_URL.to_string(),
        issuer: None,
        client_id: "test-client-id".to_string(),
        token_response: WrappedOAuthTokenResponse(response),
        expires_at: Some(0),
    };
    save_oauth_tokens(
        SERVER_NAME,
        &tokens,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
    )
    .await?;

    let status = auth_status(LEGACY_REFRESHABLE_SERVER_URL).await?;
    assert_eq!(
        status,
        McpAuthState::LoggedOut(McpLoginRequirement::Reauthentication)
    );

    let response = OAuthTokenResponse::new(
        AccessToken::new("unexpired-access-token".to_string()),
        BasicTokenType::Bearer,
        VendorExtraTokenFields::default(),
    );
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_else(|_| Duration::from_secs(0))
        .as_millis() as u64;
    let tokens = StoredOAuthTokens {
        server_name: SERVER_NAME.to_string(),
        url: UNEXPIRED_SERVER_URL.to_string(),
        issuer: None,
        client_id: "test-client-id".to_string(),
        token_response: WrappedOAuthTokenResponse(response),
        expires_at: Some(now.saturating_add(/*rhs*/ 60_000)),
    };
    save_oauth_tokens(
        SERVER_NAME,
        &tokens,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
    )
    .await?;

    let status = auth_status(UNEXPIRED_SERVER_URL).await?;
    assert_eq!(status, McpAuthState::OAuth);

    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())));
    let tokens = StoredOAuthTokens {
        server_name: SERVER_NAME.to_string(),
        url: REFRESHABLE_SERVER_URL.to_string(),
        issuer: Some("https://issuer.example.test".to_string()),
        client_id: "test-client-id".to_string(),
        token_response: WrappedOAuthTokenResponse(response),
        expires_at: Some(0),
    };
    save_oauth_tokens(
        SERVER_NAME,
        &tokens,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
    )
    .await?;

    let status = auth_status(REFRESHABLE_SERVER_URL).await?;
    assert_eq!(status, McpAuthState::OAuth);
    Ok(())
}

async fn auth_status(server_url: &str) -> anyhow::Result<McpAuthState> {
    determine_streamable_http_auth_status(
        SERVER_NAME,
        server_url,
        /*bearer_token_env_var*/ None,
        /*http_headers*/ None,
        /*env_http_headers*/ None,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
        Environment::default_for_tests().get_http_client(),
        OAuthDiscoveryTimeout::LOCAL,
        StreamableHttpRedirectMode::Legacy,
    )
    .await
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"]
async fn oauth_startup_child() -> anyhow::Result<()> {
    let server_url = std::env::var(CHILD_SERVER_URL_ENV)?;
    let refresh_mode = if std::env::var(CHILD_REFRESH_MODE_ENV).as_deref() == Ok("coordinated") {
        McpOAuthRefreshMode::Coordinated
    } else {
        McpOAuthRefreshMode::Legacy
    };

    // Save an expired access token with a valid refresh token so startup must
    // refresh before sending the initialize request.
    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(),
        issuer: Some(std::env::var(CHILD_STORED_ISSUER_ENV)?),
        client_id: "test-client-id".to_string(),
        token_response: WrappedOAuthTokenResponse(response),
        expires_at: Some(0),
    };
    save_oauth_tokens(
        SERVER_NAME,
        &tokens,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
    )
    .await?;

    // This mirrors create_client's transport and initialization setup, except
    // it omits the direct bearer token. Supplying that token would bypass the
    // persisted OAuth credentials and the startup refresh under test.
    let mut http_client = Environment::default_for_tests().get_http_client();
    if let Ok(helper_command) = std::env::var(CHILD_HELPER_COMMAND_ENV) {
        http_client = with_http_headers_helper(
            http_client,
            &server_url,
            &helper_command,
            std::env::current_dir()?,
        )?;
    }
    let client = RmcpClient::new_streamable_http_client_with_protocol_mode_and_redirect_mode(
        SERVER_NAME,
        &server_url,
        /*bearer_token*/ None,
        Some(HashMap::from([(
            "User-Agent".to_string(),
            RESOURCE_USER_AGENT.to_string(),
        )])),
        Some(HashMap::from([(
            "X-Api-Key".to_string(),
            CHILD_RESOURCE_API_KEY_ENV.to_string(),
        )])),
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
        http_client,
        /*auth_provider*/ None,
        McpProtocolMode::Legacy,
        StreamableHttpRedirectMode::Legacy,
        refresh_mode,
    )
    .await?;

    if refresh_mode == McpOAuthRefreshMode::Coordinated {
        initialize_client_with_timeout(&client, Duration::from_secs(/*secs*/ 1)).await?;
    } else {
        initialize_client(&client).await?;
    }
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[ignore = "spawned by identifies_expired_unrefreshable_token_startup_error"]
async fn expired_unrefreshable_startup_child() -> anyhow::Result<()> {
    let server_url = std::env::var(CHILD_SERVER_URL_ENV)?;
    let response = OAuthTokenResponse::new(
        AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()),
        BasicTokenType::Bearer,
        VendorExtraTokenFields::default(),
    );
    let tokens = StoredOAuthTokens {
        server_name: SERVER_NAME.to_string(),
        url: server_url.clone(),
        issuer: None,
        client_id: "test-client-id".to_string(),
        token_response: WrappedOAuthTokenResponse(response),
        expires_at: Some(0),
    };
    save_oauth_tokens(
        SERVER_NAME,
        &tokens,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
    )
    .await?;

    let client = RmcpClient::new_streamable_http_client(
        SERVER_NAME,
        &server_url,
        /*bearer_token*/ None,
        /*http_headers*/ None,
        /*env_http_headers*/ None,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
        Environment::default_for_tests().get_http_client(),
        /*auth_provider*/ None,
    )
    .await?;

    let error = initialize_client(&client)
        .await
        .expect_err("expired token without a refresh token should fail startup");
    assert!(is_authentication_required_error(&error));
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[ignore = "spawned by issuer startup tests"]
async fn issuer_startup_child() -> anyhow::Result<()> {
    let server_url = std::env::var(CHILD_SERVER_URL_ENV)?;
    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 expires_at = match std::env::var(CHILD_ACCESS_TOKEN_EXPIRY_ENV).as_deref() {
        Ok("expired") | Err(_) => Some(0),
        Ok("unexpired") => {
            let now = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_else(|_| Duration::from_secs(0))
                .as_millis() as u64;
            Some(now.saturating_add(/*rhs*/ 60_000))
        }
        Ok(value) => anyhow::bail!("unexpected access token expiry fixture: {value}"),
    };
    let tokens = StoredOAuthTokens {
        server_name: SERVER_NAME.to_string(),
        url: server_url.clone(),
        issuer: Some(std::env::var(CHILD_STORED_ISSUER_ENV)?),
        client_id: "test-client-id".to_string(),
        token_response: WrappedOAuthTokenResponse(response),
        expires_at,
    };
    save_oauth_tokens(
        SERVER_NAME,
        &tokens,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
    )
    .await?;

    let client = RmcpClient::new_streamable_http_client(
        SERVER_NAME,
        &server_url,
        /*bearer_token*/ None,
        /*http_headers*/ None,
        /*env_http_headers*/ None,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
        Environment::default_for_tests().get_http_client(),
        /*auth_provider*/ None,
    )
    .await;

    let error = match client {
        Ok(client) => initialize_client(&client)
            .await
            .expect_err("stored refresh token must fail when startup cannot refresh safely"),
        Err(error) => error,
    };
    assert!(
        is_authentication_required_error(&error),
        "unexpected issuer startup error: {error:#}"
    );
    Ok(())
}