agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(core, mcp): cache codex_apps tools in memory (#29003)

Browse Source 
## 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).
This commit is contained in:
Owen Lin
2026-06-25 13:54:48 -07:00
committed by GitHub
Unverified
parent 62c7f506d9
commit 703793c22e
15 changed files with 1109 additions and 621 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/request_processors/mcp_processor.rs
+12 -11
View File
@@ -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<CodexAuth>,
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<CodexAuth>,
runtime_context: McpRuntimeContext,
codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache,
) -> Result<ListMcpServerStatusResponse, JSONRPCErrorError> {
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,
)
codex-rs/codex-mcp/src/codex_apps.rs
+1 -211
View File
@@ -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<String>,
pub(crate) chatgpt_user_id: Option<String>,
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<ToolInfo>),
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<Vec<ToolInfo>> {
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<McpServerInfo> {
load_cached_codex_apps_server_info(cache_context?)
}
#[cfg(test)]
pub(crate) fn read_cached_codex_apps_tools(
cache_context: &CodexAppsToolsCacheContext,
) -> Option<Vec<ToolInfo>> {
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<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();
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<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";
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}")
}
codex-rs/codex-mcp/src/codex_apps_cache.rs
+378
View File
@@ -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;
codex-rs/codex-mcp/src/codex_apps_cache_tests.rs
+464
View File
@@ -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<ToolName> {
tools
.iter()
.map(ToolInfo::canonical_tool_name)
.collect::<HashSet<_>>()
}
#[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<_>>(),
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::<HashSet<_>>().len(),
cache_paths.len()
);
}
codex-rs/codex-mcp/src/connection_manager.rs
+53 -37
View File
@@ -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<ToolInfo> {
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<Vec<ToolInfo>> {
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<CodexAppsToolsCacheContext> {
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<CodexAppsToolsCacheContext> {
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<Event>,
codex-rs/codex-mcp/src/connection_manager_tests.rs
+116 -309
View File
@@ -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<DuplexStream>> {
async {
let (client_stream, _server_stream) = tokio::io::duplex(1);
Ok(client_stream)
}
.boxed()
}
}
async fn create_ready_async_managed_client(tools: Vec<ToolInfo>) -> 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::<Result<ManagedClient, StartupOutcomeError>>(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<ToolName> {
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<_>>(),
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::<Result<ManagedClient, StartupOutcomeError>>()
.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::<Result<ManagedClient, StartupOutcomeError>>()
.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::<Result<ManagedClient, StartupOutcomeError>>()
.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::<Result<ManagedClient, StartupOutcomeError>>()
.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::<Result<ManagedClient, StartupOutcomeError>>()
.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::<Result<ManagedClient, StartupOutcomeError>>(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::<Result<ManagedClient, StartupOutcomeError>>()
.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,
codex-rs/codex-mcp/src/lib.rs
+4 -2
View File
@@ -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;
codex-rs/codex-mcp/src/mcp/mod.rs
+6 -1
View File
@@ -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<ReadResourceResult> {
@@ -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(),
codex-rs/codex-mcp/src/rmcp_client.rs
+52 -44
View File
@@ -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<ToolInfo> {
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<BoxFuture<'static, Result<ManagedClient, StartupOutcomeError>>>,
pub(crate) is_codex_apps_mcp_server: bool,
pub(crate) cached_tool_info_snapshot: Option<Vec<ToolInfo>>,
pub(crate) cached_server_info: Option<McpServerInfo>,
pub(crate) codex_apps_tools_cache_context: Option<CodexAppsToolsCacheContext>,
pub(crate) tool_filter: ToolFilter,
pub(crate) startup_complete: Arc<AtomicBool>,
pub(crate) tool_plugin_provenance: Arc<ToolPluginProvenance>,
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<Vec<ToolInfo>> {
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<Vec<ToolInfo>> {
// 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<Vec<ToolInfo>> {
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 {
codex-rs/core/src/connectors.rs
+1
View File
@@ -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,
codex-rs/core/src/mcp.rs
+10 -3
View File
@@ -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<PluginsManager>,
extensions: Arc<ExtensionRegistry<Config>>,
codex_apps_tools_cache: CodexAppsToolsCache,
}
impl McpManager {
pub fn new(plugins_manager: Arc<PluginsManager>) -> 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 {
codex-rs/core/src/mcp_tool_call.rs
+9 -3
View File
@@ -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| {
codex-rs/core/src/mcp_tool_call_tests.rs
+1
View File
@@ -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(),
codex-rs/core/src/session/mcp.rs
+1
View File
@@ -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,
codex-rs/core/src/session/session.rs
+1
View File
@@ -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,