1use std::time::Duration;
2
3use serde::{Deserialize, Serialize};
4
5use crate::control_plane_client::{ControlPlaneClient, ControlPlaneEndpoint, ControlPlaneError};
6
7use super::model::{FleetNodeHealth, FleetNodeSnapshot, FleetSnapshot, now_ms};
8use super::registry::FleetNodeConfig;
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
11pub struct FleetPollResult {
12 pub node_id: String,
13 pub label: String,
14 pub health: FleetNodeHealth,
15 #[serde(default, skip_serializing_if = "Option::is_none")]
16 pub snapshot: Option<FleetNodeSnapshot>,
17 #[serde(default, skip_serializing_if = "Option::is_none")]
18 pub error: Option<String>,
19}
20
21pub async fn poll_fleet_node(node_id: &str, node: &FleetNodeConfig) -> FleetPollResult {
22 let label = node.display_label(node_id);
23 let mut last_failure: Option<(FleetNodeHealth, String)> = None;
24
25 for endpoint_url in node.endpoints() {
26 let endpoint =
27 match ControlPlaneEndpoint::new(endpoint_url, node.admin_token_env.as_deref()) {
28 Ok(endpoint) => endpoint,
29 Err(err) => {
30 record_fallback_failure(
31 &mut last_failure,
32 FleetNodeHealth::Unsupported,
33 err.to_string(),
34 );
35 continue;
36 }
37 };
38 let client = match ControlPlaneClient::new(endpoint) {
39 Ok(client) => client,
40 Err(err) => {
41 record_fallback_failure(
42 &mut last_failure,
43 FleetNodeHealth::Unreachable,
44 err.to_string(),
45 );
46 continue;
47 }
48 };
49
50 match client.fleet_snapshot().await {
51 Ok(snapshot) => {
52 let mut node_snapshot = snapshot.nodes.into_iter().next().unwrap_or_else(|| {
53 empty_remote_node(
54 node_id,
55 &label,
56 FleetNodeHealth::Unsupported,
57 Some("remote fleet snapshot contained no nodes".to_string()),
58 )
59 });
60 node_snapshot.node_id = node_id.to_string();
61 node_snapshot.label = label.clone();
62 node_snapshot.kind = super::model::FleetNodeKind::Remote;
63 node_snapshot.health = FleetNodeHealth::Fresh;
64 node_snapshot.active_endpoint = Some(client.endpoint().admin_base_url.clone());
65 return FleetPollResult {
66 node_id: node_id.to_string(),
67 label,
68 health: FleetNodeHealth::Fresh,
69 snapshot: Some(node_snapshot),
70 error: None,
71 };
72 }
73 Err(err) => {
74 record_fallback_failure(
75 &mut last_failure,
76 health_from_control_plane_error(&err),
77 err.to_string(),
78 );
79 continue;
80 }
81 }
82 }
83
84 let (health, error) = last_failure.unwrap_or_else(|| {
85 (
86 FleetNodeHealth::Unreachable,
87 "no configured fleet endpoint could be reached".to_string(),
88 )
89 });
90 FleetPollResult {
91 node_id: node_id.to_string(),
92 label: label.clone(),
93 health,
94 snapshot: Some(empty_remote_node(
95 node_id,
96 &label,
97 health,
98 Some(error.clone()),
99 )),
100 error: Some(error),
101 }
102}
103
104pub fn stale_node_from_previous(
105 previous: &FleetNodeSnapshot,
106 health: FleetNodeHealth,
107 error: impl Into<String>,
108) -> FleetNodeSnapshot {
109 let now = now_ms();
110 let mut snapshot = previous.clone();
111 snapshot.health = health;
112 snapshot.snapshot_age_ms = Some(now.saturating_sub(previous.refreshed_at_ms));
113 snapshot.stale_since_ms = snapshot.stale_since_ms.or(Some(now));
114 snapshot.last_error = Some(error.into());
115 snapshot
116}
117
118pub fn node_snapshot_from_poll_result(
119 node_id: &str,
120 result: FleetPollResult,
121 previous: Option<&FleetNodeSnapshot>,
122) -> FleetNodeSnapshot {
123 let error = result
124 .error
125 .clone()
126 .unwrap_or_else(|| "fleet node refresh failed".to_string());
127 if result.health.is_current() {
128 return result.snapshot.unwrap_or_else(|| {
129 empty_remote_node(node_id, result.label.as_str(), result.health, Some(error))
130 });
131 }
132
133 previous
134 .map(|previous| stale_node_from_previous(previous, result.health, error.clone()))
135 .or(result.snapshot)
136 .unwrap_or_else(|| {
137 empty_remote_node(node_id, result.label.as_str(), result.health, Some(error))
138 })
139}
140
141pub fn health_from_control_plane_error(err: &ControlPlaneError) -> FleetNodeHealth {
142 match err {
143 ControlPlaneError::HttpStatus { status, .. } if *status == 401 || *status == 403 => {
144 FleetNodeHealth::AuthFailed
145 }
146 ControlPlaneError::HttpStatus { status, .. } if *status == 429 => {
147 FleetNodeHealth::RateLimited
148 }
149 ControlPlaneError::HttpStatus { status, .. } if *status == 404 || *status == 501 => {
150 FleetNodeHealth::Unsupported
151 }
152 ControlPlaneError::HttpStatus { .. } => FleetNodeHealth::Unreachable,
153 ControlPlaneError::Decode { .. } => FleetNodeHealth::ParseFailed,
154 ControlPlaneError::Transport { .. } => FleetNodeHealth::Unreachable,
155 }
156}
157
158fn empty_remote_node(
159 node_id: &str,
160 label: &str,
161 health: FleetNodeHealth,
162 error: Option<String>,
163) -> FleetNodeSnapshot {
164 FleetNodeSnapshot {
165 node_id: node_id.to_string(),
166 label: label.to_string(),
167 kind: super::model::FleetNodeKind::Remote,
168 health,
169 refreshed_at_ms: now_ms(),
170 stale_since_ms: None,
171 snapshot_age_ms: None,
172 active_endpoint: None,
173 last_error: error,
174 processes: super::model::FleetProcessSummary::default(),
175 topology: super::model::FleetTopology::default(),
176 work_units: Vec::new(),
177 }
178}
179
180fn record_fallback_failure(
181 last_failure: &mut Option<(FleetNodeHealth, String)>,
182 health: FleetNodeHealth,
183 error: String,
184) {
185 let should_replace = match last_failure {
186 Some((current_health, _)) => health_priority(health) >= health_priority(*current_health),
187 None => true,
188 };
189 if should_replace {
190 *last_failure = Some((health, error));
191 }
192}
193
194fn health_priority(health: FleetNodeHealth) -> u8 {
195 match health {
196 FleetNodeHealth::Fresh => 5,
197 FleetNodeHealth::AuthFailed | FleetNodeHealth::RateLimited => 4,
198 FleetNodeHealth::Unsupported => 3,
199 FleetNodeHealth::ParseFailed => 2,
200 FleetNodeHealth::Stale => 1,
201 FleetNodeHealth::Unreachable => 0,
202 }
203}
204
205pub async fn poll_fleet_registry(
206 service_name: &str,
207 registry: &super::registry::FleetRegistryConfig,
208) -> FleetSnapshot {
209 poll_fleet_registry_with_previous(service_name, registry, None).await
210}
211
212pub async fn poll_fleet_registry_with_previous(
213 service_name: &str,
214 registry: &super::registry::FleetRegistryConfig,
215 previous: Option<&FleetSnapshot>,
216) -> FleetSnapshot {
217 let mut nodes = Vec::new();
218 for (node_id, node) in registry.enabled_nodes() {
219 let result = poll_fleet_node(node_id, node).await;
220 let previous_node = previous.and_then(|snapshot| {
221 snapshot
222 .nodes
223 .iter()
224 .find(|snapshot| snapshot.node_id.as_str() == node_id.as_str())
225 });
226 nodes.push(node_snapshot_from_poll_result(
227 node_id,
228 result,
229 previous_node,
230 ));
231 }
232 super::merge::merge_fleet_nodes(service_name, nodes)
233}
234
235pub fn default_poll_interval() -> Duration {
236 Duration::from_millis(1_500)
237}
238
239#[cfg(test)]
240mod tests {
241 use super::*;
242 use crate::fleet::{FleetNodeKind, FleetProcessSummary, FleetTopology};
243 use axum::{Json, http::StatusCode, routing::get};
244 use std::collections::BTreeMap;
245 use std::net::SocketAddr;
246 use std::sync::Arc;
247 use std::sync::atomic::{AtomicUsize, Ordering};
248 use tokio::net::TcpListener;
249
250 async fn spawn_axum_server(app: axum::Router) -> (SocketAddr, tokio::task::JoinHandle<()>) {
251 let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
252 let addr = listener.local_addr().expect("local_addr");
253 let handle = tokio::spawn(async move {
254 axum::serve(listener, app).await.expect("serve");
255 });
256 (addr, handle)
257 }
258
259 #[test]
260 fn classifies_http_errors_for_fleet_degradation() {
261 let err = |status| ControlPlaneError::HttpStatus {
262 status,
263 body: String::new(),
264 };
265
266 assert_eq!(
267 health_from_control_plane_error(&err(403)),
268 FleetNodeHealth::AuthFailed
269 );
270 assert_eq!(
271 health_from_control_plane_error(&err(429)),
272 FleetNodeHealth::RateLimited
273 );
274 assert_eq!(
275 health_from_control_plane_error(&err(404)),
276 FleetNodeHealth::Unsupported
277 );
278 }
279
280 #[tokio::test]
281 async fn poll_fleet_node_falls_back_to_second_endpoint_after_http_failure() {
282 let first_hits = Arc::new(AtomicUsize::new(0));
283 let second_hits = Arc::new(AtomicUsize::new(0));
284
285 let first_count = first_hits.clone();
286 let first_app = axum::Router::new().route(
287 "/__codex_helper/api/v1/fleet/snapshot",
288 get(move || {
289 let first_count = first_count.clone();
290 async move {
291 first_count.fetch_add(1, Ordering::SeqCst);
292 (StatusCode::FORBIDDEN, "forbidden")
293 }
294 }),
295 );
296 let (first_addr, first_handle) = spawn_axum_server(first_app).await;
297
298 let second_count = second_hits.clone();
299 let success_snapshot = FleetSnapshot {
300 api_version: 1,
301 service_name: "codex".to_string(),
302 refreshed_at_ms: 42,
303 nodes: vec![FleetNodeSnapshot {
304 node_id: "placeholder".to_string(),
305 label: "placeholder".to_string(),
306 kind: FleetNodeKind::Remote,
307 health: FleetNodeHealth::Fresh,
308 refreshed_at_ms: 42,
309 stale_since_ms: None,
310 snapshot_age_ms: None,
311 active_endpoint: None,
312 last_error: None,
313 processes: FleetProcessSummary::default(),
314 topology: FleetTopology::default(),
315 work_units: Vec::new(),
316 }],
317 };
318 let second_snapshot = success_snapshot.clone();
319 let second_app = axum::Router::new().route(
320 "/__codex_helper/api/v1/fleet/snapshot",
321 get(move || {
322 let second_count = second_count.clone();
323 let second_snapshot = second_snapshot.clone();
324 async move {
325 second_count.fetch_add(1, Ordering::SeqCst);
326 Json(second_snapshot)
327 }
328 }),
329 );
330 let (second_addr, second_handle) = spawn_axum_server(second_app).await;
331
332 let first_url = format!("http://{}", first_addr);
333 let second_url = format!("http://{}", second_addr);
334 let node = FleetNodeConfig {
335 label: Some("remote-node".to_string()),
336 admin_url: Some(first_url.clone()),
337 admin_urls: vec![second_url.clone()],
338 admin_token_env: None,
339 enabled: true,
340 };
341
342 let result = poll_fleet_node("node-a", &node).await;
343 assert_eq!(result.health, FleetNodeHealth::Fresh);
344 assert!(result.error.is_none());
345
346 let snapshot = result.snapshot.expect("snapshot");
347 assert_eq!(snapshot.node_id, "node-a");
348 assert_eq!(snapshot.label, "remote-node");
349 assert_eq!(snapshot.health, FleetNodeHealth::Fresh);
350 assert_eq!(
351 snapshot.active_endpoint.as_deref(),
352 Some(second_url.as_str())
353 );
354 assert_eq!(first_hits.load(Ordering::SeqCst), 1);
355 assert_eq!(second_hits.load(Ordering::SeqCst), 1);
356
357 first_handle.abort();
358 second_handle.abort();
359 }
360
361 #[tokio::test]
362 async fn poll_fleet_registry_uses_previous_snapshot_when_node_refresh_fails() {
363 let hits = Arc::new(AtomicUsize::new(0));
364 let hit_count = hits.clone();
365 let app = axum::Router::new().route(
366 "/__codex_helper/api/v1/fleet/snapshot",
367 get(move || {
368 let hit_count = hit_count.clone();
369 async move {
370 hit_count.fetch_add(1, Ordering::SeqCst);
371 (StatusCode::FORBIDDEN, "forbidden")
372 }
373 }),
374 );
375 let (addr, handle) = spawn_axum_server(app).await;
376
377 let previous = FleetSnapshot {
378 api_version: 1,
379 service_name: "codex".to_string(),
380 refreshed_at_ms: 200,
381 nodes: vec![FleetNodeSnapshot {
382 node_id: "node-a".to_string(),
383 label: "remote-node".to_string(),
384 kind: FleetNodeKind::Remote,
385 health: FleetNodeHealth::Fresh,
386 refreshed_at_ms: 100,
387 stale_since_ms: None,
388 snapshot_age_ms: None,
389 active_endpoint: Some("http://previous.example".to_string()),
390 last_error: None,
391 processes: FleetProcessSummary {
392 scan_available: true,
393 codex_like_processes: 2,
394 error: None,
395 },
396 topology: FleetTopology::default(),
397 work_units: vec![super::super::model::FleetWorkUnit {
398 node_id: "node-a".to_string(),
399 id: "session:keep".to_string(),
400 parent_id: None,
401 kind: super::super::model::FleetWorkUnitKind::Root,
402 state: super::super::model::FleetWorkUnitState::Running,
403 evidence: super::super::model::FleetEvidence::high(
404 super::super::model::FleetEvidenceSource::RuntimeStatus,
405 ),
406 session_id: Some("keep".to_string()),
407 local_thread_id: Some("keep".to_string()),
408 task_name: Some("preserved".to_string()),
409 cwd: None,
410 model: None,
411 station_name: None,
412 provider_id: None,
413 last_status: None,
414 active_started_at_ms: None,
415 last_activity_ms: Some(100),
416 last_error: None,
417 usage: Default::default(),
418 }],
419 }],
420 };
421 let registry = super::super::registry::FleetRegistryConfig {
422 nodes: BTreeMap::from([(
423 "node-a".to_string(),
424 FleetNodeConfig {
425 label: Some("remote-node".to_string()),
426 admin_url: Some(format!("http://{}", addr)),
427 admin_urls: Vec::new(),
428 admin_token_env: None,
429 enabled: true,
430 },
431 )]),
432 };
433
434 let snapshot = poll_fleet_registry_with_previous("codex", ®istry, Some(&previous)).await;
435
436 assert_eq!(hits.load(Ordering::SeqCst), 1);
437 assert_eq!(snapshot.nodes.len(), 1);
438 let node = &snapshot.nodes[0];
439 assert_eq!(node.node_id, "node-a");
440 assert_eq!(node.health, FleetNodeHealth::AuthFailed);
441 assert_eq!(node.work_units.len(), 1);
442 assert_eq!(node.work_units[0].task_name.as_deref(), Some("preserved"));
443 assert_eq!(
444 node.active_endpoint.as_deref(),
445 Some("http://previous.example")
446 );
447 assert!(node.snapshot_age_ms.is_some());
448 assert!(node.stale_since_ms.is_some());
449 assert!(
450 node.last_error
451 .as_deref()
452 .is_some_and(|err| err.contains("403"))
453 );
454
455 handle.abort();
456 }
457}