From 4914889cdd3a0549beffdb09002b71dfc66fefe2 Mon Sep 17 00:00:00 2001 From: Cooper Gamble Date: Wed, 17 Jun 2026 15:46:17 +0000 Subject: [PATCH] [codex] close remaining WIF readiness gaps [ci changed_files] --- codex-rs/Cargo.lock | 1 + .../src/request_processors/mcp_processor.rs | 16 +++- codex-rs/app-server/tests/common/lib.rs | 1 + .../tests/common/workload_identity.rs | 29 +++++-- .../tests/suite/v2/mcp_server_status.rs | 84 ++++++++++++++++++- codex-rs/cloud-config/Cargo.toml | 1 + codex-rs/cloud-config/src/service.rs | 3 +- codex-rs/cloud-config/src/service_tests.rs | 54 ++++++++++++ codex-rs/exec/Cargo.toml | 4 + codex-rs/tui/Cargo.toml | 4 + 10 files changed, 188 insertions(+), 9 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index fe4f747d1e..4cad1ab480 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2418,6 +2418,7 @@ version = "0.0.0" dependencies = [ "base64 0.22.1", "chrono", + "codex-app-server-protocol", "codex-backend-client", "codex-config", "codex-core", diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index bbf05e8606..95fc7e891a 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -125,6 +125,20 @@ impl McpRequestProcessor { } } + async fn resolve_optional_mcp_auth(&self, operation: &'static str) -> Option { + match self.auth_manager.auth().await { + Ok(auth) => auth, + Err(err) => { + warn!( + error = %err, + operation, + "failed to resolve optional MCP auth; using cached auth" + ); + self.auth_manager.auth_cached() + } + } + } + async fn mcp_server_oauth_login_response( &self, params: McpServerOauthLoginParams, @@ -245,7 +259,7 @@ impl McpRequestProcessor { .await } }; - let auth = self.auth_manager.auth_cached(); + let auth = self.resolve_optional_mcp_auth("mcpServerStatus/list").await; let environment_manager = self.thread_manager.environment_manager(); // This status path has no turn-selected environment. Use config cwd // as the local stdio fallback; named environment stdio MCPs must diff --git a/codex-rs/app-server/tests/common/lib.rs b/codex-rs/app-server/tests/common/lib.rs index 09d48e2590..74ba37c85e 100644 --- a/codex-rs/app-server/tests/common/lib.rs +++ b/codex-rs/app-server/tests/common/lib.rs @@ -50,6 +50,7 @@ pub use test_app_server::DISABLE_PLUGIN_STARTUP_TASKS_ARG; pub use test_app_server::TestAppServer; pub use workload_identity::ExpiringWorkloadIdentityFixture; pub use workload_identity::configure_expiring_workload_identity; +pub use workload_identity::configure_expiring_workload_identity_without_cloud_config_mock; pub fn to_response(response: JSONRPCResponse) -> anyhow::Result { let value = serde_json::to_value(response.result)?; diff --git a/codex-rs/app-server/tests/common/workload_identity.rs b/codex-rs/app-server/tests/common/workload_identity.rs index a820a949a7..f9eb6c2c8c 100644 --- a/codex-rs/app-server/tests/common/workload_identity.rs +++ b/codex-rs/app-server/tests/common/workload_identity.rs @@ -28,6 +28,21 @@ impl ExpiringWorkloadIdentityFixture { pub async fn configure_expiring_workload_identity( codex_home: &Path, server: &MockServer, +) -> Result { + configure_expiring_workload_identity_inner(codex_home, server, true).await +} + +pub async fn configure_expiring_workload_identity_without_cloud_config_mock( + codex_home: &Path, + server: &MockServer, +) -> Result { + configure_expiring_workload_identity_inner(codex_home, server, false).await +} + +async fn configure_expiring_workload_identity_inner( + codex_home: &Path, + server: &MockServer, + mock_cloud_config: bool, ) -> Result { let token_path = codex_home.join("projected-workload-token"); tokio::fs::write(&token_path, "external-subject-token\n").await?; @@ -72,12 +87,14 @@ path = {token_path_toml} .expect(1..) .mount(server) .await; - Mock::given(method("GET")) - .and(path("/backend-api/wham/config/bundle")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) - .expect(1) - .mount(server) - .await; + if mock_cloud_config { + Mock::given(method("GET")) + .and(path("/backend-api/wham/config/bundle")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .expect(1) + .mount(server) + .await; + } Mock::given(method("GET")) .and(path("/backend-api/accounts/workspace_test/settings")) .and(header("authorization", format!("Bearer {access_token}"))) diff --git a/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs b/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs index aa2d5f6868..36c5fdd218 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs @@ -7,6 +7,7 @@ use std::time::Duration; use anyhow::Result; use app_test_support::TestAppServer; use app_test_support::configure_expiring_workload_identity; +use app_test_support::configure_expiring_workload_identity_without_cloud_config_mock; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::to_response; use app_test_support::write_mock_responses_config_toml; @@ -131,6 +132,10 @@ async fn mcp_server_status_list_succeeds_when_workload_identity_is_unavailable() )?; let config_path = codex_home.path().join("config.toml"); let mut config_toml = std::fs::read_to_string(&config_path)?; + config_toml = config_toml.replace( + "supports_websockets = false\n", + "supports_websockets = false\nrequires_openai_auth = false\n", + ); config_toml.insert_str( 0, &format!("chatgpt_base_url = \"{}/backend-api\"\n", server.uri()), @@ -177,6 +182,74 @@ url = "{mcp_server_url}/mcp" Ok(()) } +#[tokio::test] +async fn mcp_server_status_list_resolves_fresh_wif_for_codex_apps() -> Result<()> { + let auth_server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + let (apps_server_url, apps_server_handle) = + start_mcp_server_at_path("apps.lookup", "/api/codex/ps/mcp").await?; + let codex_home = TempDir::new()?; + let config_path = codex_home.path().join("config.toml"); + std::fs::write( + config_path, + format!( + r#" +model = "mock-model" +model_provider = "mock_provider" +chatgpt_base_url = "{apps_server_url}" + +[features] +apps = true + +[model_providers.mock_provider] +name = "Local provider" +base_url = "{}/v1" +wire_api = "responses" +requires_openai_auth = false +"#, + auth_server.uri(), + ), + )?; + let _workload_identity = configure_expiring_workload_identity_without_cloud_config_mock( + codex_home.path(), + &auth_server, + ) + .await?; + + let mut mcp = TestAppServer::new_without_managed_config_with_env( + codex_home.path(), + &[("OPENAI_API_KEY", None)], + ) + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_list_mcp_server_status_request(ListMcpServerStatusParams { + cursor: None, + limit: None, + detail: None, + thread_id: None, + }) + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let response: ListMcpServerStatusResponse = to_response(response)?; + + let codex_apps = response + .data + .iter() + .find(|status| status.name == "codex_apps") + .expect("fresh WIF auth should enable Codex Apps status"); + assert!(codex_apps.tools.contains_key("apps.lookup")); + auth_server.verify().await; + + apps_server_handle.abort(); + let _ = apps_server_handle.await; + Ok(()) +} + #[tokio::test] async fn mcp_server_status_list_uses_thread_project_local_config() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; @@ -514,6 +587,10 @@ url = "{underscore_server_url}/mcp" } async fn start_mcp_server(tool_name: &str) -> Result<(String, JoinHandle<()>)> { + start_mcp_server_at_path(tool_name, "/mcp").await +} + +async fn start_mcp_server_at_path(tool_name: &str, path: &str) -> Result<(String, JoinHandle<()>)> { let listener = TcpListener::bind("127.0.0.1:0").await?; let addr = listener.local_addr()?; let tool_name = Arc::new(tool_name.to_string()); @@ -526,7 +603,12 @@ async fn start_mcp_server(tool_name: &str) -> Result<(String, JoinHandle<()>)> { Arc::new(LocalSessionManager::default()), StreamableHttpServerConfig::default(), ); - let router = Router::new().nest_service("/mcp", mcp_service); + let router = Router::new() + .route( + "/api/codex/config/bundle", + axum::routing::get(|| async { axum::Json(json!({})) }), + ) + .nest_service(path, mcp_service); let handle = tokio::spawn(async move { let _ = axum::serve(listener, router).await; diff --git a/codex-rs/cloud-config/Cargo.toml b/codex-rs/cloud-config/Cargo.toml index 6bf58c8399..eef1114cef 100644 --- a/codex-rs/cloud-config/Cargo.toml +++ b/codex-rs/cloud-config/Cargo.toml @@ -25,6 +25,7 @@ tokio = { workspace = true, features = ["fs", "rt", "sync", "time"] } tracing = { workspace = true } [dev-dependencies] +codex-app-server-protocol = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "test-util", "time"] } diff --git a/codex-rs/cloud-config/src/service.rs b/codex-rs/cloud-config/src/service.rs index 457514c9fe..47dcd11e81 100644 --- a/codex-rs/cloud-config/src/service.rs +++ b/codex-rs/cloud-config/src/service.rs @@ -494,7 +494,8 @@ where Ok(auth) => auth, Err(error) => { tracing::error!(error = %error, "Failed to resolve auth while refreshing cloud config bundle"); - return false; + emit_load_metric("refresh", "error", /*bundle*/ None); + return true; } }; let Some(auth) = auth else { diff --git a/codex-rs/cloud-config/src/service_tests.rs b/codex-rs/cloud-config/src/service_tests.rs index d63fef4c54..041ee80ee7 100644 --- a/codex-rs/cloud-config/src/service_tests.rs +++ b/codex-rs/cloud-config/src/service_tests.rs @@ -8,6 +8,7 @@ use crate::cache::CloudConfigBundleCache; use crate::metrics::bundle_shape_tag; use base64::Engine; use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use codex_app_server_protocol::AuthMode; use codex_backend_client::ConfigBundleResponse; use codex_backend_client::DeliveredTomlFragment; use codex_config::AbsolutePathBuf; @@ -17,6 +18,10 @@ use codex_config::CloudRequirementsFragment; use codex_config::CloudRequirementsTomlBundle; use codex_config::types::AuthCredentialsStoreMode; use codex_login::AuthKeyringBackendKind; +use codex_login::auth::ExternalAuth; +use codex_login::auth::ExternalAuthFuture; +use codex_login::auth::ExternalAuthRefreshContext; +use codex_login::auth::ExternalAuthTokens; use pretty_assertions::assert_eq; use serde_json::json; use std::collections::VecDeque; @@ -238,6 +243,29 @@ impl BundleClient for PendingBundleClient { } } +struct FailingRequiredExternalAuth; + +impl ExternalAuth for FailingRequiredExternalAuth { + fn auth_mode(&self) -> AuthMode { + AuthMode::Chatgpt + } + + fn requires_successful_resolution(&self) -> bool { + true + } + + fn resolve(&self) -> ExternalAuthFuture<'_, Option> { + Box::pin(async { Err(std::io::Error::other("transient WIF failure")) }) + } + + fn refresh( + &self, + _context: ExternalAuthRefreshContext, + ) -> ExternalAuthFuture<'_, ExternalAuthTokens> { + Box::pin(async { Err(std::io::Error::other("transient WIF failure")) }) + } +} + struct SequenceBundleClient { responses: tokio::sync::Mutex>>, request_count: AtomicUsize, @@ -962,6 +990,32 @@ async fn refresh_from_remote_updates_cached_bundle() { assert_eq!(signed_payload.bundle, replacement_bundle); } +#[tokio::test] +async fn refresh_continues_after_required_external_auth_resolution_error() { + let codex_home = tempdir().expect("tempdir"); + let auth_manager = Arc::new( + AuthManager::new( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + ) + .await, + ); + auth_manager.set_external_auth(Arc::new(FailingRequiredExternalAuth)); + let fetcher = Arc::new(StaticBundleClient::new(test_bundle())); + let service = CloudConfigBundleService::new( + auth_manager, + fetcher.clone(), + codex_home.path().to_path_buf(), + CLOUD_CONFIG_BUNDLE_TIMEOUT, + ); + + assert!(service.refresh_cache_once().await); + assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 0); +} + #[test] fn bundle_response_conversion_preserves_fragment_order() { let response = ConfigBundleResponse { diff --git a/codex-rs/exec/Cargo.toml b/codex-rs/exec/Cargo.toml index cdfc569b68..1e49059238 100644 --- a/codex-rs/exec/Cargo.toml +++ b/codex-rs/exec/Cargo.toml @@ -18,6 +18,10 @@ doctest = false name = "all" path = "tests/all.rs" +[features] +default = ["workload-identity-all-providers"] +workload-identity-all-providers = ["codex-login/workload-identity-all-providers"] + [lints] workspace = true diff --git a/codex-rs/tui/Cargo.toml b/codex-rs/tui/Cargo.toml index b000cba124..d05eaa25ac 100644 --- a/codex-rs/tui/Cargo.toml +++ b/codex-rs/tui/Cargo.toml @@ -18,6 +18,10 @@ name = "codex_tui" path = "src/lib.rs" doctest = false +[features] +default = ["workload-identity-all-providers"] +workload-identity-all-providers = ["codex-login/workload-identity-all-providers"] + [lints] workspace = true