codex/codex-rs/exec-server/src/capability_discovery_cache.rs

use std::collections::BTreeMap;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;

use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use tokio::sync::Mutex;

use crate::CapabilityRootDiscoverRequest;
use crate::CapabilityRootDiscovery;
use crate::CapabilityRootsDiscoverParams;
use crate::EnvironmentManager;
use crate::ExecutorCapabilityDiscoverySnapshot;
use crate::FileSystemSandboxContext;

/// Thread-scoped cache shared by capability consumers using the high-level executor API.
///
/// A single miss batches every requested root by environment. Successful discoveries and
/// permanent failures remain cached by root and sandbox; transient failures are retried on the
/// next request. Recovery is reported so dependent MCP projections can be invalidated.
pub struct ExecutorCapabilityDiscoveryCache {
    environment_manager: Arc<EnvironmentManager>,
    entries: Mutex<Vec<CachedRoot>>,
    recovered_discovery: AtomicBool,
}

impl std::fmt::Debug for ExecutorCapabilityDiscoveryCache {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("ExecutorCapabilityDiscoveryCache")
            .finish_non_exhaustive()
    }
}

struct CachedRoot {
    selected_root: SelectedCapabilityRoot,
    sandbox: Option<FileSystemSandboxContext>,
    result: Result<Arc<CapabilityRootDiscovery>, String>,
    // Preserve transport classification after the public snapshot reduces errors to strings.
    retryable: bool,
}

impl ExecutorCapabilityDiscoveryCache {
    pub fn new(environment_manager: Arc<EnvironmentManager>) -> Self {
        Self {
            environment_manager,
            entries: Mutex::new(Vec::new()),
            recovered_discovery: AtomicBool::new(false),
        }
    }

    /// Reports whether a previously failed root has recovered since the last observation.
    pub fn take_recovered_discovery(&self) -> bool {
        self.recovered_discovery.swap(false, Ordering::AcqRel)
    }

