mod streamable_http_test_support;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::McpAuthStatus;use codex_rmcp_client::RmcpClient;use codex_rmcp_client::StoredOAuthTokens;use codex_rmcp_client::WrappedOAuthTokenResponse;use codex_rmcp_client::determine_streamable_http_auth_status;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;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;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 CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL";const UNREFRESHABLE_SERVER_URL: &str = "https://unrefreshable.example/mcp";const UNEXPIRED_SERVER_URL: &str = "https://unexpired.example/mcp";const REFRESHABLE_SERVER_URL: &str = "https://refreshable.example/mcp";#[tokio::test(flavor = "multi_thread", worker_threads = 1)]async fn refreshes_expired_persisted_token_before_initialize() -> 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={REFRESH_TOKEN}"        )))        .respond_with(ResponseTemplate::new(200).set_body_json(json!({            "access_token": REFRESHED_ACCESS_TOKEN,            "token_type": "Bearer",            "expires_in": 7200,            "refresh_token": REFRESH_TOKEN,        })))        .expect(1)        .mount(&server)        .await;    Mock::given(method("POST"))        .and(path("/mcp"))        .and(header(            "authorization",            format!("Bearer {REFRESHED_ACCESS_TOKEN}"),        ))        .respond_with(|request: &Request| {            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-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 server_url = format!("{}/mcp", server.uri());    // 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 status = Command::new(std::env::current_exe()?)        .args(["oauth_startup_child", "--exact", "--ignored", "--nocapture"])        .env("CODEX_HOME", codex_home.path())        .env(CHILD_SERVER_URL_ENV, server_url)        .status()        .await?;    assert!(status.success(), "OAuth startup child failed: {status}");    server.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)]#[ignore = "spawned by reports_auth_status_for_persisted_credentials"]async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> {    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: UNREFRESHABLE_SERVER_URL.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(),    )?;    let status = auth_status(UNREFRESHABLE_SERVER_URL).await?;    assert_eq!(status, McpAuthStatus::NotLoggedIn);    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(),        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(),    )?;    let status = auth_status(UNEXPIRED_SERVER_URL).await?;    assert_eq!(status, McpAuthStatus::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(),        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(),    )?;    let status = auth_status(REFRESHABLE_SERVER_URL).await?;    assert_eq!(status, McpAuthStatus::OAuth);    Ok(())}async fn auth_status(server_url: &str) -> anyhow::Result<McpAuthStatus> {    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(),    )    .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)?;    // 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(),        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(),    )?;    // 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 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?;    Ok(())}