1use std::collections::{HashMap, HashSet};
2
3use anyhow::{bail, Result};
4use serde::{Deserialize, Serialize};
5
6use crate::{
7 config::{HarnessConfig, Profile},
8 transport::StreamTransportBackend,
9};
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
12#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
13#[serde(rename_all = "kebab-case")]
14pub enum ModuleState {
15 Planned,
16 Starting,
17 Running,
18 Stopped,
19 Failed,
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
23#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
24#[serde(rename_all = "kebab-case")]
25pub enum StreamDirection {
26 Input,
27 Output,
28}
29
30#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
31#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
32pub struct StreamDescriptor {
33 pub name: String,
34 pub direction: StreamDirection,
35 pub message_type: String,
36 pub transport: String,
37}
38
39#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
40#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
41pub struct ModuleHealth {
42 pub ok: bool,
43 pub state: ModuleState,
44 pub message: Option<String>,
45}
46
47#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
48#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
49pub struct ModuleInfo {
50 pub id: usize,
51 pub name: String,
52 pub module_type: String,
53 pub state: ModuleState,
54 pub health: ModuleHealth,
55 pub dependencies: Vec<String>,
56 pub inputs: Vec<StreamDescriptor>,
57 pub outputs: Vec<StreamDescriptor>,
58 pub capabilities: Vec<String>,
59 pub physical: bool,
60 pub required: bool,
61}
62
63#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
64#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
65pub struct ModuleGraph {
66 pub modules: Vec<ModuleInfo>,
67}
68
69impl ModuleGraph {
70 pub fn new(modules: Vec<ModuleInfo>) -> Result<Self> {
71 let mut names = HashSet::with_capacity(modules.len());
72 let mut ids = HashSet::with_capacity(modules.len());
73 for module in &modules {
74 if !ids.insert(module.id) {
75 bail!("duplicate module id '{}'", module.id);
76 }
77 if module.name.trim().is_empty() {
78 bail!("module name cannot be empty");
79 }
80 if !names.insert(module.name.clone()) {
81 bail!("duplicate module name '{}'", module.name);
82 }
83 }
84 for module in &modules {
85 for dependency in &module.dependencies {
86 if dependency == &module.name {
87 bail!("module '{}' cannot depend on itself", module.name);
88 }
89 if !names.contains(dependency) {
90 bail!(
91 "module '{}' depends on unknown module '{}'",
92 module.name,
93 dependency
94 );
95 }
96 }
97 }
98 Ok(Self { modules })
99 }
100
101 pub fn modules(&self) -> &[ModuleInfo] {
102 &self.modules
103 }
104
105 pub fn stream_count(&self) -> usize {
106 self.modules
107 .iter()
108 .map(|module| module.inputs.len() + module.outputs.len())
109 .sum()
110 }
111
112 pub fn to_dot(&self) -> String {
113 let mut dot = String::from("digraph leash_module_graph {\n");
114 dot.push_str(" graph [rankdir=LR];\n");
115 dot.push_str(" node [shape=box];\n");
116
117 for module in &self.modules {
118 let risk = if module.physical {
119 "physical"
120 } else {
121 "non-physical"
122 };
123 dot.push_str(&format!(
124 " {} [label=\"{}\\n{}\\n{}\", color=\"{}\"];\n",
125 dot_id(module.id),
126 escape_dot(&module.name),
127 escape_dot(&module.module_type),
128 risk,
129 if module.physical { "red" } else { "black" }
130 ));
131 }
132
133 for from in &self.modules {
134 for output in &from.outputs {
135 for to in &self.modules {
136 if from.id == to.id {
137 continue;
138 }
139 for input in &to.inputs {
140 if output.name == input.name && output.message_type == input.message_type {
141 dot.push_str(&format!(
142 " {} -> {} [label=\"{}:{} via {}\"];\n",
143 dot_id(from.id),
144 dot_id(to.id),
145 escape_dot(&output.name),
146 escape_dot(&output.message_type),
147 escape_dot(&output.transport)
148 ));
149 }
150 }
151 }
152 }
153 }
154
155 dot.push_str("}\n");
156 dot
157 }
158}
159
160#[derive(Debug, Clone)]
161pub struct ModuleCoordinator {
162 modules: Vec<ModuleInfo>,
163 start_failures: HashMap<String, String>,
164}
165
166impl ModuleCoordinator {
167 pub fn new(graph: ModuleGraph) -> Self {
168 Self {
169 modules: graph.modules,
170 start_failures: HashMap::new(),
171 }
172 }
173
174 pub fn with_start_failure(mut self, module_name: &str, message: impl Into<String>) -> Self {
175 self.start_failures
176 .insert(module_name.to_string(), message.into());
177 self
178 }
179
180 pub fn start(&mut self) -> Result<()> {
181 let order = self.dependency_order()?;
182 let mut started = Vec::new();
183
184 for index in order {
185 self.set_health(index, ModuleState::Starting, true, "starting");
186 let module_name = self.modules[index].name.clone();
187 if let Some(message) = self.start_failures.get(&module_name).cloned() {
188 self.set_health(index, ModuleState::Failed, false, message.clone());
189 self.stop_started_after_failure(&started);
190 bail!("module '{}' failed to start: {}", module_name, message);
191 }
192 self.set_health(index, ModuleState::Running, true, "running");
193 started.push(index);
194 }
195
196 Ok(())
197 }
198
199 pub fn stop(&mut self) {
200 let Ok(order) = self.dependency_order() else {
201 for index in (0..self.modules.len()).rev() {
202 if matches!(
203 self.modules[index].state,
204 ModuleState::Starting | ModuleState::Running
205 ) {
206 self.set_health(index, ModuleState::Stopped, true, "stopped");
207 }
208 }
209 return;
210 };
211 for index in order.into_iter().rev() {
212 if matches!(
213 self.modules[index].state,
214 ModuleState::Starting | ModuleState::Running
215 ) {
216 self.set_health(index, ModuleState::Stopped, true, "stopped");
217 }
218 }
219 }
220
221 pub fn graph(&self) -> ModuleGraph {
222 ModuleGraph::new(self.modules.clone()).expect("coordinator preserves valid module graph")
223 }
224
225 pub fn modules(&self) -> &[ModuleInfo] {
226 &self.modules
227 }
228
229 pub fn is_healthy(&self) -> bool {
230 self.modules
231 .iter()
232 .all(|module| module.health.ok && module.state == ModuleState::Running)
233 }
234
235 fn dependency_order(&self) -> Result<Vec<usize>> {
236 #[derive(Clone, Copy, PartialEq, Eq)]
237 enum Visit {
238 New,
239 Active,
240 Done,
241 }
242
243 let by_name: HashMap<&str, usize> = self
244 .modules
245 .iter()
246 .enumerate()
247 .map(|(index, module)| (module.name.as_str(), index))
248 .collect();
249 let mut visits = vec![Visit::New; self.modules.len()];
250 let mut order = Vec::with_capacity(self.modules.len());
251
252 fn visit(
253 index: usize,
254 modules: &[ModuleInfo],
255 by_name: &HashMap<&str, usize>,
256 visits: &mut [Visit],
257 order: &mut Vec<usize>,
258 ) -> Result<()> {
259 match visits[index] {
260 Visit::Done => return Ok(()),
261 Visit::Active => bail!("module dependency cycle at '{}'", modules[index].name),
262 Visit::New => {}
263 }
264
265 visits[index] = Visit::Active;
266 for dependency in &modules[index].dependencies {
267 let Some(dependency_index) = by_name.get(dependency.as_str()).copied() else {
268 bail!(
269 "module '{}' depends on unknown module '{}'",
270 modules[index].name,
271 dependency
272 );
273 };
274 visit(dependency_index, modules, by_name, visits, order)?;
275 }
276 visits[index] = Visit::Done;
277 order.push(index);
278 Ok(())
279 }
280
281 for index in 0..self.modules.len() {
282 visit(index, &self.modules, &by_name, &mut visits, &mut order)?;
283 }
284 Ok(order)
285 }
286
287 fn stop_started_after_failure(&mut self, started: &[usize]) {
288 for index in started.iter().copied().rev() {
289 self.set_health(
290 index,
291 ModuleState::Stopped,
292 true,
293 "stopped after startup failure",
294 );
295 }
296 }
297
298 fn set_health(
299 &mut self,
300 index: usize,
301 state: ModuleState,
302 ok: bool,
303 message: impl Into<String>,
304 ) {
305 let module = &mut self.modules[index];
306 module.state = state;
307 module.health = ModuleHealth {
308 ok,
309 state,
310 message: Some(message.into()),
311 };
312 }
313}
314
315pub fn default_module_graph(config: &HarnessConfig, capabilities: Vec<String>) -> ModuleGraph {
316 let driver_name = match config.profile {
317 Profile::Sim => "sim-driver",
318 Profile::Replay => "replay-driver",
319 Profile::WaveshareUgv => "waveshare-ugv-driver",
320 };
321 let transport = config.stream_transport;
322 ModuleGraph::new(vec![
323 ModuleInfo {
324 id: 0,
325 name: "harness-runtime".to_string(),
326 module_type: "runtime".to_string(),
327 state: ModuleState::Planned,
328 health: planned_health(),
329 dependencies: Vec::new(),
330 inputs: vec![stream(
331 "capability_request",
332 StreamDirection::Input,
333 "json",
334 transport,
335 )],
336 outputs: vec![
337 stream("health", StreamDirection::Output, "Health", transport),
338 stream(
339 "capabilities",
340 StreamDirection::Output,
341 "Capabilities",
342 transport,
343 ),
344 ],
345 capabilities,
346 physical: false,
347 required: true,
348 },
349 ModuleInfo {
350 id: 1,
351 name: driver_name.to_string(),
352 module_type: "driver".to_string(),
353 state: ModuleState::Planned,
354 health: planned_health(),
355 dependencies: vec!["harness-runtime".to_string()],
356 inputs: vec![stream(
357 "drive_command",
358 StreamDirection::Input,
359 "DriveReq",
360 transport,
361 )],
362 outputs: vec![stream(
363 "odometry",
364 StreamDirection::Output,
365 "OdometryStatus",
366 transport,
367 )],
368 capabilities: vec!["drive".to_string(), "stop".to_string(), "estop".to_string()],
369 physical: config.profile.is_physical(),
370 required: true,
371 },
372 ModuleInfo {
373 id: 2,
374 name: "telemetry".to_string(),
375 module_type: "telemetry".to_string(),
376 state: ModuleState::Planned,
377 health: planned_health(),
378 dependencies: vec![driver_name.to_string()],
379 inputs: vec![stream(
380 "odometry",
381 StreamDirection::Input,
382 "OdometryStatus",
383 transport,
384 )],
385 outputs: vec![
386 stream(
387 "telemetry",
388 StreamDirection::Output,
389 "TelemetryFrame",
390 transport,
391 ),
392 stream(
393 "sensors",
394 StreamDirection::Output,
395 "SensorSnapshot",
396 transport,
397 ),
398 ],
399 capabilities: vec!["observe".to_string(), "capture".to_string()],
400 physical: false,
401 required: true,
402 },
403 ])
404 .expect("default module graph uses unique non-empty names")
405}
406
407fn dot_id(id: usize) -> String {
408 format!("module_{id}")
409}
410
411fn escape_dot(value: &str) -> String {
412 value.replace('\\', "\\\\").replace('"', "\\\"")
413}
414
415fn stream(
416 name: &str,
417 direction: StreamDirection,
418 message_type: &str,
419 transport: StreamTransportBackend,
420) -> StreamDescriptor {
421 StreamDescriptor {
422 name: name.to_string(),
423 direction,
424 message_type: message_type.to_string(),
425 transport: transport.as_str().to_string(),
426 }
427}
428
429fn planned_health() -> ModuleHealth {
430 ModuleHealth {
431 ok: true,
432 state: ModuleState::Planned,
433 message: Some("planned".to_string()),
434 }
435}
436
437#[cfg(test)]
438mod tests {
439 use super::*;
440
441 fn module(name: &str) -> ModuleInfo {
442 ModuleInfo {
443 id: 0,
444 name: name.to_string(),
445 module_type: "test".to_string(),
446 state: ModuleState::Planned,
447 health: planned_health(),
448 dependencies: Vec::new(),
449 inputs: Vec::new(),
450 outputs: Vec::new(),
451 capabilities: Vec::new(),
452 physical: false,
453 required: true,
454 }
455 }
456
457 fn module_with_id(id: usize, name: &str) -> ModuleInfo {
458 ModuleInfo { id, ..module(name) }
459 }
460
461 #[test]
462 fn rejects_duplicate_module_names() {
463 let err = ModuleGraph::new(vec![module_with_id(0, "a"), module_with_id(1, "a")])
464 .unwrap_err()
465 .to_string();
466 assert!(err.contains("duplicate module name"));
467 }
468
469 #[test]
470 fn rejects_duplicate_module_ids() {
471 let err = ModuleGraph::new(vec![module_with_id(0, "a"), module_with_id(0, "b")])
472 .unwrap_err()
473 .to_string();
474 assert!(err.contains("duplicate module id"));
475 }
476
477 #[test]
478 fn counts_declared_streams() {
479 let graph = default_module_graph(&HarnessConfig::default(), vec!["health".to_string()]);
480 assert_eq!(graph.modules().len(), 3);
481 assert_eq!(graph.stream_count(), 8);
482 }
483
484 #[test]
485 fn exports_dot_with_edges_and_risk_metadata() {
486 let graph = default_module_graph(
487 &HarnessConfig {
488 profile: Profile::WaveshareUgv,
489 ..HarnessConfig::default()
490 },
491 vec!["drive".to_string()],
492 );
493 let dot = graph.to_dot();
494 assert!(dot.contains("digraph leash_module_graph"));
495 assert!(dot.contains("waveshare-ugv-driver"));
496 assert!(dot.contains("physical"));
497 assert!(dot.contains("module_1 -> module_2"));
498 assert!(dot.contains("odometry:OdometryStatus via local-pubsub"));
499 }
500
501 #[test]
502 fn graph_stream_transport_can_switch_backends() {
503 let graph = default_module_graph(
504 &HarnessConfig {
505 stream_transport: StreamTransportBackend::Memory,
506 ..HarnessConfig::default()
507 },
508 vec!["observe".to_string()],
509 );
510
511 assert!(graph
512 .modules()
513 .iter()
514 .flat_map(|module| module.inputs.iter().chain(module.outputs.iter()))
515 .all(|stream| stream.transport == "memory"));
516 }
517
518 #[test]
519 fn coordinator_starts_in_dependency_order() {
520 let graph = default_module_graph(&HarnessConfig::default(), vec!["health".to_string()]);
521 let mut coordinator = ModuleCoordinator::new(graph);
522 coordinator.start().unwrap();
523
524 let modules = coordinator.modules();
525 assert_eq!(modules[0].state, ModuleState::Running);
526 assert_eq!(modules[1].state, ModuleState::Running);
527 assert_eq!(modules[2].state, ModuleState::Running);
528 assert!(coordinator.is_healthy());
529 }
530
531 #[test]
532 fn coordinator_cleans_up_started_modules_after_start_failure() {
533 let graph = default_module_graph(&HarnessConfig::default(), vec!["drive".to_string()]);
534 let mut coordinator =
535 ModuleCoordinator::new(graph).with_start_failure("telemetry", "loop refused");
536
537 let err = coordinator.start().unwrap_err().to_string();
538 assert!(err.contains("telemetry"));
539
540 let modules = coordinator.modules();
541 assert_eq!(modules[0].state, ModuleState::Stopped);
542 assert_eq!(modules[1].state, ModuleState::Stopped);
543 assert_eq!(modules[2].state, ModuleState::Failed);
544 assert!(!coordinator.is_healthy());
545 }
546
547 #[test]
548 fn coordinator_stop_is_idempotent_during_partial_startup() {
549 let mut runtime = module_with_id(0, "runtime");
550 runtime.state = ModuleState::Running;
551 runtime.health = ModuleHealth {
552 ok: true,
553 state: ModuleState::Running,
554 message: Some("running".to_string()),
555 };
556 let mut driver = module_with_id(1, "driver");
557 driver.dependencies = vec!["runtime".to_string()];
558 driver.state = ModuleState::Starting;
559 driver.health = ModuleHealth {
560 ok: true,
561 state: ModuleState::Starting,
562 message: Some("starting".to_string()),
563 };
564 let graph = ModuleGraph::new(vec![runtime, driver]).unwrap();
565 let mut coordinator = ModuleCoordinator::new(graph);
566
567 coordinator.stop();
568 coordinator.stop();
569
570 assert!(coordinator
571 .modules()
572 .iter()
573 .all(|module| module.state == ModuleState::Stopped));
574 }
575
576 #[test]
577 fn failed_optional_module_does_not_leave_runtime_starting() {
578 let mut optional = module_with_id(1, "optional-viewer");
579 optional.required = false;
580 optional.dependencies = vec!["runtime".to_string()];
581 let graph = ModuleGraph::new(vec![module_with_id(0, "runtime"), optional]).unwrap();
582 let mut coordinator =
583 ModuleCoordinator::new(graph).with_start_failure("optional-viewer", "not available");
584
585 let _ = coordinator.start().unwrap_err();
586
587 assert!(coordinator
588 .modules()
589 .iter()
590 .all(|module| { matches!(module.state, ModuleState::Stopped | ModuleState::Failed) }));
591 }
592}