codex_exec_server/
resolved_capability.rs1use 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#[derive(Clone)]
17pub struct ResolvedSelectedCapabilityRoot {
18 selected_root: SelectedCapabilityRoot,
19 environment: Arc<Environment>,
20}
21
22#[derive(Clone, Debug, Default, PartialEq, Eq)]
24pub struct SelectedCapabilityRootsStatus {
25 pub ready_roots: Vec<SelectedCapabilityRoot>,
27 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 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 #[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}