openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 19:08:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(apps): persist codex mcp affinity cookies

Browse Source
This commit is contained in:
Casey Chow
2026-03-11 16:33:36 -04:00
parent 402ca6cc54
commit 36351c5efa
17 changed files with 691 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+56
View File
@@ -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",
codex-rs/app-server/src/codex_message_processor.rs
+11 -7
View File
@@ -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<codex_core::AppsMcpCookieStore>,
) {
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));
});
codex-rs/app-server/tests/suite/v2/app_list.rs
+150 -3
View File
@@ -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<AppListUpdatedNotification> {
@@ -1222,6 +1295,12 @@ struct AppsServerState {
expected_account_id: String,
response: Arc<StdMutex<serde_json::Value>>,
directory_delay: Duration,
affinity_cookie_state: Option<Arc<AffinityCookieState>>,
}
#[derive(Default)]
struct AffinityCookieState {
request_count: AtomicUsize,
}
#[derive(Clone)]
@@ -1310,6 +1389,32 @@ async fn start_apps_server_with_delays_and_control(
tools: Vec<Tool>,
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<AppInfo>,
tools: Vec<Tool>,
) -> Result<(String, Arc<AffinityCookieState>, 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<AppInfo>,
tools: Vec<Tool>,
directory_delay: Duration,
tools_delay: Duration,
affinity_cookie_state: Option<Arc<AffinityCookieState>>,
) -> 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<Arc<AppsServerState>>,
request: Request<Body>,
next: Next,
) -> Result<Response, StatusCode> {
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<Arc<AppsServerState>>,
headers: HeaderMap,
codex-rs/chatgpt/src/connectors.rs
+1
View File
@@ -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;
codex-rs/core/Cargo.toml
+1 -1
View File
@@ -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",
codex-rs/core/src/apps_mcp_cookie_store.rs
+122
View File
@@ -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<String>,
chatgpt_user_id: Option<String>,
is_workspace_account: bool,
}
#[derive(Clone, Default)]
pub struct AppsMcpCookieStore {
jars: Arc<StdMutex<HashMap<AppsMcpCookieStoreKey, Arc<Jar>>>>,
}
impl AppsMcpCookieStore {
pub fn jar_for(&self, config: &Config, auth: Option<&CodexAuth>) -> Arc<Jar> {
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<String>,
chatgpt_user_id: Option<String>,
is_workspace_account: bool,
) -> Arc<Jar> {
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));
}
}
codex-rs/core/src/codex.rs
+17
View File
@@ -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<AuthManager>,
apps_mcp_cookie_store: Arc<AppsMcpCookieStore>,
models_manager: Arc<ModelsManager>,
skills_manager: Arc<SkillsManager>,
plugins_manager: Arc<PluginsManager>,
@@ -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<Config>,
auth_manager: Arc<AuthManager>,
apps_mcp_cookie_store: Arc<AppsMcpCookieStore>,
models_manager: Arc<ModelsManager>,
exec_policy: ExecPolicyManager,
tx_event: Sender<Event>,
@@ -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;
{
codex-rs/core/src/codex_delegate.rs
+1
View File
@@ -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),
codex-rs/core/src/connectors.rs
+43 -3
View File
@@ -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<Vec<AppInfo>> {
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<AccessibleConnectorsStatus> {
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<Arc<AppsMcpCookieStore>>,
) -> anyhow::Result<Vec<AppInfo>> {
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<Arc<AppsMcpCookieStore>>,
) -> anyhow::Result<AccessibleConnectorsStatus> {
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;
codex-rs/core/src/lib.rs
+3
View File
@@ -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")]
codex-rs/core/src/mcp/mod.rs
+1
View File
@@ -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;
codex-rs/core/src/mcp_connection_manager.rs
+20 -2
View File
@@ -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<Arc<Jar>>,
cancel_token: CancellationToken,
tx_event: Sender<Event>,
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<Arc<Jar>>,
) -> (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<Arc<Jar>>,
) -> Result<RmcpClient, StartupOutcomeError> {
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)
codex-rs/core/src/state/service.rs
+2
View File
@@ -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<RwLock<McpConnectionManager>>,
pub(crate) mcp_startup_cancellation_token: Mutex<CancellationToken>,
pub(crate) apps_mcp_cookie_store: Arc<AppsMcpCookieStore>,
pub(crate) unified_exec_manager: UnifiedExecProcessManager,
#[cfg_attr(not(unix), allow(dead_code))]
pub(crate) shell_zsh_path: Option<PathBuf>,
codex-rs/core/src/thread_manager.rs
+9
View File
@@ -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<RwLock<HashMap<ThreadId, Arc<CodexThread>>>>,
thread_created_tx: broadcast::Sender<ThreadId>,
apps_mcp_cookie_store: Arc<AppsMcpCookieStore>,
auth_manager: Arc<AuthManager>,
models_manager: Arc<ModelsManager>,
skills_manager: Arc<SkillsManager>,
@@ -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<AppsMcpCookieStore> {
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),
codex-rs/rmcp-client/Cargo.toml
+1
View File
@@ -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",
codex-rs/rmcp-client/src/rmcp_client.rs
+22 -7
View File
@@ -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<HashMap<String, String>>,
env_http_headers: Option<HashMap<String, String>>,
store_mode: OAuthCredentialsStoreMode,
cookie_jar: Option<Arc<Jar>>,
},
}
@@ -504,6 +506,7 @@ impl RmcpClient {
http_headers: Option<HashMap<String, String>>,
env_http_headers: Option<HashMap<String, String>>,
store_mode: OAuthCredentialsStoreMode,
cookie_jar: Option<Arc<Jar>>,
) -> Result<Self> {
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<Arc<Jar>>,
) -> Result<(
StreamableHttpClientTransport<AuthClient<StreamableHttpResponseClient>>,
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<Arc<Jar>>,
) -> Result<reqwest::Client> {
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()?)
}
codex-rs/rmcp-client/tests/cookies.rs
+231
View File
@@ -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<Vec<Tool>>,
}
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<rmcp::model::PaginatedRequestParams>,
_context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
) -> impl std::future::Future<Output = Result<ListToolsResult, rmcp::ErrorData>> + 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<Arc<AffinityCookieState>>,
request: Request<Body>,
next: Next,
) -> Result<Response, StatusCode> {
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)
}