diff --git a/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs b/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs index 6808e75883..5a1ba15d78 100644 --- a/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs +++ b/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs @@ -210,6 +210,7 @@ async fn enterprise_login_attempt_lifecycle() -> Result<()> { keyring::set_default_credential_builder(Box::new(keyring)); let server = MockServer::builder().start().await; let idp_server = MockServer::start().await; + app_test_support::mount_workspace_routing(&server).await; let origin = server.uri(); let issuer = format!("{}/idp", idp_server.uri()); let resource = format!("{origin}/api/codex/ps/mcp"); @@ -250,6 +251,7 @@ async fn enterprise_login_attempt_lifecycle() -> Result<()> { .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "issuer": format!("{origin}/as"), "token_endpoint": format!("{origin}/as/token"), "grant_types_supported": ["urn:ietf:params:oauth:grant-type:jwt-bearer"], + "authorization_grant_profiles_supported": ["urn:ietf:params:oauth:grant-profile:id-jag"], "token_endpoint_auth_methods_supported": ["none"], }))) .mount(&server) @@ -272,10 +274,25 @@ async fn enterprise_login_attempt_lifecycle() -> Result<()> { let expected_jag = jag.clone(); let resource_exchanges = Arc::new(std::sync::atomic::AtomicUsize::new(0)); let exchanges = Arc::clone(&resource_exchanges); + let pending_account_refresh = Arc::new(Mutex::new(std::collections::VecDeque::<( + std::path::PathBuf, + Vec, + )>::new())); + let account_refresh = Arc::clone(&pending_account_refresh); Mock::given(method("POST")).and(path("/as/token")) .respond_with(move |request: &wiremock::Request| { let fields = url::form_urlencoded::parse(&request.body).into_owned().collect::>(); - assert_eq!(fields.get("assertion"), Some(&expected_jag)); + // EMA section 5: the resource AS gets the ID-JAG, never the IdP credential. + assert!(request.headers.get("authorization").is_none()); + assert_eq!(serde_json::to_value(fields).unwrap(), json!({ + "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", + "assertion": expected_jag, + "client_id": "mcp-client", + })); + if let Some((path, auth)) = account_refresh.lock().unwrap().pop_front() { + // Another process refreshes the account while the resource exchange runs. + std::fs::write(path, auth).expect("publish account refresh"); + } ResponseTemplate::new(200).set_body_json(json!({ "access_token": format!("enterprise-resource-{}", exchanges.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1), "token_type": "Bearer", "expires_in": 3600, @@ -291,6 +308,7 @@ async fn enterprise_login_attempt_lifecycle() -> Result<()> { .mount(&server) .await; let token_resource = resource.clone(); + let token_audience = format!("{origin}/as"); let pending_exchange = Arc::new(Mutex::new( None::<( tokio::sync::oneshot::Sender<()>, @@ -307,10 +325,29 @@ async fn enterprise_login_attempt_lifecycle() -> Result<()> { let fields = url::form_urlencoded::parse(&request.body) .into_owned() .collect::>(); - assert_eq!( - fields.get("subject_token").map(String::as_str), - Some("enterprise-refresh") - ); + let subject = fields.get("subject_token").map(String::as_str); + assert!(matches!( + subject, + Some( + "enterprise-refresh" + | "discovered-account-refresh" + | "rotated-account-refresh" + | "exchanged-account-refresh" + | "retried-account-refresh" + ) + )); + // EMA section 4: audience is the authorization server; resource is the MCP server. + assert!(request.headers.get("authorization").is_none()); + assert_eq!(serde_json::to_value(&fields).unwrap(), json!({ + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "requested_token_type": "urn:ietf:params:oauth:token-type:id-jag", + "subject_token_type": "urn:ietf:params:oauth:token-type:refresh_token", + "subject_token": subject, + "client_id": if subject == Some("enterprise-refresh") { "idp-client" } else { "account-client" }, + "audience": token_audience, + "resource": token_resource, + "scope": "files.read", + })); if let Some((arrived, release)) = exchange_gate.lock().unwrap().take() { let _ = arrived.send(()); let _ = release.recv(); @@ -331,14 +368,28 @@ async fn enterprise_login_attempt_lifecycle() -> Result<()> { ) .mount(&server) .await; + let pending_discovery = Arc::new(Mutex::new( + None::<( + tokio::sync::oneshot::Sender<()>, + std::sync::mpsc::Receiver<()>, + )>, + )); + let discovery_gate = Arc::clone(&pending_discovery); + let idp_metadata = json!({ + "issuer": issuer, + "authorization_endpoint": format!("{issuer}/authorize"), + "token_endpoint": format!("{issuer}/token"), + "token_endpoint_auth_methods_supported": ["none"], + }); Mock::given(method("GET")) .and(path("/.well-known/oauth-authorization-server/idp")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "issuer": issuer, - "authorization_endpoint": format!("{issuer}/authorize"), - "token_endpoint": format!("{issuer}/token"), - "token_endpoint_auth_methods_supported": ["none"], - }))) + .respond_with(move |_: &wiremock::Request| { + if let Some((arrived, release)) = discovery_gate.lock().unwrap().take() { + let _ = arrived.send(()); + let _ = release.recv(); + } + ResponseTemplate::new(200).set_body_json(idp_metadata.clone()) + }) .mount(&idp_server) .await; let assertion = format!( @@ -727,6 +778,44 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" ); } std::fs::write(config_path, saved_config)?; + let handshakes = server.received_requests().await.context("MCP handshakes")?; + let enterprise_discovery = handshakes + .iter() + .find_map(|request| { + let body: Value = serde_json::from_slice(&request.body).ok()?; + (request.headers.get("authorization")?.to_str().ok()? == "Bearer enterprise-resource-1" + && body["method"] == "server/discover") + .then_some(body) + }) + .context("authenticated EMA connection negotiates modern MCP")?; + let metadata = &enterprise_discovery["params"]["_meta"]; + assert_eq!( + metadata["io.modelcontextprotocol/protocolVersion"], + "2026-07-28" + ); + assert_eq!( + metadata["io.modelcontextprotocol/clientCapabilities"]["extensions"] + [codex_protocol::mcp::ENTERPRISE_MANAGED_AUTHORIZATION_EXTENSION_ID], + json!({}) + ); + let ordinary_initialize = handshakes + .iter() + .find_map(|request| { + let body: Value = serde_json::from_slice(&request.body).ok()?; + (request.headers.get("authorization")?.to_str().ok()? == "Bearer ordinary-access" + && body["method"] == "initialize") + .then_some(body) + }) + .context("ordinary MCP keeps its legacy handshake")?; + assert_eq!( + ordinary_initialize["params"]["protocolVersion"], + "2025-06-18" + ); + assert!( + ordinary_initialize["params"]["capabilities"]["extensions"] + .get(codex_protocol::mcp::ENTERPRISE_MANAGED_AUTHORIZATION_EXTENSION_ID) + .is_none() + ); let authorization = server .received_requests() .await @@ -912,6 +1001,312 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" json!({"name":"enterprise","threadId":thread_id,"loginId":recovery["loginId"],"success":true}) ); + // An eligible account login uses its own OAuth client, without a second sign-in. + let auth_path = home.path().join("auth.json"); + let mut account_auth: Value = serde_json::from_slice(&std::fs::read(&auth_path)?)?; + account_auth["tokens"]["oauth_provenance"] = + json!({"issuer": issuer, "client_id": "account-client"}); + account_auth["tokens"]["login_owner_id"] = json!("completed-account-login"); + account_auth["tokens"]["refresh_token"] = json!("account-refresh"); + account_auth["tokens"]["id_token"] = json!(format!("{}.{}.c2lnbmF0dXJl", + URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#), + URL_SAFE_NO_PAD.encode(json!({ + "iss": issuer, "sub": "enterprise-user", "aud": "account-client", "exp": now + 3600, + "https://api.openai.com/auth": {"chatgpt_user_id": "enterprise-user", "chatgpt_account_id": WORKSPACE_ID_REFRESHED, "chatgpt_plan_type": "business"} + }).to_string()))); + std::fs::write(&auth_path, serde_json::to_vec(&account_auth)?)?; + // Startup does not adopt a different login. Explicitly reload the account + // before starting a thread under the new completed login. + enterprise_rpc(&client, "account/read", json!({"refreshToken":true})).await?; + let (discovery_arrived, discovery_wait) = tokio::sync::oneshot::channel(); + let (release_discovery, discovery_release) = std::sync::mpsc::channel(); + *pending_discovery.lock().unwrap() = Some((discovery_arrived, discovery_release)); + let account_thread = enterprise_rpc(&client, "thread/start", json!({})).await?; + let account_id = account_thread["thread"]["id"] + .as_str() + .context("account thread ID")?; + let account_call = json!({"threadId":account_id,"server":"enterprise","tool":"calendar_list_events","arguments":{"query":"hello"}}); + let mut initial_call = Box::pin(enterprise_rpc( + &client, + "mcpServer/tool/call", + account_call.clone(), + )); + tokio::select! { + arrived = discovery_wait => arrived?, + result = &mut initial_call => anyhow::bail!( + "fresh EMA startup failed before IdP discovery: {result:?}" + ), + } + // Account refresh can publish a newer grant while IdP metadata is still pending. + // EMA must acquire that newer credential before sending anything to the IdP. + account_auth["tokens"]["access_token"] = json!("account-after-discovery"); + account_auth["tokens"]["refresh_token"] = json!("discovered-account-refresh"); + std::fs::write(&auth_path, serde_json::to_vec(&account_auth)?)?; + let refreshed = enterprise_rpc(&client, "account/read", json!({"refreshToken":true})).await; + account_auth["tokens"]["access_token"] = json!("account-after-resource-exchange"); + account_auth["tokens"]["id_token"] = json!(format!("{}.{}.c2lnbmF0dXJl", + URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#), + URL_SAFE_NO_PAD.encode(json!({ + "iss": issuer, "sub": "enterprise-user", "aud": "account-client", "iat": now, "exp": now + 7200, + "https://api.openai.com/auth": {"chatgpt_user_id": "enterprise-user", "chatgpt_account_id": WORKSPACE_ID_REFRESHED, "chatgpt_plan_type": "business"} + }).to_string()))); + pending_account_refresh + .lock() + .unwrap() + .push_back((auth_path.clone(), serde_json::to_vec(&account_auth)?)); + release_discovery.send(()).expect("release IdP discovery"); + refreshed.context("IdP discovery must not block account refresh")?; + assert_eq!(initial_call.await?, ordinary_result); + assert_eq!( + resource_exchanges.load(std::sync::atomic::Ordering::SeqCst), + 4 + ); + + // A no-op refresh keeps the established connection and its cached bearer. + enterprise_rpc(&client, "config/mcpServer/reload", json!({})).await?; + assert_eq!( + enterprise_rpc(&client, "mcpServer/tool/call", account_call.clone()).await?, + ordinary_result + ); + assert_eq!( + resource_exchanges.load(std::sync::atomic::Ordering::SeqCst), + 4 + ); + + let request_count = server + .received_requests() + .await + .context("MCP requests")? + .len(); + account_auth["tokens"]["refresh_token"] = json!("rotated-account-refresh"); + std::fs::write(&auth_path, serde_json::to_vec(&account_auth)?)?; + Mock::given(method("POST")) + .and(path("/api/codex/ps/mcp")) + .and(header("authorization", "Bearer enterprise-resource-4")) + .and(body_partial_json(json!({"method":"tools/call"}))) + .respond_with(ResponseTemplate::new(401)) + .with_priority(1) + .expect(1) + .up_to_n_times(1) + .mount(&server) + .await; + assert_eq!( + enterprise_rpc(&client, "mcpServer/tool/call", account_call.clone()).await?["isError"], + true + ); + assert_eq!( + resource_exchanges.load(std::sync::atomic::Ordering::SeqCst), + 4 + ); + // Rotate during both resource exchanges; discard stale bearers without revoking the owner. + for refresh_token in ["exchanged-account-refresh", "retried-account-refresh"] { + account_auth["tokens"]["refresh_token"] = json!(refresh_token); + pending_account_refresh + .lock() + .unwrap() + .push_back((auth_path.clone(), serde_json::to_vec(&account_auth)?)); + } + assert!( + enterprise_rpc(&client, "mcpServer/tool/call", account_call.clone()) + .await + .is_err() + ); + assert_eq!( + resource_exchanges.load(std::sync::atomic::Ordering::SeqCst), + 6 + ); + // Retry exhaustion fails only this call. The next explicit call can still renew. + assert_eq!( + enterprise_rpc(&client, "mcpServer/tool/call", account_call.clone()).await?, + ordinary_result + ); + let account_subjects = idp_server + .received_requests() + .await + .context("IdP requests")? + .iter() + .filter_map(|request| { + url::form_urlencoded::parse(&request.body) + .find(|(name, _)| name == "subject_token") + .map(|(_, token)| token.into_owned()) + }) + .filter(|token| token != "enterprise-refresh") + .collect::>(); + assert_eq!( + account_subjects, + [ + "discovered-account-refresh", + "rotated-account-refresh", + "exchanged-account-refresh", + "retried-account-refresh" + ] + ); + let renewed_calls = server + .received_requests() + .await + .context("MCP requests")? + .into_iter() + .skip(request_count) + .filter_map(|request| { + let body: Value = serde_json::from_slice(&request.body).ok()?; + (body["method"] == "tools/call").then(|| { + request.headers["authorization"] + .to_str() + .expect("valid bearer header") + .to_owned() + }) + }) + .collect::>(); + assert_eq!( + renewed_calls, + [ + "Bearer enterprise-resource-4", + "Bearer enterprise-resource-7" + ], + "failed calls are not replayed and stale exchange bearers never reach MCP" + ); + + // One-shot responses avoid scoped-mock IDs being reordered by wiremock priorities. + // Resource-AS step-up is an administrator error, not an account sign-in failure. + Mock::given(method("POST")) + .and(path("/as/token")) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(json!({"error":"insufficient_user_authentication"})), + ) + .with_priority(1) + .expect(1) + .up_to_n_times(1) + .mount(&server) + .await; + let rejected_thread = enterprise_rpc(&client, "thread/start", json!({})).await?; + let rejected_id = &rejected_thread["thread"]["id"]; + let status = enterprise_rpc( + &client, + "mcpServerStatus/list", + json!({"threadId":rejected_id,"serverName":"enterprise"}), + ) + .await?; + assert_eq!(status["data"][0]["authStatus"], "bearerToken"); + let error = enterprise_rpc(&client, "mcpServer/tool/call", json!({"threadId":rejected_id,"server":"enterprise","tool":"calendar_list_events","arguments":{}})).await.unwrap_err(); + assert!( + error.to_string().contains("contact your administrator"), + "{error}" + ); + assert_eq!( + serde_json::from_slice::(&std::fs::read(&auth_path)?)?, + account_auth + ); + + // IdP step-up preserves the account grant and affects this registration across threads. + Mock::given(method("POST")) + .and(path("/idp/token")) + .and(wiremock::matchers::body_string_contains( + "subject_token=retried-account-refresh", + )) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(json!({"error":"insufficient_user_authentication"})), + ) + .with_priority(1) + .expect(1) + .up_to_n_times(1) + .mount(&idp_server) + .await; + Mock::given(method("POST")) + .and(path("/api/codex/ps/mcp")) + .and(header("authorization", "Bearer enterprise-resource-7")) + .and(body_partial_json(json!({"method":"tools/call"}))) + .respond_with(ResponseTemplate::new(401)) + .with_priority(1) + .expect(1) + .up_to_n_times(1) + .mount(&server) + .await; + assert_eq!( + enterprise_rpc(&client, "mcpServer/tool/call", account_call.clone()).await?["isError"], + true + ); + assert_eq!( + enterprise_rpc(&client, "mcpServer/tool/call", account_call).await?["isError"], + true + ); + assert_eq!( + serde_json::from_slice::(&std::fs::read(&auth_path)?)?, + account_auth + ); + // Access-token rotation must not forget the unchanged grant's sign-in requirement. + account_auth["tokens"]["access_token"] = json!("account-after-step-up"); + std::fs::write(&auth_path, serde_json::to_vec(&account_auth)?)?; + enterprise_rpc(&client, "account/read", json!({"refreshToken":true})).await?; + let recovered = enterprise_rpc(&client, "thread/start", json!({})).await?; + let recovered_call = json!({"threadId":recovered["thread"]["id"],"server":"enterprise","tool":"calendar_list_events","arguments":{"query":"hello"}}); + assert_eq!( + enterprise_rpc(&client, "mcpServer/tool/call", recovered_call).await?, + ordinary_result + ); + assert_eq!( + resource_exchanges.load(std::sync::atomic::Ordering::SeqCst), + 8 + ); + let recovered_subject = idp_server + .received_requests() + .await + .context("IdP requests")? + .iter() + .filter_map(|request| { + url::form_urlencoded::parse(&request.body) + .find(|(name, _)| name == "subject_token") + .map(|(_, token)| token.into_owned()) + }) + .next_back(); + assert_eq!(recovered_subject.as_deref(), Some("enterprise-refresh")); + + // Both failure stages preserve a dedicated grant. Status must finish startup first. + for (endpoint, identity_provider) in [("/idp/token", true), ("/as/token", false)] { + let target = if identity_provider { + &idp_server + } else { + &server + }; + Mock::given(method("POST")) + .and(path(endpoint)) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(json!({"error":"insufficient_user_authentication"})), + ) + .with_priority(1) + .expect(1) + .up_to_n_times(1) + .mount(target) + .await; + let rejected = enterprise_rpc(&client, "thread/start", json!({})).await?; + let rejected_id = &rejected["thread"]["id"]; + let status = enterprise_rpc( + &client, + "mcpServerStatus/list", + json!({"threadId":rejected_id,"serverName":"enterprise"}), + ) + .await?; + assert_eq!( + status["data"][0]["authStatus"], + if identity_provider { + "notLoggedIn" + } else { + "oAuth" + } + ); + let error = enterprise_rpc(&client, "mcpServer/tool/call", json!({"threadId":rejected_id,"server":"enterprise","tool":"calendar_list_events","arguments":{}})).await.unwrap_err(); + assert!( + error.to_string().contains(if identity_provider { + "Complete enterprise sign-in" + } else { + "contact your administrator" + }), + "{error}" + ); + } + let stalled = enterprise_rpc( &client, "mcpServer/oauth/login", @@ -973,7 +1368,7 @@ oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" ); assert_eq!( resource_exchanges.load(std::sync::atomic::Ordering::SeqCst), - 3 + 8 ); assert!( enterprise_rpc( diff --git a/codex-rs/core/tests/suite/scenarios_mcp_ema.rs b/codex-rs/core/tests/suite/scenarios_mcp_ema.rs index 1fd38f05b1..30599c8cef 100644 --- a/codex-rs/core/tests/suite/scenarios_mcp_ema.rs +++ b/codex-rs/core/tests/suite/scenarios_mcp_ema.rs @@ -1,4 +1,4 @@ -//! Failed enterprise authentication leaves ordinary MCP calls and model history unchanged. +//! Enterprise authentication and revocation preserve ordinary tools and accumulated history. use anyhow::Result; use codex_config::test_support::CloudConfigBundleFixture; @@ -13,6 +13,7 @@ use core_test_support::responses; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event_match; +use core_test_support::wait_for_event_with_timeout; use pretty_assertions::assert_eq; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -67,11 +68,16 @@ http_headers = {{ Authorization = "Bearer ordinary-access" }} .await; insta::assert_snapshot!("enterprise_sign_in_required", error); } - let startup = wait_for_event_match(&test.codex, |event| match event { - EventMsg::McpStartupComplete(summary) => Some(summary.clone()), - _ => None, - }) + // Allow MCP's 30-second startup deadline before checking the result. + let startup = wait_for_event_with_timeout( + &test.codex, + |event| matches!(event, EventMsg::McpStartupComplete(_)), + std::time::Duration::from_secs(35), + ) .await; + let EventMsg::McpStartupComplete(startup) = startup else { + unreachable!("waited for MCP startup completion"); + }; assert_eq!(startup.ready, ["ordinary"]); assert_eq!( startup @@ -131,3 +137,593 @@ http_headers = {{ Authorization = "Bearer ordinary-access" }} ); Ok(()) } + +#[derive(Clone, Copy)] +enum EnterpriseLifecycle { + Policy, + Removal, + RefreshPermissions, +} + +#[test_case::test_case(EnterpriseLifecycle::Policy; "policy")] +#[test_case::test_case(EnterpriseLifecycle::Removal; "removal")] +#[test_case::test_case(EnterpriseLifecycle::RefreshPermissions; "permission_refresh")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn enterprise_lifecycle_preserves_history_and_ordinary_tools( + scenario: EnterpriseLifecycle, +) -> Result<()> { + use base64::Engine; + use base64::engine::general_purpose::URL_SAFE_NO_PAD; + use codex_http_client::HttpClientFactory; + use codex_http_client::OutboundProxyPolicy; + use codex_login::AuthCredentialsStoreMode; + use codex_login::AuthKeyringBackendKind; + use codex_login::AuthManager; + use codex_login::AuthRouteConfig; + use serde_json::json; + use std::sync::Arc; + use wiremock::Mock; + use wiremock::ResponseTemplate; + use wiremock::matchers::method; + use wiremock::matchers::path; + + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let mcp = responses::start_mock_server().await; + AppsTestServer::mount(&mcp).await?; + if matches!(scenario, EnterpriseLifecycle::RefreshPermissions) { + // Keep the enterprise client on the protocol that carries input requests. + Mock::given(method("POST")) + .and(path("/api/codex/ps/mcp")) + .and(wiremock::matchers::header("authorization", "Bearer enterprise-access")) + .and(wiremock::matchers::body_partial_json(json!({"method": "server/discover"}))) + .respond_with(|request: &wiremock::Request| { + let body: serde_json::Value = request.body_json().expect("discovery request"); + ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", "id": body["id"], "result": { + "resultType": "complete", "supportedVersions": ["2026-07-28"], + "capabilities": {"tools": {}}, "ttlMs": 0, "cacheScope": "private", + "_meta": {"io.modelcontextprotocol/serverInfo": {"name": "enterprise", "version": "1"}} + } + })) + }) + .with_priority(1) + .mount(&mcp) + .await; + } + let origin = mcp.uri(); + let resource = format!("{origin}/api/codex/ps/mcp"); + let issuer = format!("{origin}/idp"); + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH)? + .as_secs(); + for (endpoint, metadata) in [ + ( + "/.well-known/oauth-protected-resource/api/codex/ps/mcp", + json!({ + "resource": resource, "authorization_servers": [format!("{origin}/as")], + }), + ), + ( + "/.well-known/oauth-authorization-server/as", + json!({ + "issuer": format!("{origin}/as"), "token_endpoint": format!("{origin}/as/token"), + "grant_types_supported": ["urn:ietf:params:oauth:grant-type:jwt-bearer"], + "token_endpoint_auth_methods_supported": ["none"], + }), + ), + ( + "/.well-known/oauth-authorization-server/idp", + json!({ + "issuer": issuer, "token_endpoint": format!("{issuer}/token"), + "authorization_endpoint": format!("{issuer}/authorize"), + "token_endpoint_auth_methods_supported": ["none"], + }), + ), + ] { + Mock::given(method("GET")) + .and(path(endpoint)) + .respond_with(ResponseTemplate::new(200).set_body_json(metadata)) + .with_priority(1) + .mount(&mcp) + .await; + } + let assertion = format!( + "{}.{}.c2lnbmF0dXJl", + URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256","typ":"oauth-id-jag+jwt"}"#), + URL_SAFE_NO_PAD.encode( + json!({ + "iss": issuer, "aud": format!("{origin}/as"), "sub": "user", + "client_id": "mcp-client", "jti": "grant", "iat": now, + "exp": now + 3600, "resource": resource, + }) + .to_string() + ) + ); + Mock::given(method("POST")).and(path("/idp/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": assertion, "issued_token_type": "urn:ietf:params:oauth:token-type:id-jag", + "token_type": "N_A", "resource": resource, + }))).expect(1).mount(&mcp).await; + Mock::given(method("POST")) + .and(path("/as/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "enterprise-access", "token_type": "Bearer", "expires_in": 3600, + }))) + .expect(1) + .mount(&mcp) + .await; + let home = Arc::new(tempfile::TempDir::new()?); + let id_token = format!("{}.{}.c2lnbmF0dXJl", + URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#), + URL_SAFE_NO_PAD.encode(json!({ + "iss": issuer, "sub": "user", "aud": "account-client", "exp": now + 3600, + "https://api.openai.com/auth": {"chatgpt_user_id":"user","chatgpt_account_id":"workspace","chatgpt_plan_type":"business"}, + }).to_string())); + std::fs::write( + home.path().join("auth.json"), + serde_json::to_vec(&json!({ + "auth_mode": "chatgpt", "last_refresh": chrono::Utc::now(), + "tokens": {"id_token": id_token, "access_token":"account-access", "refresh_token":"account-refresh", + "account_id":"workspace", "login_owner_id":"completed-login", + "oauth_provenance":{"issuer":issuer,"client_id":"account-client"}}, + }))?, + )?; + let auth = AuthManager::shared( + home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + AuthRouteConfig::from_http_client_factory(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + )), + ) + .await; + let policy = format!( + r#" +[features] +use_xaa = true +apps = false +secret_auth_storage = false +[mcp_enterprise_managed_auth.idp] +issuer = "{issuer}" +client_id = "enterprise-client" +[mcp_servers.enterprise] +url = "{resource}" +auth = "ema_auth" +oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" }} +[mcp_servers.ordinary] +url = "{resource}" +http_headers = {{ Authorization = "Bearer ordinary-access" }} +"# + ); + let delivered_policy = codex_config::CloudConfigBundlePolicy::default(); + let mut snapshot = codex_config::CloudConfigBundleSnapshot { + bundle: Ok(Some( + CloudConfigBundleFixture::enterprise_config(policy).into_bundle(), + )), + binding: None, + }; + delivered_policy.publish_snapshot(&mut snapshot); + let test = test_codex() + .with_home(home) + .with_auth_manager(auth) + .with_cloud_config_bundle(codex_config::CloudConfigBundleLoader::from_snapshot( + snapshot, + )) + .with_config(move |config| { + super::configure_scenario_catalog(config); + if matches!(scenario, EnterpriseLifecycle::RefreshPermissions) { + config.permissions.approval_policy = codex_config::Constrained::allow_any( + codex_protocol::protocol::AskForApproval::Never, + ); + config + .permissions + .set_permission_profile(codex_protocol::models::PermissionProfile::Disabled) + .expect("full access"); + config.approvals_reviewer = codex_protocol::config_types::ApprovalsReviewer::User; + config + .features + .enable(codex_features::Feature::ToolCallMcpElicitation) + .expect("MCP approval prompts"); + } + }) + .with_model_info_override("gpt-6-astra", |model| { + model.tool_mode = Some(ToolMode::Direct); + model.use_responses_lite = false; + model.supports_search_tool = false; + }) + .build_with_auto_env(&server) + .await?; + let startup = wait_for_event_match(&test.codex, |event| match event { + EventMsg::McpStartupComplete(summary) => Some(summary.clone()), + _ => None, + }) + .await; + assert!(startup.failed.is_empty(), "{startup:?}"); + if matches!(scenario, EnterpriseLifecycle::RefreshPermissions) { + use codex_protocol::protocol::AskForApproval; + use codex_protocol::protocol::ElicitationAction; + use codex_protocol::protocol::Op; + use codex_protocol::protocol::ThreadSettingsOverrides; + use core_test_support::wait_for_event; + + // Force a new manager while still in full access. A queued settings update + // is a barrier before status reads flush the hard refresh. + let requests_before_refresh = mcp.received_requests().await.expect("MCP requests").len(); + test.codex.submit(Op::RefreshMcpServers).await?; + test.codex + .update_thread_settings(ThreadSettingsOverrides::default()) + .await?; + test.codex + .mcp_server_status_snapshot( + "ordinary", + codex_mcp::McpSnapshotDetail::ToolsAndAuthOnly, + /*expected_registration*/ None, + ) + .await?; + let discovery_requests = mcp + .received_requests() + .await + .expect("MCP requests") + .into_iter() + .skip(requests_before_refresh) + .filter(|request| { + request.body_json::().is_ok_and(|body| { + matches!( + body["method"].as_str(), + Some("initialize" | "server/discover" | "tools/list") + ) + }) + }) + .map(|request| { + request + .headers + .get("authorization") + .and_then(|header| header.to_str().ok()) + .map(str::to_owned) + }) + .collect::>(); + assert!( + !discovery_requests.is_empty(), + "hard refresh must reconnect ordinary MCP" + ); + assert!( + discovery_requests + .iter() + .all(|authorization| authorization.as_deref() == Some("Bearer ordinary-access")), + "hard refresh must retain EMA without rediscovery or relisting: {discovery_requests:?}" + ); + test.codex + .update_thread_settings(ThreadSettingsOverrides { + approval_policy: Some(AskForApproval::OnRequest), + ..Default::default() + }) + .await?; + + Mock::given(method("POST")) + .and(path("/api/codex/ps/mcp")) + .and(wiremock::matchers::header("authorization", "Bearer enterprise-access")) + .and(wiremock::matchers::body_partial_json(json!({ + "method": "tools/call", "params": {"arguments": {"query": "permission"}} + }))) + .respond_with(|request: &wiremock::Request| { + let body: serde_json::Value = request.body_json().expect("MCP request"); + let result = if body["params"].get("inputResponses").is_none() { + json!({ + "resultType": "input_required", "content": [], "isError": false, + "inputRequests": {"approval": {"method": "elicitation/create", "params": { + "mode": "form", "message": "Approve the calendar action", + "requestedSchema": {"type": "object", "properties": {}}, + "_meta": {"codex_request_type": "approval_request", "codex_approval_kind": "mcp_tool_call", + "codex_sensitive_action": true, "tool_name": "calendar_list_events", "tool_params": {}} + }}}, "requestState": "awaiting-approval" + }) + } else { + assert_eq!(body["params"]["inputResponses"]["approval"]["action"], "accept"); + json!({"resultType": "complete", "content": [{"type": "text", "text": "approved"}], "isError": false}) + }; + ResponseTemplate::new(200).set_body_json(json!({"jsonrpc": "2.0", "id": body["id"], "result": result})) + }) + .with_priority(1) + .expect(2) + .mount(&mcp) + .await; + let mut call = Box::pin(test.codex.call_mcp_tool( + "enterprise", + "calendar_list_events", + Some(json!({"query": "permission"})), + /*meta*/ None, + )); + let request = tokio::select! { + result = &mut call => anyhow::bail!("EMA bypassed the updated approval policy: {result:?}"), + request = wait_for_event_match(&test.codex, |event| match event { + EventMsg::ElicitationRequest(request) => Some(request.clone()), + _ => None, + }) => request, + }; + assert_eq!(request.server_name, "enterprise"); + test.codex + .submit(Op::ResolveElicitation { + server_name: request.server_name, + request_id: request.id, + decision: ElicitationAction::Accept, + content: Some(json!({})), + meta: None, + }) + .await?; + assert_eq!(call.await?.is_error, Some(false)); + + // Tool policy changes retain this EMA authority, while replacing its tool view. + let current = test.codex.config().await; + let mut next = current.as_ref().clone(); + let mut servers = next.mcp_servers.get().clone(); + let enterprise = servers.get_mut("enterprise").expect("enterprise server"); + enterprise.disabled_tools = Some(vec!["calendar_create_event".into()]); + enterprise.tools.insert( + "calendar_list_events".into(), + codex_config::McpServerToolConfig { + approval_mode: Some(codex_config::AppToolApproval::Prompt), + ..Default::default() + }, + ); + next.mcp_servers.set(servers)?; + assert_eq!( + test.codex.refresh_mcp_config(current, next).await, + codex_core::ConfigRefreshOutcome::Published + ); + let error = test + .codex + .call_mcp_tool( + "enterprise", + "calendar_create_event", + /*arguments*/ None, + /*meta*/ None, + ) + .await + .expect_err("the disabled EMA tool must not dispatch"); + assert_eq!( + error.to_string(), + "tool 'calendar_create_event' is disabled for MCP server 'enterprise'" + ); + + // A read-only tool now explicitly requires approval. Missing updated metadata + // would let this turn complete without asking. + let mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_function_call_with_namespace( + "restricted-call", + "mcp__enterprise", + "calendar_list_events", + "{}", + ), + responses::ev_completed("restricted-response"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("done", "Cancelled."), + responses::ev_completed("done-response"), + ]), + ], + ) + .await; + test.codex + .start_or_steer_turn(codex_core::TurnInputRequest::user_input(vec![ + codex_protocol::user_input::UserInput::Text { + text: "List my work events.".into(), + text_elements: Vec::new(), + }, + ])) + .await?; + let event = wait_for_event(&test.codex, |event| { + matches!( + event, + EventMsg::ElicitationRequest(_) | EventMsg::TurnComplete(_) + ) + }) + .await; + let EventMsg::ElicitationRequest(request) = event else { + anyhow::bail!("EMA ignored the refreshed per-tool approval setting"); + }; + assert_eq!(request.server_name, "enterprise"); + assert_eq!(mock.requests().len(), 1, "the tool must wait for approval"); + let request_to_model = mock.single_request(); + assert_eq!( + ( + request_to_model + .tool_by_name("mcp__enterprise", "calendar_create_event") + .is_some(), + request_to_model + .tool_by_name("mcp__enterprise", "calendar_list_events") + .is_some(), + request_to_model + .tool_by_name("mcp__ordinary", "calendar_create_event") + .is_some(), + ), + (false, true, true), + ); + test.codex + .submit(Op::ResolveElicitation { + server_name: request.server_name, + request_id: request.id, + decision: ElicitationAction::Decline, + content: None, + meta: None, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let ordinary = test + .codex + .call_mcp_tool( + "ordinary", + "calendar_create_event", + Some(json!({"title": "control", "starts_at": "2026-06-18T12:00:00Z"})), + /*meta*/ None, + ) + .await?; + assert_eq!(ordinary.is_error, Some(false)); + let calls = core_test_support::apps_test_server::recorded_apps_tool_calls(&mcp).await; + assert_eq!( + calls + .iter() + .map(|call| call["params"]["name"].as_str().expect("tool name")) + .collect::>(), + [ + "calendar_list_events", + "calendar_list_events", + "calendar_create_event" + ] + ); + test.codex.shutdown_and_wait().await?; + return Ok(()); + } + let mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_function_call_with_namespace( + "enterprise-call", + "mcp__enterprise", + "calendar_list_events", + "{}", + ), + responses::ev_completed("enterprise-response"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("answer-1", "Here are your work events."), + responses::ev_completed("answer-1-response"), + ]), + responses::sse(vec![ + responses::ev_function_call_with_namespace( + "ordinary-call", + "mcp__ordinary", + "calendar_list_events", + "{}", + ), + responses::ev_completed("ordinary-response"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("answer-2", "Here are your other events."), + responses::ev_completed("answer-2-response"), + ]), + ], + ) + .await; + test.submit_text_turn("List my work events.").await?; + let pending_request = Mock::given(method("POST")) + .and(path("/api/codex/ps/mcp")) + .and(wiremock::matchers::body_partial_json(json!({ + "method": "tools/call", "params": {"arguments": {"query": "blocked"}} + }))) + .respond_with(ResponseTemplate::new(200).set_delay(std::time::Duration::from_secs(60))) + .with_priority(1) + .expect(1) + .mount_as_scoped(&mcp) + .await; + let thread = Arc::clone(&test.codex); + let request_started = std::time::Instant::now(); + let pending = tokio::spawn(async move { + thread + .call_mcp_tool( + "enterprise", + "calendar_list_events", + Some(json!({"query":"blocked"})), + /*meta*/ None, + ) + .await + }); + pending_request.wait_until_satisfied().await; + match scenario { + EnterpriseLifecycle::Policy => { + let denied = + CloudConfigBundleFixture::enterprise_requirement("[features]\nuse_xaa = false\n") + .into_bundle(); + delivered_policy.observe_remote_bundle(&denied); + let mut snapshot = codex_config::CloudConfigBundleSnapshot { + bundle: Ok(Some(denied)), + binding: None, + }; + delivered_policy.publish_snapshot(&mut snapshot); + } + EnterpriseLifecycle::Removal => { + // Archive/delete removes the registry entry before asynchronous shutdown. + // A retained caller must lose EMA access at removal, even before shutdown. + let removed = test + .thread_manager + .remove_thread_for_client(&test.session_configured.thread_id) + .await? + .expect("registered thread"); + assert!(Arc::ptr_eq(&removed, &test.codex)); + let result = removed + .call_mcp_tool( + "enterprise", + "calendar_list_events", + /*arguments*/ None, + /*meta*/ None, + ) + .await?; + assert_eq!(serde_json::to_value(result)?["isError"], true); + } + EnterpriseLifecycle::RefreshPermissions => unreachable!("handled before revocation"), + } + // No config reload or shutdown: the existing owner cancels admitted work. + let result = serde_json::to_value(pending.await??)?; + assert!( + request_started.elapsed() < std::time::Duration::from_secs(10), + "revocation must finish while the 60-second response is still pending" + ); + assert_eq!(result["isError"], true); + assert_eq!( + result["_meta"]["mcp/www_authenticate"], + "Bearer error=\"invalid_token\"" + ); + // The pending response is already cancelled. Apply the config update before + // building the next turn's tool catalog, preserving the original history check. + test.codex.disable_mcp_enterprise_auth().await; + test.submit_text_turn("List my other calendar events.") + .await?; + let requests = mock.requests(); + assert_eq!(requests.len(), 4); + let calls = mcp + .received_requests() + .await + .expect("MCP requests") + .into_iter() + .filter(|request| { + serde_json::from_slice::(&request.body) + .is_ok_and(|body| body["method"] == "tools/call") + }) + .map(|request| { + request.headers["authorization"] + .to_str() + .expect("valid authorization header") + .to_owned() + }) + .collect::>(); + assert_eq!( + calls, + [ + "Bearer enterprise-access", + "Bearer enterprise-access", + "Bearer ordinary-access" + ] + ); + insta::assert_snapshot!( + "enterprise_revocation_preserves_history_and_ordinary_tools", + context_snapshot::format_request_history_snapshot( + "Enterprise access disappears after revocation. Prior calls remain in history and ordinary tools still work.", + &requests, + &ContextSnapshotOptions::default() + .rewrite_known_segments() + .include_request_settings(), + ) + ); + test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mcp_ema__enterprise_revocation_preserves_history_and_ordinary_tools.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mcp_ema__enterprise_revocation_preserves_history_and_ordinary_tools.snap new file mode 100644 index 0000000000..abd3e86398 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mcp_ema__enterprise_revocation_preserves_history_and_ordinary_tools.snap @@ -0,0 +1,98 @@ +--- +source: core/tests/suite/scenarios_mcp_ema.rs +assertion_line: 257 +expression: "context_snapshot::format_request_history_snapshot(\"Enterprise access disappears after revocation. Prior calls remain in history and ordinary tools still work.\",\n&requests,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)" +--- +Scenario: Enterprise access disappears after revocation. Prior calls remain in history and ordinary tools still work. + +## Window 1 +Settings: + include: ["reasoning.encrypted_content"] + model: "gpt-6-astra" + parallel_tool_calls: true + prompt_cache_key: "" + reasoning: {"effort":"low"} + store: false + stream: true + text: {"verbosity":"low"} + tool_choice: "auto" + tools (14; hash=c23674e45d0ff686): + - function/exec_command: Runs a command in a PTY, returning output or a session ID for ongoing interaction.; args=[cmd, justification, login, max_output_tokens, prefix_rule, sandbox_permissions, shell, tty, workdir, yield_time_ms]; hash=c04108e51a8548c5 + - function/write_stdin: Writes characters to an existing unified exec session and returns recent output.; args=[chars, max_output_tokens, session_id, yield_time_ms]; hash=7fb62d8a3e2df4e2 + - function/list_mcp_resources: Lists resources provided by MCP servers. Resources allow servers to share data that provides cont...; args=[cursor, server]; hash=6792edaff13561bd + - function/list_mcp_resource_templates: Lists resource templates provided by MCP servers. Parameterized resource templates allow servers ...; args=[cursor, server]; hash=60f17c90ef7a4ebb + - function/read_mcp_resource: Read a specific resource from an MCP server given the server name and resource URI.; args=[server, uri]; hash=24246896b39008ae + - function/request_user_input: Request user input for one to three short questions and wait for the response. This tool is only ...; args=[questions]; hash=64a729b25c424fc6 + - function/request_user_input_async: Ask the user one or more questions during ongoing work. Use this tool only to request missing inf...; args=[questions]; hash=d6d74f33db330162 + - namespace/clock: Tools for reading and waiting on time.; hash=a80435c7a3a3943f + - function/curr_time + - function/sleep + - custom/apply_patch: The `apply_patch` tool can be used to edit files. This is a FREEFORM tool, so do not wrap the pat...; hash=dea9f12b0111a4ba + - function/view_image: View a local image file from the filesystem when visual inspection is needed. Use this for images...; args=[detail, path]; hash=3b766d5bdec5d50e + - namespace/collaboration: Tools for spawning and managing sub-agents.; hash=2c8d6295978dbe63 + - function/followup_task + - function/interrupt_agent + - function/list_agents + - function/send_message + - function/spawn_agent + - function/wait_agent + - namespace/mcp__enterprise: Tools in the mcp__enterprise namespace.; hash=b12dd185b8ee8c88 + - function/calendar_create_event + - function/calendar_extract_text + - function/calendar_list_events + - namespace/mcp__ordinary: Tools in the mcp__ordinary namespace.; hash=101b30f5cdc3cae0 + - function/calendar_create_event + - function/calendar_extract_text + - function/calendar_list_events + - web_search; hash=f103e327c6870a46 +-- request 1 (turn) -- +00:message/developer: + +01:message/developer[2]: + [01] + [02] +02:message/developer: + +03:message/developer: + +04:message/user: + +05:message/user: + List my work events. +-- request 2 (turn) -- +06:function_call/mcp__enterprise.calendar_list_events:{} +07:function_call_output:Wall time: seconds + Output: + {"_codex_apps":null} + +## Window 2 (after request 2) +Settings: relative to window 1 + tools (13; hash=75ef17d8adc1b3db): + - removed namespace/mcp__enterprise +-- request 3 (turn) -- +00:message/developer: + +01:message/developer[2]: + [01] + [02] +02:message/developer: + +03:message/developer: + +04:message/user: + +05:message/user: + List my work events. +06:function_call/mcp__enterprise.calendar_list_events:{} +07:function_call_output:Wall time: seconds + Output: + {"_codex_apps":null} +08:message/assistant: + Here are your work events. +09:message/user: + List my other calendar events. +-- request 4 (turn) -- +10:function_call/mcp__ordinary.calendar_list_events:{} +11:function_call_output:Wall time: seconds + Output: + {"_codex_apps":null}