mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
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:
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user