Preserve orchestrator skill caches across MCP runtime updates (#46033)

## Why

Orchestrator skill caches were tied to the entire MCP connection set, so runtime updates could discard cached catalogs and contents even when the Apps connection was reused. Plugin changes still need to refresh those resources without requiring a reconnect.

## What changed

Key orchestrator skill caches by the Apps server connection and an explicit invalidation generation. Preserve cached resources while that identity stays unchanged, and invalidate them on MCP runtime refresh requests or configuration refreshes. Replacing the connection also changes the cache identity.

## Testing

Extend the production-turn integration test to verify cache reuse across ordinary turns and approval-policy changes, refreshed skill contents after plugin invalidation without reconnecting, and fresh resource requests after a forced reconnect.

GitOrigin-RevId: 43ce059972369c4c735742e8224b60a692244b02
This commit is contained in:
Henry Levy
2026-09-16 21:41:23 +00:00
committed by copyberry
parent 73bf181272
commit 29e6bc814e
8 changed files with 130 additions and 9 deletions
@@ -17,9 +17,23 @@ use tracing::warn;
use super::McpConnectionSet;
use crate::pagination::collect_paginated;
use crate::resource_client::McpResourceServerCacheKey;
use crate::rmcp_client::ManagedClient;
impl McpConnectionSet {
pub(crate) fn resource_cache_key(
&self,
server: &str,
generation: u64,
) -> Option<McpResourceServerCacheKey> {
self.servers
.get(server)
.map(|view| McpResourceServerCacheKey {
connection: Arc::downgrade(&view.connection),
generation,
})
}
/// Returns resources from servers selected by `include_server`.
pub async fn list_all_resources(
&self,
+1
View File
@@ -18,6 +18,7 @@ pub use resource_client::McpResourceClient;
pub use resource_client::McpResourceClientCacheKey;
pub use resource_client::McpResourcePage;
pub use resource_client::McpResourceReadResult;
pub use resource_client::McpResourceServerCacheKey;
pub use rmcp::model::ReadResourceRequestParams;
pub use rmcp_client::MCP_SANDBOX_STATE_META_CAPABILITY;
pub use runtime::McpRuntime;
+24
View File
@@ -25,6 +25,7 @@ use tokio::sync::watch;
use crate::McpEventStreamOpener;
use crate::McpRuntime;
use crate::connection_manager::McpConnectionSet;
use crate::connection_manager::McpServerConnection;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
/// One page of resources returned by an MCP server.
@@ -200,6 +201,24 @@ impl PartialEq for McpResourceClientCacheKey {
impl Eq for McpResourceClientCacheKey {}
/// Identity of one server's resource cache across MCP runtime publications.
///
/// Survives publications that reuse the same connection until resource caches are explicitly
/// invalidated. The key does not keep superseded connections alive.
#[derive(Clone)]
pub struct McpResourceServerCacheKey {
pub(crate) connection: Weak<McpServerConnection>,
pub(crate) generation: u64,
}
impl PartialEq for McpResourceServerCacheKey {
fn eq(&self, other: &Self) -> bool {
self.connection.ptr_eq(&other.connection) && self.generation == other.generation
}
}
impl Eq for McpResourceServerCacheKey {}
impl std::fmt::Debug for McpResourceClient {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
@@ -219,6 +238,11 @@ impl McpResourceClient {
McpResourceClientCacheKey(Arc::downgrade(&self.runtime.latest_connections()))
}
/// Returns a server's resource cache identity without starting its connection.
pub fn server_cache_key(&self, server: &str) -> Option<McpResourceServerCacheKey> {
self.runtime.resource_cache_key(server)
}
/// Returns whether this client can address the named server.
///
/// This does not wait for server startup.
+16
View File
@@ -12,6 +12,7 @@ use std::sync::Arc;
use std::sync::Mutex;
use std::sync::Weak;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::time::Duration;
@@ -55,6 +56,7 @@ use crate::elicitation::ElicitationRequestRouter;
use crate::elicitation::ElicitationReviewerHandle;
use crate::event_stream::McpEventStreamOpener;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
use crate::resource_client::McpResourceServerCacheKey;
use crate::resource_origin::ResourceOrigins;
use crate::server::EffectiveMcpServer;
use crate::tool_catalog_cache::McpToolCatalogCache;
@@ -98,6 +100,7 @@ pub struct McpRuntime {
current: ArcSwap<PublishedMcpRuntime>,
event_stream_cancellation: Mutex<EventStreamCancellation>,
reconnect_pending: AtomicBool,
resource_cache_generation: AtomicU64,
elicitation_router: ElicitationRequestRouter,
resource_origins: Mutex<ResourceOrigins>,
}
@@ -215,6 +218,7 @@ impl McpRuntime {
retained_subscription_cancellation: None,
}),
reconnect_pending: AtomicBool::new(false),
resource_cache_generation: AtomicU64::new(0),
elicitation_router: ElicitationRequestRouter::default(),
resource_origins: Mutex::default(),
}
@@ -357,6 +361,18 @@ impl McpRuntime {
self.reconnect_pending.store(true, Ordering::Release);
}
/// Invalidates resource catalogs and contents even when connections can be reused.
pub fn invalidate_resource_caches(&self) {
self.resource_cache_generation
.fetch_add(1, Ordering::AcqRel);
}
pub(crate) fn resource_cache_key(&self, server: &str) -> Option<McpResourceServerCacheKey> {
let generation = self.resource_cache_generation.load(Ordering::Acquire);
self.latest_connections()
.resource_cache_key(server, generation)
}
/// Captures the latest published configuration and live client handles.
pub async fn current_binding(&self) -> Option<Arc<McpBinding>> {
self.current_binding_with_requirements(&[], &HashSet::new())
+2
View File
@@ -7,6 +7,8 @@ use super::*;
impl Session {
pub(crate) fn request_mcp_runtime_refresh(&self) {
// Plugin changes can reuse connections but still change their skill resources.
self.services.mcp_runtime.invalidate_resource_caches();
self.mark_mcp_runtime_dirty();
self.schedule_mcp_prewarm();
}
+1
View File
@@ -2160,6 +2160,7 @@ impl Session {
warn!("failed to refresh MCP OAuth coordination config: {err}");
}
state.session_configuration.original_config_do_not_use = Arc::new(config);
self.services.mcp_runtime.invalidate_resource_caches();
self.mark_mcp_runtime_dirty();
drop(state);
self.schedule_mcp_prewarm();
+67 -6
View File
@@ -1,6 +1,7 @@
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use anyhow::Result;
@@ -43,6 +44,7 @@ use codex_protocol::permissions::FileSystemSandboxEntry;
use codex_protocol::permissions::FileSystemSandboxPolicy;
use codex_protocol::permissions::FileSystemSpecialPath;
use codex_protocol::permissions::NetworkSandboxPolicy;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
@@ -1334,19 +1336,31 @@ text({ names: result.skills.map(skill => skill.name), warnings: result.warnings,
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn production_turn_aliases_discovered_singleton_orchestrator_root() -> Result<()> {
async fn production_turn_reuses_orchestrator_skills_until_mcp_invalidation() -> Result<()> {
skip_if_no_network!(Ok(()));
const SKILL_ROOT: &str = "skill://plugin_connector_1p_2330815c823c8191941e5dc465bb899f";
const SKILL_BODY: &str = "ORCHESTRATOR_SKILL_REMAINS_AVAILABLE_WITHOUT_HOST_DISCOVERY";
const UPDATED_SKILL_BODY: &str = "UPDATED_ORCHESTRATOR_SKILL";
let server = responses::start_mock_server().await;
let apps_server = AppsTestServer::mount(&server).await?;
let response = mount_sse_once(
let (apps_server, startup) = AppsTestServer::mount_with_startup_control(&server).await?;
let response = responses::mount_sse_sequence(
&server,
sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
(1..=5)
.map(|index| {
let id = format!("resp-{index}");
sse(vec![ev_response_created(&id), ev_completed(&id)])
})
.collect(),
)
.await;
let list_calls = Arc::new(AtomicUsize::new(0));
let read_calls = Arc::new(AtomicUsize::new(0));
let updated = Arc::new(AtomicBool::new(false));
let resource_list_calls = Arc::clone(&list_calls);
let resource_read_calls = Arc::clone(&read_calls);
let resource_updated = Arc::clone(&updated);
Mock::given(method("POST"))
.and(path_regex("^/api/codex/ps/mcp/?$"))
@@ -1358,20 +1372,27 @@ async fn production_turn_aliases_discovered_singleton_orchestrator_root() -> Res
)
})
})
.respond_with(|request: &Request| {
.respond_with(move |request: &Request| {
let body: Value = serde_json::from_slice(&request.body)
.expect("MCP resource request should be valid JSON");
let result = if body["method"] == "resources/read" {
resource_read_calls.fetch_add(1, Ordering::SeqCst);
let uri = format!("{SKILL_ROOT}/search/SKILL.md");
assert_eq!(body["params"]["uri"], uri);
let contents = if resource_updated.load(Ordering::SeqCst) {
UPDATED_SKILL_BODY
} else {
SKILL_BODY
};
json!({
"contents": [{
"uri": uri,
"mimeType": "text/markdown",
"text": SKILL_BODY,
"text": contents,
}],
})
} else {
resource_list_calls.fetch_add(1, Ordering::SeqCst);
json!({
"resources": [{
"name": "search",
@@ -1452,6 +1473,46 @@ async fn production_turn_aliases_discovered_singleton_orchestrator_root() -> Res
"orchestrator instruction reads must remain available without host discovery: {user_text}"
);
// A normal turn and a policy change must reuse the same Apps skill snapshot.
test.submit_text_turn("Use $demo:search again.").await?;
test.submit_turn_with_approval_and_permission_profile(
"Use $demo:search after updating approval policy.",
AskForApproval::OnRequest,
PermissionProfile::Disabled,
)
.await?;
assert_eq!(
startup.initialize_attempts(),
1,
"Apps connection was reused"
);
assert_eq!(list_calls.load(Ordering::SeqCst), 1);
assert_eq!(read_calls.load(Ordering::SeqCst), 1);
// Plugin invalidation must refresh both the catalog and contents, even when
// the underlying Apps connection is reused.
updated.store(true, Ordering::SeqCst);
test.thread_manager.invalidate_mcp_runtimes().await;
test.submit_text_turn("Use $demo:search after its plugin changed.")
.await?;
assert_eq!(startup.initialize_attempts(), 1);
assert_eq!(list_calls.load(Ordering::SeqCst), 2);
assert_eq!(read_calls.load(Ordering::SeqCst), 2);
assert!(
response.requests()[3]
.message_input_texts("user")
.join("\n")
.contains(UPDATED_SKILL_BODY)
);
// A forced reconnect must also refresh the skills snapshot.
test.codex.submit(Op::RefreshMcpServers).await?;
test.submit_text_turn("Use $demo:search after reconnecting.")
.await?;
assert_eq!(startup.initialize_attempts(), 2);
assert_eq!(list_calls.load(Ordering::SeqCst), 3);
assert_eq!(read_calls.load(Ordering::SeqCst), 3);
Ok(())
}
+5 -3
View File
@@ -7,8 +7,9 @@ use std::sync::Weak;
use codex_exec_server::Environment;
use codex_exec_server::FileSystemSandboxContext;
use codex_extension_api::ExtensionMetrics;
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
use codex_mcp::McpResourceClient;
use codex_mcp::McpResourceClientCacheKey;
use codex_mcp::McpResourceServerCacheKey;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use tokio::sync::OnceCell;
@@ -259,7 +260,8 @@ impl SkillsThreadState {
.orchestrator_cache
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let cache_key = mcp_resources.map(McpResourceClient::cache_key);
let cache_key =
mcp_resources.and_then(|client| client.server_cache_key(CODEX_APPS_MCP_SERVER_NAME));
if let Some(cache) = cache
.as_ref()
.filter(|cache| cache.mcp_cache_key == cache_key)
@@ -337,7 +339,7 @@ struct CachedExecutorDiscoveryCatalog {
}
struct OrchestratorGenerationCache {
mcp_cache_key: Option<McpResourceClientCacheKey>,
mcp_cache_key: Option<McpResourceServerCacheKey>,
catalog: OnceCell<SkillCatalog>,
resources: Mutex<OrchestratorResourceCache>,
}