mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00:00
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).
This commit is contained in:
1 parent
62c7f506d9
commit
703793c22e
15 files changed
+1109
-621
No files matched your search
@@ -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>,
|
||||
|
||||
Reference in new issue
Block a user