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;pub(crate) struct SkillsThreadState { config: Mutex<SkillsExtensionConfig>, selected_roots: Vec<SelectedCapabilityRoot>, orchestrator_skills_enabled: bool, orchestrator_cache: Mutex<Option<Arc<OrchestratorGenerationCache>>>,}impl SkillsThreadState { pub(crate) fn new( config: SkillsExtensionConfig, selected_roots: Vec<SelectedCapabilityRoot>, orchestrator_skills_enabled: bool, ) -> Self { Self { config: Mutex::new(config), selected_roots, orchestrator_skills_enabled, orchestrator_cache: Mutex::new(None), } } pub(crate) fn config(&self) -> SkillsExtensionConfig { self.config .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .clone() } pub(crate) fn set_config(&self, config: SkillsExtensionConfig) { *self .config .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) = config; } pub(crate) fn selected_roots(&self) -> &[SelectedCapabilityRoot] { &self.selected_roots } pub(crate) fn orchestrator_skills_enabled(&self) -> bool { self.orchestrator_skills_enabled } pub(crate) async fn orchestrator_catalog_snapshot( &self, mcp_resources: Option<&McpResourceClient>, initialize: impl Future<Output = Result<SkillCatalog, SkillProviderError>> + Send, ) -> SkillCatalog { self.orchestrator_cache(mcp_resources) .catalog .get_or_init(|| async { initialize.await.unwrap_or_else(|err| SkillCatalog { warnings: vec![err.message], ..Default::default() }) }) .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)]pub(crate) struct SkillsTurnState { pub(crate) catalog: SkillCatalog, pub(crate) selected_entries: Vec<SkillCatalogEntry>, pub(crate) warnings: Vec<String>, pub(crate) main_prompts_injected: bool,}