From 2f7d89b1419bf7064346855b0acde23514b1ebc5 Mon Sep 17 00:00:00 2001 From: Matthew Zeng Date: Mon, 13 Jul 2026 02:17:14 +0000 Subject: [PATCH] Extract connector runtime snapshot management (#32698) ## What changed - Move the Codex Apps tool cache into a reusable `codex-connectors` runtime manager keyed by account and workspace. - Represent cached tools as atomically published snapshots with refresh timestamps while preserving the newest accepted fetch generation. - Harden disk persistence with bounded reads, atomic file replacement, and serialized writes so an older fetch cannot overwrite newer state. ## Testing Add coverage for identity isolation, snapshot timestamps, oversized cache files, atomic replacement, and concurrent persistence ordering. GitOrigin-RevId: 5ea2234469daae3abf54b030244c3251de62ca5a --- codex-rs/Cargo.lock | 4 + .../src/request_processors/mcp_processor.rs | 6 +- codex-rs/codex-mcp/src/codex_apps_cache.rs | 378 ------------------ codex-rs/codex-mcp/src/connection_manager.rs | 21 +- .../codex-mcp/src/connection_manager_tests.rs | 79 ++-- codex-rs/codex-mcp/src/lib.rs | 11 +- codex-rs/codex-mcp/src/mcp/mod.rs | 13 +- codex-rs/codex-mcp/src/rmcp_client.rs | 27 +- codex-rs/codex-mcp/src/tools.rs | 3 - codex-rs/connectors/Cargo.toml | 6 +- .../connectors/src/connector_runtime/mod.rs | 335 ++++++++++++++++ .../src/connector_runtime/persistence.rs | 268 +++++++++++++ .../src/connector_runtime/tests.rs} | 359 ++++++++++++++--- codex-rs/connectors/src/lib.rs | 9 + codex-rs/core/src/connectors.rs | 4 +- codex-rs/core/src/mcp.rs | 9 +- codex-rs/core/src/mcp_tool_call_tests.rs | 2 +- codex-rs/core/src/session/mcp.rs | 2 +- codex-rs/core/src/session/mod.rs | 2 +- codex-rs/core/src/session/session.rs | 2 +- 20 files changed, 1028 insertions(+), 512 deletions(-) delete mode 100644 codex-rs/codex-mcp/src/codex_apps_cache.rs create mode 100644 codex-rs/connectors/src/connector_runtime/mod.rs create mode 100644 codex-rs/connectors/src/connector_runtime/persistence.rs rename codex-rs/{codex-mcp/src/codex_apps_cache_tests.rs => connectors/src/connector_runtime/tests.rs} (54%) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 0c689b10a3..6ea0242251 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2588,8 +2588,12 @@ name = "codex-connectors" version = "0.0.0" dependencies = [ "anyhow", + "arc-swap", "codex-config", + "codex-login", + "codex-otel", "codex-plugin", + "codex-protocol", "indexmap 2.14.0", "pretty_assertions", "serde", 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 ac2125c2cf..34ec3887ce 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -1,4 +1,6 @@ use super::*; +use codex_connectors::ConnectorRuntimeManager; +use codex_mcp::ToolInfo; const MCP_TOOL_THREAD_ID_META_KEY: &str = "threadId"; @@ -284,7 +286,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, - codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache, + codex_apps_tools_cache: ConnectorRuntimeManager, ) { let result = Self::list_mcp_server_status_response( request_id.request_id.to_string(), @@ -304,7 +306,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, - codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache, + codex_apps_tools_cache: ConnectorRuntimeManager, ) -> Result { let detail = match params.detail.unwrap_or(McpServerStatusDetail::Full) { McpServerStatusDetail::Full => McpSnapshotDetail::Full, diff --git a/codex-rs/codex-mcp/src/codex_apps_cache.rs b/codex-rs/codex-mcp/src/codex_apps_cache.rs deleted file mode 100644 index ea39444628..0000000000 --- a/codex-rs/codex-mcp/src/codex_apps_cache.rs +++ /dev/null @@ -1,378 +0,0 @@ -//! 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/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 2e6c518d11..ef5cb45599 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -15,9 +15,6 @@ use std::time::Instant; use crate::McpAuthStatusEntry; use crate::codex_apps::prepare_openai_file_params_for_model; -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::ElicitationRequestRouter; use crate::elicitation::ElicitationReviewerHandle; @@ -47,6 +44,9 @@ use codex_config::McpServerAuth; use codex_config::McpServerTransportConfig; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; +use codex_connectors::ConnectorRuntimeContextKey; +use codex_connectors::ConnectorRuntimeFetchSource; +use codex_connectors::ConnectorRuntimeManager; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_protocol::mcp::CallToolResult; @@ -135,8 +135,8 @@ impl McpConnectionManager { initial_permission_profile: PermissionProfile, runtime_context: McpRuntimeContext, codex_home: PathBuf, - codex_apps_tools_cache: CodexAppsToolsCache, - codex_apps_tools_cache_key: CodexAppsToolsCacheKey, + codex_apps_tools_cache: ConnectorRuntimeManager, + codex_apps_tools_cache_key: ConnectorRuntimeContextKey, prefix_mcp_tool_names: bool, client_elicitation_capability: ElicitationCapability, supports_openai_form_elicitation: bool, @@ -564,10 +564,13 @@ impl McpConnectionManager { .context("failed to get client")?; let list_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 fetch_ticket = + managed_client + .codex_apps_tools_cache_context + .as_ref() + .map(|cache_context| { + cache_context.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh) + }); let tools = list_tools_for_client_uncached( CODEX_APPS_MCP_SERVER_NAME, /*is_codex_apps_mcp_server*/ true, diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index faaa198b68..ff11ef596f 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1,6 +1,4 @@ use super::*; -use crate::codex_apps_cache::CodexAppsToolsCache; -use crate::codex_apps_cache::CodexAppsToolsCacheContext; use crate::elicitation::ElicitationLifecycle; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationRequestRouter; @@ -24,6 +22,10 @@ use codex_config::McpServerConfig; use codex_config::McpServerToolConfig; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; +use codex_connectors::ConnectorRuntimeContext; +use codex_connectors::ConnectorRuntimeContextKey; +use codex_connectors::ConnectorRuntimeFetchSource; +use codex_connectors::ConnectorRuntimeManager; use codex_exec_server::EnvironmentManager; use codex_protocol::ToolName; use codex_protocol::mcp::McpServerInfo; @@ -72,17 +74,24 @@ fn create_codex_apps_tools_cache_context( codex_home: PathBuf, account_id: Option<&str>, chatgpt_user_id: Option<&str>, -) -> CodexAppsToolsCacheContext { - CodexAppsToolsCache::default().context( +) -> ConnectorRuntimeContext { + ConnectorRuntimeManager::::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, - }, + ConnectorRuntimeContextKey::personal( + account_id.map(ToOwned::to_owned), + chatgpt_user_id.map(ToOwned::to_owned), + ), ) } +fn store_current_tools(cache_context: &ConnectorRuntimeContext, tools: Vec) { + let _ = cache_context.publish_if_newest_accepted( + cache_context.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh), + &create_test_server_info("Codex Apps"), + tools, + ); +} + fn create_test_server_info(title: &str) -> McpServerInfo { McpServerInfo { name: "codex-apps".to_string(), @@ -157,7 +166,7 @@ fn create_test_manager_with_failed_apps_startup( Some("reconnect-test-account"), Some("reconnect-test-user"), ); - cache_context.store_current_tools_for_test(cached_tools); + store_current_tools(&cache_context, cached_tools); let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionManager::new_uninitialized( @@ -689,10 +698,13 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { 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", - )]); + store_current_tools( + &cache_context, + vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_create_event", + )], + ); let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -864,6 +876,22 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { assert!(timeout_result.is_err()); } +#[tokio::test] +async fn cancelling_startup_does_not_disable_a_ready_client() { + let client = create_ready_async_managed_client(vec![create_test_tool("ready", "search")]).await; + + client.cancel_token.cancel(); + + let managed = client + .client() + .await + .expect("startup cancellation should not disable a ready client"); + assert_eq!( + model_tool_names(&managed.tools), + HashSet::from([ToolName::namespaced("ready", "search")]) + ); +} + #[tokio::test] async fn shutdown_cancels_pending_tool_listing() { let cancel_token = CancellationToken::new(); @@ -972,7 +1000,7 @@ async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() { Some("account-one"), Some("user-one"), ); - cache_context.store_current_tools_for_test(Vec::new()); + store_current_tools(&cache_context, Vec::new()); let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -1012,10 +1040,13 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() 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", - )]); + store_current_tools( + &cache_context, + 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 { @@ -1472,12 +1503,10 @@ 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, - is_workspace_account: false, - }, + ConnectorRuntimeManager::::default(), + ConnectorRuntimeContextKey::personal( + /*account_id*/ None, /*chatgpt_user_id*/ None, + ), /*prefix_mcp_tool_names*/ true, ElicitationCapability::default(), /*supports_openai_form_elicitation*/ false, diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 8451f87c25..079444070a 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -14,6 +14,11 @@ pub use runtime::McpRuntimeContext; pub use runtime::SandboxState; pub use tools::ToolInfo; +/// Backward-compatible name for the shared Codex Apps tools runtime. +pub type CodexAppsToolsCache = codex_connectors::ConnectorRuntimeManager; +/// Backward-compatible name for the Codex Apps runtime context key. +pub type CodexAppsToolsCacheKey = codex_connectors::ConnectorRuntimeContextKey; + pub use catalog::McpCatalogBuilder; pub use catalog::McpPluginAttribution; pub use catalog::McpServerConflict; @@ -37,9 +42,8 @@ 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_cache::CodexAppsToolsCache; -pub use codex_apps_cache::CodexAppsToolsCacheKey; -pub use codex_apps_cache::codex_apps_tools_cache_key; +/// Backward-compatible name for the Codex Apps runtime context key builder. +pub use codex_connectors::connector_runtime_context_key as 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; @@ -78,7 +82,6 @@ pub use mcp::qualified_mcp_tool_name_prefix; 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 8a1847bc96..ad2245e029 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -27,7 +27,9 @@ use codex_config::McpServerTransportConfig; use codex_config::types::AppToolApproval; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; +use codex_connectors::ConnectorRuntimeManager; use codex_connectors::ConnectorSnapshot; +use codex_connectors::connector_runtime_context_key; use codex_login::CodexAuth; use codex_model_provider::CHATGPT_CODEX_BASE_URL; use codex_protocol::mcp::McpServerInfo; @@ -44,11 +46,10 @@ use serde_json::Value; use tokio_util::sync::CancellationToken; use crate::ResolvedMcpCatalog; -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; +use crate::tools::ToolInfo; pub const CODEX_APPS_MCP_SERVER_NAME: &str = "codex_apps"; const DEFAULT_CODEX_APPS_MCP_PRODUCT_SKU: &str = "codex"; @@ -303,7 +304,7 @@ pub async fn read_mcp_resource( config: &McpConfig, auth: Option<&CodexAuth>, runtime_context: McpRuntimeContext, - codex_apps_tools_cache: CodexAppsToolsCache, + codex_apps_tools_cache: ConnectorRuntimeManager, server: &str, uri: &str, ) -> anyhow::Result { @@ -333,7 +334,7 @@ pub async fn read_mcp_resource( runtime_context, config.codex_home.clone(), codex_apps_tools_cache, - codex_apps_tools_cache_key(auth), + connector_runtime_context_key(auth), config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), /*supports_openai_form_elicitation*/ false, @@ -368,7 +369,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, + codex_apps_tools_cache: ConnectorRuntimeManager, detail: McpSnapshotDetail, ) -> McpServerStatusSnapshot { let mcp_servers = effective_mcp_servers(config, auth); @@ -412,7 +413,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( runtime_context, config.codex_home.clone(), codex_apps_tools_cache, - codex_apps_tools_cache_key(auth), + connector_runtime_context_key(auth), config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), /*supports_openai_form_elicitation*/ false, diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index f9210aea3e..52c4a5d462 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -22,9 +22,6 @@ 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::prepare_openai_file_params_for_model; -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; @@ -45,6 +42,8 @@ use codex_config::McpServerConfig; use codex_config::McpServerTransportConfig; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; +use codex_connectors::ConnectorRuntimeContext; +use codex_connectors::ConnectorRuntimeFetchSource; use codex_exec_server::HttpClient; use codex_exec_server::ReqwestHttpClient; use codex_protocol::mcp::McpServerInfo; @@ -106,7 +105,7 @@ pub(crate) struct ManagedClient { pub(crate) tool_timeout: Option, pub(crate) server_instructions: Option, pub(crate) server_supports_sandbox_state_meta_capability: bool, - pub(crate) codex_apps_tools_cache_context: Option, + pub(crate) codex_apps_tools_cache_context: Option>, } impl ManagedClient { @@ -115,7 +114,7 @@ impl ManagedClient { if let Some(tools) = self .codex_apps_tools_cache_context .as_ref() - .and_then(CodexAppsToolsCacheContext::current_tools) + .and_then(ConnectorRuntimeContext::current_tools) { emit_duration( MCP_TOOLS_LIST_DURATION_METRIC, @@ -272,7 +271,7 @@ struct ManagedClientStartup { keyring_backend_kind: AuthKeyringBackendKind, tx_event: Sender, elicitation_requests: ElicitationRequestManager, - codex_apps_tools_cache_context: Option, + codex_apps_tools_cache_context: Option>, runtime_context: McpRuntimeContext, runtime_auth_provider: Option, client_elicitation_capability: ElicitationCapability, @@ -375,7 +374,7 @@ pub(crate) struct AsyncManagedClient { pub(crate) client: ManagedClientFuture, pub(crate) is_codex_apps_mcp_server: bool, pub(crate) cached_server_info: Option, - pub(crate) codex_apps_tools_cache_context: Option, + pub(crate) codex_apps_tools_cache_context: Option>, pub(crate) tool_filter: ToolFilter, pub(crate) startup_complete: Arc, pub(crate) startup_reconnect: Option>, @@ -397,7 +396,7 @@ impl AsyncManagedClient { cancel_token: CancellationToken, tx_event: Sender, elicitation_requests: ElicitationRequestManager, - codex_apps_tools_cache_context: Option, + codex_apps_tools_cache_context: Option>, tool_plugin_provenance: Arc, runtime_context: McpRuntimeContext, runtime_auth_provider: Option, @@ -414,7 +413,7 @@ impl AsyncManagedClient { 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) + .and_then(ConnectorRuntimeContext::cached_server_info) } else { None }; @@ -448,7 +447,7 @@ impl AsyncManagedClient { }); if codex_apps_tools_cache_context .as_ref() - .is_some_and(CodexAppsToolsCacheContext::has_current_tools) + .is_some_and(ConnectorRuntimeContext::has_current_tools) { let startup_task = client.clone(); tokio::spawn(async move { @@ -506,13 +505,13 @@ 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) + .is_some_and(ConnectorRuntimeContext::has_current_tools) } fn cached_tools(&self) -> Option> { self.codex_apps_tools_cache_context .as_ref() - .and_then(CodexAppsToolsCacheContext::current_tools) + .and_then(ConnectorRuntimeContext::current_tools) .map(|tools| filter_tools(tools, &self.tool_filter)) } @@ -850,7 +849,7 @@ async fn start_server_task( let list_start = Instant::now(); let fetch_ticket = codex_apps_tools_cache_context .as_ref() - .map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::Startup)); + .map(|cache_context| cache_context.begin_fetch(ConnectorRuntimeFetchSource::Startup)); let tools = list_tools_for_client_uncached( &server_name, is_codex_apps_mcp_server, @@ -934,7 +933,7 @@ struct StartServerTaskParams { tool_filter: ToolFilter, tx_event: Sender, elicitation_requests: ElicitationRequestManager, - codex_apps_tools_cache_context: Option, + codex_apps_tools_cache_context: Option>, client_elicitation_capability: ElicitationCapability, supports_openai_form_elicitation: bool, } diff --git a/codex-rs/codex-mcp/src/tools.rs b/codex-rs/codex-mcp/src/tools.rs index 8f79741edd..1fb0c28afc 100644 --- a/codex-rs/codex-mcp/src/tools.rs +++ b/codex-rs/codex-mcp/src/tools.rs @@ -19,9 +19,6 @@ use tracing::warn; use crate::mcp::sanitize_responses_api_tool_name; -pub(crate) const MCP_TOOLS_CACHE_WRITE_DURATION_METRIC: &str = - "codex.mcp.tools.cache_write.duration_ms"; - const LEGACY_MCP_TOOL_NAME_PREFIX: &str = "mcp__"; #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/codex-rs/connectors/Cargo.toml b/codex-rs/connectors/Cargo.toml index 6890408ea3..eb2306495b 100644 --- a/codex-rs/connectors/Cargo.toml +++ b/codex-rs/connectors/Cargo.toml @@ -9,18 +9,22 @@ workspace = true [dependencies] anyhow = { workspace = true } +arc-swap = { workspace = true } codex-config = { workspace = true } +codex-login = { workspace = true } +codex-otel = { workspace = true } codex-plugin = { workspace = true } +codex-protocol = { workspace = true } indexmap = { workspace = true, features = ["serde"] } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } sha1 = { workspace = true } +tempfile = { workspace = true } tracing = { workspace = true } urlencoding = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } -tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } [lib] diff --git a/codex-rs/connectors/src/connector_runtime/mod.rs b/codex-rs/connectors/src/connector_runtime/mod.rs new file mode 100644 index 0000000000..be2949e9f7 --- /dev/null +++ b/codex-rs/connectors/src/connector_runtime/mod.rs @@ -0,0 +1,335 @@ +//! Shared runtime snapshot for connector-backed MCP tools. +//! +//! Runtime snapshots are process-local live state scoped by account and +//! workspace. Disk is best-effort cold-start persistence; a context reads it +//! once when created and never rereads it. Full connector metadata is +//! owned by the connector metadata store, not by this module. + +use std::collections::HashMap; +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::Duration; +use std::time::Instant; +use std::time::SystemTime; + +use arc_swap::ArcSwapOption; +use codex_login::CodexAuth; +use codex_protocol::mcp::McpServerInfo; +use serde::Deserialize; +use serde::Serialize; +use serde::de::DeserializeOwned; + +use self::persistence::load_cached_codex_apps_server_info; +use self::persistence::load_cached_connector_runtime_for_identity; +use self::persistence::persist_codex_apps_cache; +use self::persistence::server_info_cache_path; +use self::persistence::tools_cache_path; + +const MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC: &str = "codex.mcp.tools.cache_publish.duration_ms"; + +/// Values stored in the connector runtime's persisted tool snapshot. +/// +/// The runtime uses the connector-owned Codex Apps cache layout for every +/// serializable, cloneable payload. +pub trait ConnectorRuntimePayload: Clone + Serialize + DeserializeOwned {} + +impl ConnectorRuntimePayload for T where T: Clone + Serialize + DeserializeOwned {} + +/// The account and workspace identity of a connector runtime catalog. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct ConnectorRuntimeContextKey { + account_id: Option, + chatgpt_user_id: Option, + is_workspace_account: bool, +} + +impl ConnectorRuntimeContextKey { + pub fn personal(account_id: Option, chatgpt_user_id: Option) -> Self { + Self { + account_id, + chatgpt_user_id, + is_workspace_account: false, + } + } + + pub fn workspace(account_id: Option, chatgpt_user_id: Option) -> Self { + Self { + account_id, + chatgpt_user_id, + is_workspace_account: true, + } + } +} + +/// Builds the connector runtime context key for the active Codex auth. +pub fn connector_runtime_context_key(auth: Option<&CodexAuth>) -> ConnectorRuntimeContextKey { + let account_id = auth.and_then(CodexAuth::get_account_id); + let chatgpt_user_id = auth.and_then(CodexAuth::get_chatgpt_user_id); + if auth.is_some_and(CodexAuth::is_workspace_account) { + ConnectorRuntimeContextKey::workspace(account_id, chatgpt_user_id) + } else { + ConnectorRuntimeContextKey::personal(account_id, chatgpt_user_id) + } +} + +/// One atomically published connector runtime state. +/// +/// Tools remain raw and in response order. Local and managed configuration is +/// intentionally applied by readers rather than persisted in this snapshot. +#[derive(Debug, Clone)] +pub struct ConnectorRuntimeSnapshot { + tools: Vec, + refreshed_at: SystemTime, +} + +impl ConnectorRuntimeSnapshot { + pub fn tools(&self) -> &[T] { + &self.tools + } + + pub fn refreshed_at(&self) -> SystemTime { + self.refreshed_at + } + + pub fn age(&self) -> Duration { + SystemTime::now() + .duration_since(self.refreshed_at) + .unwrap_or_default() + } +} + +/// Process-scoped registry of connector runtime state by account and workspace. +/// +/// Contexts with the same identity share one live entry. Different identities +/// remain independently available for clients that already hold their context. +pub struct ConnectorRuntimeManager { + entries: Arc>>>>, +} + +impl Clone for ConnectorRuntimeManager { + fn clone(&self) -> Self { + Self { + entries: Arc::clone(&self.entries), + } + } +} + +impl Default for ConnectorRuntimeManager { + fn default() -> Self { + Self { + entries: Arc::new(Mutex::new(HashMap::new())), + } + } +} + +impl ConnectorRuntimeManager { + pub fn current_snapshot( + &self, + codex_home: PathBuf, + key: ConnectorRuntimeContextKey, + ) -> Option>> { + self.context(codex_home, key).current_snapshot() + } + + pub fn context( + &self, + codex_home: PathBuf, + key: ConnectorRuntimeContextKey, + ) -> ConnectorRuntimeContext { + let identity = ConnectorRuntimeIdentity { codex_home, key }; + let mut entries = lock_unpoisoned(&self.entries); + let entry = entries + .entry(identity.clone()) + .or_insert_with(|| Arc::new(ConnectorRuntimeEntry::new(identity))) + .clone(); + ConnectorRuntimeContext { entry } + } +} + +/// Handle to one shared account/workspace connector runtime. +pub struct ConnectorRuntimeContext { + entry: Arc>, +} + +impl Clone for ConnectorRuntimeContext { + fn clone(&self) -> Self { + Self { + entry: Arc::clone(&self.entry), + } + } +} + +impl ConnectorRuntimeContext { + pub fn current_snapshot(&self) -> Option>> { + self.entry.current_snapshot.load_full() + } + + pub fn has_current_tools(&self) -> bool { + self.current_snapshot().is_some() + } + + pub fn begin_fetch(&self, source: ConnectorRuntimeFetchSource) -> ConnectorRuntimeFetchTicket { + ConnectorRuntimeFetchTicket { + generation: self + .entry + .next_fetch_generation + .fetch_add(1, Ordering::Relaxed) + + 1, + source, + } + } + + pub fn cached_server_info(&self) -> Option { + load_cached_codex_apps_server_info(self) + } + + fn tools_cache_path(&self) -> PathBuf { + tools_cache_path(&self.entry.identity) + } + + fn server_info_cache_path(&self) -> PathBuf { + server_info_cache_path(&self.entry.identity) + } + + pub fn current_tools(&self) -> Option> { + self.current_snapshot() + .map(|snapshot| snapshot.tools.clone()) + } + + pub fn publish_runtime_if_newest_accepted( + &self, + ticket: ConnectorRuntimeFetchTicket, + server_info: &McpServerInfo, + tools: Vec, + ) -> Arc> { + self.publish_runtime_if_newest_accepted_with( + ticket, + server_info, + tools, + persist_codex_apps_cache, + ) + } + + fn publish_runtime_if_newest_accepted_with( + &self, + ticket: ConnectorRuntimeFetchTicket, + server_info: &McpServerInfo, + tools: Vec, + persist: impl FnOnce(&ConnectorRuntimeContext, &McpServerInfo, &ConnectorRuntimeSnapshot), + ) -> Arc> { + let publish_start = Instant::now(); + let mut last_accepted_generation = lock_unpoisoned(&self.entry.last_accepted_generation); + if ticket.generation <= *last_accepted_generation + && let Some(snapshot) = self.current_snapshot() + { + drop(last_accepted_generation); + emit_duration( + MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC, + publish_start.elapsed(), + &[("source", ticket.source.as_str()), ("result", "stale")], + ); + return snapshot; + } + + let snapshot = Arc::new(ConnectorRuntimeSnapshot { + tools, + refreshed_at: SystemTime::now(), + }); + + *last_accepted_generation = ticket.generation; + self.entry + .current_snapshot + .store(Some(Arc::clone(&snapshot))); + // Keep the generation guard through persistence so accepted generations cannot reach disk + // out of order. + persist(self, server_info, snapshot.as_ref()); + drop(last_accepted_generation); + emit_duration( + MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC, + publish_start.elapsed(), + &[("source", ticket.source.as_str()), ("result", "published")], + ); + snapshot + } + + pub fn publish_if_newest_accepted( + &self, + ticket: ConnectorRuntimeFetchTicket, + server_info: &McpServerInfo, + tools: Vec, + ) -> Vec { + self.publish_runtime_if_newest_accepted(ticket, server_info, tools) + .tools + .clone() + } +} + +#[derive(Debug, Clone, Copy)] +pub enum ConnectorRuntimeFetchSource { + Startup, + HardRefresh, +} + +impl ConnectorRuntimeFetchSource { + fn as_str(self) -> &'static str { + match self { + Self::Startup => "startup", + Self::HardRefresh => "hard_refresh", + } + } +} + +pub struct ConnectorRuntimeFetchTicket { + generation: u64, + source: ConnectorRuntimeFetchSource, +} + +/// All live state owned by one connector identity. +struct ConnectorRuntimeEntry { + identity: ConnectorRuntimeIdentity, + current_snapshot: ArcSwapOption>, + next_fetch_generation: AtomicU64, + last_accepted_generation: Mutex, +} + +impl ConnectorRuntimeEntry { + fn new(identity: ConnectorRuntimeIdentity) -> Self { + let current_snapshot = load_cached_connector_runtime_for_identity(&identity).map(Arc::new); + Self { + identity, + current_snapshot: ArcSwapOption::from(current_snapshot), + next_fetch_generation: AtomicU64::new(0), + last_accepted_generation: Mutex::new(0), + } + } +} + +/// Everything that decides whether two connector runtime clients can share a snapshot. +/// +/// The auth key says whose runtime catalog we are reading. `codex_home` keeps +/// the persisted cache under the right home directory. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct ConnectorRuntimeIdentity { + codex_home: PathBuf, + key: ConnectorRuntimeContextKey, +} + +fn emit_duration(metric: &str, duration: Duration, tags: &[(&str, &str)]) { + if let Some(metrics) = codex_otel::global() { + let _ = metrics.record_duration(metric, duration, tags); + } +} + +fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +mod persistence; + +#[cfg(test)] +mod tests; diff --git a/codex-rs/connectors/src/connector_runtime/persistence.rs b/codex-rs/connectors/src/connector_runtime/persistence.rs new file mode 100644 index 0000000000..349577b310 --- /dev/null +++ b/codex-rs/connectors/src/connector_runtime/persistence.rs @@ -0,0 +1,268 @@ +//! Bounded, atomic persistence for connector runtime snapshots. + +use std::fs::File; +use std::io::Read; +use std::io::Write; +use std::path::Path; +use std::path::PathBuf; +#[cfg(test)] +use std::sync::Arc; +use std::time::Instant; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; + +use anyhow::Context; +use anyhow::anyhow; +use codex_protocol::mcp::McpServerInfo; +use serde::Deserialize; +use serde::Serialize; +use sha1::Digest; +use sha1::Sha1; +use tempfile::NamedTempFile; +use tracing::instrument; + +use super::ConnectorRuntimeContext; +use super::ConnectorRuntimeIdentity; +use super::ConnectorRuntimePayload; +use super::ConnectorRuntimeSnapshot; +use super::emit_duration; + +const MCP_TOOLS_CACHE_WRITE_DURATION_METRIC: &str = "codex.mcp.tools.cache_write.duration_ms"; +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; +pub(crate) const CODEX_APPS_TOOLS_CACHE_MAX_BYTES: u64 = 32 * 1024 * 1024; + +pub(crate) fn tools_cache_path(identity: &ConnectorRuntimeIdentity) -> PathBuf { + cache_path_in(identity, CODEX_APPS_TOOLS_CACHE_DIR) +} + +pub(crate) fn server_info_cache_path(identity: &ConnectorRuntimeIdentity) -> PathBuf { + cache_path_in(identity, CODEX_APPS_SERVER_INFO_CACHE_DIR) +} + +fn cache_path_in(identity: &ConnectorRuntimeIdentity, 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. + let identity_json = serde_json::to_string(&identity.key).unwrap_or_default(); + let identity_hash = sha1_hex(&identity_json); + identity + .codex_home + .join(cache_dir) + .join(format!("{identity_hash}.json")) +} + +#[instrument(level = "trace", skip_all)] +pub(crate) fn load_cached_connector_runtime_for_identity( + identity: &ConnectorRuntimeIdentity, +) -> Option> { + let cache_path = tools_cache_path(identity); + let (bytes, modified_at) = read_bounded_cache_file(&cache_path).ok()?; + let cache: CodexAppsToolsDiskCache = serde_json::from_slice(&bytes).ok()?; + (cache.schema_version == CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION).then_some( + ConnectorRuntimeSnapshot { + tools: cache.tools, + refreshed_at: modified_at, + }, + ) +} + +pub(crate) fn write_cached_connector_runtime( + cache_context: &ConnectorRuntimeContext, + snapshot: &ConnectorRuntimeSnapshot, +) -> anyhow::Result<()> +where + T: ConnectorRuntimePayload, +{ + 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: snapshot.tools.clone(), + }) + .context("failed to serialize connector runtime cache")?; + write_codex_apps_cache_file(&cache_path, "runtime", bytes) +} + +#[instrument(level = "trace", skip_all)] +pub(crate) fn load_cached_codex_apps_server_info( + cache_context: &ConnectorRuntimeContext, +) -> Option { + let (bytes, _) = read_bounded_cache_file(&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: &ConnectorRuntimeContext, + 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) +} + +pub(crate) fn persist_codex_apps_cache( + cache_context: &ConnectorRuntimeContext, + server_info: &McpServerInfo, + snapshot: &ConnectorRuntimeSnapshot, +) where + T: ConnectorRuntimePayload, +{ + let cache_write_start = Instant::now(); + let tools_result = write_cached_connector_runtime(cache_context, snapshot); + if let Err(err) = &tools_result { + tracing::warn!("failed to write connector runtime 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)], + ); +} + +fn read_bounded_cache_file(cache_path: &Path) -> anyhow::Result<(Vec, SystemTime)> { + let mut file = File::open(cache_path) + .with_context(|| format!("failed to open cache `{}`", cache_path.display()))?; + let metadata = file + .metadata() + .with_context(|| format!("failed to stat cache `{}`", cache_path.display()))?; + if metadata.len() > CODEX_APPS_TOOLS_CACHE_MAX_BYTES { + return Err(anyhow!( + "cache `{}` is {} bytes, exceeding the {} byte limit", + cache_path.display(), + metadata.len(), + CODEX_APPS_TOOLS_CACHE_MAX_BYTES + )); + } + let mut bytes = Vec::with_capacity(metadata.len() as usize); + std::io::Read::by_ref(&mut file) + .take(CODEX_APPS_TOOLS_CACHE_MAX_BYTES + 1) + .read_to_end(&mut bytes) + .with_context(|| format!("failed to read cache `{}`", cache_path.display()))?; + if bytes.len() as u64 > CODEX_APPS_TOOLS_CACHE_MAX_BYTES { + return Err(anyhow!( + "cache `{}` grew beyond the {} byte limit while reading", + cache_path.display(), + CODEX_APPS_TOOLS_CACHE_MAX_BYTES + )); + } + Ok((bytes, metadata.modified().unwrap_or(UNIX_EPOCH))) +} + +fn write_codex_apps_cache_file( + cache_path: &Path, + cache_name: &str, + bytes: Vec, +) -> anyhow::Result<()> { + let parent = cache_path.parent().ok_or_else(|| { + anyhow!( + "Codex Apps {cache_name} cache path `{}` has no parent", + cache_path.display() + ) + })?; + std::fs::create_dir_all(parent).with_context(|| { + format!( + "failed to create Codex Apps {cache_name} cache directory `{}`", + parent.display() + ) + })?; + let mut temporary = NamedTempFile::new_in(parent).with_context(|| { + format!( + "failed to create temporary Codex Apps {cache_name} cache in `{}`", + parent.display() + ) + })?; + temporary.write_all(&bytes).with_context(|| { + format!( + "failed to write temporary Codex Apps {cache_name} cache for `{}`", + cache_path.display() + ) + })?; + temporary.persist(cache_path).map_err(|error| { + anyhow!( + "failed to atomically replace Codex Apps {cache_name} cache `{}`: {}", + cache_path.display(), + error.error + ) + })?; + 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, +} + +fn sha1_hex(s: &str) -> String { + let mut hasher = Sha1::new(); + hasher.update(s.as_bytes()); + let sha1 = hasher.finalize(); + format!("{sha1:x}") +} + +#[cfg(test)] +pub(crate) fn write_cached_codex_apps_tools_for_test( + cache_context: &ConnectorRuntimeContext, + server_info: &McpServerInfo, + tools: &[T], +) where + T: ConnectorRuntimePayload, +{ + let snapshot = ConnectorRuntimeSnapshot { + tools: tools.to_vec(), + refreshed_at: SystemTime::now(), + }; + cache_context + .entry + .current_snapshot + .store(Some(Arc::new(snapshot.clone()))); + persist_codex_apps_cache(cache_context, server_info, &snapshot); +} + +#[cfg(test)] +pub(crate) fn read_cached_codex_apps_tools( + cache_context: &ConnectorRuntimeContext, +) -> Option> +where + T: ConnectorRuntimePayload, +{ + load_cached_connector_runtime_for_identity(&cache_context.entry.identity) + .map(|snapshot| snapshot.tools) +} + +#[cfg(test)] +pub(crate) fn write_cached_codex_apps_tools( + cache_context: &ConnectorRuntimeContext, + tools: &[T], +) -> anyhow::Result<()> +where + T: ConnectorRuntimePayload, +{ + let snapshot = ConnectorRuntimeSnapshot { + tools: tools.to_vec(), + refreshed_at: SystemTime::now(), + }; + write_cached_connector_runtime(cache_context, &snapshot) +} diff --git a/codex-rs/codex-mcp/src/codex_apps_cache_tests.rs b/codex-rs/connectors/src/connector_runtime/tests.rs similarity index 54% rename from codex-rs/codex-mcp/src/codex_apps_cache_tests.rs rename to codex-rs/connectors/src/connector_runtime/tests.rs index 0d1e2ffdc5..fb88695718 100644 --- a/codex-rs/codex-mcp/src/codex_apps_cache_tests.rs +++ b/codex-rs/connectors/src/connector_runtime/tests.rs @@ -1,35 +1,35 @@ +use super::persistence::CODEX_APPS_TOOLS_CACHE_MAX_BYTES; +use super::persistence::CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION; +use super::persistence::read_cached_codex_apps_tools; +use super::persistence::write_cached_codex_apps_tools; +use super::persistence::write_cached_codex_apps_tools_for_test; 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; +use serde::Deserialize; +use serde::Serialize; #[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 { +const CODEX_APPS_MCP_SERVER_NAME: &str = "codex_apps"; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct TestTool { + server_name: String, + callable_name: String, + connector_id: Option, + connector_name: Option, +} + +fn create_test_tool(server_name: &str, tool_name: &str) -> TestTool { + TestTool { 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()), - ), - openai_file_input_optional_fields: Default::default(), connector_id: None, connector_name: None, - plugin_display_names: Vec::new(), } } @@ -38,7 +38,7 @@ fn create_test_tool_with_connector( tool_name: &str, connector_id: &str, connector_name: Option<&str>, -) -> ToolInfo { +) -> TestTool { 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); @@ -49,10 +49,10 @@ fn create_codex_apps_tools_cache_context( codex_home: PathBuf, account_id: Option<&str>, chatgpt_user_id: Option<&str>, -) -> CodexAppsToolsCacheContext { - CodexAppsToolsCache::default().context( +) -> ConnectorRuntimeContext { + ConnectorRuntimeManager::::default().context( codex_home, - CodexAppsToolsCacheKey { + ConnectorRuntimeContextKey { account_id: account_id.map(ToOwned::to_owned), chatgpt_user_id: chatgpt_user_id.map(ToOwned::to_owned), is_workspace_account: false, @@ -71,13 +71,6 @@ fn create_test_server_info(title: &str) -> McpServerInfo { } } -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"); @@ -238,7 +231,7 @@ fn startup_cached_codex_apps_tools_loads_from_disk_cache() { 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); + let cached_server_info = cache_context.cached_server_info(); assert_eq!(startup_tools.len(), 1); assert_eq!(startup_tools[0].server_name, CODEX_APPS_MCP_SERVER_NAME); @@ -273,7 +266,7 @@ fn startup_cached_codex_apps_tools_loads_without_server_info_cache() { 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); + let cached_server_info = cache_context.cached_server_info(); assert_eq!(startup_tools.len(), 1); assert_eq!(startup_tools[0].callable_name, "calendar_search"); @@ -315,7 +308,7 @@ fn codex_apps_server_info_cache_survives_legacy_tools_cache_write() { ); assert_eq!( - load_startup_cached_codex_apps_server_info(&startup_cache_context), + startup_cache_context.cached_server_info(), Some(server_info) ); assert!(startup_cache_context.current_tools().is_none()); @@ -355,10 +348,10 @@ fn codex_apps_tools_cache_context_does_not_reread_disk_after_creation() { #[test] fn codex_apps_tools_cache_publishes_newest_shared_snapshot() { let codex_home = tempdir().expect("tempdir"); - let cache = CodexAppsToolsCache::default(); + let cache = ConnectorRuntimeManager::::default(); let cache_context_1 = cache.context( codex_home.path().to_path_buf(), - CodexAppsToolsCacheKey { + ConnectorRuntimeContextKey { account_id: Some("account-one".to_string()), chatgpt_user_id: Some("user-one".to_string()), is_workspace_account: false, @@ -366,28 +359,21 @@ fn codex_apps_tools_cache_publishes_newest_shared_snapshot() { ); let cache_context_2 = cache.context( codex_home.path().to_path_buf(), - CodexAppsToolsCacheKey { + ConnectorRuntimeContextKey { 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 older_ticket = cache_context_1.begin_fetch(ConnectorRuntimeFetchSource::Startup); + let newer_ticket = cache_context_2.begin_fetch(ConnectorRuntimeFetchSource::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") - ) - ); + assert_eq!(cache_context_1.current_tools(), Some(published_tools)); let current_tools = cache_context_1.publish_if_newest_accepted(older_ticket, &server_info, older_tools); @@ -408,9 +394,9 @@ 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( + let cache_context = ConnectorRuntimeManager::::default().context( codex_home_file, - CodexAppsToolsCacheKey { + ConnectorRuntimeContextKey { account_id: Some("account-one".to_string()), chatgpt_user_id: Some("user-one".to_string()), is_workspace_account: false, @@ -418,16 +404,13 @@ fn codex_apps_tools_cache_keeps_live_publish_when_disk_persistence_fails() { ); 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), + cache_context.begin_fetch(ConnectorRuntimeFetchSource::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) - ); + assert_eq!(published_tools, tools); + assert_eq!(cache_context.current_tools(), Some(tools)); } #[cfg(unix)] @@ -436,10 +419,10 @@ 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 cache = ConnectorRuntimeManager::::default(); let user_one_context = cache.context( codex_home.clone(), - CodexAppsToolsCacheKey { + ConnectorRuntimeContextKey { account_id: Some("account-one".to_string()), chatgpt_user_id: Some("user-one".to_string()), is_workspace_account: false, @@ -447,7 +430,7 @@ fn codex_apps_tools_cache_scopes_non_utf8_home_disk_paths() { ); let user_two_context = cache.context( codex_home, - CodexAppsToolsCacheKey { + ConnectorRuntimeContextKey { account_id: Some("account-two".to_string()), chatgpt_user_id: Some("user-two".to_string()), is_workspace_account: false, @@ -458,8 +441,264 @@ fn codex_apps_tools_cache_scopes_non_utf8_home_disk_paths() { user_two_context.tools_cache_path(), ]; - assert_eq!( - cache_paths.iter().collect::>().len(), - cache_paths.len() + assert_ne!(cache_paths[0], cache_paths[1]); +} + +#[test] +fn contexts_for_different_identities_keep_isolated_snapshots() { + let codex_home = tempdir().expect("tempdir"); + let manager = ConnectorRuntimeManager::::default(); + let context_a = manager.context( + codex_home.path().to_path_buf(), + ConnectorRuntimeContextKey { + account_id: Some("account-a".to_string()), + chatgpt_user_id: Some("user-a".to_string()), + is_workspace_account: false, + }, + ); + let tools_a = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "tool-a")]; + let snapshot_a = context_a.publish_runtime_if_newest_accepted( + context_a.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh), + &create_test_server_info("Codex Apps"), + tools_a.clone(), + ); + let older_ticket_a = context_a.begin_fetch(ConnectorRuntimeFetchSource::Startup); + let context_b = manager.context( + codex_home.path().to_path_buf(), + ConnectorRuntimeContextKey { + account_id: Some("account-b".to_string()), + chatgpt_user_id: Some("user-b".to_string()), + is_workspace_account: false, + }, + ); + let same_context_a = manager.context( + codex_home.path().to_path_buf(), + ConnectorRuntimeContextKey { + account_id: Some("account-a".to_string()), + chatgpt_user_id: Some("user-a".to_string()), + is_workspace_account: false, + }, + ); + + assert!(Arc::ptr_eq( + &snapshot_a, + &same_context_a + .current_snapshot() + .expect("context A snapshot") + )); + assert!(context_b.current_snapshot().is_none()); + + let tools_b = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "tool-b")]; + let snapshot_b = context_b.publish_runtime_if_newest_accepted( + context_b.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh), + &create_test_server_info("Codex Apps"), + tools_b.clone(), + ); + let newer_tools_a = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "newer-a")]; + let newer_snapshot_a = same_context_a.publish_runtime_if_newest_accepted( + same_context_a.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh), + &create_test_server_info("Codex Apps"), + newer_tools_a.clone(), + ); + let stale_snapshot_a = context_a.publish_runtime_if_newest_accepted( + older_ticket_a, + &create_test_server_info("Codex Apps"), + vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "stale-a")], + ); + + assert_eq!(snapshot_a.tools(), &tools_a); + assert_eq!(snapshot_b.tools(), &tools_b); + assert_eq!(newer_snapshot_a.tools(), &newer_tools_a); + assert!(Arc::ptr_eq(&newer_snapshot_a, &stale_snapshot_a)); + assert!(Arc::ptr_eq( + &newer_snapshot_a, + &context_a.current_snapshot().expect("context A snapshot") + )); + assert!(Arc::ptr_eq( + &snapshot_b, + &context_b.current_snapshot().expect("context B snapshot") + )); +} + +#[test] +fn oversized_tools_cache_is_ignored_during_initial_load() { + let codex_home = tempdir().expect("tempdir"); + let context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let cache_path = context.tools_cache_path(); + std::fs::create_dir_all(cache_path.parent().expect("cache parent")) + .expect("create cache parent"); + let file = std::fs::File::create(cache_path).expect("create oversized cache"); + file.set_len(CODEX_APPS_TOOLS_CACHE_MAX_BYTES + 1) + .expect("size oversized cache"); + + let reloaded = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + + assert!(reloaded.current_snapshot().is_none()); +} + +#[test] +fn cold_loaded_snapshot_uses_cache_modification_time() { + let codex_home = tempdir().expect("tempdir"); + let writer = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "cached")]; + write_cached_codex_apps_tools(&writer, &tools).expect("write tools cache"); + let modified_at = std::fs::metadata(writer.tools_cache_path()) + .and_then(|metadata| metadata.modified()) + .expect("cache modification time"); + + let reloaded = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let snapshot = reloaded.current_snapshot().expect("cold-loaded snapshot"); + + assert_eq!(snapshot.tools(), &tools); + assert_eq!(snapshot.refreshed_at(), modified_at); +} +#[test] +fn accepted_generations_finish_persistence_in_order() { + let codex_home = tempdir().expect("tempdir"); + let context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let older_ticket = context.begin_fetch(ConnectorRuntimeFetchSource::Startup); + let newer_ticket = context.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh); + let (older_persisting_tx, older_persisting_rx) = std::sync::mpsc::channel(); + let (release_older_tx, release_older_rx) = std::sync::mpsc::channel(); + let older_context = context.clone(); + let older_publish = std::thread::spawn(move || { + older_context.publish_runtime_if_newest_accepted_with( + older_ticket, + &create_test_server_info("Codex Apps"), + vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "older")], + move |_, _, _| { + older_persisting_tx + .send(()) + .expect("signal older persistence"); + release_older_rx.recv().expect("release older persistence"); + }, + ) + }); + older_persisting_rx + .recv_timeout(Duration::from_secs(1)) + .expect("older generation should enter persistence"); + + let (newer_persisting_tx, newer_persisting_rx) = std::sync::mpsc::channel(); + let newer_context = context; + let newer_publish = std::thread::spawn(move || { + newer_context.publish_runtime_if_newest_accepted_with( + newer_ticket, + &create_test_server_info("Codex Apps"), + vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "newer")], + move |_, _, _| { + newer_persisting_tx + .send(()) + .expect("signal newer persistence"); + }, + ) + }); + + assert!( + newer_persisting_rx + .recv_timeout(Duration::from_millis(20)) + .is_err() + ); + release_older_tx + .send(()) + .expect("allow older persistence to finish"); + newer_persisting_rx + .recv_timeout(Duration::from_secs(1)) + .expect("newer generation should persist after older generation"); + + older_publish.join().expect("join older publish"); + let newer_snapshot = newer_publish.join().expect("join newer publish"); + assert_eq!(newer_snapshot.tools()[0].callable_name, "newer"); +} + +#[test] +fn personal_and_workspace_contexts_are_distinct_even_with_matching_ids() { + let codex_home = tempdir().expect("tempdir"); + let manager = ConnectorRuntimeManager::::default(); + let personal_context = manager.context( + codex_home.path().to_path_buf(), + ConnectorRuntimeContextKey { + account_id: Some("account".to_string()), + chatgpt_user_id: Some("user".to_string()), + is_workspace_account: false, + }, + ); + let personal_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "personal")]; + let _ = personal_context.publish_runtime_if_newest_accepted( + personal_context.begin_fetch(ConnectorRuntimeFetchSource::Startup), + &create_test_server_info("Codex Apps"), + personal_tools.clone(), + ); + + let workspace_context = manager.context( + codex_home.path().to_path_buf(), + ConnectorRuntimeContextKey { + account_id: Some("account".to_string()), + chatgpt_user_id: Some("user".to_string()), + is_workspace_account: true, + }, + ); + + let workspace_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "workspace")]; + let _ = workspace_context.publish_runtime_if_newest_accepted( + workspace_context.begin_fetch(ConnectorRuntimeFetchSource::Startup), + &create_test_server_info("Codex Apps"), + workspace_tools.clone(), + ); + + assert_eq!(personal_context.current_tools(), Some(personal_tools)); + assert_eq!(workspace_context.current_tools(), Some(workspace_tools)); + assert_ne!( + personal_context.tools_cache_path(), + workspace_context.tools_cache_path() ); } + +#[test] +fn live_publish_sets_timestamp_and_stale_publish_preserves_it() { + let codex_home = tempdir().expect("tempdir"); + let context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let stale_ticket = context.begin_fetch(ConnectorRuntimeFetchSource::Startup); + let current_ticket = context.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh); + let before = SystemTime::now(); + let current = context.publish_runtime_if_newest_accepted( + current_ticket, + &create_test_server_info("Codex Apps"), + vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "current")], + ); + let after = SystemTime::now(); + + assert!(current.refreshed_at() >= before); + assert!(current.refreshed_at() <= after); + + let stale = context.publish_runtime_if_newest_accepted( + stale_ticket, + &create_test_server_info("Codex Apps"), + vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "stale")], + ); + assert!(Arc::ptr_eq(¤t, &stale)); + assert_eq!(stale.refreshed_at(), current.refreshed_at()); +} diff --git a/codex-rs/connectors/src/lib.rs b/codex-rs/connectors/src/lib.rs index f77936e88f..8d072c6d67 100644 --- a/codex-rs/connectors/src/lib.rs +++ b/codex-rs/connectors/src/lib.rs @@ -11,6 +11,7 @@ use serde::Serialize; pub mod accessible; mod app_info; mod app_tool_policy; +mod connector_runtime; mod directory_cache; pub mod filter; pub mod merge; @@ -28,6 +29,14 @@ pub use app_tool_policy::AppToolPolicyEvaluator; pub use app_tool_policy::AppToolPolicyInput; pub use app_tool_policy::app_is_enabled; pub use app_tool_policy::apps_config_from_layer_stack; +pub use connector_runtime::ConnectorRuntimeContext; +pub use connector_runtime::ConnectorRuntimeContextKey; +pub use connector_runtime::ConnectorRuntimeFetchSource; +pub use connector_runtime::ConnectorRuntimeFetchTicket; +pub use connector_runtime::ConnectorRuntimeManager; +pub use connector_runtime::ConnectorRuntimePayload; +pub use connector_runtime::ConnectorRuntimeSnapshot; +pub use connector_runtime::connector_runtime_context_key; pub use directory_cache::ConnectorDirectoryCacheContext; pub use plugin_config::parse_plugin_app_config; pub use plugin_config::parse_plugin_app_config_value; diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index a35b717cad..bd12cf3053 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -13,6 +13,7 @@ use codex_connectors::ConnectorDirectoryCacheContext; use codex_connectors::ConnectorDirectoryCacheKey; use codex_connectors::app_is_enabled; use codex_connectors::apps_config_from_layer_stack; +use codex_connectors::connector_runtime_context_key; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecServerRuntimePaths; use codex_protocol::models::PermissionProfile; @@ -37,7 +38,6 @@ use codex_mcp::McpConnectionManager; use codex_mcp::McpRuntimeContext; use codex_mcp::ToolInfo; use codex_mcp::ToolPluginProvenance; -use codex_mcp::codex_apps_tools_cache_key; use codex_mcp::compute_auth_statuses; use codex_mcp::effective_mcp_servers; use codex_mcp::tool_plugin_provenance; @@ -266,7 +266,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( runtime_context, config.codex_home.to_path_buf(), mcp_manager.codex_apps_tools_cache(), - codex_apps_tools_cache_key(auth.as_ref()), + connector_runtime_context_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability, /*supports_openai_form_elicitation*/ false, diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index bd09ad8564..e51aa36385 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use crate::config::Config; use codex_config::McpServerConfig; +use codex_connectors::ConnectorRuntimeManager; use codex_connectors::ConnectorSnapshot; use codex_connectors::PluginConnectorSource; use codex_core_plugins::PluginsManager; @@ -13,11 +14,11 @@ 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; use codex_mcp::McpServerRegistration; +use codex_mcp::ToolInfo; use codex_mcp::codex_apps_mcp_server_config; use codex_mcp::configured_mcp_servers; use codex_mcp::effective_mcp_servers; @@ -49,7 +50,7 @@ enum OrderedMcpOverlay { pub struct McpManager { plugins_manager: Arc, extensions: Arc>, - codex_apps_tools_cache: CodexAppsToolsCache, + codex_apps_tools_cache: ConnectorRuntimeManager, } impl McpManager { @@ -68,11 +69,11 @@ impl McpManager { Self { plugins_manager, extensions, - codex_apps_tools_cache: CodexAppsToolsCache::default(), + codex_apps_tools_cache: ConnectorRuntimeManager::default(), } } - pub fn codex_apps_tools_cache(&self) -> CodexAppsToolsCache { + pub fn codex_apps_tools_cache(&self) -> ConnectorRuntimeManager { self.codex_apps_tools_cache.clone() } diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 8c76c60334..3356880edc 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1454,7 +1454,7 @@ async fn host_owned_codex_apps_manager( ), 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()), + codex_connectors::connector_runtime_context_key(auth.as_ref()), turn_context.config.prefix_mcp_tool_names(), rmcp::model::ElicitationCapability::default(), /*supports_openai_form_elicitation*/ false, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 2c33cc041b..a057267900 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -380,7 +380,7 @@ impl Session { mcp_runtime_context.clone(), mcp_config.codex_home.clone(), self.services.mcp_manager.codex_apps_tools_cache(), - codex_apps_tools_cache_key(auth.as_ref()), + connector_runtime_context_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability.clone(), self.services diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 97e5ef3865..e61d59b2e2 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -55,6 +55,7 @@ use codex_analytics::SubAgentThreadStartedInput; use codex_analytics::TurnCodexErrorFact; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; +use codex_connectors::connector_runtime_context_key; use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionDataInit; @@ -73,7 +74,6 @@ use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_mcp::McpConnectionManager; use codex_mcp::McpResourceClient; use codex_mcp::McpRuntimeContext; -use codex_mcp::codex_apps_tools_cache_key; use codex_models_manager::manager::RefreshStrategy; use codex_models_manager::manager::SharedModelsManager; use codex_network_proxy::NetworkProxy; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index af3e4c512b..4ba1503ea3 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1221,7 +1221,7 @@ impl Session { mcp_runtime_context.clone(), config.codex_home.to_path_buf(), sess.services.mcp_manager.codex_apps_tools_cache(), - codex_apps_tools_cache_key(auth), + connector_runtime_context_key(auth), config.prefix_mcp_tool_names(), mcp_projection .config