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

Add files using upload-large-folder tool

ea39c0e verified about 1 month ago

10.3 kB

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()
    }
}

	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()
	}
	}