From 36351c5efa59fc41b8ec35bb46f3d44eb267c241 Mon Sep 17 00:00:00 2001 From: Casey Chow Date: Wed, 4 Mar 2026 19:06:44 -0500 Subject: [PATCH] fix(apps): persist codex mcp affinity cookies --- codex-rs/Cargo.lock | 56 +++++ .../app-server/src/codex_message_processor.rs | 18 +- .../app-server/tests/suite/v2/app_list.rs | 153 +++++++++++- codex-rs/chatgpt/src/connectors.rs | 1 + codex-rs/core/Cargo.toml | 2 +- codex-rs/core/src/apps_mcp_cookie_store.rs | 122 +++++++++ codex-rs/core/src/codex.rs | 17 ++ codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/connectors.rs | 46 +++- codex-rs/core/src/lib.rs | 3 + codex-rs/core/src/mcp/mod.rs | 1 + codex-rs/core/src/mcp_connection_manager.rs | 22 +- codex-rs/core/src/state/service.rs | 2 + codex-rs/core/src/thread_manager.rs | 9 + codex-rs/rmcp-client/Cargo.toml | 1 + codex-rs/rmcp-client/src/rmcp_client.rs | 29 ++- codex-rs/rmcp-client/tests/cookies.rs | 231 ++++++++++++++++++ 17 files changed, 691 insertions(+), 23 deletions(-) create mode 100644 codex-rs/core/src/apps_mcp_cookie_store.rs create mode 100644 codex-rs/rmcp-client/tests/cookies.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 64abfaac3c..2520e3c0ae 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2881,6 +2881,17 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "cookie" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747" +dependencies = [ + "percent-encoding", + "time", + "version_check", +] + [[package]] name = "cookie-factory" version = "0.3.3" @@ -2890,6 +2901,24 @@ dependencies = [ "futures", ] +[[package]] +name = "cookie_store" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b2c103cf610ec6cae3da84a766285b42fd16aad564758459e6ecf128c75206" +dependencies = [ + "cookie", + "document-features", + "idna", + "log", + "publicsuffix", + "serde", + "serde_derive", + "serde_json", + "time", + "url", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -3586,6 +3615,15 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "document-features" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" +dependencies = [ + "litrs", +] + [[package]] name = "dotenvy" version = "0.15.7" @@ -5497,6 +5535,12 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" +[[package]] +name = "litrs" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" + [[package]] name = "local-waker" version = "0.1.4" @@ -7069,6 +7113,16 @@ version = "2.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33cb294fe86a74cbcf50d4445b37da762029549ebeea341421c7c70370f86cac" +[[package]] +name = "publicsuffix" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f42ea446cab60335f76979ec15e12619a2165b5ae2c12166bef27d283a9fadf" +dependencies = [ + "idna", + "psl-types", +] + [[package]] name = "pulldown-cmark" version = "0.10.3" @@ -7752,6 +7806,8 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ "base64 0.22.1", "bytes", + "cookie", + "cookie_store", "encoding_rs", "futures-channel", "futures-core", diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index a4909c7745..047d8ca5ea 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -4993,8 +4993,9 @@ impl CodexMessageProcessor { let request = request_id.clone(); let outgoing = Arc::clone(&self.outgoing); + let apps_mcp_cookie_store = self.thread_manager.apps_mcp_cookie_store(); tokio::spawn(async move { - Self::apps_list_task(outgoing, request, params, config).await; + Self::apps_list_task(outgoing, request, params, config, apps_mcp_cookie_store).await; }); } @@ -5003,6 +5004,7 @@ impl CodexMessageProcessor { request_id: ConnectionRequestId, params: AppsListParams, config: Config, + apps_mcp_cookie_store: Arc, ) { let AppsListParams { cursor, @@ -5037,12 +5039,14 @@ impl CodexMessageProcessor { let accessible_config = config.clone(); let accessible_tx = tx.clone(); tokio::spawn(async move { - let result = connectors::list_accessible_connectors_from_mcp_tools_with_options( - &accessible_config, - force_refetch, - ) - .await - .map_err(|err| format!("failed to load accessible apps: {err}")); + let result = + connectors::list_accessible_connectors_from_mcp_tools_with_options_and_cookie_store( + &accessible_config, + force_refetch, + Some(apps_mcp_cookie_store), + ) + .await + .map_err(|err| format!("failed to load accessible apps: {err}")); let _ = accessible_tx.send(AppListLoadResult::Accessible(result)); }); diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index a19cdad857..a4890e276f 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -2,6 +2,8 @@ use std::borrow::Cow; use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use anyhow::Result; @@ -12,11 +14,19 @@ use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use axum::Json; use axum::Router; +use axum::body::Body; use axum::extract::State; use axum::http::HeaderMap; +use axum::http::HeaderValue; +use axum::http::Request; use axum::http::StatusCode; use axum::http::Uri; use axum::http::header::AUTHORIZATION; +use axum::http::header::COOKIE; +use axum::http::header::SET_COOKIE; +use axum::middleware; +use axum::middleware::Next; +use axum::response::Response; use axum::routing::get; use codex_app_server_protocol::AppBranding; use codex_app_server_protocol::AppInfo; @@ -1016,7 +1026,6 @@ async fn list_apps_force_refetch_patches_updates_from_cached_snapshots() -> Resu let mut mcp = McpProcess::new(codex_home.path()).await?; timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let warm_request = mcp .send_apps_list_request(AppsListParams { limit: None, @@ -1201,6 +1210,70 @@ async fn list_apps_force_refetch_patches_updates_from_cached_snapshots() -> Resu Ok(()) } +#[tokio::test] +async fn list_apps_force_refetch_reuses_affinity_cookie_across_temp_mcp_clients() -> Result<()> { + let connectors = vec![AppInfo { + id: "beta".to_string(), + name: "Beta App".to_string(), + description: Some("Beta connector".to_string()), + logo_url: None, + logo_url_dark: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: None, + is_accessible: false, + is_enabled: true, + }]; + let tools = vec![connector_tool("beta", "Beta App")?]; + let (server_url, affinity_cookie_state, server_handle) = + start_apps_server_with_affinity_cookie(connectors, tools).await?; + + let codex_home = TempDir::new()?; + write_connectors_config(codex_home.path(), &server_url)?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("chatgpt-token") + .account_id("account-123") + .chatgpt_user_id("user-123") + .chatgpt_account_id("account-123"), + AuthCredentialsStoreMode::File, + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; + + for _ in 0..2 { + let request = mcp + .send_apps_list_request(AppsListParams { + limit: None, + cursor: None, + thread_id: None, + force_refetch: true, + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request)), + ) + .await??; + let AppsListResponse { data, next_cursor } = to_response(response)?; + assert_eq!(data.len(), 1); + assert!(data.iter().all(|app| app.is_accessible)); + assert!(next_cursor.is_none()); + } + + assert!( + affinity_cookie_state.request_count.load(Ordering::SeqCst) >= 5, + "expected the initial refresh plus a later cookie-gated second-manager request", + ); + + server_handle.abort(); + let _ = server_handle.await; + Ok(()) +} + async fn read_app_list_updated_notification( mcp: &mut McpProcess, ) -> Result { @@ -1222,6 +1295,12 @@ struct AppsServerState { expected_account_id: String, response: Arc>, directory_delay: Duration, + affinity_cookie_state: Option>, +} + +#[derive(Default)] +struct AffinityCookieState { + request_count: AtomicUsize, } #[derive(Clone)] @@ -1310,6 +1389,32 @@ async fn start_apps_server_with_delays_and_control( tools: Vec, directory_delay: Duration, tools_delay: Duration, +) -> Result<(String, JoinHandle<()>, AppsServerControl)> { + start_apps_server(connectors, tools, directory_delay, tools_delay, None).await +} + +async fn start_apps_server_with_affinity_cookie( + connectors: Vec, + tools: Vec, +) -> Result<(String, Arc, JoinHandle<()>)> { + let affinity_cookie_state = Arc::new(AffinityCookieState::default()); + let (server_url, handle, _server_control) = start_apps_server( + connectors, + tools, + Duration::ZERO, + Duration::ZERO, + Some(Arc::clone(&affinity_cookie_state)), + ) + .await?; + Ok((server_url, affinity_cookie_state, handle)) +} + +async fn start_apps_server( + connectors: Vec, + tools: Vec, + directory_delay: Duration, + tools_delay: Duration, + affinity_cookie_state: Option>, ) -> Result<(String, JoinHandle<()>, AppsServerControl)> { let response = Arc::new(StdMutex::new( json!({ "apps": connectors, "next_token": null }), @@ -1320,6 +1425,7 @@ async fn start_apps_server_with_delays_and_control( expected_account_id: "account-123".to_string(), response: response.clone(), directory_delay, + affinity_cookie_state, }; let state = Arc::new(state); let server_control = AppsServerControl { @@ -1345,8 +1451,12 @@ async fn start_apps_server_with_delays_and_control( "/connectors/directory/list_workspace", get(list_directory_connectors), ) - .with_state(state) - .nest_service("/api/codex/apps", mcp_service); + .nest_service("/api/codex/apps", mcp_service) + .layer(middleware::from_fn_with_state( + Arc::clone(&state), + require_affinity_cookie, + )) + .with_state(state); let handle = tokio::spawn(async move { let _ = axum::serve(listener, router).await; @@ -1355,6 +1465,43 @@ async fn start_apps_server_with_delays_and_control( Ok((format!("http://{addr}"), handle, server_control)) } +async fn require_affinity_cookie( + State(state): State>, + request: Request, + next: Next, +) -> Result { + if !request.uri().path().starts_with("/api/codex/apps") { + return Ok(next.run(request).await); + } + + let Some(affinity_cookie_state) = state.affinity_cookie_state.as_ref() else { + return Ok(next.run(request).await); + }; + + let request_number = affinity_cookie_state + .request_count + .fetch_add(1, Ordering::SeqCst) + + 1; + if request_number > 4 + && !request + .headers() + .get(COOKIE) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.contains("__cflb=sticky")) + { + return Err(StatusCode::UNAUTHORIZED); + } + + let mut response = next.run(request).await; + if request_number <= 4 { + response.headers_mut().insert( + SET_COOKIE, + HeaderValue::from_static("__cflb=sticky; Path=/"), + ); + } + Ok(response) +} + async fn list_directory_connectors( State(state): State>, headers: HeaderMap, diff --git a/codex-rs/chatgpt/src/connectors.rs b/codex-rs/chatgpt/src/connectors.rs index dfc05fe317..863bf3f01b 100644 --- a/codex-rs/chatgpt/src/connectors.rs +++ b/codex-rs/chatgpt/src/connectors.rs @@ -25,6 +25,7 @@ pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options_and_status; pub use codex_core::connectors::list_cached_accessible_connectors_from_mcp_tools; +pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options_and_cookie_store; use codex_core::connectors::merge_connectors; use codex_core::connectors::merge_plugin_apps; pub use codex_core::connectors::with_app_enabled_state; diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index e406586524..137a792b18 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -74,7 +74,7 @@ once_cell = { workspace = true } os_info = { workspace = true } rand = { workspace = true } regex-lite = { workspace = true } -reqwest = { workspace = true, features = ["json", "stream"] } +reqwest = { workspace = true, features = ["cookies", "json", "stream"] } rmcp = { workspace = true, default-features = false, features = [ "base64", "macros", diff --git a/codex-rs/core/src/apps_mcp_cookie_store.rs b/codex-rs/core/src/apps_mcp_cookie_store.rs new file mode 100644 index 0000000000..2a8d624d76 --- /dev/null +++ b/codex-rs/core/src/apps_mcp_cookie_store.rs @@ -0,0 +1,122 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; + +use reqwest::cookie::Jar; + +use crate::CodexAuth; +use crate::config::Config; +use crate::mcp::codex_apps_mcp_url; + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct AppsMcpCookieStoreKey { + url: String, + account_id: Option, + chatgpt_user_id: Option, + is_workspace_account: bool, +} + +#[derive(Clone, Default)] +pub struct AppsMcpCookieStore { + jars: Arc>>>, +} + +impl AppsMcpCookieStore { + pub fn jar_for(&self, config: &Config, auth: Option<&CodexAuth>) -> Arc { + let token_data = auth.and_then(|auth| auth.get_token_data().ok()); + self.jar_for_identity( + codex_apps_mcp_url(config), + token_data + .as_ref() + .and_then(|token_data| token_data.account_id.clone()), + token_data + .as_ref() + .and_then(|token_data| token_data.id_token.chatgpt_user_id.clone()), + token_data + .as_ref() + .is_some_and(|token_data| token_data.id_token.is_workspace_account()), + ) + } + + fn jar_for_identity( + &self, + url: String, + account_id: Option, + chatgpt_user_id: Option, + is_workspace_account: bool, + ) -> Arc { + let key = AppsMcpCookieStoreKey { + url, + account_id, + chatgpt_user_id, + is_workspace_account, + }; + let mut jars = self + .jars + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Arc::clone(jars.entry(key).or_insert_with(|| Arc::new(Jar::default()))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn jar_for_identity_reuses_matching_key() { + let store = AppsMcpCookieStore::default(); + let left = store.jar_for_identity( + "https://example.test/apps".to_string(), + Some("account-1".to_string()), + Some("user-1".to_string()), + false, + ); + let right = store.jar_for_identity( + "https://example.test/apps".to_string(), + Some("account-1".to_string()), + Some("user-1".to_string()), + false, + ); + + assert!(Arc::ptr_eq(&left, &right)); + } + + #[test] + fn jar_for_identity_separates_account_boundaries() { + let store = AppsMcpCookieStore::default(); + let left = store.jar_for_identity( + "https://example.test/apps".to_string(), + Some("account-1".to_string()), + Some("user-1".to_string()), + false, + ); + let right = store.jar_for_identity( + "https://example.test/apps".to_string(), + Some("account-2".to_string()), + Some("user-1".to_string()), + false, + ); + + assert!(!Arc::ptr_eq(&left, &right)); + } + + #[test] + fn jar_for_identity_separates_workspace_boundaries() { + let store = AppsMcpCookieStore::default(); + let left = store.jar_for_identity( + "https://example.test/apps".to_string(), + Some("account-1".to_string()), + Some("user-1".to_string()), + false, + ); + let right = store.jar_for_identity( + "https://example.test/apps".to_string(), + Some("account-1".to_string()), + Some("user-1".to_string()), + true, + ); + + assert!(!Arc::ptr_eq(&left, &right)); + } +} diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index b2bbae4a78..893804a2d6 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -6,6 +6,7 @@ use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::AtomicU64; +use crate::AppsMcpCookieStore; use crate::AuthManager; use crate::CodexAuth; use crate::SandboxState; @@ -356,6 +357,7 @@ impl Codex { pub(crate) async fn spawn( mut config: Config, auth_manager: Arc, + apps_mcp_cookie_store: Arc, models_manager: Arc, skills_manager: Arc, plugins_manager: Arc, @@ -539,6 +541,7 @@ impl Codex { session_configuration, config.clone(), auth_manager.clone(), + apps_mcp_cookie_store, models_manager.clone(), exec_policy, tx_event.clone(), @@ -1257,6 +1260,7 @@ impl Session { mut session_configuration: SessionConfiguration, config: Arc, auth_manager: Arc, + apps_mcp_cookie_store: Arc, models_manager: Arc, exec_policy: ExecPolicyManager, tx_event: Sender, @@ -1427,6 +1431,9 @@ impl Session { } let auth = auth.as_ref(); + let apps_mcp_cookie_jar = mcp_servers + .contains_key(CODEX_APPS_MCP_SERVER_NAME) + .then(|| apps_mcp_cookie_store.jar_for(config.as_ref(), auth)); let auth_mode = auth.map(CodexAuth::auth_mode).map(TelemetryAuthMode::from); let account_id = auth.and_then(CodexAuth::get_account_id); let account_email = auth.and_then(CodexAuth::get_account_email); @@ -1612,6 +1619,7 @@ impl Session { &config.permissions.approval_policy, ))), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), + apps_mcp_cookie_store, unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), @@ -1738,6 +1746,7 @@ impl Session { config.codex_home.clone(), codex_apps_tools_cache_key(auth), tool_plugin_provenance, + apps_mcp_cookie_jar, ) .await; { @@ -3912,6 +3921,13 @@ impl Session { auth.as_ref(), config.as_ref(), ); + let apps_mcp_cookie_jar = mcp_servers + .contains_key(CODEX_APPS_MCP_SERVER_NAME) + .then(|| { + self.services + .apps_mcp_cookie_store + .jar_for(config.as_ref(), auth.as_ref()) + }); let auth_statuses = compute_auth_statuses(mcp_servers.iter(), store_mode).await; let sandbox_state = SandboxState { sandbox_policy: turn_context.sandbox_policy.get().clone(), @@ -3934,6 +3950,7 @@ impl Session { config.codex_home.clone(), codex_apps_tools_cache_key(auth.as_ref()), tool_plugin_provenance, + apps_mcp_cookie_jar, ) .await; { diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index f353febd04..416c3ac0fb 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -58,6 +58,7 @@ pub(crate) async fn run_codex_thread_interactive( let CodexSpawnOk { codex, .. } = Codex::spawn( config, auth_manager, + Arc::clone(&parent_session.services.apps_mcp_cookie_store), models_manager, Arc::clone(&parent_session.services.skills_manager), Arc::clone(&parent_session.services.plugins_manager), diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index b2b8ac109e..cfa00bc27d 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -18,6 +18,7 @@ use rmcp::model::ToolAnnotations; use serde::Deserialize; use tracing::warn; +use crate::AppsMcpCookieStore; use crate::AuthManager; use crate::CodexAuth; use crate::SandboxState; @@ -107,15 +108,48 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_options( force_refetch: bool, ) -> anyhow::Result> { Ok( - list_accessible_connectors_from_mcp_tools_with_options_and_status(config, force_refetch) - .await? - .connectors, + list_accessible_connectors_from_mcp_tools_with_options_and_status_and_cookie_store( + config, + force_refetch, + None, + ) + .await? + .connectors, ) } pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_status( config: &Config, force_refetch: bool, +) -> anyhow::Result { + list_accessible_connectors_from_mcp_tools_with_options_and_status_and_cookie_store( + config, + force_refetch, + None, + ) + .await +} + +pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_cookie_store( + config: &Config, + force_refetch: bool, + apps_mcp_cookie_store: Option>, +) -> anyhow::Result> { + Ok( + list_accessible_connectors_from_mcp_tools_with_options_and_status_and_cookie_store( + config, + force_refetch, + apps_mcp_cookie_store, + ) + .await? + .connectors, + ) +} + +async fn list_accessible_connectors_from_mcp_tools_with_options_and_status_and_cookie_store( + config: &Config, + force_refetch: bool, + apps_mcp_cookie_store: Option>, ) -> anyhow::Result { let auth_manager = auth_manager_from_config(config); let auth = auth_manager.auth().await; @@ -145,6 +179,11 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_status( codex_apps_ready: true, }); } + let apps_mcp_cookie_store = + apps_mcp_cookie_store.unwrap_or_else(|| Arc::new(AppsMcpCookieStore::default())); + let apps_mcp_cookie_jar = mcp_servers + .contains_key(CODEX_APPS_MCP_SERVER_NAME) + .then(|| apps_mcp_cookie_store.jar_for(config, auth.as_ref())); let auth_status_entries = compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await; @@ -169,6 +208,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_status( config.codex_home.clone(), codex_apps_tools_cache_key(auth.as_ref()), ToolPluginProvenance::default(), + apps_mcp_cookie_jar, ) .await; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 7e84577ee6..c0b14338c1 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -10,6 +10,7 @@ pub mod api_bridge; mod apply_patch; mod apps; mod arc_monitor; +mod apps_mcp_cookie_store; pub mod auth; mod client; mod client_common; @@ -90,6 +91,8 @@ mod seatbelt_permissions; mod thread_manager; pub mod web_search; pub mod windows_sandbox_read_grants; +pub use apps_mcp_cookie_store::AppsMcpCookieStore; +pub use codex_protocol::protocol::InitialHistory; pub use thread_manager::NewThread; pub use thread_manager::ThreadManager; #[deprecated(note = "use ThreadManager")] diff --git a/codex-rs/core/src/mcp/mod.rs b/codex-rs/core/src/mcp/mod.rs index 39eca4b29b..105cb63121 100644 --- a/codex-rs/core/src/mcp/mod.rs +++ b/codex-rs/core/src/mcp/mod.rs @@ -317,6 +317,7 @@ pub async fn collect_mcp_snapshot(config: &Config) -> McpListToolsResponseEvent config.codex_home.clone(), codex_apps_tools_cache_key(auth.as_ref()), tool_plugin_provenance, + None, ) .await; diff --git a/codex-rs/core/src/mcp_connection_manager.rs b/codex-rs/core/src/mcp_connection_manager.rs index 442e7e0c6e..7cbe9b853d 100644 --- a/codex-rs/core/src/mcp_connection_manager.rs +++ b/codex-rs/core/src/mcp_connection_manager.rs @@ -48,6 +48,7 @@ use codex_rmcp_client::SendElicitation; use futures::future::BoxFuture; use futures::future::FutureExt; use futures::future::Shared; +use reqwest::cookie::Jar; use rmcp::model::ClientCapabilities; use rmcp::model::CreateElicitationRequestParams; use rmcp::model::ElicitationAction; @@ -428,6 +429,7 @@ impl AsyncManagedClient { server_name: String, config: McpServerConfig, store_mode: OAuthCredentialsStoreMode, + codex_apps_cookie_jar: Option>, cancel_token: CancellationToken, tx_event: Sender, elicitation_requests: ElicitationRequestManager, @@ -449,8 +451,15 @@ impl AsyncManagedClient { return Err(error.into()); } - let client = - Arc::new(make_rmcp_client(&server_name, config.transport, store_mode).await?); + let client = Arc::new( + make_rmcp_client( + &server_name, + config.transport, + store_mode, + codex_apps_cookie_jar, + ) + .await?, + ); match start_server_task( server_name, client, @@ -634,6 +643,7 @@ impl McpConnectionManager { codex_home: PathBuf, codex_apps_tools_cache_key: CodexAppsToolsCacheKey, tool_plugin_provenance: ToolPluginProvenance, + codex_apps_cookie_jar: Option>, ) -> (Self, CancellationToken) { let cancel_token = CancellationToken::new(); let mut clients = HashMap::new(); @@ -647,6 +657,11 @@ impl McpConnectionManager { server_origins.insert(server_name.clone(), origin); } let cancel_token = cancel_token.child_token(); + let server_cookie_jar = if server_name == CODEX_APPS_MCP_SERVER_NAME { + codex_apps_cookie_jar.clone() + } else { + None + }; let _ = emit_update( &tx_event, McpStartupUpdateEvent { @@ -667,6 +682,7 @@ impl McpConnectionManager { server_name.clone(), cfg, store_mode, + server_cookie_jar, cancel_token.clone(), tx_event.clone(), elicitation_requests.clone(), @@ -1402,6 +1418,7 @@ async fn make_rmcp_client( server_name: &str, transport: McpServerTransportConfig, store_mode: OAuthCredentialsStoreMode, + cookie_jar: Option>, ) -> Result { match transport { McpServerTransportConfig::Stdio { @@ -1435,6 +1452,7 @@ async fn make_rmcp_client( http_headers, env_http_headers, store_mode, + cookie_jar, ) .await .map_err(StartupOutcomeError::from) diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 5c0a741a12..b9145fa466 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::sync::Arc; +use crate::AppsMcpCookieStore; use crate::AuthManager; use crate::RolloutRecorder; use crate::agent::AgentControl; @@ -54,6 +55,7 @@ impl CodeModeStoreService { pub(crate) struct SessionServices { pub(crate) mcp_connection_manager: Arc>, pub(crate) mcp_startup_cancellation_token: Mutex, + pub(crate) apps_mcp_cookie_store: Arc, pub(crate) unified_exec_manager: UnifiedExecProcessManager, #[cfg_attr(not(unix), allow(dead_code))] pub(crate) shell_zsh_path: Option, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index cc6de47104..0e0a530a27 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1,3 +1,4 @@ +use crate::AppsMcpCookieStore; use crate::AuthManager; use crate::CodexAuth; use crate::ModelProviderInfo; @@ -131,6 +132,7 @@ pub struct ThreadManager { pub(crate) struct ThreadManagerState { threads: Arc>>>, thread_created_tx: broadcast::Sender, + apps_mcp_cookie_store: Arc, auth_manager: Arc, models_manager: Arc, skills_manager: Arc, @@ -163,6 +165,7 @@ impl ThreadManager { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), thread_created_tx, + apps_mcp_cookie_store: Arc::new(AppsMcpCookieStore::default()), models_manager: Arc::new(ModelsManager::new( codex_home, auth_manager.clone(), @@ -223,6 +226,7 @@ impl ThreadManager { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), thread_created_tx, + apps_mcp_cookie_store: Arc::new(AppsMcpCookieStore::default()), models_manager: Arc::new(ModelsManager::with_provider_for_tests( codex_home, auth_manager.clone(), @@ -265,6 +269,10 @@ impl ThreadManager { self.state.models_manager.clone() } + pub fn apps_mcp_cookie_store(&self) -> Arc { + Arc::clone(&self.state.apps_mcp_cookie_store) + } + pub async fn list_models( &self, refresh_strategy: crate::models_manager::manager::RefreshStrategy, @@ -600,6 +608,7 @@ impl ThreadManagerState { } = Codex::spawn( config, auth_manager, + Arc::clone(&self.apps_mcp_cookie_store), Arc::clone(&self.models_manager), Arc::clone(&self.skills_manager), Arc::clone(&self.plugins_manager), diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index 7393368b32..76607ed810 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -21,6 +21,7 @@ futures = { workspace = true, default-features = false, features = ["std"] } keyring = { workspace = true, features = ["crypto-rust"] } oauth2 = "5" reqwest = { version = "0.12", default-features = false, features = [ + "cookies", "json", "stream", "rustls-tls", diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index c3bb6be0e6..950ff848ee 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -15,6 +15,7 @@ use futures::future::BoxFuture; use futures::stream::BoxStream; use oauth2::TokenResponse; use reqwest::header::ACCEPT; +use reqwest::cookie::Jar; use reqwest::header::AUTHORIZATION; use reqwest::header::CONTENT_TYPE; use reqwest::header::HeaderMap; @@ -395,6 +396,7 @@ enum TransportRecipe { http_headers: Option>, env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, + cookie_jar: Option>, }, } @@ -504,6 +506,7 @@ impl RmcpClient { http_headers: Option>, env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, + cookie_jar: Option>, ) -> Result { let transport_recipe = TransportRecipe::StreamableHttp { server_name: server_name.to_string(), @@ -512,6 +515,7 @@ impl RmcpClient { http_headers, env_http_headers, store_mode, + cookie_jar, }; let transport = Self::create_pending_transport(&transport_recipe).await?; Ok(Self { @@ -868,6 +872,7 @@ impl RmcpClient { http_headers, env_http_headers, store_mode, + cookie_jar, } => { let default_headers = build_default_headers(http_headers.clone(), env_http_headers.clone())?; @@ -892,6 +897,7 @@ impl RmcpClient { initial_tokens.clone(), *store_mode, default_headers.clone(), + cookie_jar.clone(), ) .await { @@ -919,8 +925,7 @@ impl RmcpClient { StreamableHttpClientTransportConfig::with_uri(url.clone()) .auth_header(access_token); let http_client = - apply_default_headers(reqwest::Client::builder(), &default_headers) - .build()?; + build_http_client(default_headers.clone(), cookie_jar.clone())?; let transport = StreamableHttpClientTransport::with_client( StreamableHttpResponseClient::new(http_client), http_config, @@ -936,9 +941,7 @@ impl RmcpClient { http_config = http_config.auth_header(bearer_token); } - let http_client = - apply_default_headers(reqwest::Client::builder(), &default_headers) - .build()?; + let http_client = build_http_client(default_headers, cookie_jar.clone())?; let transport = StreamableHttpClientTransport::with_client( StreamableHttpResponseClient::new(http_client), @@ -1122,12 +1125,12 @@ async fn create_oauth_transport_and_runtime( initial_tokens: StoredOAuthTokens, credentials_store: OAuthCredentialsStoreMode, default_headers: HeaderMap, + cookie_jar: Option>, ) -> Result<( StreamableHttpClientTransport>, OAuthPersistor, )> { - let http_client = - apply_default_headers(reqwest::Client::builder(), &default_headers).build()?; + let http_client = build_http_client(default_headers, cookie_jar)?; let mut oauth_state = OAuthState::new(url.to_string(), Some(http_client.clone())).await?; oauth_state @@ -1163,3 +1166,15 @@ async fn create_oauth_transport_and_runtime( Ok((transport, runtime)) } + +fn build_http_client( + default_headers: HeaderMap, + cookie_jar: Option>, +) -> Result { + let builder = apply_default_headers(reqwest::Client::builder(), &default_headers); + let builder = match cookie_jar { + Some(cookie_jar) => builder.cookie_provider(cookie_jar), + None => builder, + }; + Ok(builder.build()?) +} diff --git a/codex-rs/rmcp-client/tests/cookies.rs b/codex-rs/rmcp-client/tests/cookies.rs new file mode 100644 index 0000000000..f995b6af7f --- /dev/null +++ b/codex-rs/rmcp-client/tests/cookies.rs @@ -0,0 +1,231 @@ +use std::borrow::Cow; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use anyhow::Result; +use axum::Router; +use axum::body::Body; +use axum::extract::State; +use axum::http::HeaderValue; +use axum::http::Request; +use axum::http::StatusCode; +use axum::http::header::COOKIE; +use axum::http::header::SET_COOKIE; +use axum::middleware; +use axum::middleware::Next; +use axum::response::Response; +use codex_rmcp_client::ElicitationAction; +use codex_rmcp_client::ElicitationResponse; +use codex_rmcp_client::OAuthCredentialsStoreMode; +use codex_rmcp_client::RmcpClient; +use futures::FutureExt as _; +use pretty_assertions::assert_eq; +use reqwest::cookie::Jar; +use rmcp::handler::server::ServerHandler; +use rmcp::model::ClientCapabilities; +use rmcp::model::ElicitationCapability; +use rmcp::model::FormElicitationCapability; +use rmcp::model::Implementation; +use rmcp::model::InitializeRequestParams; +use rmcp::model::JsonObject; +use rmcp::model::ListToolsResult; +use rmcp::model::ProtocolVersion; +use rmcp::model::ServerCapabilities; +use rmcp::model::ServerInfo; +use rmcp::model::Tool; +use rmcp::transport::StreamableHttpServerConfig; +use rmcp::transport::StreamableHttpService; +use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; +use serde_json::json; +use tokio::net::TcpListener; + +const AFFINITY_COOKIE: &str = "__cflb=sticky"; + +#[derive(Default)] +struct AffinityCookieState { + request_count: AtomicUsize, +} + +#[derive(Clone)] +struct CookieToolServer { + tools: Arc>, +} + +impl CookieToolServer { + fn new() -> Self { + Self { + tools: Arc::new(vec![echo_tool()]), + } + } +} + +impl ServerHandler for CookieToolServer { + fn get_info(&self) -> ServerInfo { + ServerInfo { + capabilities: ServerCapabilities::builder().enable_tools().build(), + ..ServerInfo::default() + } + } + + fn list_tools( + &self, + _request: Option, + _context: rmcp::service::RequestContext, + ) -> impl std::future::Future> + Send + '_ + { + let tools = Arc::clone(&self.tools); + async move { + Ok(ListToolsResult { + tools: (*tools).clone(), + next_cursor: None, + meta: None, + }) + } + } +} + +fn echo_tool() -> Tool { + let schema: JsonObject = serde_json::from_value(json!({ + "type": "object", + "additionalProperties": false, + })) + .unwrap_or_else(|err| panic!("cookie test tool schema should deserialize: {err}")); + + Tool::new( + Cow::Borrowed("echo"), + Cow::Borrowed("Echo back the provided message."), + Arc::new(schema), + ) +} + +fn init_params() -> InitializeRequestParams { + InitializeRequestParams { + meta: None, + capabilities: ClientCapabilities { + experimental: None, + extensions: None, + roots: None, + sampling: None, + elicitation: Some(ElicitationCapability { + form: Some(FormElicitationCapability { + schema_validation: None, + }), + url: None, + }), + tasks: None, + }, + client_info: Implementation { + name: "codex-test".into(), + version: "0.0.0-test".into(), + title: Some("Codex rmcp cookie test".into()), + description: None, + icons: None, + website_url: None, + }, + protocol_version: ProtocolVersion::V_2025_06_18, + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_client_replays_affinity_cookie_from_initialize_response() -> Result<()> { + let affinity_cookie_state = Arc::new(AffinityCookieState::default()); + let listener = TcpListener::bind("127.0.0.1:0").await?; + let addr = listener.local_addr()?; + + let mcp_service = StreamableHttpService::new( + || Ok(CookieToolServer::new()), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default(), + ); + + let router = + Router::new() + .nest_service("/mcp", mcp_service) + .layer(middleware::from_fn_with_state( + Arc::clone(&affinity_cookie_state), + require_affinity_cookie, + )); + + let handle = tokio::spawn(async move { + let _ = axum::serve(listener, router).await; + }); + + let cookie_jar = Arc::new(Jar::default()); + let client = RmcpClient::new_streamable_http_client( + "codex_apps", + &format!("http://{addr}/mcp"), + Some("bearer-token".to_string()), + None, + None, + OAuthCredentialsStoreMode::File, + Some(cookie_jar), + ) + .await?; + + client + .initialize( + init_params(), + Some(Duration::from_secs(5)), + Box::new(|_, _| { + async { + Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({})), + }) + } + .boxed() + }), + ) + .await?; + + let tools = client + .list_tools(None, Some(Duration::from_secs(5))) + .await?; + assert_eq!( + tools, + ListToolsResult { + tools: vec![echo_tool()], + next_cursor: None, + meta: None, + } + ); + assert!( + affinity_cookie_state.request_count.load(Ordering::SeqCst) >= 2, + "expected initialize plus at least one follow-up request", + ); + + handle.abort(); + let _ = handle.await; + Ok(()) +} + +async fn require_affinity_cookie( + State(affinity_cookie_state): State>, + request: Request, + next: Next, +) -> Result { + let request_number = affinity_cookie_state + .request_count + .fetch_add(1, Ordering::SeqCst) + + 1; + if request_number > 1 + && !request + .headers() + .get(COOKIE) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.contains(AFFINITY_COOKIE)) + { + return Err(StatusCode::UNAUTHORIZED); + } + + let mut response = next.run(request).await; + if request_number == 1 { + response.headers_mut().insert( + SET_COOKIE, + HeaderValue::from_static("__cflb=sticky; Path=/"), + ); + } + Ok(response) +}