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).
2026-07-01 00:31:56 +08:00 · 2026-06-25 20:54:48 +00:00
parent 62c7f506d9
commit 703793c22e
15 changed files with 1109 additions and 621 deletions
@@ -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<String>,
+    pub(crate) chatgpt_user_id: Option<String>,
+    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<Mutex<HashMap<CodexAppsToolsCacheIdentity, Arc<CodexAppsToolsCacheEntry>>>>,
+}
+
+/// 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<CodexAppsToolsCacheEntry>,
+}
+
+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<Vec<ToolInfo>> {
+        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<ToolInfo>,
+    ) -> Vec<ToolInfo> {
+        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<ToolInfo>) {
+        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<Vec<ToolInfo>>,
+    next_fetch_generation: AtomicU64,
+    last_accepted_generation: Mutex<u64>,
+}
+
+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<McpServerInfo> {
+    load_cached_codex_apps_server_info(cache_context)
+}
+
+#[cfg(test)]
+fn read_cached_codex_apps_tools(
+    cache_context: &CodexAppsToolsCacheContext,
+) -> Option<Vec<ToolInfo>> {
+    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<Vec<ToolInfo>> {
+    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<McpServerInfo> {
+    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<u8>,
+) -> 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<ToolInfo>,
+}
+
+#[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<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
+    mutex
+        .lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner)
+}
+
+#[cfg(test)]
+#[path = "codex_apps_cache_tests.rs"]
+mod tests;