From 703793c22efc1e47de51703150f5cf16976ff456 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Thu, 25 Jun 2026 13:54:48 -0700 Subject: [PATCH] feat(core, mcp): cache codex_apps tools in memory (#29003) ## Description This makes Codex Apps tool reads use a shared in-memory snapshot instead of rereading the disk cache every time `list_all_tools()` runs. Disk still seeds the cache on startup and gets updated after successful fetches, but it is no longer the live read path. The core change is that `McpManager` now owns a process-scoped `CodexAppsToolsCache`. Codex threads in the same app-server process now share this Codex Apps in-memory tools snapshot. The snapshot is keyed by the Codex home plus the Codex Apps identity: the active Codex auth user/workspace and the effective Codex Apps MCP source config. There's already code to hard-refresh the cache, so we respect it in this PR. ## Local benchmark I ran a local steady-state microbenchmark of the exact repeated Codex Apps cached-tools read this PR removes, using the same real local cache payload in both trees: `3,678,138` bytes and `381` tools. The cache file was already warm in the OS page cache, so this measures same-process reread/deserialization work rather than cold-disk latency or full turn latency. Each run is 25 iterations (mimicking a turn that makes 25 inference calls). | Version | Run 1 | Run 2 | Avg | |---|---:|---:|---:| | `origin/main` disk read + JSON deserialize + `filter_tools` | `50.755 ms` | `52.894 ms` | `51.825 ms` | | This branch in-memory `current_tools` + `filter_tools` | `0.740 ms` | `0.778 ms` | `0.759 ms` | That removes about `51 ms` from each repeated Codex Apps cached-tools read on this machine, roughly `68x` faster for that subpath. It is useful evidence for the hot path this PR changes, but not a claim that every production turn gets `51 ms` faster; end-to-end impact also depends on the rest of `list_all_tools()` and tool-payload construction. This is on my M2 Max macbook, so with a slower disk this would be much worse (and indeed we did see this really blew up turn runtime with a slow disk). --- .../src/request_processors/mcp_processor.rs | 23 +- codex-rs/codex-mcp/src/codex_apps.rs | 212 +------- codex-rs/codex-mcp/src/codex_apps_cache.rs | 378 ++++++++++++++ .../codex-mcp/src/codex_apps_cache_tests.rs | 464 ++++++++++++++++++ codex-rs/codex-mcp/src/connection_manager.rs | 90 ++-- .../codex-mcp/src/connection_manager_tests.rs | 425 +++++----------- codex-rs/codex-mcp/src/lib.rs | 6 +- codex-rs/codex-mcp/src/mcp/mod.rs | 7 +- codex-rs/codex-mcp/src/rmcp_client.rs | 96 ++-- codex-rs/core/src/connectors.rs | 1 + codex-rs/core/src/mcp.rs | 13 +- codex-rs/core/src/mcp_tool_call.rs | 12 +- codex-rs/core/src/mcp_tool_call_tests.rs | 1 + codex-rs/core/src/session/mcp.rs | 1 + codex-rs/core/src/session/session.rs | 1 + 15 files changed, 1109 insertions(+), 621 deletions(-) create mode 100644 codex-rs/codex-mcp/src/codex_apps_cache.rs create mode 100644 codex-rs/codex-mcp/src/codex_apps_cache_tests.rs diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index 18b49ede3..a3d219178 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -249,15 +249,12 @@ impl McpRequestProcessor { } None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None), }; + let mcp_manager = self.thread_manager.mcp_manager(); let mcp_config = match thread { Some(thread) => thread.runtime_mcp_config(&config).await, - None => { - self.thread_manager - .mcp_manager() - .runtime_config(&config) - .await - } + None => mcp_manager.runtime_config(&config).await, }; + let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); let auth = self.auth_manager.auth().await; let environment_manager = self.thread_manager.environment_manager(); // This status path has no turn-selected environment. Use config cwd @@ -274,6 +271,7 @@ impl McpRequestProcessor { mcp_config, auth, runtime_context, + codex_apps_tools_cache, ) .await; }); @@ -287,6 +285,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache, ) { let result = Self::list_mcp_server_status_response( request_id.request_id.to_string(), @@ -294,6 +293,7 @@ impl McpRequestProcessor { mcp_config, auth, runtime_context, + codex_apps_tools_cache, ) .await; outgoing.send_result(request_id, result).await; @@ -305,6 +305,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache, ) -> Result { let detail = match params.detail.unwrap_or(McpServerStatusDetail::Full) { McpServerStatusDetail::Full => McpSnapshotDetail::Full, @@ -316,6 +317,7 @@ impl McpRequestProcessor { auth.as_ref(), request_id, runtime_context, + codex_apps_tools_cache, detail, ) .await; @@ -406,11 +408,9 @@ impl McpRequestProcessor { } let config = self.load_latest_config(/*fallback_cwd*/ None).await?; - let mcp_config = self - .thread_manager - .mcp_manager() - .runtime_config(&config) - .await; + let mcp_manager = self.thread_manager.mcp_manager(); + let mcp_config = mcp_manager.runtime_config(&config).await; + let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); let auth = self.auth_manager.auth().await; let environment_manager = self.thread_manager.environment_manager(); // This threadless resource-read path has no turn cwd or turn-selected @@ -425,6 +425,7 @@ impl McpRequestProcessor { &mcp_config, auth.as_ref(), runtime_context, + codex_apps_tools_cache, &server, &uri, ) diff --git a/codex-rs/codex-mcp/src/codex_apps.rs b/codex-rs/codex-mcp/src/codex_apps.rs index c1b5fd5fa..a33f73536 100644 --- a/codex-rs/codex-mcp/src/codex_apps.rs +++ b/codex-rs/codex-mcp/src/codex_apps.rs @@ -1,70 +1,9 @@ //! Codex Apps support for the host-owned apps MCP server. //! -//! This module owns the pieces that are unique to ChatGPT-hosted app -//! connectors: cache scoping by authenticated user, disk cache reads/writes, -//! connector allow-list filtering, and the normalization that turns app +//! This module owns the normalization that turns ChatGPT-hosted app //! connector/tool metadata into model-visible MCP callable names. -use std::path::PathBuf; -use std::time::Instant; - -use crate::runtime::emit_duration; -use crate::tools::MCP_TOOLS_CACHE_WRITE_DURATION_METRIC; -use crate::tools::ToolInfo; -use anyhow::Context; -use codex_login::CodexAuth; -use codex_protocol::mcp::McpServerInfo; use codex_utils_plugins::mcp_connector::sanitize_name; -use serde::Deserialize; -use serde::Serialize; -use sha1::Digest; -use sha1::Sha1; -use tracing::instrument; - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct CodexAppsToolsCacheKey { - pub(crate) account_id: Option, - pub(crate) chatgpt_user_id: Option, - pub(crate) is_workspace_account: bool, -} - -pub fn codex_apps_tools_cache_key(auth: Option<&CodexAuth>) -> CodexAppsToolsCacheKey { - CodexAppsToolsCacheKey { - account_id: auth.and_then(CodexAuth::get_account_id), - chatgpt_user_id: auth.and_then(CodexAuth::get_chatgpt_user_id), - is_workspace_account: auth.is_some_and(CodexAuth::is_workspace_account), - } -} - -#[derive(Clone)] -pub(crate) struct CodexAppsToolsCacheContext { - pub(crate) codex_home: PathBuf, - pub(crate) user_key: CodexAppsToolsCacheKey, -} - -impl CodexAppsToolsCacheContext { - pub(crate) fn tools_cache_path(&self) -> PathBuf { - self.cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR) - } - - pub(crate) fn server_info_cache_path(&self) -> PathBuf { - self.cache_path_in(CODEX_APPS_SERVER_INFO_CACHE_DIR) - } - - fn cache_path_in(&self, cache_dir: &str) -> PathBuf { - let user_key_json = serde_json::to_string(&self.user_key).unwrap_or_default(); - let user_key_hash = sha1_hex(&user_key_json); - self.codex_home - .join(cache_dir) - .join(format!("{user_key_hash}.json")) - } -} - -pub(crate) enum CachedCodexAppsToolsLoad { - Hit(Vec), - Missing, - Invalid, -} pub(crate) fn normalize_codex_apps_tool_title(connector_name: Option<&str>, value: &str) -> String { let Some(connector_name) = connector_name @@ -124,152 +63,3 @@ pub(crate) fn normalize_codex_apps_callable_namespace( server_name.to_string() } } - -pub(crate) fn write_codex_apps_tools_cache( - cache_context: Option<&CodexAppsToolsCacheContext>, - server_info: &McpServerInfo, - tools: &[ToolInfo], -) { - if let Some(cache_context) = cache_context { - let cache_write_start = Instant::now(); - write_cached_codex_apps_tools(cache_context, tools); - if let Err(err) = write_cached_codex_apps_server_info(cache_context, server_info) { - tracing::warn!("failed to write Codex Apps server info cache: {err:#}"); - } - emit_duration( - MCP_TOOLS_CACHE_WRITE_DURATION_METRIC, - cache_write_start.elapsed(), - &[], - ); - } -} - -pub(crate) fn load_startup_cached_codex_apps_tools_snapshot( - cache_context: Option<&CodexAppsToolsCacheContext>, -) -> Option> { - let cache_context = cache_context?; - - match load_cached_codex_apps_tools(cache_context) { - CachedCodexAppsToolsLoad::Hit(tools) => Some(tools), - CachedCodexAppsToolsLoad::Missing | CachedCodexAppsToolsLoad::Invalid => None, - } -} - -pub(crate) fn load_startup_cached_codex_apps_server_info( - cache_context: Option<&CodexAppsToolsCacheContext>, -) -> Option { - load_cached_codex_apps_server_info(cache_context?) -} - -#[cfg(test)] -pub(crate) fn read_cached_codex_apps_tools( - cache_context: &CodexAppsToolsCacheContext, -) -> Option> { - match load_cached_codex_apps_tools(cache_context) { - CachedCodexAppsToolsLoad::Hit(tools) => Some(tools), - CachedCodexAppsToolsLoad::Missing | CachedCodexAppsToolsLoad::Invalid => None, - } -} - -#[instrument(level = "trace", skip_all)] -pub(crate) fn load_cached_codex_apps_tools( - cache_context: &CodexAppsToolsCacheContext, -) -> CachedCodexAppsToolsLoad { - let cache_path = cache_context.tools_cache_path(); - let bytes = match std::fs::read(cache_path) { - Ok(bytes) => bytes, - Err(err) if err.kind() == std::io::ErrorKind::NotFound => { - return CachedCodexAppsToolsLoad::Missing; - } - Err(_) => return CachedCodexAppsToolsLoad::Invalid, - }; - let cache: CodexAppsToolsDiskCache = match serde_json::from_slice(&bytes) { - Ok(cache) => cache, - Err(_) => return CachedCodexAppsToolsLoad::Invalid, - }; - if cache.schema_version != CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION { - return CachedCodexAppsToolsLoad::Invalid; - } - CachedCodexAppsToolsLoad::Hit(cache.tools) -} - -pub(crate) fn write_cached_codex_apps_tools( - cache_context: &CodexAppsToolsCacheContext, - tools: &[ToolInfo], -) { - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() - && std::fs::create_dir_all(parent).is_err() - { - return; - } - let Ok(bytes) = serde_json::to_vec_pretty(&CodexAppsToolsDiskCache { - schema_version: CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, - tools: tools.to_vec(), - }) else { - return; - }; - let _ = std::fs::write(cache_path, bytes); -} - -#[instrument(level = "trace", skip_all)] -pub(crate) fn load_cached_codex_apps_server_info( - cache_context: &CodexAppsToolsCacheContext, -) -> Option { - let bytes = std::fs::read(cache_context.server_info_cache_path()).ok()?; - let cache: CodexAppsServerInfoDiskCache = serde_json::from_slice(&bytes).ok()?; - (cache.schema_version == CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION) - .then_some(cache.server_info) -} - -fn write_cached_codex_apps_server_info( - cache_context: &CodexAppsToolsCacheContext, - server_info: &McpServerInfo, -) -> anyhow::Result<()> { - let cache_path = cache_context.server_info_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).with_context(|| { - format!( - "failed to create Codex Apps server info cache directory `{}`", - parent.display() - ) - })?; - } - let bytes = serde_json::to_vec_pretty(&CodexAppsServerInfoDiskCache { - schema_version: CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION, - server_info: server_info.clone(), - }) - .context("failed to serialize Codex Apps server info cache")?; - std::fs::write(&cache_path, bytes).with_context(|| { - format!( - "failed to write Codex Apps server info cache `{}`", - cache_path.display() - ) - })?; - Ok(()) -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct CodexAppsToolsDiskCache { - schema_version: u8, - tools: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct CodexAppsServerInfoDiskCache { - schema_version: u8, - server_info: McpServerInfo, -} - -const CODEX_APPS_TOOLS_CACHE_DIR: &str = "cache/codex_apps_tools"; -pub(crate) const CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION: u8 = 4; - -const CODEX_APPS_SERVER_INFO_CACHE_DIR: &str = "cache/codex_apps_server_info"; -const CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION: u8 = 1; - -fn sha1_hex(s: &str) -> String { - let mut hasher = Sha1::new(); - hasher.update(s.as_bytes()); - let sha1 = hasher.finalize(); - format!("{sha1:x}") -} diff --git a/codex-rs/codex-mcp/src/codex_apps_cache.rs b/codex-rs/codex-mcp/src/codex_apps_cache.rs new file mode 100644 index 000000000..ea3944462 --- /dev/null +++ b/codex-rs/codex-mcp/src/codex_apps_cache.rs @@ -0,0 +1,378 @@ +//! Shared raw tool cache for the host-owned Codex Apps MCP server. +//! +//! Cache entries are process-local live state scoped by the active Codex auth +//! key. Disk is best-effort cold-start persistence; entries do not reread disk +//! after creation. + +use std::collections::HashMap; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Instant; + +use anyhow::Context; +use arc_swap::ArcSwapOption; +use codex_login::CodexAuth; +use codex_protocol::mcp::McpServerInfo; +use serde::Deserialize; +use serde::Serialize; +use sha1::Digest; +use sha1::Sha1; +use tracing::instrument; + +use crate::runtime::emit_duration; +use crate::tools::MCP_TOOLS_CACHE_WRITE_DURATION_METRIC; +use crate::tools::ToolInfo; + +const MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC: &str = "codex.mcp.tools.cache_publish.duration_ms"; + +/// The CodexAuth bits that identify a Codex Apps catalog. +/// +/// Debug bearer-token overrides bypass the shared cache, so shared entries only +/// need the CodexAuth-backed identity. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct CodexAppsToolsCacheKey { + pub(crate) account_id: Option, + pub(crate) chatgpt_user_id: Option, + pub(crate) is_workspace_account: bool, +} + +/// Builds the CodexAuth-backed Codex Apps cache key. +pub fn codex_apps_tools_cache_key(auth: Option<&CodexAuth>) -> CodexAppsToolsCacheKey { + CodexAppsToolsCacheKey { + account_id: auth.and_then(CodexAuth::get_account_id), + chatgpt_user_id: auth.and_then(CodexAuth::get_chatgpt_user_id), + is_workspace_account: auth.is_some_and(CodexAuth::is_workspace_account), + } +} + +/// Process-scoped registry for shared Codex Apps raw tool snapshots. +/// +/// Two clients share an entry only when they would read the same Codex Apps +/// catalog. New entries may seed from disk; live entries read from memory only. +#[derive(Clone, Default)] +pub struct CodexAppsToolsCache { + entries: Arc>>>, +} + +/// Handle to one shared Codex Apps tools cache entry. +/// +/// The connection manager creates this from the auth key, then tool +/// reads and refreshes for that managed client use the same entry. +#[derive(Clone)] +pub(crate) struct CodexAppsToolsCacheContext { + entry: Arc, +} + +impl CodexAppsToolsCacheContext { + pub(crate) fn tools_cache_path(&self) -> PathBuf { + self.entry + .identity + .cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR) + } + + pub(crate) fn server_info_cache_path(&self) -> PathBuf { + self.entry + .identity + .cache_path_in(CODEX_APPS_SERVER_INFO_CACHE_DIR) + } + + pub(crate) fn current_tools(&self) -> Option> { + self.entry + .current_tools + .load_full() + .map(|tools| tools.as_ref().clone()) + } + + pub(crate) fn has_current_tools(&self) -> bool { + self.entry.current_tools.load_full().is_some() + } + + pub(crate) fn begin_fetch( + &self, + source: CodexAppsToolsFetchSource, + ) -> CodexAppsToolsFetchTicket { + CodexAppsToolsFetchTicket { + generation: self + .entry + .next_fetch_generation + .fetch_add(1, Ordering::Relaxed) + + 1, + source, + } + } + + pub(crate) fn publish_if_newest_accepted( + &self, + ticket: CodexAppsToolsFetchTicket, + server_info: &McpServerInfo, + tools: Vec, + ) -> Vec { + let publish_start = Instant::now(); + let mut last_accepted_generation = lock_unpoisoned(&self.entry.last_accepted_generation); + if ticket.generation <= *last_accepted_generation { + emit_duration( + MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC, + publish_start.elapsed(), + &[("source", ticket.source.as_str()), ("result", "stale")], + ); + return self.current_tools().unwrap_or(tools); + } + + *last_accepted_generation = ticket.generation; + self.entry + .current_tools + .store(Some(Arc::new(tools.clone()))); + persist_codex_apps_cache(self, server_info, &tools); + emit_duration( + MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC, + publish_start.elapsed(), + &[("source", ticket.source.as_str()), ("result", "published")], + ); + tools + } + + #[cfg(test)] + pub(crate) fn store_current_tools_for_test(&self, tools: Vec) { + self.entry.current_tools.store(Some(Arc::new(tools))); + } +} + +impl CodexAppsToolsCache { + pub(crate) fn context( + &self, + codex_home: PathBuf, + auth_key: CodexAppsToolsCacheKey, + ) -> CodexAppsToolsCacheContext { + let identity = CodexAppsToolsCacheIdentity { + codex_home, + auth_key, + }; + let mut entries = lock_unpoisoned(&self.entries); + let entry = entries + .entry(identity.clone()) + .or_insert_with(|| Arc::new(CodexAppsToolsCacheEntry::new(identity))) + .clone(); + CodexAppsToolsCacheContext { entry } + } +} + +#[derive(Debug, Clone, Copy)] +pub(crate) enum CodexAppsToolsFetchSource { + Startup, + HardRefresh, +} + +impl CodexAppsToolsFetchSource { + fn as_str(self) -> &'static str { + match self { + Self::Startup => "startup", + Self::HardRefresh => "hard_refresh", + } + } +} + +pub(crate) struct CodexAppsToolsFetchTicket { + generation: u64, + source: CodexAppsToolsFetchSource, +} + +struct CodexAppsToolsCacheEntry { + identity: CodexAppsToolsCacheIdentity, + current_tools: ArcSwapOption>, + next_fetch_generation: AtomicU64, + last_accepted_generation: Mutex, +} + +impl CodexAppsToolsCacheEntry { + fn new(identity: CodexAppsToolsCacheIdentity) -> Self { + let current_tools = load_cached_codex_apps_tools_for_identity(&identity).map(Arc::new); + Self { + identity, + current_tools: ArcSwapOption::from(current_tools), + next_fetch_generation: AtomicU64::new(0), + last_accepted_generation: Mutex::new(0), + } + } +} + +/// Everything that decides whether two Codex Apps clients can share tools. +/// +/// The auth key says whose catalog we are reading. `codex_home` keeps the +/// persisted cache under the right home directory. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct CodexAppsToolsCacheIdentity { + codex_home: PathBuf, + auth_key: CodexAppsToolsCacheKey, +} + +impl CodexAppsToolsCacheIdentity { + fn cache_path_in(&self, cache_dir: &str) -> PathBuf { + // `codex_home` is already the parent directory. Keep it out of the + // filename hash so non-UTF-8 Unix paths cannot collapse distinct auth + // keys onto the same disk cache file. + let identity_json = serde_json::to_string(&self.auth_key).unwrap_or_default(); + let identity_hash = sha1_hex(&identity_json); + self.codex_home + .join(cache_dir) + .join(format!("{identity_hash}.json")) + } +} + +#[cfg(test)] +fn write_cached_codex_apps_tools_for_test( + cache_context: &CodexAppsToolsCacheContext, + server_info: &McpServerInfo, + tools: &[ToolInfo], +) { + cache_context + .entry + .current_tools + .store(Some(Arc::new(tools.to_vec()))); + persist_codex_apps_cache(cache_context, server_info, tools); +} + +pub(crate) fn load_startup_cached_codex_apps_server_info( + cache_context: &CodexAppsToolsCacheContext, +) -> Option { + load_cached_codex_apps_server_info(cache_context) +} + +#[cfg(test)] +fn read_cached_codex_apps_tools( + cache_context: &CodexAppsToolsCacheContext, +) -> Option> { + load_cached_codex_apps_tools_for_identity(&cache_context.entry.identity) +} + +#[instrument(level = "trace", skip_all)] +fn load_cached_codex_apps_tools_for_identity( + identity: &CodexAppsToolsCacheIdentity, +) -> Option> { + let cache_path = identity.cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR); + let bytes = std::fs::read(cache_path).ok()?; + let cache: CodexAppsToolsDiskCache = serde_json::from_slice(&bytes).ok()?; + (cache.schema_version == CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION).then_some(cache.tools) +} + +fn write_cached_codex_apps_tools( + cache_context: &CodexAppsToolsCacheContext, + tools: &[ToolInfo], +) -> anyhow::Result<()> { + let cache_path = cache_context.tools_cache_path(); + let bytes = serde_json::to_vec_pretty(&CodexAppsToolsDiskCache { + schema_version: CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, + tools: tools.to_vec(), + }) + .context("failed to serialize Codex Apps tools cache")?; + write_codex_apps_cache_file(&cache_path, "tools", bytes) +} + +#[instrument(level = "trace", skip_all)] +fn load_cached_codex_apps_server_info( + cache_context: &CodexAppsToolsCacheContext, +) -> Option { + let bytes = std::fs::read(cache_context.server_info_cache_path()).ok()?; + let cache: CodexAppsServerInfoDiskCache = serde_json::from_slice(&bytes).ok()?; + (cache.schema_version == CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION) + .then_some(cache.server_info) +} + +fn write_cached_codex_apps_server_info( + cache_context: &CodexAppsToolsCacheContext, + server_info: &McpServerInfo, +) -> anyhow::Result<()> { + let cache_path = cache_context.server_info_cache_path(); + let bytes = serde_json::to_vec_pretty(&CodexAppsServerInfoDiskCache { + schema_version: CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION, + server_info: server_info.clone(), + }) + .context("failed to serialize Codex Apps server info cache")?; + write_codex_apps_cache_file(&cache_path, "server info", bytes) +} + +fn write_codex_apps_cache_file( + cache_path: &Path, + cache_name: &str, + bytes: Vec, +) -> anyhow::Result<()> { + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).with_context(|| { + format!( + "failed to create Codex Apps {cache_name} cache directory `{}`", + parent.display() + ) + })?; + } + std::fs::write(cache_path, bytes).with_context(|| { + format!( + "failed to write Codex Apps {cache_name} cache `{}`", + cache_path.display() + ) + })?; + Ok(()) +} + +fn persist_codex_apps_cache( + cache_context: &CodexAppsToolsCacheContext, + server_info: &McpServerInfo, + tools: &[ToolInfo], +) { + let cache_write_start = Instant::now(); + let tools_result = write_cached_codex_apps_tools(cache_context, tools); + if let Err(err) = &tools_result { + tracing::warn!("failed to write Codex Apps tools cache: {err:#}"); + } + let server_info_result = write_cached_codex_apps_server_info(cache_context, server_info); + if let Err(err) = &server_info_result { + tracing::warn!("failed to write Codex Apps server info cache: {err:#}"); + } + let status = if tools_result.is_ok() && server_info_result.is_ok() { + "success" + } else { + "failure" + }; + emit_duration( + MCP_TOOLS_CACHE_WRITE_DURATION_METRIC, + cache_write_start.elapsed(), + &[("status", status)], + ); +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct CodexAppsToolsDiskCache { + schema_version: u8, + tools: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct CodexAppsServerInfoDiskCache { + schema_version: u8, + server_info: McpServerInfo, +} + +const CODEX_APPS_TOOLS_CACHE_DIR: &str = "cache/codex_apps_tools"; +const CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION: u8 = 4; + +const CODEX_APPS_SERVER_INFO_CACHE_DIR: &str = "cache/codex_apps_server_info"; +const CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION: u8 = 1; + +fn sha1_hex(s: &str) -> String { + let mut hasher = Sha1::new(); + hasher.update(s.as_bytes()); + let sha1 = hasher.finalize(); + format!("{sha1:x}") +} + +fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +#[cfg(test)] +#[path = "codex_apps_cache_tests.rs"] +mod tests; diff --git a/codex-rs/codex-mcp/src/codex_apps_cache_tests.rs b/codex-rs/codex-mcp/src/codex_apps_cache_tests.rs new file mode 100644 index 000000000..46ad311a1 --- /dev/null +++ b/codex-rs/codex-mcp/src/codex_apps_cache_tests.rs @@ -0,0 +1,464 @@ +use super::*; +use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; +use crate::tools::ToolInfo; +use codex_protocol::ToolName; +use codex_protocol::mcp::McpServerInfo; +use pretty_assertions::assert_eq; +use rmcp::model::JsonObject; +use rmcp::model::Tool; +use std::collections::HashSet; +#[cfg(unix)] +use std::os::unix::ffi::OsStringExt; +use std::path::PathBuf; +use std::sync::Arc; +use tempfile::tempdir; + +fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { + ToolInfo { + server_name: server_name.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, + callable_name: tool_name.to_string(), + callable_namespace: server_name.to_string(), + namespace_description: None, + tool: Tool::new( + tool_name.to_string(), + format!("Test tool: {tool_name}"), + Arc::new(JsonObject::default()), + ), + connector_id: None, + connector_name: None, + plugin_display_names: Vec::new(), + } +} + +fn create_test_tool_with_connector( + server_name: &str, + tool_name: &str, + connector_id: &str, + connector_name: Option<&str>, +) -> ToolInfo { + let mut tool = create_test_tool(server_name, tool_name); + tool.connector_id = Some(connector_id.to_string()); + tool.connector_name = connector_name.map(ToOwned::to_owned); + tool +} + +fn create_codex_apps_tools_cache_context( + codex_home: PathBuf, + account_id: Option<&str>, + chatgpt_user_id: Option<&str>, +) -> CodexAppsToolsCacheContext { + CodexAppsToolsCache::default().context( + codex_home, + CodexAppsToolsCacheKey { + account_id: account_id.map(ToOwned::to_owned), + chatgpt_user_id: chatgpt_user_id.map(ToOwned::to_owned), + is_workspace_account: false, + }, + ) +} + +fn create_test_server_info(title: &str) -> McpServerInfo { + McpServerInfo { + name: "codex-apps".to_string(), + title: Some(title.to_string()), + version: "1.0.0".to_string(), + description: None, + icons: None, + website_url: None, + } +} + +fn model_tool_names(tools: &[ToolInfo]) -> HashSet { + tools + .iter() + .map(ToolInfo::canonical_tool_name) + .collect::>() +} + +#[test] +fn codex_apps_tools_cache_is_overwritten_by_last_write() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let tools_gateway_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; + let tools_gateway_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; + + write_cached_codex_apps_tools(&cache_context, &tools_gateway_1).expect("write first cache"); + let cached_gateway_1 = + read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for first write"); + assert_eq!(cached_gateway_1[0].callable_name, "one"); + + write_cached_codex_apps_tools(&cache_context, &tools_gateway_2).expect("write second cache"); + let cached_gateway_2 = + read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for second write"); + assert_eq!(cached_gateway_2[0].callable_name, "two"); +} + +#[test] +fn codex_apps_tools_cache_is_scoped_per_user() { + let codex_home = tempdir().expect("tempdir"); + let cache_context_user_1 = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_context_user_2 = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-two"), + Some("user-two"), + ); + let tools_user_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; + let tools_user_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; + + write_cached_codex_apps_tools(&cache_context_user_1, &tools_user_1) + .expect("write user one cache"); + write_cached_codex_apps_tools(&cache_context_user_2, &tools_user_2) + .expect("write user two cache"); + + let read_user_1 = + read_cached_codex_apps_tools(&cache_context_user_1).expect("cache entry for user one"); + let read_user_2 = + read_cached_codex_apps_tools(&cache_context_user_2).expect("cache entry for user two"); + + assert_eq!(read_user_1[0].callable_name, "one"); + assert_eq!(read_user_2[0].callable_name, "two"); + assert_ne!( + cache_context_user_1.tools_cache_path(), + cache_context_user_2.tools_cache_path(), + "each user should get an isolated cache file" + ); +} + +#[test] +fn codex_apps_tools_cache_preserves_formerly_disallowed_connectors() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let tools = vec![ + create_test_tool_with_connector( + CODEX_APPS_MCP_SERVER_NAME, + "formerly_blocked_tool", + "connector_2b0a9009c9c64bf9933a3dae3f2b1254", + Some("Formerly Blocked"), + ), + create_test_tool_with_connector( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_tool", + "calendar", + Some("Calendar"), + ), + ]; + + write_cached_codex_apps_tools(&cache_context, &tools).expect("write cache"); + let cached = read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for user"); + + assert_eq!( + cached + .iter() + .map(|tool| (tool.callable_name.as_str(), tool.connector_id.as_deref())) + .collect::>(), + vec![ + ( + "formerly_blocked_tool", + Some("connector_2b0a9009c9c64bf9933a3dae3f2b1254") + ), + ("calendar_tool", Some("calendar")), + ] + ); +} + +#[test] +fn codex_apps_tools_cache_is_ignored_when_schema_version_mismatches() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_path = cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + let bytes = serde_json::to_vec_pretty(&serde_json::json!({ + "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION + 1, + "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")], + })) + .expect("serialize"); + std::fs::write(cache_path, bytes).expect("write"); + + assert!(read_cached_codex_apps_tools(&cache_context).is_none()); +} + +#[test] +fn codex_apps_tools_cache_is_ignored_when_json_is_invalid() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_path = cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + std::fs::write(cache_path, b"{not json").expect("write"); + + assert!(read_cached_codex_apps_tools(&cache_context).is_none()); +} + +#[test] +fn startup_cached_codex_apps_tools_loads_from_disk_cache() { + let codex_home = tempdir().expect("tempdir"); + let writer_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cached_tools = vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_search", + )]; + let server_info = create_test_server_info("Codex Apps"); + write_cached_codex_apps_tools_for_test(&writer_cache_context, &server_info, &cached_tools); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + + let startup_tools = cache_context + .current_tools() + .expect("expected startup snapshot to load from cache"); + let cached_server_info = load_startup_cached_codex_apps_server_info(&cache_context); + + assert_eq!(startup_tools.len(), 1); + assert_eq!(startup_tools[0].server_name, CODEX_APPS_MCP_SERVER_NAME); + assert_eq!(startup_tools[0].callable_name, "calendar_search"); + assert_eq!(cached_server_info, Some(server_info)); +} + +#[test] +fn startup_cached_codex_apps_tools_loads_without_server_info_cache() { + let codex_home = tempdir().expect("tempdir"); + let writer_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_path = writer_cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + let bytes = serde_json::to_vec_pretty(&serde_json::json!({ + "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, + "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], + })) + .expect("serialize"); + std::fs::write(cache_path, bytes).expect("write"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + + let startup_tools = cache_context + .current_tools() + .expect("legacy startup snapshot should remain available"); + let cached_server_info = load_startup_cached_codex_apps_server_info(&cache_context); + + assert_eq!(startup_tools.len(), 1); + assert_eq!(startup_tools[0].callable_name, "calendar_search"); + assert_eq!(cached_server_info, None); +} + +#[test] +fn codex_apps_server_info_cache_survives_legacy_tools_cache_write() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let server_info = create_test_server_info("Codex Apps"); + write_cached_codex_apps_tools_for_test( + &cache_context, + &server_info, + &[create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_search", + )], + ); + + let cache_path = cache_context.tools_cache_path(); + if let Some(parent) = cache_path.parent() { + std::fs::create_dir_all(parent).expect("create parent"); + } + let bytes = serde_json::to_vec_pretty(&serde_json::json!({ + "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION - 1, + "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], + })) + .expect("serialize"); + std::fs::write(cache_path, bytes).expect("write legacy tools cache"); + let startup_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + + assert_eq!( + load_startup_cached_codex_apps_server_info(&startup_cache_context), + Some(server_info) + ); + assert!(startup_cache_context.current_tools().is_none()); +} + +#[test] +fn codex_apps_tools_cache_context_does_not_reread_disk_after_creation() { + let codex_home = tempdir().expect("tempdir"); + let writer_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cached_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "cached")]; + write_cached_codex_apps_tools(&writer_cache_context, &cached_tools).expect("write cache"); + let reader_cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let updated_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "updated")]; + write_cached_codex_apps_tools(&writer_cache_context, &updated_tools).expect("rewrite cache"); + + assert_eq!( + reader_cache_context + .current_tools() + .expect("in-memory tools")[0] + .callable_name, + "cached" + ); + assert_eq!( + read_cached_codex_apps_tools(&writer_cache_context).expect("disk tools")[0].callable_name, + "updated" + ); +} + +#[test] +fn codex_apps_tools_cache_publishes_newest_shared_snapshot() { + let codex_home = tempdir().expect("tempdir"); + let cache = CodexAppsToolsCache::default(); + let cache_context_1 = cache.context( + codex_home.path().to_path_buf(), + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let cache_context_2 = cache.context( + codex_home.path().to_path_buf(), + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let older_ticket = cache_context_1.begin_fetch(CodexAppsToolsFetchSource::Startup); + let newer_ticket = cache_context_2.begin_fetch(CodexAppsToolsFetchSource::HardRefresh); + let server_info = create_test_server_info("Codex Apps"); + let newer_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "newer")]; + let older_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "older")]; + + let published_tools = + cache_context_2.publish_if_newest_accepted(newer_ticket, &server_info, newer_tools); + assert_eq!( + model_tool_names(&published_tools), + model_tool_names( + &cache_context_1 + .current_tools() + .expect("new snapshot should publish") + ) + ); + let current_tools = + cache_context_1.publish_if_newest_accepted(older_ticket, &server_info, older_tools); + + assert_eq!(current_tools[0].callable_name, "newer"); + assert_eq!( + cache_context_2.current_tools().expect("shared snapshot")[0].callable_name, + "newer" + ); + assert_eq!( + read_cached_codex_apps_tools(&cache_context_1).expect("persisted snapshot")[0] + .callable_name, + "newer" + ); +} + +#[test] +fn codex_apps_tools_cache_keeps_live_publish_when_disk_persistence_fails() { + let codex_home = tempdir().expect("tempdir"); + let codex_home_file = codex_home.path().join("not-a-directory"); + std::fs::write(&codex_home_file, b"occupied").expect("create codex home file"); + let cache_context = CodexAppsToolsCache::default().context( + codex_home_file, + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "live")]; + let published_tools = cache_context.publish_if_newest_accepted( + cache_context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh), + &create_test_server_info("Codex Apps"), + tools.clone(), + ); + + assert_eq!(model_tool_names(&published_tools), model_tool_names(&tools)); + assert_eq!( + model_tool_names(&cache_context.current_tools().expect("live snapshot")), + model_tool_names(&tools) + ); +} + +#[cfg(unix)] +#[test] +fn codex_apps_tools_cache_scopes_non_utf8_home_disk_paths() { + let codex_home = PathBuf::from(std::ffi::OsString::from_vec( + b"/tmp/codex-home-\xff".to_vec(), + )); + let cache = CodexAppsToolsCache::default(); + let user_one_context = cache.context( + codex_home.clone(), + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let user_two_context = cache.context( + codex_home, + CodexAppsToolsCacheKey { + account_id: Some("account-two".to_string()), + chatgpt_user_id: Some("user-two".to_string()), + is_workspace_account: false, + }, + ); + let cache_paths = [ + user_one_context.tools_cache_path(), + user_two_context.tools_cache_path(), + ]; + + assert_eq!( + cache_paths.iter().collect::>().len(), + cache_paths.len() + ); +} diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 55c72b046..0e8154955 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -7,7 +7,6 @@ //! `codex-core`. use std::collections::HashMap; -use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::Ordering; @@ -15,9 +14,9 @@ use std::time::Duration; use std::time::Instant; use crate::McpAuthStatusEntry; -use crate::codex_apps::CodexAppsToolsCacheContext; -use crate::codex_apps::CodexAppsToolsCacheKey; -use crate::codex_apps::write_codex_apps_tools_cache; +use crate::codex_apps_cache::CodexAppsToolsCache; +use crate::codex_apps_cache::CodexAppsToolsCacheKey; +use crate::codex_apps_cache::CodexAppsToolsFetchSource; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationReviewerHandle; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; @@ -131,6 +130,7 @@ impl McpConnectionManager { initial_permission_profile: PermissionProfile, runtime_context: McpRuntimeContext, codex_home: PathBuf, + codex_apps_tools_cache: CodexAppsToolsCache, codex_apps_tools_cache_key: CodexAppsToolsCacheKey, prefix_mcp_tool_names: bool, client_elicitation_capability: ElicitationCapability, @@ -174,13 +174,32 @@ impl McpConnectionManager { }, ) .await; - let codex_apps_tools_cache_context = if server_name == CODEX_APPS_MCP_SERVER_NAME { - codex_apps_tools_cache_context(&codex_home, &codex_apps_tools_cache_key) - } else { - regular_mcp_tools_cache_context() - }; + let configured_config = server.configured_config(); + // For built-in Codex Apps, `CODEX_CONNECTORS_TOKEN` is a debug + // override: it supplies runtime auth but bypasses the shared tools + // cache. + let uses_env_bearer_token = + configured_config.is_some_and(|config| match &config.transport { + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var, + .. + } => bearer_token_env_var.is_some(), + McpServerTransportConfig::Stdio { .. } => false, + }); + let shares_codex_apps_tools_cache = + should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token); + let codex_apps_tools_cache_context = shares_codex_apps_tools_cache.then(|| { + codex_apps_tools_cache + .context(codex_home.clone(), codex_apps_tools_cache_key.clone()) + }); + // If Codex Apps has an env bearer token, that is its auth path. Do + // not also attach the ambient CodexAuth provider. let runtime_auth_provider = - chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider.clone()); + if server_name == CODEX_APPS_MCP_SERVER_NAME && uses_env_bearer_token { + None + } else { + chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider.clone()) + }; let async_managed_client = AsyncManagedClient::new( server_name.clone(), server, @@ -445,13 +464,13 @@ impl McpConnectionManager { pub async fn list_all_tools(&self) -> Vec { let mut tools = Vec::new(); for (server_name, managed_client) in &self.clients { - let has_cached_tool_info_snapshot = managed_client.cached_tool_info_snapshot.is_some(); + let has_cached_tools = managed_client.has_cached_tools(); let startup_complete = managed_client .startup_complete .load(std::sync::atomic::Ordering::Acquire); trace!( server_name = %server_name, - has_cached_tool_info_snapshot, + has_cached_tools, startup_complete, "waiting for MCP server tools while building tool list" ); @@ -460,7 +479,7 @@ impl McpConnectionManager { .instrument(trace_span!( "list_tools_for_server", server_name = %server_name, - has_cached_tool_info_snapshot, + has_cached_tools, startup_complete )) .await @@ -483,9 +502,9 @@ impl McpConnectionManager { /// Force-refresh codex apps tools by bypassing the in-process cache. /// - /// On success, the refreshed tools replace the cache contents and the - /// latest filtered tools are returned directly to the caller. On - /// failure, the existing cache remains unchanged. + /// On success, the refreshed tools replace shared cache contents when the + /// cache is enabled and the latest filtered tools are returned directly to + /// the caller. On failure, existing shared cache contents remain unchanged. pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { let managed_client = self .clients @@ -497,6 +516,10 @@ impl McpConnectionManager { let list_start = Instant::now(); let fetch_start = Instant::now(); + let fetch_ticket = managed_client + .codex_apps_tools_cache_context + .as_ref() + .map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh)); let tools = list_tools_for_client_uncached( CODEX_APPS_MCP_SERVER_NAME, /*is_codex_apps_mcp_server*/ true, @@ -514,11 +537,16 @@ impl McpConnectionManager { &[], ); - write_codex_apps_tools_cache( - managed_client.codex_apps_tools_cache_context.as_ref(), - &managed_client.server_info, - &tools, - ); + let tools = + match ( + managed_client.codex_apps_tools_cache_context.as_ref(), + fetch_ticket, + ) { + (Some(cache_context), Some(fetch_ticket)) => cache_context + .publish_if_newest_accepted(fetch_ticket, &managed_client.server_info, tools), + (None, None) => tools, + _ => unreachable!("Codex Apps fetch ticket requires cache context"), + }; emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, list_start.elapsed(), @@ -849,22 +877,6 @@ impl Drop for McpConnectionManager { } } -/// Creates the per-user tools cache context used only by the Codex Apps server. -fn codex_apps_tools_cache_context( - codex_home: &Path, - codex_apps_tools_cache_key: &CodexAppsToolsCacheKey, -) -> Option { - Some(CodexAppsToolsCacheContext { - codex_home: codex_home.to_path_buf(), - user_key: codex_apps_tools_cache_key.clone(), - }) -} - -/// Keeps regular MCP servers isolated from the Codex Apps tools cache. -fn regular_mcp_tools_cache_context() -> Option { - None -} - /// Makes ChatGPT authentication available to servers that explicitly opt in. /// The HTTP transport applies it only when no configured authorization resolves. fn chatgpt_auth_provider_for_server( @@ -880,6 +892,10 @@ fn chatgpt_auth_provider_for_server( chatgpt_auth_provider } +fn should_share_codex_apps_tools_cache(server_name: &str, uses_env_bearer_token: bool) -> bool { + server_name == CODEX_APPS_MCP_SERVER_NAME && !uses_env_bearer_token +} + async fn emit_update( submit_id: &str, tx_event: &Sender, diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index d66474c7b..05caf0489 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1,11 +1,6 @@ use super::*; -use crate::codex_apps::CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION; -use crate::codex_apps::CodexAppsToolsCacheContext; -use crate::codex_apps::load_startup_cached_codex_apps_server_info; -use crate::codex_apps::load_startup_cached_codex_apps_tools_snapshot; -use crate::codex_apps::read_cached_codex_apps_tools; -use crate::codex_apps::write_cached_codex_apps_tools; -use crate::codex_apps::write_codex_apps_tools_cache; +use crate::codex_apps_cache::CodexAppsToolsCache; +use crate::codex_apps_cache::CodexAppsToolsCacheContext; use crate::declared_openai_file_input_param_names; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::elicitation_is_rejected_by_policy; @@ -31,7 +26,10 @@ use codex_protocol::mcp::McpServerInfo; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::GranularApprovalConfig; use codex_protocol::protocol::McpAuthStatus; +use codex_rmcp_client::InProcessTransportFactory; +use codex_rmcp_client::RmcpClient; use futures::FutureExt; +use futures::future::BoxFuture; use pretty_assertions::assert_eq; use rmcp::model::CreateElicitationRequestParams; use rmcp::model::ElicitationAction; @@ -41,8 +39,10 @@ use rmcp::model::Meta; use rmcp::model::NumberOrString; use rmcp::model::Tool; use std::collections::HashSet; +use std::io; use std::sync::Arc; use tempfile::tempdir; +use tokio::io::DuplexStream; fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { ToolInfo { @@ -63,31 +63,19 @@ fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { } } -fn create_test_tool_with_connector( - server_name: &str, - tool_name: &str, - connector_id: &str, - connector_name: Option<&str>, -) -> ToolInfo { - let mut tool = create_test_tool(server_name, tool_name); - tool.connector_id = Some(connector_id.to_string()); - tool.connector_name = connector_name.map(ToOwned::to_owned); - tool -} - fn create_codex_apps_tools_cache_context( codex_home: PathBuf, account_id: Option<&str>, chatgpt_user_id: Option<&str>, ) -> CodexAppsToolsCacheContext { - CodexAppsToolsCacheContext { + CodexAppsToolsCache::default().context( codex_home, - user_key: CodexAppsToolsCacheKey { + CodexAppsToolsCacheKey { account_id: account_id.map(ToOwned::to_owned), chatgpt_user_id: chatgpt_user_id.map(ToOwned::to_owned), is_workspace_account: false, }, - } + ) } fn create_test_server_info(title: &str) -> McpServerInfo { @@ -101,6 +89,51 @@ fn create_test_server_info(title: &str) -> McpServerInfo { } } +struct TestInProcessTransportFactory; + +impl InProcessTransportFactory for TestInProcessTransportFactory { + fn open(&self) -> BoxFuture<'static, io::Result> { + async { + let (client_stream, _server_stream) = tokio::io::duplex(1); + Ok(client_stream) + } + .boxed() + } +} + +async fn create_ready_async_managed_client(tools: Vec) -> AsyncManagedClient { + let tool_filter = ToolFilter::default(); + let managed_client = ManagedClient { + client: Arc::new( + RmcpClient::new_in_process_client(Arc::new(TestInProcessTransportFactory)) + .await + .expect("create in-process RMCP client"), + ), + server_info: create_test_server_info("Ready"), + tools, + tool_filter: tool_filter.clone(), + tool_timeout: None, + server_instructions: None, + server_supports_sandbox_state_meta_capability: false, + codex_apps_tools_cache_context: None, + }; + + AsyncManagedClient { + client: futures::future::ready::>(Ok( + managed_client, + )) + .boxed() + .shared(), + is_codex_apps_mcp_server: false, + cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter, + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + } +} + fn model_tool_names(tools: &[ToolInfo]) -> HashSet { tools .iter() @@ -554,235 +587,25 @@ fn filter_tools_applies_per_server_filters() { } #[test] -fn codex_apps_tools_cache_is_overwritten_by_last_write() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let tools_gateway_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; - let tools_gateway_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; - - write_cached_codex_apps_tools(&cache_context, &tools_gateway_1); - let cached_gateway_1 = - read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for first write"); - assert_eq!(cached_gateway_1[0].callable_name, "one"); - - write_cached_codex_apps_tools(&cache_context, &tools_gateway_2); - let cached_gateway_2 = - read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for second write"); - assert_eq!(cached_gateway_2[0].callable_name, "two"); -} - -#[test] -fn codex_apps_tools_cache_is_scoped_per_user() { - let codex_home = tempdir().expect("tempdir"); - let cache_context_user_1 = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_context_user_2 = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-two"), - Some("user-two"), - ); - let tools_user_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")]; - let tools_user_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")]; - - write_cached_codex_apps_tools(&cache_context_user_1, &tools_user_1); - write_cached_codex_apps_tools(&cache_context_user_2, &tools_user_2); - - let read_user_1 = - read_cached_codex_apps_tools(&cache_context_user_1).expect("cache entry for user one"); - let read_user_2 = - read_cached_codex_apps_tools(&cache_context_user_2).expect("cache entry for user two"); - - assert_eq!(read_user_1[0].callable_name, "one"); - assert_eq!(read_user_2[0].callable_name, "two"); - assert_ne!( - cache_context_user_1.tools_cache_path(), - cache_context_user_2.tools_cache_path(), - "each user should get an isolated cache file" - ); -} - -#[test] -fn codex_apps_tools_cache_preserves_formerly_disallowed_connectors() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let tools = vec![ - create_test_tool_with_connector( - CODEX_APPS_MCP_SERVER_NAME, - "formerly_blocked_tool", - "connector_2b0a9009c9c64bf9933a3dae3f2b1254", - Some("Formerly Blocked"), - ), - create_test_tool_with_connector( - CODEX_APPS_MCP_SERVER_NAME, - "calendar_tool", - "calendar", - Some("Calendar"), - ), - ]; - - write_cached_codex_apps_tools(&cache_context, &tools); - let cached = read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for user"); - - assert_eq!( - cached - .iter() - .map(|tool| (tool.callable_name.as_str(), tool.connector_id.as_deref())) - .collect::>(), - vec![ - ( - "formerly_blocked_tool", - Some("connector_2b0a9009c9c64bf9933a3dae3f2b1254") - ), - ("calendar_tool", Some("calendar")), - ] - ); -} - -#[test] -fn codex_apps_tools_cache_is_ignored_when_schema_version_mismatches() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - let bytes = serde_json::to_vec_pretty(&serde_json::json!({ - "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION + 1, - "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")], - })) - .expect("serialize"); - std::fs::write(cache_path, bytes).expect("write"); - - assert!(read_cached_codex_apps_tools(&cache_context).is_none()); -} - -#[test] -fn codex_apps_tools_cache_is_ignored_when_json_is_invalid() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - std::fs::write(cache_path, b"{not json").expect("write"); - - assert!(read_cached_codex_apps_tools(&cache_context).is_none()); -} - -#[test] -fn startup_cached_codex_apps_tools_loads_from_disk_cache() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cached_tools = vec![create_test_tool( +fn codex_apps_env_bearer_token_bypasses_shared_tools_cache() { + assert!(!should_share_codex_apps_tools_cache( CODEX_APPS_MCP_SERVER_NAME, - "calendar_search", - )]; - let server_info = create_test_server_info("Codex Apps"); - write_codex_apps_tools_cache(Some(&cache_context), &server_info, &cached_tools); - - let startup_tools = load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context)) - .expect("expected startup snapshot to load from cache"); - let cached_server_info = load_startup_cached_codex_apps_server_info(Some(&cache_context)); - - assert_eq!(startup_tools.len(), 1); - assert_eq!(startup_tools[0].server_name, CODEX_APPS_MCP_SERVER_NAME); - assert_eq!(startup_tools[0].callable_name, "calendar_search"); - assert_eq!(cached_server_info, Some(server_info)); -} - -#[test] -fn startup_cached_codex_apps_tools_loads_without_server_info_cache() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - let bytes = serde_json::to_vec_pretty(&serde_json::json!({ - "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION, - "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], - })) - .expect("serialize"); - std::fs::write(cache_path, bytes).expect("write"); - - let startup_tools = load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context)) - .expect("legacy startup snapshot should remain available"); - let cached_server_info = load_startup_cached_codex_apps_server_info(Some(&cache_context)); - - assert_eq!(startup_tools.len(), 1); - assert_eq!(startup_tools[0].callable_name, "calendar_search"); - assert_eq!(cached_server_info, None); -} - -#[test] -fn codex_apps_server_info_cache_survives_legacy_tools_cache_write() { - let codex_home = tempdir().expect("tempdir"); - let cache_context = create_codex_apps_tools_cache_context( - codex_home.path().to_path_buf(), - Some("account-one"), - Some("user-one"), - ); - let server_info = create_test_server_info("Codex Apps"); - write_codex_apps_tools_cache( - Some(&cache_context), - &server_info, - &[create_test_tool( - CODEX_APPS_MCP_SERVER_NAME, - "calendar_search", - )], - ); - - let cache_path = cache_context.tools_cache_path(); - if let Some(parent) = cache_path.parent() { - std::fs::create_dir_all(parent).expect("create parent"); - } - let bytes = serde_json::to_vec_pretty(&serde_json::json!({ - "schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION - 1, - "tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")], - })) - .expect("serialize"); - std::fs::write(cache_path, bytes).expect("write legacy tools cache"); - - assert_eq!( - load_startup_cached_codex_apps_server_info(Some(&cache_context)), - Some(server_info) - ); - assert!(load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context)).is_none()); + /*uses_env_bearer_token*/ true, + )); } #[tokio::test] -async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() { - let startup_tools = vec![create_test_tool( +async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + cache_context.store_current_tools_for_test(vec![create_test_tool( CODEX_APPS_MCP_SERVER_NAME, "calendar_create_event", - )]; + )]); let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -798,8 +621,9 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(startup_tools), cached_server_info: None, + codex_apps_tools_cache_context: Some(cache_context), + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -813,7 +637,7 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() tool.canonical_tool_name() == ToolName::namespaced("mcp__codex_apps", "calendar_create_event") }) - .expect("tool from startup cache"); + .expect("tool from shared cache"); assert_eq!(tool.server_name, CODEX_APPS_MCP_SERVER_NAME); assert_eq!(tool.callable_name, "calendar_create_event"); } @@ -836,8 +660,9 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(Vec::new()), cached_server_info: Some(server_info.clone()), + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -858,10 +683,8 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { #[tokio::test] async fn list_all_tools_accepts_canonical_namespaced_tool_names() { - let startup_tools = vec![create_test_tool("rmcp", "echo")]; - let pending_client = futures::future::pending::>() - .boxed() - .shared(); + let managed_client = + create_ready_async_managed_client(vec![create_test_tool("rmcp", "echo")]).await; let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionManager::new_uninitialized( @@ -869,18 +692,7 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() { &permission_profile, /*prefix_mcp_tool_names*/ false, ); - manager.clients.insert( - "rmcp".to_string(), - AsyncManagedClient { - client: pending_client, - is_codex_apps_mcp_server: false, - cached_tool_info_snapshot: Some(startup_tools), - cached_server_info: None, - startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - cancel_token: CancellationToken::new(), - }, - ); + manager.clients.insert("rmcp".to_string(), managed_client); let tools = manager.list_all_tools().await; let tool = tools @@ -902,10 +714,8 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() { #[tokio::test] async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { - let startup_tools = vec![create_test_tool("rmcp", "echo")]; - let pending_client = futures::future::pending::>() - .boxed() - .shared(); + let managed_client = + create_ready_async_managed_client(vec![create_test_tool("rmcp", "echo")]).await; let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionManager::new_uninitialized( @@ -913,18 +723,7 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( - "rmcp".to_string(), - AsyncManagedClient { - client: pending_client, - is_codex_apps_mcp_server: false, - cached_tool_info_snapshot: Some(startup_tools), - cached_server_info: None, - startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - cancel_token: CancellationToken::new(), - }, - ); + manager.clients.insert("rmcp".to_string(), managed_client); let tools = manager.list_all_tools().await; let tool = tools @@ -945,7 +744,7 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { } #[tokio::test] -async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_snapshot() { +async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -961,8 +760,9 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_ AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: None, cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -998,8 +798,9 @@ async fn shutdown_cancels_pending_tool_listing() { AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: None, cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token, @@ -1043,8 +844,9 @@ async fn shutdown_continues_after_caller_is_aborted() { AsyncManagedClient { client: blocking_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: None, cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -1071,7 +873,14 @@ async fn shutdown_continues_after_caller_is_aborted() { } #[tokio::test] -async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() { +async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + cache_context.store_current_tools_for_test(Vec::new()); let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -1087,8 +896,9 @@ async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() AsyncManagedClient { client: pending_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(Vec::new()), cached_server_info: None, + codex_apps_tools_cache_context: Some(cache_context), + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -1097,16 +907,22 @@ async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() let timeout_result = tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()).await; - let tools = timeout_result.expect("cache-hit startup snapshot should not block"); + let tools = timeout_result.expect("shared empty cache should not block"); assert!(tools.is_empty()); } #[tokio::test] -async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails() { - let startup_tools = vec![create_test_tool( +async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() { + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + cache_context.store_current_tools_for_test(vec![create_test_tool( CODEX_APPS_MCP_SERVER_NAME, "calendar_create_event", - )]; + )]); let server_info = create_test_server_info("Codex Apps"); let failed_client = futures::future::ready::>(Err( StartupOutcomeError::Failed { @@ -1128,8 +944,9 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails AsyncManagedClient { client: failed_client, is_codex_apps_mcp_server: true, - cached_tool_info_snapshot: Some(startup_tools), cached_server_info: Some(server_info.clone()), + codex_apps_tools_cache_context: Some(cache_context), + tool_filter: ToolFilter::default(), startup_complete, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), @@ -1143,7 +960,7 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails tool.canonical_tool_name() == ToolName::namespaced("mcp__codex_apps", "calendar_create_event") }) - .expect("tool from startup cache"); + .expect("tool from shared cache"); assert_eq!(tool.server_name, CODEX_APPS_MCP_SERVER_NAME); assert_eq!(tool.callable_name, "calendar_create_event"); assert_eq!( @@ -1156,12 +973,10 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails } #[tokio::test] -async fn list_all_tools_adds_server_metadata_to_cached_tools() { +async fn list_all_tools_adds_server_metadata_to_tools() { let server_name = "docs"; - let startup_tools = vec![create_test_tool(server_name, "search")]; - let pending_client = futures::future::pending::>() - .boxed() - .shared(); + let managed_client = + create_ready_async_managed_client(vec![create_test_tool(server_name, "search")]).await; let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionManager::new_uninitialized( @@ -1182,18 +997,9 @@ async fn list_all_tools_adds_server_metadata_to_cached_tools() { tool_approval_modes: HashMap::new(), }, ); - manager.clients.insert( - server_name.to_string(), - AsyncManagedClient { - client: pending_client, - is_codex_apps_mcp_server: false, - cached_tool_info_snapshot: Some(startup_tools), - cached_server_info: None, - startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - cancel_token: CancellationToken::new(), - }, - ); + manager + .clients + .insert(server_name.to_string(), managed_client); let tools = manager.list_all_tools().await; assert_eq!(tools.len(), 1); @@ -1340,6 +1146,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { PathBuf::from("/tmp"), ), codex_home.path().to_path_buf(), + CodexAppsToolsCache::default(), CodexAppsToolsCacheKey { account_id: None, chatgpt_user_id: None, diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 98efdc1b5..e611623d4 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -35,8 +35,9 @@ pub use auth_elicitation::auth_elicitation_id; pub use auth_elicitation::build_auth_elicitation; pub use auth_elicitation::build_auth_elicitation_plan; pub use auth_elicitation::connector_auth_failure_from_tool_result; -pub use codex_apps::CodexAppsToolsCacheKey; -pub use codex_apps::codex_apps_tools_cache_key; +pub use codex_apps_cache::CodexAppsToolsCache; +pub use codex_apps_cache::CodexAppsToolsCacheKey; +pub use codex_apps_cache::codex_apps_tools_cache_key; pub use mcp::codex_apps_mcp_server_config; pub use mcp::configured_mcp_servers; pub use mcp::effective_mcp_servers; @@ -75,6 +76,7 @@ pub use tools::declared_openai_file_input_param_names; pub(crate) mod auth_elicitation; mod catalog; pub(crate) mod codex_apps; +pub(crate) mod codex_apps_cache; pub(crate) mod connection_manager; pub(crate) mod elicitation; pub(crate) mod mcp; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 4550afa79..7532c6466 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -44,7 +44,8 @@ use serde_json::Value; use tokio_util::sync::CancellationToken; use crate::ResolvedMcpCatalog; -use crate::codex_apps::codex_apps_tools_cache_key; +use crate::codex_apps_cache::CodexAppsToolsCache; +use crate::codex_apps_cache::codex_apps_tools_cache_key; use crate::connection_manager::McpConnectionManager; use crate::runtime::McpRuntimeContext; use crate::server::EffectiveMcpServer; @@ -301,6 +302,7 @@ pub async fn read_mcp_resource( config: &McpConfig, auth: Option<&CodexAuth>, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: CodexAppsToolsCache, server: &str, uri: &str, ) -> anyhow::Result { @@ -329,6 +331,7 @@ pub async fn read_mcp_resource( PermissionProfile::default(), runtime_context, config.codex_home.clone(), + codex_apps_tools_cache, codex_apps_tools_cache_key(auth), config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), @@ -361,6 +364,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( auth: Option<&CodexAuth>, submit_id: String, runtime_context: McpRuntimeContext, + codex_apps_tools_cache: CodexAppsToolsCache, detail: McpSnapshotDetail, ) -> McpServerStatusSnapshot { let mcp_servers = effective_mcp_servers(config, auth); @@ -403,6 +407,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( PermissionProfile::default(), runtime_context, config.codex_home.clone(), + codex_apps_tools_cache, codex_apps_tools_cache_key(auth), config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 1e8f5a331..73ad33df7 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -2,7 +2,7 @@ //! //! This module owns startup of individual RMCP clients: building the transport, //! initializing the server, listing raw tools, applying per-server tool filters, -//! and exposing cached startup snapshots while a client is still connecting. +//! and exposing cached Codex Apps tools while a client is still connecting. //! Higher-level aggregation and resource/tool APIs live in //! [`crate::connection_manager`]. @@ -17,15 +17,12 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; -use crate::codex_apps::CachedCodexAppsToolsLoad; -use crate::codex_apps::CodexAppsToolsCacheContext; -use crate::codex_apps::load_cached_codex_apps_tools; -use crate::codex_apps::load_startup_cached_codex_apps_server_info; -use crate::codex_apps::load_startup_cached_codex_apps_tools_snapshot; use crate::codex_apps::normalize_codex_apps_callable_name; use crate::codex_apps::normalize_codex_apps_callable_namespace; use crate::codex_apps::normalize_codex_apps_tool_title; -use crate::codex_apps::write_codex_apps_tools_cache; +use crate::codex_apps_cache::CodexAppsToolsCacheContext; +use crate::codex_apps_cache::CodexAppsToolsFetchSource; +use crate::codex_apps_cache::load_startup_cached_codex_apps_server_info; use crate::elicitation::ElicitationRequestManager; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::ToolPluginProvenance; @@ -105,9 +102,10 @@ pub(crate) struct ManagedClient { impl ManagedClient { fn listed_tools(&self) -> Vec { let total_start = Instant::now(); - if let Some(cache_context) = self.codex_apps_tools_cache_context.as_ref() - && let CachedCodexAppsToolsLoad::Hit(tools) = - load_cached_codex_apps_tools(cache_context) + if let Some(tools) = self + .codex_apps_tools_cache_context + .as_ref() + .and_then(CodexAppsToolsCacheContext::current_tools) { emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, @@ -133,8 +131,9 @@ impl ManagedClient { pub(crate) struct AsyncManagedClient { pub(crate) client: Shared>>, pub(crate) is_codex_apps_mcp_server: bool, - pub(crate) cached_tool_info_snapshot: Option>, pub(crate) cached_server_info: Option, + pub(crate) codex_apps_tools_cache_context: Option, + pub(crate) tool_filter: ToolFilter, pub(crate) startup_complete: Arc, pub(crate) tool_plugin_provenance: Arc, pub(crate) cancel_token: CancellationToken, @@ -165,19 +164,15 @@ impl AsyncManagedClient { .configured_config() .map(ToolFilter::from_config) .unwrap_or_default(); - let (cached_tool_info_snapshot, cached_server_info) = if is_codex_apps_mcp_server { - ( - load_startup_cached_codex_apps_tools_snapshot( - codex_apps_tools_cache_context.as_ref(), - ), - load_startup_cached_codex_apps_server_info(codex_apps_tools_cache_context.as_ref()), - ) + let cached_server_info = if is_codex_apps_mcp_server { + codex_apps_tools_cache_context + .as_ref() + .and_then(load_startup_cached_codex_apps_server_info) } else { - (None, None) + None }; - let cached_tool_info_snapshot = - cached_tool_info_snapshot.map(|tools| filter_tools(tools, &tool_filter)); - let startup_tool_filter = tool_filter; + let startup_tool_filter = tool_filter.clone(); + let codex_apps_tools_cache_context_for_fut = codex_apps_tools_cache_context.clone(); let startup_complete = Arc::new(AtomicBool::new(false)); let startup_complete_for_fut = Arc::clone(&startup_complete); let cancel_token_for_fut = cancel_token.clone(); @@ -214,7 +209,7 @@ impl AsyncManagedClient { tool_filter: startup_tool_filter, tx_event, elicitation_requests, - codex_apps_tools_cache_context, + codex_apps_tools_cache_context: codex_apps_tools_cache_context_for_fut, client_elicitation_capability, supports_openai_form_elicitation, }, @@ -232,7 +227,10 @@ impl AsyncManagedClient { outcome }; let client = fut.in_current_span().boxed().shared(); - if cached_tool_info_snapshot.is_some() { + if codex_apps_tools_cache_context + .as_ref() + .is_some_and(CodexAppsToolsCacheContext::has_current_tools) + { let startup_task = client.clone(); tokio::spawn(async move { let _ = startup_task.await; @@ -242,8 +240,9 @@ impl AsyncManagedClient { Self { client, is_codex_apps_mcp_server, - cached_tool_info_snapshot, cached_server_info, + codex_apps_tools_cache_context, + tool_filter, startup_complete, tool_plugin_provenance, cancel_token, @@ -265,15 +264,29 @@ impl AsyncManagedClient { } } + pub(crate) fn has_cached_tools(&self) -> bool { + self.codex_apps_tools_cache_context + .as_ref() + .is_some_and(CodexAppsToolsCacheContext::has_current_tools) + } + + fn cached_tools(&self) -> Option> { + self.codex_apps_tools_cache_context + .as_ref() + .and_then(CodexAppsToolsCacheContext::current_tools) + .map(|tools| filter_tools(tools, &self.tool_filter)) + } + pub(crate) async fn listed_tools(&self) -> Option> { // Keep cache payloads raw; plugin provenance is resolved per-session at read time. - let tools = if let Some(startup_tools) = self.cached_tool_info_snapshot_while_initializing() + let tools = if !self.startup_complete.load(Ordering::Acquire) + && let Some(startup_tools) = self.cached_tools() { Some(startup_tools) } else { match self.client().await { Ok(client) => Some(client.listed_tools()), - Err(_) => self.cached_tool_info_snapshot.clone(), + Err(_) => self.cached_tools(), } }?; Some(if self.is_codex_apps_mcp_server { @@ -282,13 +295,6 @@ impl AsyncManagedClient { prepare_regular_mcp_tools_for_model(tools, &self.tool_plugin_provenance) }) } - - fn cached_tool_info_snapshot_while_initializing(&self) -> Option> { - if !self.startup_complete.load(Ordering::Acquire) { - return self.cached_tool_info_snapshot.clone(); - } - None - } } #[derive(Debug, Clone, thiserror::Error)] @@ -570,6 +576,9 @@ async fn start_server_task( .is_some(); let list_start = Instant::now(); let fetch_start = Instant::now(); + let fetch_ticket = codex_apps_tools_cache_context + .as_ref() + .map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::Startup)); let tools = list_tools_for_client_uncached( &server_name, is_codex_apps_mcp_server, @@ -585,21 +594,20 @@ async fn start_server_task( &[], ); let server_info = mcp_server_info_from_implementation(initialize_result.server_info); - let codex_apps_tools_cache_context = if is_codex_apps_mcp_server { - write_codex_apps_tools_cache( - codex_apps_tools_cache_context.as_ref(), - &server_info, - &tools, - ); + let tools = match (codex_apps_tools_cache_context.as_ref(), fetch_ticket) { + (Some(cache_context), Some(fetch_ticket)) => { + cache_context.publish_if_newest_accepted(fetch_ticket, &server_info, tools) + } + (None, None) => tools, + _ => unreachable!("Codex Apps fetch ticket requires cache context"), + }; + if is_codex_apps_mcp_server { emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, list_start.elapsed(), &[("cache", "miss")], ); - codex_apps_tools_cache_context - } else { - None - }; + } let tools = filter_tools(tools, &tool_filter); let managed = ManagedClient { diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 76b4a8483..97f24e043 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -275,6 +275,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( // one exists, but do not reintroduce the old hidden-local fallback. runtime_context, config.codex_home.to_path_buf(), + mcp_manager.codex_apps_tools_cache(), codex_apps_tools_cache_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability, diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 90e6b01e7..658c5471e 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -10,6 +10,7 @@ use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; use codex_login::CodexAuth; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::CodexAppsToolsCache; use codex_mcp::EffectiveMcpServer; use codex_mcp::McpConfig; use codex_mcp::McpPluginAttribution; @@ -38,14 +39,15 @@ enum OrderedMcpOverlay { pub struct McpManager { plugins_manager: Arc, extensions: Arc>, + codex_apps_tools_cache: CodexAppsToolsCache, } impl McpManager { pub fn new(plugins_manager: Arc) -> Self { - Self { + Self::new_with_extensions( plugins_manager, - extensions: codex_extension_api::empty_extension_registry(), - } + codex_extension_api::empty_extension_registry(), + ) } /// Creates a manager that resolves host-installed MCP contributions. @@ -56,9 +58,14 @@ impl McpManager { Self { plugins_manager, extensions, + codex_apps_tools_cache: CodexAppsToolsCache::default(), } } + pub fn codex_apps_tools_cache(&self) -> CodexAppsToolsCache { + self.codex_apps_tools_cache.clone() + } + /// Returns the MCP config after applying compatibility built-ins and /// runtime-only extension overlays. pub async fn runtime_config(&self, config: &Config) -> McpConfig { diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index df2a375e3..a270af785 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -1471,9 +1471,15 @@ pub(crate) async fn lookup_mcp_tool_metadata( { Some(connectors) => Some(connectors), None => { - connectors::list_accessible_connectors_from_mcp_tools(turn_context.config.as_ref()) - .await - .ok() + connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager( + turn_context.config.as_ref(), + /*force_refetch*/ false, + sess.services.turn_environments.environment_manager(), + Arc::clone(&sess.services.mcp_manager), + ) + .await + .ok() + .map(|status| status.connectors) } }; connectors.and_then(|connectors| { diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 387987c9f..41701d557 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1400,6 +1400,7 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: }, ), turn_context.config.codex_home.to_path_buf(), + session.services.mcp_manager.codex_apps_tools_cache(), codex_mcp::codex_apps_tools_cache_key(auth.as_ref()), turn_context.config.prefix_mcp_tool_names(), rmcp::model::ElicitationCapability::default(), diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index c060ef090..249a6994b 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -311,6 +311,7 @@ impl Session { turn_context.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), + self.services.mcp_manager.codex_apps_tools_cache(), codex_apps_tools_cache_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 96dffed88..57a66b998 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1193,6 +1193,7 @@ impl Session { session_configuration.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), + sess.services.mcp_manager.codex_apps_tools_cache(), codex_apps_tools_cache_key(auth), config.prefix_mcp_tool_names(), client_elicitation_capability,