agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

skills: cache orchestrator resources per thread (#28336)

Browse Source 
## 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 20:20:19 +02:00
committed by GitHub
parent ee40dddbf6
commit 0afe559318
7 changed files with 340 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/ext/skills/src/extension.rs
+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)
}
codex-rs/ext/skills/src/state.rs
+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)]
codex-rs/ext/skills/src/tools/mod.rs
+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
}
}
}
}
codex-rs/ext/skills/src/tools/read.rs
+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!(