    /// Returns discoveries in the same order as `selected_roots`.
    #[tracing::instrument(
        name = "capability_roots.discovery_cache.resolve",
        skip_all,
        fields(root_count = selected_roots.len())
    )]
    pub async fn discover(
        &self,
        selected_roots: &[SelectedCapabilityRoot],
        sandbox_contexts: &HashMap<String, FileSystemSandboxContext>,
    ) -> Vec<Result<Arc<CapabilityRootDiscovery>, String>> {
        let missing = {
            let entries = self.entries.lock().await;
            selected_roots
                .iter()
                .filter(|selected_root| {
                    let CapabilityRootLocation::Environment { environment_id, .. } =
                        &selected_root.location;
                    let sandbox = sandbox_contexts.get(environment_id);
                    !entries.iter().any(|cached| {
                        cached.selected_root == **selected_root
                            && cached.sandbox.as_ref() == sandbox
                            && !cached.retryable
                    })
                })
                .cloned()
                .collect::<Vec<_>>()
        };
        let discovered = self.discover_missing(missing, sandbox_contexts).await;
        let mut entries = self.entries.lock().await;
        for discovered_root in discovered {
            if let Some(cached) = entries
                .iter_mut()
                .find(|cached| cached.selected_root == discovered_root.selected_root)
            {
                if cached.sandbox != discovered_root.sandbox || cached.result.is_err() {
                    if cached.result.is_err() && discovered_root.result.is_ok() {
                        self.recovered_discovery.store(true, Ordering::Release);
                    }
                    *cached = discovered_root;
                }
            } else {
                entries.push(discovered_root);
            }
        }
        selected_roots
            .iter()
            .map(|selected_root| {
                let CapabilityRootLocation::Environment { environment_id, .. } =
                    &selected_root.location;
                let sandbox = sandbox_contexts.get(environment_id);
                match entries.iter().find(|cached| {
                    cached.selected_root == *selected_root && cached.sandbox.as_ref() == sandbox
                }) {
                    Some(cached) => cached.result.clone(),
                    None => Err(format!(
                        "selected capability root `{}` was not discovered",
                        selected_root.id
                    )),
                }
            })
            .collect()
    }

    /// Resolves the selected roots once and freezes their results for one model step.
    pub async fn snapshot(
        &self,
        selected_roots: &[SelectedCapabilityRoot],
        sandbox_contexts: &HashMap<String, FileSystemSandboxContext>,
    ) -> ExecutorCapabilityDiscoverySnapshot {
        ExecutorCapabilityDiscoverySnapshot::new(
            selected_roots,
            self.discover(selected_roots, sandbox_contexts).await,
            sandbox_contexts.clone(),
        )
    }

    async fn discover_missing(
        &self,
        missing: Vec<SelectedCapabilityRoot>,
        sandbox_contexts: &HashMap<String, FileSystemSandboxContext>,
    ) -> Vec<CachedRoot> {
        let mut grouped = BTreeMap::<String, Vec<SelectedCapabilityRoot>>::new();
        for selected_root in missing {
            let CapabilityRootLocation::Environment { environment_id, .. } =
                &selected_root.location;
            grouped
                .entry(environment_id.clone())
                .or_default()
                .push(selected_root);
        }

        let batches = grouped.into_iter().flat_map(|(environment_id, roots)| {
            roots
                .chunks(crate::capability_discovery::MAX_ROOTS_PER_REQUEST)
                .map(|batch| (environment_id.clone(), batch.to_vec()))
                .collect::<Vec<_>>()
        });
        let discoveries =
            futures::future::join_all(batches.map(|(environment_id, selected_roots)| async move {
                let sandbox = sandbox_contexts.get(&environment_id).cloned();
                let Some(environment) = self.environment_manager.get_environment(&environment_id)
                else {
                    let error = format!("environment `{environment_id}` is unavailable");
                    return selected_roots
                        .into_iter()
                        .map(|selected_root| CachedRoot {
                            selected_root,
                            sandbox: sandbox.clone(),
                            result: Err(error.clone()),
                            retryable: true,
                        })
                        .collect::<Vec<_>>();
                };
                let params = CapabilityRootsDiscoverParams {
                    roots: selected_roots
                        .iter()
                        .map(|selected_root| {
                            let CapabilityRootLocation::Environment { path, .. } =
                                &selected_root.location;
                            CapabilityRootDiscoverRequest {
                                id: selected_root.id.clone(),
                                path: path.clone(),
                                sandbox: sandbox.clone(),
                            }
                        })
                        .collect(),
                };
                let response = match environment.discover_capability_roots(params).await {
                    Ok(response) => response,
                    Err(error) => {
                        let retryable = crate::client::is_retryable_recovery_error(&error);
                        let error = error.to_string();
                        return selected_roots
                            .into_iter()
                            .map(|selected_root| CachedRoot {
                                selected_root,
                                sandbox: sandbox.clone(),
                                result: Err(error.clone()),
                                retryable,
                            })
                            .collect();
                    }
                };
                if response.roots.len() != selected_roots.len() {
                    let error = format!(
                        "exec-server returned {} capability roots for {} requests",
                        response.roots.len(),
                        selected_roots.len()
                    );
                    return selected_roots
                        .into_iter()
                        .map(|selected_root| CachedRoot {
                            selected_root,
                            sandbox: sandbox.clone(),
                            result: Err(error.clone()),
                            retryable: false,
                        })
                        .collect();
                }
                selected_roots
                    .into_iter()
                    .zip(response.roots)
                    .map(|(selected_root, discovery)| {
                        let CapabilityRootLocation::Environment { path, .. } =
                            &selected_root.location;
                        let result = if discovery.id == selected_root.id && discovery.path == *path
                        {
                            Ok(Arc::new(discovery))
                        } else {
                            Err(format!(
                                "exec-server returned mismatched capability root `{}` at {}",
                                discovery.id, discovery.path
                            ))
                        };
                        CachedRoot {
                            selected_root,
                            sandbox: sandbox.clone(),
                            result,
                            retryable: false,
                        }
                    })
                    .collect()
            }))
            .await;
        discoveries.into_iter().flatten().collect()
    }
}