diff --git a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs index 85933bd6b..dc3b6d577 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs @@ -1,4 +1,6 @@ use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use anyhow::Result; @@ -79,11 +81,13 @@ const SKILL_REFERENCE_CONTENTS: &str = "# Deploy reference\n\nUse the orchestrator deployment API.\n"; const SKILLS_LIST_CALL_ID: &str = "skills-list"; const SKILLS_READ_CALL_ID: &str = "skills-read"; +const SKILLS_READ_AGAIN_CALL_ID: &str = "skills-read-again"; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn mcp_resource_read_returns_resource_contents() -> Result<()> { let responses_server = responses::start_mock_server().await; - let (apps_server_url, apps_server_handle) = start_resource_apps_mcp_server().await?; + let (apps_server_url, _apps_server_calls, apps_server_handle) = + start_resource_apps_mcp_server().await?; let responses_server_uri = responses_server.uri(); let (_codex_home, mut mcp) = start_resource_test_app_server(&apps_server_url, &responses_server_uri).await?; @@ -126,7 +130,8 @@ async fn mcp_resource_read_returns_resource_contents() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn orchestrator_skill_can_read_referenced_resource_without_an_executor() -> Result<()> { let responses_server = responses::start_mock_server().await; - let (apps_server_url, apps_server_handle) = start_resource_apps_mcp_server().await?; + let (apps_server_url, apps_server_calls, apps_server_handle) = + start_resource_apps_mcp_server().await?; let responses_server_uri = responses_server.uri(); let (_codex_home, mut mcp) = start_resource_test_app_server(&apps_server_url, &responses_server_uri).await?; @@ -180,17 +185,39 @@ async fn orchestrator_skill_can_read_referenced_resource_without_an_executor() - ), responses::ev_completed("resp-skills-read"), ]), + responses::sse(vec![ + responses::ev_response_created("resp-skills-read-again"), + responses::ev_function_call_with_namespace( + SKILLS_READ_AGAIN_CALL_ID, + "skills", + "read", + &json!({ + "authority": { + "kind": "orchestrator", + }, + "package": SKILL_RESOURCE_URI, + "resource": SKILL_REFERENCE_URI, + }) + .to_string(), + ), + responses::ev_completed("resp-skills-read-again"), + ]), responses::sse(vec![ responses::ev_response_created("resp-orchestrator-skill"), responses::ev_assistant_message("msg-orchestrator-skill", "Done"), responses::ev_completed("resp-orchestrator-skill"), ]), + responses::sse(vec![ + responses::ev_response_created("resp-orchestrator-skill-after-refresh"), + responses::ev_assistant_message("msg-orchestrator-skill-after-refresh", "Done"), + responses::ev_completed("resp-orchestrator-skill-after-refresh"), + ]), ], ) .await; let turn_start_id = mcp .send_turn_start_request(TurnStartParams { - thread_id: thread.id, + thread_id: thread.id.clone(), input: vec![UserInput::Text { text: format!("Use ${SKILL_NAME}"), text_elements: Vec::new(), @@ -210,7 +237,7 @@ async fn orchestrator_skill_can_read_referenced_resource_without_an_executor() - .await??; let requests = response_mock.requests(); - assert_eq!(requests.len(), 3); + assert_eq!(requests.len(), 4); let first_request = &requests[0]; assert!(first_request.tool_by_name("skills", "list").is_some()); assert!(first_request.tool_by_name("skills", "read").is_some()); @@ -276,6 +303,61 @@ async fn orchestrator_skill_can_read_referenced_resource_without_an_executor() - "contents": SKILL_REFERENCE_CONTENTS, }) ); + let repeated_read_output = requests[3] + .function_call_output_text(SKILLS_READ_AGAIN_CALL_ID) + .ok_or_else(|| { + anyhow::anyhow!("repeated skills.read output should be sent to the model") + })?; + assert_eq!(read_output, repeated_read_output); + assert_eq!( + ResourceAppsMcpCallCounts { + list_resources: 3, + main_prompt_reads: 1, + reference_reads: 1, + }, + apps_server_calls.snapshot() + ); + + let refresh_request_id = mcp + .send_raw_request("config/mcpServer/reload", /*params*/ None) + .await?; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(refresh_request_id)), + ) + .await??; + + let refreshed_turn_start_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id, + input: vec![UserInput::Text { + text: format!("Use ${SKILL_NAME} after refreshing MCP"), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(refreshed_turn_start_id)), + ) + .await??; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 5); + assert_eq!( + ResourceAppsMcpCallCounts { + list_resources: 6, + main_prompt_reads: 2, + reference_reads: 1, + }, + apps_server_calls.snapshot() + ); apps_server_handle.abort(); let _ = apps_server_handle.await; Ok(()) @@ -284,7 +366,8 @@ async fn orchestrator_skill_can_read_referenced_resource_without_an_executor() - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn local_executor_does_not_expose_orchestrator_skills() -> Result<()> { let responses_server = responses::start_mock_server().await; - let (apps_server_url, apps_server_handle) = start_resource_apps_mcp_server().await?; + let (apps_server_url, _apps_server_calls, apps_server_handle) = + start_resource_apps_mcp_server().await?; let responses_server_uri = responses_server.uri(); let (_codex_home, mut mcp) = start_resource_test_app_server(&apps_server_url, &responses_server_uri).await?; @@ -355,7 +438,8 @@ async fn local_executor_does_not_expose_orchestrator_skills() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn mcp_resource_read_returns_resource_contents_without_thread() -> Result<()> { - let (apps_server_url, apps_server_handle) = start_resource_apps_mcp_server().await?; + let (apps_server_url, _apps_server_calls, apps_server_handle) = + start_resource_apps_mcp_server().await?; let codex_home = TempDir::new()?; std::fs::write( @@ -514,13 +598,20 @@ stream_max_retries = 0 Ok((codex_home, mcp)) } -async fn start_resource_apps_mcp_server() -> Result<(String, JoinHandle<()>)> { +async fn start_resource_apps_mcp_server() +-> Result<(String, Arc, JoinHandle<()>)> { let listener = TcpListener::bind("127.0.0.1:0").await?; let addr = listener.local_addr()?; let apps_server_url = format!("http://{addr}"); + let calls = Arc::new(ResourceAppsMcpCalls::default()); + let server_calls = Arc::clone(&calls); let mcp_service = StreamableHttpService::new( - move || Ok(ResourceAppsMcpServer), + move || { + Ok(ResourceAppsMcpServer { + calls: Arc::clone(&server_calls), + }) + }, Arc::new(LocalSessionManager::default()), StreamableHttpServerConfig::default(), ); @@ -529,7 +620,7 @@ async fn start_resource_apps_mcp_server() -> Result<(String, JoinHandle<()>)> { let _ = axum::serve(listener, router).await; }); - Ok((apps_server_url, apps_server_handle)) + Ok((apps_server_url, calls, apps_server_handle)) } fn expected_resource_read_response() -> McpResourceReadResponse { @@ -551,8 +642,34 @@ fn expected_resource_read_response() -> McpResourceReadResponse { } } -#[derive(Clone, Default)] -struct ResourceAppsMcpServer; +#[derive(Debug, Default)] +struct ResourceAppsMcpCalls { + list_resources: AtomicUsize, + main_prompt_reads: AtomicUsize, + reference_reads: AtomicUsize, +} + +impl ResourceAppsMcpCalls { + fn snapshot(&self) -> ResourceAppsMcpCallCounts { + ResourceAppsMcpCallCounts { + list_resources: self.list_resources.load(Ordering::Relaxed), + main_prompt_reads: self.main_prompt_reads.load(Ordering::Relaxed), + reference_reads: self.reference_reads.load(Ordering::Relaxed), + } + } +} + +#[derive(Debug, PartialEq, Eq)] +struct ResourceAppsMcpCallCounts { + list_resources: usize, + main_prompt_reads: usize, + reference_reads: usize, +} + +#[derive(Clone)] +struct ResourceAppsMcpServer { + calls: Arc, +} impl ServerHandler for ResourceAppsMcpServer { fn get_info(&self) -> ServerInfo { @@ -565,6 +682,7 @@ impl ServerHandler for ResourceAppsMcpServer { request: Option, _context: RequestContext, ) -> Result { + self.calls.list_resources.fetch_add(1, Ordering::Relaxed); let cursor = request.and_then(|request| request.cursor); if cursor.is_none() { return Ok(ListResourcesResult { @@ -614,6 +732,7 @@ impl ServerHandler for ResourceAppsMcpServer { ) -> Result { let uri = request.uri; if uri == SKILL_MAIN_PROMPT_URI { + self.calls.main_prompt_reads.fetch_add(1, Ordering::Relaxed); return Ok(ReadResourceResult::new(vec![ ResourceContents::TextResourceContents { uri: SKILL_MAIN_PROMPT_URI.to_string(), @@ -624,6 +743,7 @@ impl ServerHandler for ResourceAppsMcpServer { ])); } if uri == SKILL_REFERENCE_URI { + self.calls.reference_reads.fetch_add(1, Ordering::Relaxed); return Ok(ReadResourceResult::new(vec![ ResourceContents::TextResourceContents { uri: SKILL_REFERENCE_URI.to_string(), diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 598aaeac8..b56a4f907 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -4,6 +4,7 @@ pub use elicitation::ElicitationReviewRequest; pub use elicitation::ElicitationReviewer; pub use elicitation::ElicitationReviewerHandle; pub use resource_client::McpResourceClient; +pub use resource_client::McpResourceClientCacheKey; pub use resource_client::McpResourcePage; pub use resource_client::McpResourceReadResult; pub use rmcp_client::MCP_SANDBOX_STATE_META_CAPABILITY; diff --git a/codex-rs/codex-mcp/src/resource_client.rs b/codex-rs/codex-mcp/src/resource_client.rs index 85dc974f2..4406b81a5 100644 --- a/codex-rs/codex-mcp/src/resource_client.rs +++ b/codex-rs/codex-mcp/src/resource_client.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::sync::Weak; use anyhow::Context; use anyhow::Result; @@ -35,6 +36,18 @@ pub struct McpResourceClient { manager: Arc>, } +/// Opaque identity for the manager currently used by an MCP resource client. +#[derive(Clone)] +pub struct McpResourceClientCacheKey(Weak); + +impl PartialEq for McpResourceClientCacheKey { + fn eq(&self, other: &Self) -> bool { + self.0.ptr_eq(&other.0) + } +} + +impl Eq for McpResourceClientCacheKey {} + impl std::fmt::Debug for McpResourceClient { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter @@ -49,6 +62,11 @@ impl McpResourceClient { Self { manager } } + /// Returns an identity that changes whenever the published manager changes. + pub fn cache_key(&self) -> McpResourceClientCacheKey { + McpResourceClientCacheKey(Arc::downgrade(&self.manager.load_full())) + } + /// Returns whether the current manager contains the named server. /// /// This does not wait for server startup or imply that startup succeeded. diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index 2bbf14952..bf4d86d36 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -150,10 +150,11 @@ where session_store: &ExtensionData, thread_store: &ExtensionData, ) -> Vec>> { + let Some(thread_state) = thread_store.get::() else { + return Vec::new(); + }; if !self.providers.has_orchestrator_provider() - || !thread_store - .get::() - .is_some_and(|state| state.orchestrator_skills_enabled()) + || !thread_state.orchestrator_skills_enabled() { return Vec::new(); } @@ -161,6 +162,7 @@ where skill_tools( self.providers.clone(), session_store.get::(), + thread_state, ) } } @@ -215,7 +217,12 @@ where let mut injected_host_skill_prompts = InjectedHostSkillPrompts::default(); for entry in &selected_entries { match self - .read_main_prompt(entry, host_loaded_skills.clone(), session_store) + .read_main_prompt( + entry, + host_loaded_skills.clone(), + session_store, + &thread_state, + ) .await { Ok(read_result) => { @@ -292,12 +299,14 @@ impl SkillsExtension { ) -> SkillCatalog { let include_orchestrator_skills = query.include_orchestrator_skills; let orchestrator_query = query.clone(); + let mcp_resources = orchestrator_query.mcp_resources.clone(); query.include_orchestrator_skills = false; let mut catalog = self.providers.list_for_turn(query).await; if include_orchestrator_skills { let orchestrator_catalog = thread_state .orchestrator_catalog_snapshot( + mcp_resources.as_deref(), self.providers .list_orchestrator_for_turn(orchestrator_query), ) @@ -312,15 +321,19 @@ impl SkillsExtension { entry: &SkillCatalogEntry, host_loaded_skills: Option>, session_store: &ExtensionData, + thread_state: &SkillsThreadState, ) -> Result { - self.providers - .read(SkillReadRequest { - authority: entry.authority.clone(), - package: entry.id.clone(), - resource: entry.main_prompt.clone(), - host: host_loaded_skills, - mcp_resources: session_store.get::(), - }) + thread_state + .read_skill( + &self.providers, + SkillReadRequest { + authority: entry.authority.clone(), + package: entry.id.clone(), + resource: entry.main_prompt.clone(), + host: host_loaded_skills, + mcp_resources: session_store.get::(), + }, + ) .await .map_err(|err| err.message) } diff --git a/codex-rs/ext/skills/src/state.rs b/codex-rs/ext/skills/src/state.rs index dc3c7d0b8..f718f6e02 100644 --- a/codex-rs/ext/skills/src/state.rs +++ b/codex-rs/ext/skills/src/state.rs @@ -1,19 +1,34 @@ -use codex_protocol::capabilities::SelectedCapabilityRoot; +use std::collections::HashMap; use std::future::Future; +use std::sync::Arc; use std::sync::Mutex; + +use codex_mcp::McpResourceClient; +use codex_mcp::McpResourceClientCacheKey; +use codex_protocol::capabilities::SelectedCapabilityRoot; use tokio::sync::OnceCell; use crate::SkillsExtensionConfig; +use crate::catalog::SkillAuthority; use crate::catalog::SkillCatalog; use crate::catalog::SkillCatalogEntry; +use crate::catalog::SkillPackageId; use crate::catalog::SkillProviderError; +use crate::catalog::SkillProviderResult; +use crate::catalog::SkillReadResult; +use crate::catalog::SkillResourceId; +use crate::catalog::SkillSourceKind; +use crate::provider::SkillReadRequest; +use crate::sources::SkillProviders; + +const MAX_CACHED_ORCHESTRATOR_RESOURCES: usize = 100; +const MAX_CACHED_ORCHESTRATOR_CONTENT_BYTES: usize = 8 * 1024 * 1024; -#[derive(Debug)] pub(crate) struct SkillsThreadState { config: Mutex, selected_roots: Vec, orchestrator_skills_enabled: bool, - orchestrator_catalog: OnceCell, + orchestrator_cache: Mutex>>, } impl SkillsThreadState { @@ -26,7 +41,7 @@ impl SkillsThreadState { config: Mutex::new(config), selected_roots, orchestrator_skills_enabled, - orchestrator_catalog: OnceCell::new(), + orchestrator_cache: Mutex::new(None), } } @@ -54,9 +69,11 @@ impl SkillsThreadState { pub(crate) async fn orchestrator_catalog_snapshot( &self, + mcp_resources: Option<&McpResourceClient>, initialize: impl Future> + Send, ) -> SkillCatalog { - self.orchestrator_catalog + self.orchestrator_cache(mcp_resources) + .catalog .get_or_init(|| async { initialize.await.unwrap_or_else(|err| SkillCatalog { warnings: vec![err.message], @@ -66,6 +83,118 @@ impl SkillsThreadState { .await .clone() } + + pub(crate) async fn read_skill( + &self, + providers: &SkillProviders, + request: SkillReadRequest, + ) -> SkillProviderResult { + if request.authority.kind != SkillSourceKind::Orchestrator { + return providers.read(request).await; + } + + let cache = self.orchestrator_cache(request.mcp_resources.as_deref()); + let cache_key = SkillReadCacheKey::from(&request); + if let Some(result) = cache + .resources + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(&cache_key) + { + return Ok(result); + } + + let result = providers.read(request).await?; + if result.resource != cache_key.resource { + return Ok(result); + } + + Ok(cache + .resources + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(cache_key, result)) + } + + fn orchestrator_cache( + &self, + mcp_resources: Option<&McpResourceClient>, + ) -> Arc { + let mut cache = self + .orchestrator_cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let cache_key = mcp_resources.map(McpResourceClient::cache_key); + if let Some(cache) = cache + .as_ref() + .filter(|cache| cache.mcp_cache_key == cache_key) + { + return Arc::clone(cache); + } + + let next_cache = Arc::new(OrchestratorGenerationCache { + mcp_cache_key: cache_key, + catalog: OnceCell::new(), + resources: Mutex::new(OrchestratorResourceCache::default()), + }); + *cache = Some(Arc::clone(&next_cache)); + next_cache + } +} + +struct OrchestratorGenerationCache { + mcp_cache_key: Option, + catalog: OnceCell, + resources: Mutex, +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct SkillReadCacheKey { + authority: SkillAuthority, + package: SkillPackageId, + resource: SkillResourceId, +} + +impl From<&SkillReadRequest> for SkillReadCacheKey { + fn from(request: &SkillReadRequest) -> Self { + Self { + authority: request.authority.clone(), + package: request.package.clone(), + resource: request.resource.clone(), + } + } +} + +#[derive(Default)] +struct OrchestratorResourceCache { + entries: HashMap, + contents_bytes: usize, +} + +impl OrchestratorResourceCache { + fn get(&self, key: &SkillReadCacheKey) -> Option { + self.entries.get(key).cloned() + } + + fn insert(&mut self, key: SkillReadCacheKey, result: SkillReadResult) -> SkillReadResult { + if let Some(cached) = self.entries.get(&key) { + return cached.clone(); + } + + let contents_bytes = result.contents.len(); + let Some(next_contents_bytes) = self.contents_bytes.checked_add(contents_bytes) else { + return result; + }; + if self.entries.len() >= MAX_CACHED_ORCHESTRATOR_RESOURCES + || next_contents_bytes > MAX_CACHED_ORCHESTRATOR_CONTENT_BYTES + { + return result; + } + + self.contents_bytes = next_contents_bytes; + self.entries.insert(key, result.clone()); + result + } } #[derive(Clone, Debug, Default, PartialEq, Eq)] diff --git a/codex-rs/ext/skills/src/tools/mod.rs b/codex-rs/ext/skills/src/tools/mod.rs index 12b050338..a1898b1a2 100644 --- a/codex-rs/ext/skills/src/tools/mod.rs +++ b/codex-rs/ext/skills/src/tools/mod.rs @@ -24,6 +24,7 @@ use crate::catalog::SkillCatalog; use crate::catalog::SkillSourceKind; use crate::provider::SkillListQuery; use crate::sources::SkillProviders; +use crate::state::SkillsThreadState; mod list; mod read; @@ -35,10 +36,12 @@ const MAX_HANDLE_BYTES: usize = 2_048; pub(crate) fn skill_tools( providers: SkillProviders, mcp_resources: Option>, + thread_state: Arc, ) -> Vec>> { let context = SkillToolContext { providers, mcp_resources, + thread_state, }; vec![ Arc::new(list::ListTool { @@ -52,30 +55,28 @@ pub(crate) fn skill_tools( struct SkillToolContext { providers: SkillProviders, mcp_resources: Option>, + thread_state: Arc, } impl SkillToolContext { async fn catalog(&self, turn_id: &str, authority: SkillToolAuthority) -> SkillCatalog { match authority { - SkillToolAuthority::Orchestrator => match self - .providers - .list_orchestrator_for_turn(SkillListQuery { - turn_id: turn_id.to_string(), - executor_roots: Vec::new(), - host: None, - include_host_skills: false, - include_bundled_skills: false, - include_orchestrator_skills: true, - mcp_resources: self.mcp_resources.clone(), - }) - .await - { - Ok(catalog) => catalog, - Err(err) => SkillCatalog { - warnings: vec![err.message], - ..Default::default() - }, - }, + SkillToolAuthority::Orchestrator => { + self.thread_state + .orchestrator_catalog_snapshot( + self.mcp_resources.as_deref(), + self.providers.list_orchestrator_for_turn(SkillListQuery { + turn_id: turn_id.to_string(), + executor_roots: Vec::new(), + host: None, + include_host_skills: false, + include_bundled_skills: false, + include_orchestrator_skills: true, + mcp_resources: self.mcp_resources.clone(), + }), + ) + .await + } } } } diff --git a/codex-rs/ext/skills/src/tools/read.rs b/codex-rs/ext/skills/src/tools/read.rs index e8b29905b..358842aa1 100644 --- a/codex-rs/ext/skills/src/tools/read.rs +++ b/codex-rs/ext/skills/src/tools/read.rs @@ -75,14 +75,17 @@ impl ToolExecutor for ReadTool { let requested_resource = SkillResourceId::new(args.resource); let result = self .context - .providers - .read(SkillReadRequest { - authority, - package: SkillPackageId(args.package), - resource: requested_resource.clone(), - host: None, - mcp_resources: self.context.mcp_resources.clone(), - }) + .thread_state + .read_skill( + &self.context.providers, + SkillReadRequest { + authority, + package: SkillPackageId(args.package), + resource: requested_resource.clone(), + host: None, + mcp_resources: self.context.mcp_resources.clone(), + }, + ) .await .map_err(|err| { tracing::warn!(