Skip to main content

rustant_core/nodes/
mod.rs

1//! # Node System
2//!
3//! Multi-device capability nodes for the Rustant agent.
4//! Each node implements the `Node` trait, providing capabilities
5//! like shell execution, AppleScript, screenshots, etc.
6
7pub mod consent;
8pub mod discovery;
9pub mod linux;
10pub mod macos;
11pub mod manager;
12pub mod types;
13
14pub use consent::{ConsentEntry, ConsentStore};
15pub use discovery::{
16    DiscoveredNode, MDNS_MULTICAST_ADDR, MDNS_PORT, MdnsConfig, MdnsDiscovery, MdnsServiceRecord,
17    MdnsTransport, NodeDiscovery, RUSTANT_SERVICE_NAME, UdpMdnsTransport,
18};
19pub use manager::NodeManager;
20pub use types::{
21    Capability, NodeCapability, NodeHealth, NodeId, NodeInfo, NodeMessage, NodeResult, NodeTask,
22    Platform, RateLimit,
23};
24
25use crate::error::RustantError;
26use async_trait::async_trait;
27
28/// Core trait for node implementations.
29#[async_trait]
30pub trait Node: Send + Sync {
31    /// Unique identifier for this node.
32    fn node_id(&self) -> &NodeId;
33
34    /// Descriptive info about this node.
35    fn info(&self) -> &NodeInfo;
36
37    /// Capabilities this node can provide.
38    fn capabilities(&self) -> &[Capability];
39
40    /// Execute a task on this node.
41    async fn execute(&self, task: NodeTask) -> Result<NodeResult, RustantError>;
42
43    /// Current health of this node.
44    fn health(&self) -> NodeHealth;
45
46    /// Perform a heartbeat check.
47    async fn heartbeat(&self) -> Result<NodeHealth, RustantError>;
48
49    /// Rich capability descriptions. Default wraps each Capability into a basic NodeCapability.
50    fn rich_capabilities(&self) -> Vec<NodeCapability> {
51        self.capabilities()
52            .iter()
53            .map(|c| NodeCapability::basic(c.clone()))
54            .collect()
55    }
56
57    /// Handle an inter-node protocol message. Default: responds to Ping and CapabilityQuery.
58    async fn handle_message(&self, msg: NodeMessage) -> Result<Option<NodeMessage>, RustantError> {
59        match msg {
60            NodeMessage::Ping => Ok(Some(NodeMessage::Pong {
61                uptime_secs: self.info().uptime_secs,
62            })),
63            NodeMessage::CapabilityQuery => Ok(Some(NodeMessage::CapabilityResponse {
64                capabilities: self.rich_capabilities(),
65            })),
66            _ => Ok(None),
67        }
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use chrono::Utc;
75
76    struct DefaultTestNode {
77        id: NodeId,
78        info: NodeInfo,
79        capabilities: Vec<Capability>,
80    }
81
82    impl DefaultTestNode {
83        fn new() -> Self {
84            let id = NodeId::new("test-node");
85            Self {
86                id: id.clone(),
87                info: NodeInfo {
88                    node_id: id,
89                    name: "test".into(),
90                    platform: Platform::MacOS,
91                    hostname: "test-host".into(),
92                    registered_at: Utc::now(),
93                    os_version: None,
94                    agent_version: "0.1.0".into(),
95                    uptime_secs: 42,
96                },
97                capabilities: vec![Capability::Shell, Capability::FileSystem],
98            }
99        }
100    }
101
102    #[async_trait]
103    impl Node for DefaultTestNode {
104        fn node_id(&self) -> &NodeId {
105            &self.id
106        }
107        fn info(&self) -> &NodeInfo {
108            &self.info
109        }
110        fn capabilities(&self) -> &[Capability] {
111            &self.capabilities
112        }
113        async fn execute(&self, task: NodeTask) -> Result<NodeResult, RustantError> {
114            Ok(NodeResult {
115                task_id: task.task_id,
116                success: true,
117                output: "ok".into(),
118                exit_code: Some(0),
119                duration_ms: 1,
120            })
121        }
122        fn health(&self) -> NodeHealth {
123            NodeHealth::Healthy
124        }
125        async fn heartbeat(&self) -> Result<NodeHealth, RustantError> {
126            Ok(NodeHealth::Healthy)
127        }
128    }
129
130    #[test]
131    fn test_node_types_reexported() {
132        let _ = NodeHealth::Healthy;
133        let _ = Capability::Shell;
134        let _ = Platform::MacOS;
135    }
136
137    #[test]
138    fn test_node_default_rich_capabilities() {
139        let node = DefaultTestNode::new();
140        let rich = node.rich_capabilities();
141        assert_eq!(rich.len(), 2);
142        assert_eq!(rich[0].capability, Capability::Shell);
143        assert!(rich[0].requires_consent);
144        assert_eq!(rich[1].capability, Capability::FileSystem);
145    }
146
147    #[tokio::test]
148    async fn test_node_default_handle_ping() {
149        let node = DefaultTestNode::new();
150        let response = node.handle_message(NodeMessage::Ping).await.unwrap();
151        match response {
152            Some(NodeMessage::Pong { uptime_secs }) => assert_eq!(uptime_secs, 42),
153            _ => panic!("Expected Pong"),
154        }
155    }
156
157    #[tokio::test]
158    async fn test_node_default_handle_capability_query() {
159        let node = DefaultTestNode::new();
160        let response = node
161            .handle_message(NodeMessage::CapabilityQuery)
162            .await
163            .unwrap();
164        match response {
165            Some(NodeMessage::CapabilityResponse { capabilities }) => {
166                assert_eq!(capabilities.len(), 2);
167            }
168            _ => panic!("Expected CapabilityResponse"),
169        }
170    }
171
172    #[tokio::test]
173    async fn test_node_default_handle_unknown() {
174        let node = DefaultTestNode::new();
175        let response = node.handle_message(NodeMessage::Shutdown).await.unwrap();
176        assert!(response.is_none());
177    }
178}