Skip to main content

codex_exec_server/
resolved_capability.rs

1use std::collections::HashMap;
2use std::fmt;
3use std::sync::Arc;
4
5use codex_protocol::capabilities::CapabilityRootLocation;
6use codex_protocol::capabilities::SelectedCapabilityRoot;
7
8use crate::Environment;
9use crate::EnvironmentManager;
10
11/// A selected capability root paired with its currently ready environment handle.
12///
13/// Environment IDs have stable identity and contents. This process-local value must not be
14/// persisted: it only keeps the current connection handle alive while one model step uses the
15/// stable environment.
16#[derive(Clone)]
17pub struct ResolvedSelectedCapabilityRoot {
18    selected_root: SelectedCapabilityRoot,
19    environment: Arc<Environment>,
20}
21
22/// A passive view of selected capability roots and unavailable environments.
23#[derive(Clone, Debug, Default, PartialEq, Eq)]
24pub struct SelectedCapabilityRootsStatus {
25    /// Selected roots whose environments are ready.
26    pub ready_roots: Vec<SelectedCapabilityRoot>,
27    /// Missing environments and terminal connection failures.
28    pub warnings: Vec<String>,
29}
30
31impl ResolvedSelectedCapabilityRoot {
32    pub fn selected_root(&self) -> &SelectedCapabilityRoot {
33        &self.selected_root
34    }
35
36    pub fn environment(&self) -> &Arc<Environment> {
37        &self.environment
38    }
39}
40
41impl EnvironmentManager {
42    /// Inspects selected roots without starting or waiting for an environment.
43    ///
44    /// Starting or recovering environments are omitted. Missing environments and terminal
45    /// connection failures are returned as warnings so read-only catalog clients can distinguish
46    /// them from an empty catalog.
47    ///
48    /// Environment IDs are stable identities, so callers can safely resolve a returned root by
49    /// ID when they read it.
50    pub fn inspect_selected_capability_roots(
51        &self,
52        selected_roots: &[SelectedCapabilityRoot],
53    ) -> SelectedCapabilityRootsStatus {
54        let (candidates, mut warnings) = {
55            let environments = self
56                .environments
57                .read()
58                .unwrap_or_else(std::sync::PoisonError::into_inner);
59            let mut candidates = Vec::with_capacity(selected_roots.len());
60            let mut warnings = Vec::new();
61            for selected_root in selected_roots {
62                let CapabilityRootLocation::Environment { environment_id, .. } =
63                    &selected_root.location;
64                let Some(environment) = environments.get(environment_id) else {
65                    warnings.push(format!(
66                        "selected capability root `{}` references unavailable environment `{environment_id}`",
67                        selected_root.id
68                    ));
69                    continue;
70                };
71                candidates.push((selected_root.clone(), Arc::clone(environment)));
72            }
73            (candidates, warnings)
74        };
75        let mut readiness = HashMap::new();
76        for (selected_root, environment) in &candidates {
77            let CapabilityRootLocation::Environment { environment_id, .. } =
78                &selected_root.location;
79            if readiness.contains_key(environment_id) {
80                continue;
81            }
82            let ready = match environment.readiness_result() {
83                Some(Ok(())) => true,
84                Some(Err(error)) => {
85                    warnings.push(format!(
86                        "selected capability environment `{environment_id}` is unavailable: {error}"
87                    ));
88                    false
89                }
90                None => false,
91            };
92            readiness.insert(environment_id.clone(), ready);
93        }
94
95        let ready_roots = candidates
96            .into_iter()
97            .filter(|(selected_root, _)| {
98                let CapabilityRootLocation::Environment { environment_id, .. } =
99                    &selected_root.location;
100                readiness.get(environment_id).copied().unwrap_or(false)
101            })
102            .map(|(selected_root, _)| selected_root)
103            .collect();
104        SelectedCapabilityRootsStatus {
105            ready_roots,
106            warnings,
107        }
108    }
109
110    /// Resolves selected roots whose stable environments are ready for the current model step.
111    ///
112    /// Environment identity comes from the selected root's stable environment ID. A ready
113    /// environment captured for the step carries its exact process-local handle so readiness and
114    /// execution cannot come from different registry snapshots. Missing, starting, or failed
115    /// environments are omitted. A lazy environment is started for a later step.
116    #[tracing::instrument(name = "capability_roots.resolve", skip_all)]
117    pub async fn resolve_selected_capability_roots(
118        &self,
119        selected_roots: &[SelectedCapabilityRoot],
120        captured_environments: &HashMap<String, Option<Arc<Environment>>>,
121    ) -> Vec<ResolvedSelectedCapabilityRoot> {
122        let candidates = {
123            let environments = self
124                .environments
125                .read()
126                .unwrap_or_else(std::sync::PoisonError::into_inner);
127            selected_roots
128                .iter()
129                .filter_map(|selected_root| {
130                    let CapabilityRootLocation::Environment { environment_id, .. } =
131                        &selected_root.location;
132                    let (environment, already_ready) =
133                        match captured_environments.get(environment_id) {
134                            Some(Some(environment)) => (Arc::clone(environment), true),
135                            Some(None) => return None,
136                            None => (Arc::clone(environments.get(environment_id)?), false),
137                        };
138                    Some((
139                        ResolvedSelectedCapabilityRoot {
140                            selected_root: selected_root.clone(),
141                            environment,
142                        },
143                        already_ready,
144                    ))
145                })
146                .collect::<Vec<_>>()
147        };
148
149        let mut readiness = HashMap::new();
150        for (candidate, already_ready) in &candidates {
151            let CapabilityRootLocation::Environment { environment_id, .. } =
152                &candidate.selected_root().location;
153            if readiness.contains_key(environment_id) {
154                continue;
155            }
156            let environment = candidate.environment();
157            let ready = if *already_ready {
158                true
159            } else if environment.startup_finished() {
160                environment.wait_until_ready().await.is_ok()
161            } else {
162                Environment::start_connecting_for_use(environment);
163                false
164            };
165            readiness.insert(environment_id.clone(), ready);
166        }
167
168        candidates
169            .into_iter()
170            .map(|(candidate, _)| candidate)
171            .filter(|candidate| {
172                let CapabilityRootLocation::Environment { environment_id, .. } =
173                    &candidate.selected_root().location;
174                readiness.get(environment_id).copied().unwrap_or(false)
175            })
176            .collect()
177    }
178}
179
180impl fmt::Debug for ResolvedSelectedCapabilityRoot {
181    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
182        formatter
183            .debug_struct("ResolvedSelectedCapabilityRoot")
184            .field("selected_root", &self.selected_root)
185            .finish_non_exhaustive()
186    }
187}