skills: cache orchestrator resources per thread (#28336)

## Why

Hosted orchestrator skills are read through the remote MCP resource
server. Within one thread, the same catalog or skill resource can be
requested multiple times by prompt injection and the `skills.list` /
`skills.read` tools. Re-fetching adds latency and can make those
surfaces observe different remote contents during the same thread.

This is a follow-up to #28333: orchestrator skills remain limited to
threads without a local executor, and those threads now get a stable
per-thread view of the remote skill data they use.

## What changed

- Reuse the existing per-thread orchestrator catalog snapshot for
`skills.list` and `skills.read` availability checks.
- Cache successful orchestrator resource reads by authority, package,
and resource so prompt injection and tool calls share the same contents.
- Keep the cache memory-only and bounded to 100 resources and 8 MiB per
thread.
- Leave host and executor skill reads unchanged, and do not cache failed
remote reads.

## Verification

- Extended the app-server MCP resource integration test to read the same
hosted skill resource twice and verify that the remote server receives
one read.
- The same test verifies that catalog discovery and the selected skill's
main prompt are each fetched only once per thread.
This commit is contained in:
jif
2026-06-15 19:20:19 +01:00
committed by GitHub
Unverified
parent ee40dddbf6
commit 0afe559318
7 changed files with 340 additions and 55 deletions
@@ -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<ResourceAppsMcpCalls>, 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<ResourceAppsMcpCalls>,
}
impl ServerHandler for ResourceAppsMcpServer {
fn get_info(&self) -> ServerInfo {
@@ -565,6 +682,7 @@ impl ServerHandler for ResourceAppsMcpServer {
request: Option<PaginatedRequestParams>,
_context: RequestContext<RoleServer>,
) -> Result<ListResourcesResult, rmcp::ErrorData> {
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<ReadResourceResult, rmcp::ErrorData> {
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(),
+1
View File
@@ -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;
+18
View File
@@ -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<ArcSwap<McpConnectionManager>>,
}
/// Opaque identity for the manager currently used by an MCP resource client.
#[derive(Clone)]
pub struct McpResourceClientCacheKey(Weak<McpConnectionManager>);
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.
+25 -12
View File
@@ -150,10 +150,11 @@ where
session_store: &ExtensionData,
thread_store: &ExtensionData,
) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
let Some(thread_state) = thread_store.get::<SkillsThreadState>() else {
return Vec::new();
};
if !self.providers.has_orchestrator_provider()
|| !thread_store
.get::<SkillsThreadState>()
.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::<McpResourceClient>(),
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<C> SkillsExtension<C> {
) -> 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<C> SkillsExtension<C> {
entry: &SkillCatalogEntry,
host_loaded_skills: Option<Arc<HostLoadedSkills>>,
session_store: &ExtensionData,
thread_state: &SkillsThreadState,
) -> Result<SkillReadResult, String> {
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::<McpResourceClient>(),
})
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::<McpResourceClient>(),
},
)
.await
.map_err(|err| err.message)
}
+134 -5
View File
@@ -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<SkillsExtensionConfig>,
selected_roots: Vec<SelectedCapabilityRoot>,
orchestrator_skills_enabled: bool,
orchestrator_catalog: OnceCell<SkillCatalog>,
orchestrator_cache: Mutex<Option<Arc<OrchestratorGenerationCache>>>,
}
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<Output = Result<SkillCatalog, SkillProviderError>> + 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<SkillReadResult> {
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<OrchestratorGenerationCache> {
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<McpResourceClientCacheKey>,
catalog: OnceCell<SkillCatalog>,
resources: Mutex<OrchestratorResourceCache>,
}
#[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<SkillReadCacheKey, SkillReadResult>,
contents_bytes: usize,
}
impl OrchestratorResourceCache {
fn get(&self, key: &SkillReadCacheKey) -> Option<SkillReadResult> {
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)]
+20 -19
View File
@@ -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<Arc<McpResourceClient>>,
thread_state: Arc<SkillsThreadState>,
) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
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<Arc<McpResourceClient>>,
thread_state: Arc<SkillsThreadState>,
}
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
}
}
}
}
+11 -8
View File
@@ -75,14 +75,17 @@ impl ToolExecutor<ToolCall> 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!(