diff --git a/Cargo.lock b/Cargo.lock index 0c524704..1fa12238 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -157,7 +157,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -168,7 +168,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -2136,7 +2136,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2323,7 +2323,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2887,6 +2887,17 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "hostname" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "617aaa3557aef3810a6369d0a99fac8a080891b68bd9f9812a1eeda0c0730cbd" +dependencies = [ + "cfg-if", + "libc", + "windows-link", +] + [[package]] name = "html-escape" version = "0.2.13" @@ -3388,126 +3399,6 @@ dependencies = [ "serde", ] -[[package]] -name = "ironclaw" -version = "0.22.0" -dependencies = [ - "aes-gcm", - "aho-corasick", - "anyhow", - "async-trait", - "aws-config", - "aws-sdk-bedrockruntime", - "aws-smithy-types", - "axum 0.8.8", - "base64 0.22.1", - "blake3", - "bollard", - "bytes", - "chrono", - "chrono-tz", - "clap", - "clap_complete", - "criterion", - "cron", - "crossterm", - "deadpool-postgres", - "dirs 6.0.0", - "dotenvy", - "ed25519-dalek", - "eventsource-stream", - "flate2", - "fs4", - "futures", - "hex", - "hkdf", - "hmac", - "html-to-markdown-rs", - "http-body-util", - "hyper 1.8.1", - "hyper-util", - "iana-time-zone", - "insta", - "ironclaw_common", - "ironclaw_safety", - "json5", - "libsql", - "lru", - "mime_guess", - "open", - "pdf-extract", - "pgvector", - "postgres-types", - "pretty_assertions", - "pty-process", - "rand 0.8.5", - "readabilityrs", - "refinery", - "regex", - "reqwest", - "rig-core", - "rust_decimal", - "rust_decimal_macros", - "rustls 0.23.37", - "rustls-native-certs 0.8.3", - "rustyline", - "secrecy", - "secret-service", - "security-framework 3.7.0", - "semver", - "serde", - "serde_json", - "serde_yml", - "sha2", - "subtle", - "tar", - "tempfile", - "termimad", - "testcontainers-modules", - "thiserror 2.0.18", - "tokio", - "tokio-postgres", - "tokio-postgres-rustls", - "tokio-stream", - "tokio-test", - "tokio-tungstenite 0.26.2", - "toml", - "tower 0.5.3", - "tower-http 0.6.8", - "tracing", - "tracing-subscriber", - "tracing-test", - "url", - "urlencoding", - "uuid", - "wasmparser 0.220.1", - "wasmtime", - "wasmtime-wasi", - "webpki-roots 0.26.11", - "zbus", - "zip", -] - -[[package]] -name = "ironclaw_common" -version = "0.1.0" -dependencies = [ - "serde", - "serde_json", -] - -[[package]] -name = "ironclaw_safety" -version = "0.2.0" -dependencies = [ - "aho-corasick", - "regex", - "serde_json", - "thiserror 2.0.18", - "tracing", - "url", -] - [[package]] name = "is-docker" version = "0.2.0" @@ -4145,7 +4036,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4323,6 +4214,130 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "optimclaw" +version = "0.22.0" +dependencies = [ + "aes-gcm", + "aho-corasick", + "anyhow", + "async-trait", + "aws-config", + "aws-sdk-bedrockruntime", + "aws-smithy-types", + "axum 0.8.8", + "base64 0.22.1", + "blake3", + "bollard", + "bytes", + "chrono", + "chrono-tz", + "clap", + "clap_complete", + "criterion", + "cron", + "crossterm", + "deadpool-postgres", + "dirs 6.0.0", + "dotenvy", + "ed25519-dalek", + "eventsource-stream", + "flate2", + "fs4", + "futures", + "hex", + "hkdf", + "hmac", + "hostname", + "html-to-markdown-rs", + "http-body-util", + "hyper 1.8.1", + "hyper-util", + "iana-time-zone", + "insta", + "json5", + "libsql", + "lru", + "mime_guess", + "open", + "optimclaw_common", + "optimclaw_safety", + "pdf-extract", + "pgvector", + "postgres-types", + "pqcrypto-kyber", + "pqcrypto-traits", + "pretty_assertions", + "pty-process", + "rand 0.8.5", + "readabilityrs", + "refinery", + "regex", + "reqwest", + "rig-core", + "rust_decimal", + "rust_decimal_macros", + "rustls 0.23.37", + "rustls-native-certs 0.8.3", + "rustyline", + "secrecy", + "secret-service", + "security-framework 3.7.0", + "semver", + "serde", + "serde_json", + "serde_yml", + "sha2", + "subtle", + "sys-info", + "tar", + "tempfile", + "termimad", + "testcontainers-modules", + "thiserror 2.0.18", + "tokio", + "tokio-postgres", + "tokio-postgres-rustls", + "tokio-stream", + "tokio-test", + "tokio-tungstenite 0.26.2", + "toml", + "tower 0.5.3", + "tower-http 0.6.8", + "tracing", + "tracing-subscriber", + "tracing-test", + "url", + "urlencoding", + "uuid", + "wasmparser 0.220.1", + "wasmtime", + "wasmtime-wasi", + "webpki-roots 0.26.11", + "zbus", + "zip", +] + +[[package]] +name = "optimclaw_common" +version = "0.1.0" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "optimclaw_safety" +version = "0.2.0" +dependencies = [ + "aho-corasick", + "regex", + "serde_json", + "thiserror 2.0.18", + "tracing", + "url", +] + [[package]] name = "option-ext" version = "0.2.0" @@ -4810,6 +4825,37 @@ dependencies = [ "zerocopy 0.8.42", ] +[[package]] +name = "pqcrypto-internals" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4a326caf27cbf2ac291ca7fd56300497ba9e76a8cc6a7d95b7a18b57f22b61d" +dependencies = [ + "cc", + "dunce", + "getrandom 0.3.4", + "libc", +] + +[[package]] +name = "pqcrypto-kyber" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15c00293cf898859d0c771455388054fd69ab712263c73fdc7f287a39b1ba000" +dependencies = [ + "cc", + "glob", + "libc", + "pqcrypto-internals", + "pqcrypto-traits", +] + +[[package]] +name = "pqcrypto-traits" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94e851c7654eed9e68d7d27164c454961a616cf8c203d500607ef22c737b51bb" + [[package]] name = "precomputed-hash" version = "0.1.1" @@ -5493,7 +5539,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -6175,7 +6221,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -6330,6 +6376,16 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "sys-info" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b3a0d0aba8bf96a0e1ddfdc352fc53b3df7f39318c71854910c3c4b024ae52c" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "system-configuration" version = "0.7.0" @@ -6400,7 +6456,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix 1.1.4", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -7201,7 +7257,7 @@ checksum = "f2f6fb2847f6742cd76af783a2a2c49e9375d0a111c7bef6f71cd9e738c72d6e" dependencies = [ "memoffset", "tempfile", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -8051,7 +8107,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 0493a1f8..6fc818bc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -185,6 +185,13 @@ hex = "0.4.3" # OpenClaw import (feature gated) json5 = { version = "0.4", optional = true } +# Mesh cluster (feature gated) +pqcrypto-kyber = { version = "0.8", optional = true } +pqcrypto-traits = { version = "0.3", optional = true } +sys-info = { version = "0.9", optional = true } +hostname = { version = "0.4", optional = true } +tokio-tungstenite = { version = "0.26", optional = true } + # macOS keychain [target.'cfg(target_os = "macos")'.dependencies] security-framework = "3" @@ -237,6 +244,7 @@ integration = [] html-to-markdown = ["dep:html-to-markdown-rs", "dep:readabilityrs"] bedrock = ["dep:aws-config", "dep:aws-sdk-bedrockruntime", "dep:aws-smithy-types"] import = ["dep:json5", "libsql"] +cluster = ["dep:pqcrypto-kyber", "dep:pqcrypto-traits", "dep:sys-info", "dep:hostname", "dep:tokio-tungstenite"] [[test]] name = "e2e_thread_scheduling" diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 4420a450..42a5a29e 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -154,7 +154,19 @@ impl Agent { // Build system prompts once for this turn. Two variants: with tools // (normal iterations) and without (force_text final iteration). - let initial_tool_defs = self.tools().tool_definitions().await; + // + // Lazy tool loading: only send core tools to keep the prompt small. + // The model can discover additional tools via `tool_info`. + let initial_tool_defs = if std::env::var("OPTIMCLAW_LAZY_TOOLS").is_ok() { + let core = &[ + "shell", "read_file", "write_file", "list_dir", "apply_patch", + "tool_info", "memory_read", "memory_write", "memory_search", + "echo", "time", "http", + ]; + self.tools().tool_definitions_for(core).await + } else { + self.tools().tool_definitions().await + }; let initial_tool_defs = if !active_skills.is_empty() { crate::skills::attenuate_tools(&initial_tool_defs, &active_skills).tools } else { @@ -286,7 +298,16 @@ impl<'a> LoopDelegate for ChatDelegate<'a> { let force_text = iteration >= self.force_text_at; // Refresh tool definitions each iteration so newly built tools become visible - let tool_defs = self.agent.tools().tool_definitions().await; + let tool_defs = if std::env::var("OPTIMCLAW_LAZY_TOOLS").is_ok() { + let core = &[ + "shell", "read_file", "write_file", "list_dir", "apply_patch", + "tool_info", "memory_read", "memory_write", "memory_search", + "echo", "time", "http", + ]; + self.agent.tools().tool_definitions_for(core).await + } else { + self.agent.tools().tool_definitions().await + }; // Apply trust-based tool attenuation if skills are active. let tool_defs = if !self.active_skills.is_empty() { diff --git a/src/cluster/api.rs b/src/cluster/api.rs new file mode 100644 index 00000000..64df977f --- /dev/null +++ b/src/cluster/api.rs @@ -0,0 +1,105 @@ +//! REST API endpoints for mesh cluster status and task submission. + +use axum::extract::State; +use axum::http::StatusCode; +use axum::routing::get; +use axum::{Json, Router}; +use serde::Serialize; +use std::sync::Arc; + +use super::types::*; +use super::MeshNode; + +/// Build the cluster API routes. +pub fn cluster_routes() -> Router>> { + Router::new() + .route("/api/mesh/status", get(mesh_status)) + .route("/api/mesh/nodes", get(mesh_nodes)) +} + +#[derive(Serialize)] +struct MeshStatus { + enabled: bool, + node_id: String, + hostname: String, + peer_count: usize, + total_nodes: usize, + mesh_port: u16, +} + +#[derive(Serialize)] +struct MeshNodeInfo { + node_id: String, + hostname: String, + status: String, + load: f32, + gpu: Option, + loaded_model: Option, + tool_count: usize, + free_memory_mb: u64, +} + +async fn mesh_status( + State(mesh): State>>, +) -> Json { + match mesh { + Some(node) => { + let peer_count = node.overlay.peer_count().await; + let total = node.gossip.node_count().await + 1; // +1 for self + Json(MeshStatus { + enabled: true, + node_id: node_id_hex(&node.identity.node_id), + hostname: node.config.node_name.clone(), + peer_count, + total_nodes: total, + mesh_port: node.config.mesh_port, + }) + } + None => Json(MeshStatus { + enabled: false, + node_id: String::new(), + hostname: String::new(), + peer_count: 0, + total_nodes: 0, + mesh_port: 0, + }), + } +} + +async fn mesh_nodes( + State(mesh): State>>, +) -> Json> { + let Some(node) = mesh else { + return Json(vec![]); + }; + + let mut nodes = Vec::new(); + + // Add self + nodes.push(MeshNodeInfo { + node_id: node_id_hex(&node.identity.node_id), + hostname: node.config.node_name.clone(), + status: "self".into(), + load: super::compute_load(), + gpu: node.local_capabilities().gpu.as_ref().map(|g| g.name.clone()), + loaded_model: node.local_capabilities().loaded_model.clone(), + tool_count: node.local_capabilities().available_tools.len(), + free_memory_mb: node.local_capabilities().free_memory_mb, + }); + + // Add peers + for (info, status) in node.gossip.all_peers().await { + nodes.push(MeshNodeInfo { + node_id: node_id_hex(&info.id), + hostname: info.hostname, + status: format!("{:?}", status), + load: info.load, + gpu: info.capabilities.gpu.as_ref().map(|g| g.name.clone()), + loaded_model: info.capabilities.loaded_model.clone(), + tool_count: info.capabilities.available_tools.len(), + free_memory_mb: info.capabilities.free_memory_mb, + }); + } + + Json(nodes) +} diff --git a/src/cluster/beacon.rs b/src/cluster/beacon.rs new file mode 100644 index 00000000..9e1167c4 --- /dev/null +++ b/src/cluster/beacon.rs @@ -0,0 +1,136 @@ +//! UDP beacon for zero-config mesh discovery. +//! +//! Broadcasts a compact signed beacon every N seconds. +//! Listens for beacons from other nodes on the same LAN. + +use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; +use std::sync::Arc; +use tokio::net::UdpSocket; +use tokio::sync::mpsc; + +use super::config::ClusterConfig; +use super::crypto::MeshIdentity; +use super::types::*; + +/// A discovered peer from a beacon. +#[derive(Debug, Clone)] +pub struct DiscoveredPeer { + pub node_id: NodeId, + pub mesh_addr: SocketAddr, + pub flags: u8, + pub source_addr: SocketAddr, +} + +/// Start the beacon broadcaster and listener. +/// Returns a receiver that yields newly discovered peers. +pub async fn start_beacon( + config: &ClusterConfig, + identity: Arc, +) -> anyhow::Result> { + let (tx, rx) = mpsc::channel(64); + + // Bind listener + let listen_addr = SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, config.beacon_port); + let listener = UdpSocket::bind(listen_addr).await?; + listener.set_broadcast(true)?; + + // Bind sender (ephemeral port) + let sender = UdpSocket::bind("0.0.0.0:0").await?; + sender.set_broadcast(true)?; + + let broadcast_dest = SocketAddrV4::new(config.broadcast_addr, config.beacon_port); + let beacon_interval = config.beacon_interval; + let mesh_port = config.mesh_port; + let my_id = identity.node_id; + + // Broadcaster task + let id_clone = identity.clone(); + tokio::spawn(async move { + loop { + let has_gpu = super::detect_gpu().is_some(); + let flags = FLAG_ACCEPTING_TASKS | if has_gpu { FLAG_HAS_GPU } else { 0 }; + let packet = build_beacon(&id_clone, mesh_port, flags); + if let Err(e) = sender.send_to(&packet, broadcast_dest).await { + tracing::warn!("Beacon send failed: {}", e); + } + tokio::time::sleep(beacon_interval).await; + } + }); + + // Listener task + tokio::spawn(async move { + let mut buf = [0u8; 256]; + loop { + match listener.recv_from(&mut buf).await { + Ok((len, src)) => { + if let Some(peer) = parse_beacon(&buf[..len], src, &my_id) { + if tx.send(peer).await.is_err() { + break; // receiver dropped + } + } + } + Err(e) => { + tracing::warn!("Beacon recv error: {}", e); + tokio::time::sleep(std::time::Duration::from_secs(1)).await; + } + } + } + }); + + Ok(rx) +} + +/// Build a beacon packet. +/// +/// Format: MAGIC(6) + NodeId(16) + MeshPort(2) + Flags(1) + Signature(64) = 89 bytes +fn build_beacon(identity: &MeshIdentity, mesh_port: u16, flags: u8) -> Vec { + let mut payload = Vec::with_capacity(25); + payload.extend_from_slice(BEACON_MAGIC); + payload.extend_from_slice(&identity.node_id); + payload.extend_from_slice(&mesh_port.to_be_bytes()); + payload.push(flags); + + let sig = identity.sign(&payload); + + let mut packet = payload; + packet.extend_from_slice(&sig); + packet +} + +/// Parse a beacon packet. Returns None if invalid or from self. +fn parse_beacon(data: &[u8], source: SocketAddr, my_id: &NodeId) -> Option { + // Minimum: 6 + 16 + 2 + 1 + 64 = 89 bytes + if data.len() < 89 { + return None; + } + + // Check magic + if &data[..6] != BEACON_MAGIC { + return None; + } + + // Extract fields + let mut node_id = [0u8; 16]; + node_id.copy_from_slice(&data[6..22]); + + // Ignore own beacons + if &node_id == my_id { + return None; + } + + let mesh_port = u16::from_be_bytes([data[22], data[23]]); + let flags = data[24]; + + // Signature verification happens during the PQ WebSocket handshake. + // The beacon signature ensures the packet wasn't tampered with in transit, + // but full identity verification requires the PQ key exchange. + + let mesh_addr = SocketAddr::new(source.ip(), mesh_port); + + Some(DiscoveredPeer { + node_id, + mesh_addr, + flags, + source_addr: source, + }) +} diff --git a/src/cluster/config.rs b/src/cluster/config.rs new file mode 100644 index 00000000..8aeba8c5 --- /dev/null +++ b/src/cluster/config.rs @@ -0,0 +1,78 @@ +//! Cluster configuration from environment variables. + +use std::net::{IpAddr, Ipv4Addr}; +use std::time::Duration; + +/// Configuration for the mesh cluster. +#[derive(Debug, Clone)] +pub struct ClusterConfig { + /// Whether clustering is enabled. + pub enabled: bool, + /// Human-readable node name (defaults to hostname). + pub node_name: String, + /// UDP broadcast port for beacon discovery. + pub beacon_port: u16, + /// WebSocket port for the encrypted overlay mesh. + pub mesh_port: u16, + /// How often to broadcast presence beacon. + pub beacon_interval: Duration, + /// How often to send heartbeat pings to peers. + pub heartbeat_interval: Duration, + /// How long before a silent peer is marked suspect. + pub suspect_timeout: Duration, + /// How long before a suspect peer is marked dead. + pub dead_timeout: Duration, + /// How often to run gossip exchange rounds. + pub gossip_interval: Duration, + /// Maximum number of direct peer connections. + pub max_peers: usize, + /// Address to bind UDP and WS listeners. + pub bind_addr: IpAddr, + /// UDP broadcast destination address. + pub broadcast_addr: Ipv4Addr, + /// Path to persist mesh keypair. + pub keys_path: String, +} + +impl ClusterConfig { + /// Load from environment variables with sane defaults. + pub fn from_env() -> Self { + let hostname = hostname::get() + .map(|h| h.to_string_lossy().to_string()) + .unwrap_or_else(|_| "unknown".into()); + + let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".into()); + + Self { + enabled: std::env::var("CLUSTER_ENABLED") + .map(|v| v == "1" || v == "true") + .unwrap_or(false), + node_name: std::env::var("CLUSTER_NODE_NAME").unwrap_or(hostname), + beacon_port: parse_env("CLUSTER_BEACON_PORT", 9900), + mesh_port: parse_env("CLUSTER_MESH_PORT", 9901), + beacon_interval: Duration::from_secs(parse_env("CLUSTER_BEACON_INTERVAL_SECS", 5)), + heartbeat_interval: Duration::from_secs(parse_env("CLUSTER_HEARTBEAT_INTERVAL_SECS", 3)), + suspect_timeout: Duration::from_secs(parse_env("CLUSTER_SUSPECT_TIMEOUT_SECS", 10)), + dead_timeout: Duration::from_secs(parse_env("CLUSTER_DEAD_TIMEOUT_SECS", 15)), + gossip_interval: Duration::from_secs(parse_env("CLUSTER_GOSSIP_INTERVAL_SECS", 2)), + max_peers: parse_env("CLUSTER_MAX_PEERS", 5), + bind_addr: std::env::var("CLUSTER_BIND_ADDR") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED)), + broadcast_addr: std::env::var("CLUSTER_BROADCAST_ADDR") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(Ipv4Addr::BROADCAST), + keys_path: std::env::var("CLUSTER_KEYS_PATH") + .unwrap_or_else(|_| format!("{}/.optimclaw/mesh_keys.json", home)), + } + } +} + +fn parse_env(key: &str, default: T) -> T { + std::env::var(key) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) +} diff --git a/src/cluster/crypto.rs b/src/cluster/crypto.rs new file mode 100644 index 00000000..5dcb3c72 --- /dev/null +++ b/src/cluster/crypto.rs @@ -0,0 +1,269 @@ +//! Post-quantum cryptography for mesh communication. +//! +//! Uses ML-KEM-768 (Kyber) for key encapsulation and AES-256-GCM for +//! symmetric encryption. Ed25519 for signing beacons and handshakes. + +use aes_gcm::{Aes256Gcm, KeyInit, Nonce}; +use aes_gcm::aead::Aead; +use ed25519_dalek::{Signer, SigningKey, Verifier, VerifyingKey}; +use hkdf::Hkdf; +use sha2::Sha256; +use rand::rngs::OsRng; +use serde::{Deserialize, Serialize}; +use std::path::Path; + +use pqcrypto_kyber::kyber768; +use pqcrypto_traits::kem::{Ciphertext, PublicKey, SecretKey, SharedSecret}; + +use super::types::NodeId; + +/// Persistent mesh identity (keypairs). +#[derive(Serialize, Deserialize)] +struct StoredKeys { + kem_pk: Vec, + kem_sk: Vec, + sign_seed: [u8; 32], +} + +/// A node's cryptographic identity. +pub struct MeshIdentity { + pub node_id: NodeId, + pub kem_pk: kyber768::PublicKey, + kem_sk: kyber768::SecretKey, + pub sign_key: SigningKey, + pub verify_key: VerifyingKey, +} + +impl MeshIdentity { + /// Generate a new identity or load from disk. + pub fn load_or_generate(path: &str) -> anyhow::Result { + if Path::new(path).exists() { + Self::load(path) + } else { + let identity = Self::generate(); + identity.save(path)?; + Ok(identity) + } + } + + /// Generate fresh keypairs. + pub fn generate() -> Self { + let (kem_pk, kem_sk) = kyber768::keypair(); + let mut seed = [0u8; 32]; + rand::Fill::try_fill(&mut seed, &mut OsRng).expect("RNG fill"); + let sign_key = SigningKey::from_bytes(&seed); + let verify_key = sign_key.verifying_key(); + + // NodeId = first 16 bytes of blake3(kem_pk || sign_pk) + let mut hasher = blake3::Hasher::new(); + hasher.update(kem_pk.as_bytes()); + hasher.update(verify_key.as_bytes()); + let hash = hasher.finalize(); + let mut node_id = [0u8; 16]; + node_id.copy_from_slice(&hash.as_bytes()[..16]); + + Self { + node_id, + kem_pk, + kem_sk, + sign_key, + verify_key, + } + } + + fn save(&self, path: &str) -> anyhow::Result<()> { + if let Some(parent) = Path::new(path).parent() { + std::fs::create_dir_all(parent)?; + } + let stored = StoredKeys { + kem_pk: self.kem_pk.as_bytes().to_vec(), + kem_sk: self.kem_sk.as_bytes().to_vec(), + sign_seed: self.sign_key.to_bytes(), + }; + let json = serde_json::to_string_pretty(&stored)?; + std::fs::write(path, json)?; + // Restrict permissions + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))?; + } + Ok(()) + } + + fn load(path: &str) -> anyhow::Result { + let json = std::fs::read_to_string(path)?; + let stored: StoredKeys = serde_json::from_str(&json)?; + + let kem_pk = kyber768::PublicKey::from_bytes(&stored.kem_pk) + .map_err(|_| anyhow::anyhow!("Invalid KEM public key"))?; + let kem_sk = kyber768::SecretKey::from_bytes(&stored.kem_sk) + .map_err(|_| anyhow::anyhow!("Invalid KEM secret key"))?; + let sign_key = SigningKey::from_bytes(&stored.sign_seed); + let verify_key = sign_key.verifying_key(); + + let mut hasher = blake3::Hasher::new(); + hasher.update(kem_pk.as_bytes()); + hasher.update(verify_key.as_bytes()); + let hash = hasher.finalize(); + let mut node_id = [0u8; 16]; + node_id.copy_from_slice(&hash.as_bytes()[..16]); + + Ok(Self { + node_id, + kem_pk, + kem_sk, + sign_key, + verify_key, + }) + } + + /// Sign a message with ed25519. + pub fn sign(&self, msg: &[u8]) -> Vec { + self.sign_key.sign(msg).to_bytes().to_vec() + } + + /// Verify a signature against a public key. + pub fn verify(pubkey: &VerifyingKey, msg: &[u8], sig: &[u8]) -> bool { + if sig.len() != 64 { + return false; + } + let mut sig_bytes = [0u8; 64]; + sig_bytes.copy_from_slice(sig); + let signature = ed25519_dalek::Signature::from_bytes(&sig_bytes); + pubkey.verify(msg, &signature).is_ok() + } + + /// Decapsulate a shared secret from a ciphertext (responder side). + pub fn decapsulate(&self, ciphertext: &[u8]) -> anyhow::Result> { + let ct = kyber768::Ciphertext::from_bytes(ciphertext) + .map_err(|_| anyhow::anyhow!("Invalid KEM ciphertext"))?; + let ss = kyber768::decapsulate(&ct, &self.kem_sk); + Ok(ss.as_bytes().to_vec()) + } + + /// Public KEM key bytes for sharing. + pub fn kem_pk_bytes(&self) -> Vec { + self.kem_pk.as_bytes().to_vec() + } + + /// Public verify key bytes for sharing. + pub fn verify_key_bytes(&self) -> Vec { + self.verify_key.to_bytes().to_vec() + } +} + +/// Encapsulate a shared secret using a peer's public KEM key (initiator side). +pub fn encapsulate(peer_kem_pk: &[u8]) -> anyhow::Result<(Vec, Vec)> { + let pk = kyber768::PublicKey::from_bytes(peer_kem_pk) + .map_err(|_| anyhow::anyhow!("Invalid peer KEM public key"))?; + let (ss, ct) = kyber768::encapsulate(&pk); + Ok((ss.as_bytes().to_vec(), ct.as_bytes().to_vec())) +} + +/// Derive symmetric encryption keys from a shared secret. +pub fn derive_keys(shared_secret: &[u8], initiator_id: &NodeId, responder_id: &NodeId) -> (Aes256Gcm, Aes256Gcm) { + let hk = Hkdf::::new(None, shared_secret); + + let mut send_key = [0u8; 32]; + let mut recv_key = [0u8; 32]; + + // Deterministic key derivation: initiator always gets "init" key + let mut info_send = Vec::new(); + info_send.extend_from_slice(b"omesh-send-"); + info_send.extend_from_slice(initiator_id); + info_send.extend_from_slice(responder_id); + hk.expand(&info_send, &mut send_key).expect("HKDF expand"); + + let mut info_recv = Vec::new(); + info_recv.extend_from_slice(b"omesh-recv-"); + info_recv.extend_from_slice(responder_id); + info_recv.extend_from_slice(initiator_id); + hk.expand(&info_recv, &mut recv_key).expect("HKDF expand"); + + let send_cipher = Aes256Gcm::new_from_slice(&send_key).expect("AES key"); + let recv_cipher = Aes256Gcm::new_from_slice(&recv_key).expect("AES key"); + + (send_cipher, recv_cipher) +} + +/// Encrypt a message with AES-256-GCM. +pub fn encrypt(cipher: &Aes256Gcm, nonce_counter: u64, plaintext: &[u8]) -> Vec { + let mut nonce_bytes = [0u8; 12]; + nonce_bytes[4..].copy_from_slice(&nonce_counter.to_be_bytes()); + let nonce = Nonce::from_slice(&nonce_bytes); + + let ciphertext = cipher.encrypt(nonce, plaintext).expect("AES-GCM encrypt"); + + // Frame: nonce_counter(8 bytes) || ciphertext + let mut frame = Vec::with_capacity(8 + ciphertext.len()); + frame.extend_from_slice(&nonce_counter.to_be_bytes()); + frame.extend_from_slice(&ciphertext); + frame +} + +/// Decrypt a message with AES-256-GCM. +pub fn decrypt(cipher: &Aes256Gcm, frame: &[u8]) -> anyhow::Result> { + if frame.len() < 8 { + anyhow::bail!("Frame too short"); + } + let mut nonce_counter_bytes = [0u8; 8]; + nonce_counter_bytes.copy_from_slice(&frame[..8]); + let nonce_counter = u64::from_be_bytes(nonce_counter_bytes); + + let mut nonce_bytes = [0u8; 12]; + nonce_bytes[4..].copy_from_slice(&nonce_counter.to_be_bytes()); + let nonce = Nonce::from_slice(&nonce_bytes); + + let plaintext = cipher + .decrypt(nonce, &frame[8..]) + .map_err(|_| anyhow::anyhow!("AES-GCM decrypt failed"))?; + + Ok(plaintext) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_identity_generate() { + let id = MeshIdentity::generate(); + assert_ne!(id.node_id, [0u8; 16]); + } + + #[test] + fn test_sign_verify() { + let id = MeshIdentity::generate(); + let msg = b"hello mesh"; + let sig = id.sign(msg); + assert!(MeshIdentity::verify(&id.verify_key, msg, &sig)); + assert!(!MeshIdentity::verify(&id.verify_key, b"wrong", &sig)); + } + + #[test] + fn test_kem_roundtrip() { + let node_a = MeshIdentity::generate(); + let node_b = MeshIdentity::generate(); + + // A encapsulates for B + let (ss_a, ct) = encapsulate(&node_b.kem_pk_bytes()).unwrap(); + // B decapsulates + let ss_b = node_b.decapsulate(&ct).unwrap(); + assert_eq!(ss_a, ss_b); + } + + #[test] + fn test_encrypt_decrypt_roundtrip() { + let node_a = MeshIdentity::generate(); + let node_b = MeshIdentity::generate(); + + let (ss, _ct) = encapsulate(&node_b.kem_pk_bytes()).unwrap(); + let (send_cipher, recv_cipher) = derive_keys(&ss, &node_a.node_id, &node_b.node_id); + + let plaintext = b"secret mesh message"; + let frame = encrypt(&send_cipher, 1, plaintext); + let decrypted = decrypt(&recv_cipher, &frame).unwrap(); + assert_eq!(decrypted, plaintext); + } +} diff --git a/src/cluster/gossip.rs b/src/cluster/gossip.rs new file mode 100644 index 00000000..5621faa4 --- /dev/null +++ b/src/cluster/gossip.rs @@ -0,0 +1,219 @@ +//! SWIM-lite gossip protocol for membership and failure detection. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Instant; +use tokio::sync::RwLock; + +use super::types::*; + +/// State of a member in the gossip membership table. +#[derive(Debug, Clone)] +pub struct MemberState { + pub info: NodeInfo, + pub status: MemberStatus, + pub incarnation: u64, + pub last_seen: Instant, + pub suspect_since: Option, +} + +/// The gossip membership table. +pub struct GossipState { + members: RwLock>, + local_id: NodeId, + local_incarnation: RwLock, +} + +impl GossipState { + pub fn new(local_id: NodeId) -> Self { + Self { + members: RwLock::new(HashMap::new()), + local_id, + local_incarnation: RwLock::new(0), + } + } + + /// Add or update a node in the membership table. + pub async fn merge_node(&self, info: NodeInfo, incarnation: u64, status: MemberStatus) { + if info.id == self.local_id { + return; // never merge self + } + + let mut members = self.members.write().await; + let entry = members.entry(info.id).or_insert_with(|| MemberState { + info: info.clone(), + status: MemberStatus::Alive, + incarnation: 0, + last_seen: Instant::now(), + suspect_since: None, + }); + + // Only accept updates with higher incarnation + if incarnation > entry.incarnation + || (incarnation == entry.incarnation && status_priority(status) > status_priority(entry.status)) + { + entry.info = info; + entry.incarnation = incarnation; + entry.status = status; + entry.last_seen = Instant::now(); + + if status == MemberStatus::Suspect { + entry.suspect_since = Some(Instant::now()); + } else { + entry.suspect_since = None; + } + } + } + + /// Record a heartbeat from a node. + pub async fn heartbeat(&self, node_id: &NodeId) { + let mut members = self.members.write().await; + if let Some(entry) = members.get_mut(node_id) { + entry.last_seen = Instant::now(); + if entry.status == MemberStatus::Suspect { + entry.status = MemberStatus::Alive; + entry.suspect_since = None; + } + } + } + + /// Mark a node as suspect. + pub async fn mark_suspect(&self, node_id: &NodeId) { + let mut members = self.members.write().await; + if let Some(entry) = members.get_mut(node_id) { + if entry.status == MemberStatus::Alive { + entry.status = MemberStatus::Suspect; + entry.suspect_since = Some(Instant::now()); + tracing::info!("Node {} marked suspect", node_id_hex(node_id)); + } + } + } + + /// Mark a node as dead and remove it. + pub async fn mark_dead(&self, node_id: &NodeId) { + let mut members = self.members.write().await; + if let Some(entry) = members.get_mut(node_id) { + entry.status = MemberStatus::Dead; + tracing::info!("Node {} marked dead", node_id_hex(node_id)); + } + } + + /// Remove dead nodes that have been dead for more than the given duration. + pub async fn prune_dead(&self, max_age: std::time::Duration) { + let mut members = self.members.write().await; + members.retain(|id, entry| { + if entry.status == MemberStatus::Dead || entry.status == MemberStatus::Left { + if entry.last_seen.elapsed() > max_age { + tracing::debug!("Pruning dead node {}", node_id_hex(id)); + return false; + } + } + true + }); + } + + /// Get all alive peers. + pub async fn alive_peers(&self) -> Vec { + self.members + .read() + .await + .values() + .filter(|m| m.status == MemberStatus::Alive) + .map(|m| m.info.clone()) + .collect() + } + + /// Get all known peers (any status). + pub async fn all_peers(&self) -> Vec<(NodeInfo, MemberStatus)> { + self.members + .read() + .await + .values() + .map(|m| (m.info.clone(), m.status)) + .collect() + } + + /// Get nodes that should be checked for suspect/dead transitions. + pub async fn check_timeouts( + &self, + suspect_timeout: std::time::Duration, + dead_timeout: std::time::Duration, + ) -> (Vec, Vec) { + let members = self.members.read().await; + let mut new_suspects = Vec::new(); + let mut new_dead = Vec::new(); + + for (id, entry) in members.iter() { + match entry.status { + MemberStatus::Alive => { + if entry.last_seen.elapsed() > suspect_timeout { + new_suspects.push(*id); + } + } + MemberStatus::Suspect => { + if let Some(since) = entry.suspect_since { + if since.elapsed() > dead_timeout { + new_dead.push(*id); + } + } + } + _ => {} + } + } + + (new_suspects, new_dead) + } + + /// Get a random alive peer for ping selection. + pub async fn random_alive_peer(&self) -> Option { + let members = self.members.read().await; + let alive: Vec<_> = members + .values() + .filter(|m| m.status == MemberStatus::Alive) + .collect(); + + if alive.is_empty() { + return None; + } + + let idx = rand::random::() % alive.len(); + Some(alive[idx].info.clone()) + } + + /// Build piggybacked membership updates for gossip dissemination. + pub async fn membership_updates(&self) -> Vec { + self.members + .read() + .await + .values() + .map(|m| MembershipUpdate { + node_id: m.info.id, + incarnation: m.incarnation, + status: m.status, + }) + .collect() + } + + /// Node count. + pub async fn node_count(&self) -> usize { + self.members.read().await.len() + } + + /// Handle a node gracefully leaving. + pub async fn handle_leave(&self, node_id: &NodeId) { + let mut members = self.members.write().await; + if let Some(entry) = members.get_mut(node_id) { + entry.status = MemberStatus::Left; + tracing::info!("Node {} left the mesh", node_id_hex(node_id)); + } + } +} + +fn status_priority(status: MemberStatus) -> u8 { + match status { + MemberStatus::Alive => 0, + MemberStatus::Suspect => 1, + MemberStatus::Dead => 2, + MemberStatus::Left => 3, + } +} diff --git a/src/cluster/mod.rs b/src/cluster/mod.rs new file mode 100644 index 00000000..8627c37e --- /dev/null +++ b/src/cluster/mod.rs @@ -0,0 +1,598 @@ +//! Autonomous AI mesh network for OptimClaw. +//! +//! Nodes self-discover via UDP beacons, form a peer-to-peer overlay mesh +//! with post-quantum encrypted WebSocket channels, and route tasks +//! intelligently based on GPU capability, model availability, and load. + +pub mod api; +pub mod beacon; +pub mod config; +pub mod crypto; +pub mod gossip; +pub mod overlay; +pub mod router; +pub mod types; + +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::{mpsc, RwLock}; + +use config::ClusterConfig; +use crypto::MeshIdentity; +use gossip::GossipState; +use overlay::OverlayMesh; +use types::*; + +/// A node in the OptimClaw mesh network. +pub struct MeshNode { + pub config: ClusterConfig, + pub identity: Arc, + pub gossip: Arc, + pub overlay: Arc, + pub event_tx: tokio::sync::broadcast::Sender, + capabilities: NodeCapabilities, + pending_tasks: PendingTasks, + shutdown_tx: tokio::sync::broadcast::Sender<()>, +} + +impl MeshNode { + /// Start a new mesh node. Begins discovery, gossip, and overlay management. + pub async fn start(config: ClusterConfig) -> anyhow::Result> { + // Load or generate identity + let identity = Arc::new(MeshIdentity::load_or_generate(&config.keys_path)?); + tracing::info!( + "Mesh node starting: id={} name={}", + node_id_hex(&identity.node_id), + config.node_name + ); + + // Detect local capabilities + let capabilities = detect_capabilities(); + + // Initialize gossip state + let gossip = Arc::new(GossipState::new(identity.node_id)); + + // Message channel for incoming mesh messages + let (incoming_tx, mut incoming_rx) = mpsc::channel::<(NodeId, MeshMessage)>(256); + + // Initialize overlay mesh + let overlay = Arc::new(OverlayMesh::new( + identity.clone(), + gossip.clone(), + config.clone(), + incoming_tx, + )); + + let (shutdown_tx, _) = tokio::sync::broadcast::channel(1); + let (event_tx, _) = tokio::sync::broadcast::channel(128); + let pending_tasks: PendingTasks = Arc::new(RwLock::new(HashMap::new())); + + let node = Arc::new(Self { + config: config.clone(), + identity: identity.clone(), + gossip: gossip.clone(), + overlay: overlay.clone(), + event_tx: event_tx.clone(), + capabilities, + pending_tasks: pending_tasks.clone(), + shutdown_tx: shutdown_tx.clone(), + }); + + // Start beacon discovery + let mut beacon_rx = beacon::start_beacon(&config, identity.clone()).await?; + tracing::info!("Beacon broadcasting on port {}", config.beacon_port); + + // Start the overlay WebSocket listener + overlay.start_listener().await?; + + // Task: process discovered peers -- connect via PQ-encrypted WebSocket + let overlay_disc = overlay.clone(); + let gossip_disc = gossip.clone(); + tokio::spawn(async move { + while let Some(peer) = beacon_rx.recv().await { + if overlay_disc.is_connected(&peer.node_id).await { + continue; + } + tracing::info!( + "Discovered peer {} at {}", + node_id_hex(&peer.node_id), + peer.mesh_addr + ); + + // Record in gossip + let info = NodeInfo { + id: peer.node_id, + hostname: format!("{}", peer.source_addr.ip()), + mesh_addr: peer.mesh_addr, + gateway_addr: None, + capabilities: NodeCapabilities { + gpu: None, + loaded_model: None, + available_tools: vec![], + free_memory_mb: 0, + total_memory_mb: 0, + }, + load: 0.0, + version: String::new(), + started_at: 0, + last_heartbeat: chrono::Utc::now().timestamp(), + }; + gossip_disc + .merge_node(info, 0, MemberStatus::Alive) + .await; + + // Initiate PQ-encrypted WebSocket connection + let overlay_conn = overlay_disc.clone(); + let peer_id = peer.node_id; + let peer_addr = peer.mesh_addr; + tokio::spawn(async move { + if let Err(e) = overlay_conn.connect_to_peer(peer_addr, &peer_id).await { + tracing::debug!( + "Failed to connect to {}: {}", + node_id_hex(&peer_id), + e + ); + } + }); + } + }); + + // Task: process incoming mesh messages + let gossip_msg = gossip.clone(); + let overlay_msg = overlay.clone(); + let pending_msg = pending_tasks.clone(); + let event_tx_msg = Some(event_tx.clone()); + tokio::spawn(async move { + while let Some((from, msg)) = incoming_rx.recv().await { + handle_mesh_message(&gossip_msg, &overlay_msg, from, msg, &pending_msg, &event_tx_msg).await; + } + }); + + // Task: periodic gossip / failure detection + let gossip_tick = gossip.clone(); + let overlay_tick = overlay.clone(); + let suspect_timeout = config.suspect_timeout; + let dead_timeout = config.dead_timeout; + let gossip_interval = config.gossip_interval; + let mut shutdown_rx = shutdown_tx.subscribe(); + tokio::spawn(async move { + loop { + tokio::select! { + _ = tokio::time::sleep(gossip_interval) => { + // Check timeouts + let (suspects, dead) = gossip_tick + .check_timeouts(suspect_timeout, dead_timeout) + .await; + for id in suspects { + gossip_tick.mark_suspect(&id).await; + } + for id in dead { + gossip_tick.mark_dead(&id).await; + overlay_tick.remove_peer(&id).await; + } + + // Prune long-dead nodes + gossip_tick.prune_dead(std::time::Duration::from_secs(120)).await; + + // Gossip exchange with a random peer + if let Some(peer) = gossip_tick.random_alive_peer().await { + let updates = gossip_tick.membership_updates().await; + let msg = MeshMessage::Ping { + from: identity.node_id, + seq: rand::random(), + piggyback: updates, + }; + let _ = overlay_tick.send_to(&peer.id, &msg).await; + } + } + _ = shutdown_rx.recv() => break, + } + } + }); + + tracing::info!( + "Mesh node ready: {} peers=0 beacon=:{} mesh=:{}", + node_id_hex(&node.identity.node_id), + config.beacon_port, + config.mesh_port + ); + + Ok(node) + } + + /// Get local node capabilities. + pub fn local_capabilities(&self) -> &NodeCapabilities { + &self.capabilities + } + + /// Build a NodeInfo for the local node. + pub fn local_info(&self) -> NodeInfo { + NodeInfo { + id: self.identity.node_id, + hostname: self.config.node_name.clone(), + mesh_addr: std::net::SocketAddr::new( + self.config.bind_addr, + self.config.mesh_port, + ), + gateway_addr: None, + capabilities: self.capabilities.clone(), + load: compute_load(), + version: env!("CARGO_PKG_VERSION").to_string(), + started_at: chrono::Utc::now().timestamp(), + last_heartbeat: chrono::Utc::now().timestamp(), + } + } + + /// Submit a task to the mesh for intelligent routing. + /// Returns the result from whichever node handles it. + pub async fn submit_task(&self, content: String, timeout_secs: u64) -> anyhow::Result { + let task_id = uuid::Uuid::new_v4().to_string(); + let envelope = TaskEnvelope { + task_id: task_id.clone(), + content: content.clone(), + origin_node: self.identity.node_id, + required_model: None, + required_tools: vec![], + min_vram_mb: None, + priority: 0, + hop_count: 0, + }; + + let best = router::route_task(&self.gossip, &self.local_info(), &envelope).await; + + let target = match best { + Some(node) => node, + None => anyhow::bail!("No suitable node found for task"), + }; + + // If best node is self, execute locally + if target.id == self.identity.node_id { + tracing::info!("Executing task {} locally", task_id); + return Ok(TaskResult { + success: true, + response: format!("Local execution: {}", content), + node_id: self.identity.node_id, + duration_ms: 0, + }); + } + + // Send to remote peer and wait for result + tracing::info!( + "Routing task {} to {} ({})", + task_id, + node_id_hex(&target.id), + target.hostname + ); + + let (result_tx, result_rx) = tokio::sync::oneshot::channel(); + self.pending_tasks.write().await.insert(task_id.clone(), result_tx); + + let msg = MeshMessage::TaskRequest { + task_id: task_id.clone(), + envelope, + }; + self.overlay.send_to(&target.id, &msg).await?; + + // Wait for result with timeout + match tokio::time::timeout( + std::time::Duration::from_secs(timeout_secs), + result_rx, + ).await { + Ok(Ok(result)) => Ok(result), + Ok(Err(_)) => { + self.pending_tasks.write().await.remove(&task_id); + anyhow::bail!("Task result channel closed") + } + Err(_) => { + self.pending_tasks.write().await.remove(&task_id); + anyhow::bail!("Task timed out after {}s", timeout_secs) + } + } + } + + /// Subscribe to cluster events for the UI. + pub fn subscribe_events(&self) -> tokio::sync::broadcast::Receiver { + self.event_tx.subscribe() + } + + /// Graceful shutdown. + pub async fn shutdown(&self) { + // Broadcast leave + let msg = MeshMessage::Leaving { + id: self.identity.node_id, + }; + self.overlay.broadcast(&msg).await; + let _ = self.shutdown_tx.send(()); + tracing::info!("Mesh node shutting down"); + } +} + +/// Pending tasks waiting for results from remote nodes. +pub type PendingTasks = + Arc>>>; + +/// Handle an incoming mesh message from a peer. +async fn handle_mesh_message( + gossip: &GossipState, + overlay: &OverlayMesh, + from: NodeId, + msg: MeshMessage, + pending_tasks: &PendingTasks, + event_tx: &Option>, +) { + match msg { + MeshMessage::Ping { piggyback, seq, .. } => { + gossip.heartbeat(&from).await; + for update in piggyback { + let info = NodeInfo { + id: update.node_id, + hostname: String::new(), + mesh_addr: "0.0.0.0:0".parse().unwrap(), + gateway_addr: None, + capabilities: NodeCapabilities { + gpu: None, + loaded_model: None, + available_tools: vec![], + free_memory_mb: 0, + total_memory_mb: 0, + }, + load: 0.0, + version: String::new(), + started_at: 0, + last_heartbeat: chrono::Utc::now().timestamp(), + }; + gossip.merge_node(info, update.incarnation, update.status).await; + } + // Send ack + let updates = gossip.membership_updates().await; + let ack = MeshMessage::Ack { + from: overlay.identity.node_id, + seq, + piggyback: updates, + }; + let _ = overlay.send_to(&from, &ack).await; + } + MeshMessage::Ack { piggyback, .. } => { + gossip.heartbeat(&from).await; + for update in piggyback { + let info = NodeInfo { + id: update.node_id, + hostname: String::new(), + mesh_addr: "0.0.0.0:0".parse().unwrap(), + gateway_addr: None, + capabilities: NodeCapabilities { + gpu: None, loaded_model: None, available_tools: vec![], + free_memory_mb: 0, total_memory_mb: 0, + }, + load: 0.0, version: String::new(), started_at: 0, + last_heartbeat: chrono::Utc::now().timestamp(), + }; + gossip.merge_node(info, update.incarnation, update.status).await; + } + } + MeshMessage::PeerExchange { nodes } => { + for info in nodes { + gossip.merge_node(info, 0, MemberStatus::Alive).await; + } + } + MeshMessage::Leaving { id } => { + gossip.handle_leave(&id).await; + if let Some(tx) = event_tx { + let _ = tx.send(ClusterEvent::NodeLeft { + node_id: node_id_hex(&id), + }); + } + } + MeshMessage::TaskRequest { task_id, envelope } => { + tracing::info!("Received task {} from {}", task_id, node_id_hex(&from)); + // Check if we can handle this task + let local_info = NodeInfo { + id: overlay.identity.node_id, + hostname: String::new(), + mesh_addr: "0.0.0.0:0".parse().unwrap(), + gateway_addr: None, + capabilities: detect_capabilities(), + load: compute_load(), + version: String::new(), + started_at: 0, + last_heartbeat: 0, + }; + if router::score_node(&local_info, &envelope).is_none() { + let reject = MeshMessage::TaskReject { + task_id, + node_id: overlay.identity.node_id, + reason: "Capability mismatch".into(), + }; + let _ = overlay.send_to(&from, &reject).await; + return; + } + + let accept = MeshMessage::TaskAccept { + task_id: task_id.clone(), + node_id: overlay.identity.node_id, + }; + let _ = overlay.send_to(&from, &accept).await; + + if let Some(tx) = event_tx { + let _ = tx.send(ClusterEvent::TaskReceived { + task_id: task_id.clone(), + from: node_id_hex(&from), + content: envelope.content.chars().take(100).collect(), + }); + } + + // Execute the task via optimclaw single-message mode + let overlay_exec = overlay.identity.node_id; + let from_id = from; + let overlay_ref = overlay.incoming_tx.clone(); + let task_id_exec = task_id.clone(); + let content = envelope.content.clone(); + + // Run in a blocking thread since it spawns a subprocess + let result = tokio::task::spawn_blocking(move || { + let start = std::time::Instant::now(); + let output = std::process::Command::new("optimclaw") + .args(["run", "--cli-only", "--no-onboard", "-m", &content]) + .env("OPTIMCLAW_LAZY_TOOLS", "1") + .output(); + + match output { + Ok(out) => { + let stderr = String::from_utf8_lossy(&out.stderr); + let response = stderr.lines() + .skip_while(|l| !l.contains("────")) + .skip(1) + .collect::>() + .join("\n"); + let response = if response.trim().is_empty() { + String::from_utf8_lossy(&out.stdout).trim().to_string() + } else { + response.trim().to_string() + }; + TaskResult { + success: out.status.success(), + response, + node_id: overlay_exec, + duration_ms: start.elapsed().as_millis() as u64, + } + } + Err(e) => TaskResult { + success: false, + response: format!("Execution failed: {}", e), + node_id: overlay_exec, + duration_ms: start.elapsed().as_millis() as u64, + }, + } + }).await.unwrap_or_else(|e| TaskResult { + success: false, + response: format!("Task spawn error: {}", e), + node_id: overlay.identity.node_id, + duration_ms: 0, + }); + + // Send result back to origin + let complete = MeshMessage::TaskComplete { + task_id: task_id_exec, + result, + }; + let _ = overlay.send_to(&from_id, &complete).await; + } + MeshMessage::TaskAccept { task_id, node_id } => { + tracing::info!( + "Task {} accepted by {}", + task_id, + node_id_hex(&node_id) + ); + if let Some(tx) = event_tx { + let _ = tx.send(ClusterEvent::TaskAccepted { + task_id, + node_id: node_id_hex(&node_id), + }); + } + } + MeshMessage::TaskStream { task_id, chunk } => { + if let Some(tx) = event_tx { + let _ = tx.send(ClusterEvent::TaskStream { + task_id, + chunk, + }); + } + } + MeshMessage::TaskComplete { task_id, result } => { + tracing::info!( + "Task {} completed by {} (success={})", + task_id, + node_id_hex(&result.node_id), + result.success + ); + // Resolve pending task future + let mut pending = pending_tasks.write().await; + if let Some(tx) = pending.remove(&task_id) { + let _ = tx.send(result.clone()); + } + if let Some(etx) = event_tx { + let _ = etx.send(ClusterEvent::TaskCompleted { + task_id, + node_id: node_id_hex(&result.node_id), + success: result.success, + response: result.response, + }); + } + } + _ => { + tracing::debug!("Unhandled mesh message from {}", node_id_hex(&from)); + } + } +} + +/// Events emitted by the cluster for the UI/SSE layer. +#[derive(Debug, Clone, serde::Serialize)] +#[serde(tag = "type")] +pub enum ClusterEvent { + NodeJoined { node_id: String, hostname: String }, + NodeLeft { node_id: String }, + TaskReceived { task_id: String, from: String, content: String }, + TaskAccepted { task_id: String, node_id: String }, + TaskStream { task_id: String, chunk: TaskStreamChunk }, + TaskCompleted { task_id: String, node_id: String, success: bool, response: String }, + MeshStatus { node_count: usize, peer_count: usize }, +} + +/// Detect local system capabilities. +fn detect_capabilities() -> NodeCapabilities { + let gpu = detect_gpu(); + + let (free_mem, total_mem) = { + let info = sys_info::mem_info().ok(); + ( + info.as_ref().map(|i| i.avail / 1024).unwrap_or(0), + info.as_ref().map(|i| i.total / 1024).unwrap_or(0), + ) + }; + + NodeCapabilities { + gpu, + loaded_model: None, // Set by the agent after model loads + available_tools: vec![], // Set by the tool registry + free_memory_mb: free_mem, + total_memory_mb: total_mem, + } +} + +/// Compute current system load as 0.0-1.0. +fn compute_load() -> f32 { + if let Ok(info) = sys_info::loadavg() { + let cpus = sys_info::cpu_num().unwrap_or(1) as f64; + (info.one / cpus).min(1.0) as f32 + } else { + 0.0 + } +} + +/// Detect NVIDIA GPU via nvidia-smi. +fn detect_gpu() -> Option { + let output = std::process::Command::new("nvidia-smi") + .args([ + "--query-gpu=name,memory.total,memory.free", + "--format=csv,noheader,nounits", + ]) + .output() + .ok()?; + + if !output.status.success() { + return None; + } + + let line = String::from_utf8_lossy(&output.stdout); + let parts: Vec<&str> = line.trim().split(", ").collect(); + if parts.len() < 3 { + return None; + } + + Some(GpuInfo { + name: parts[0].to_string(), + vram_total_mb: parts[1].parse().unwrap_or(0), + vram_free_mb: parts[2].parse().unwrap_or(0), + compute_capability: None, + }) +} diff --git a/src/cluster/overlay.rs b/src/cluster/overlay.rs new file mode 100644 index 00000000..f56275d1 --- /dev/null +++ b/src/cluster/overlay.rs @@ -0,0 +1,389 @@ +//! WebSocket overlay mesh with post-quantum encrypted channels. +//! +//! Each peer connection is a WebSocket tunnel encrypted with AES-256-GCM +//! after an ML-KEM-768 key exchange handshake. + +use std::collections::HashMap; +use std::net::SocketAddr; +use std::sync::Arc; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::{mpsc, RwLock}; +use tokio_tungstenite::tungstenite::Message as WsMessage; +use tokio_tungstenite::{accept_async, connect_async}; +use futures::{SinkExt, StreamExt}; +use aes_gcm::Aes256Gcm; +use serde::{Deserialize, Serialize}; + +use super::config::ClusterConfig; +use super::crypto::{self, MeshIdentity}; +use super::gossip::GossipState; +use super::types::*; + +/// A connected peer with its encrypted channel. +pub struct PeerConnection { + pub node_id: NodeId, + pub addr: SocketAddr, + pub send_tx: mpsc::Sender>, + pub send_nonce: std::sync::atomic::AtomicU64, +} + +impl std::fmt::Debug for PeerConnection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PeerConnection") + .field("node_id", &node_id_hex(&self.node_id)) + .field("addr", &self.addr) + .finish() + } +} + +/// Handshake message sent during PQ key exchange. +#[derive(Serialize, Deserialize)] +struct HandshakeInit { + node_id: NodeId, + kem_pk: Vec, + verify_pk: Vec, +} + +#[derive(Serialize, Deserialize)] +struct HandshakeReply { + node_id: NodeId, + ciphertext: Vec, + verify_pk: Vec, + signature: Vec, +} + +/// Manages the overlay mesh connections. +pub struct OverlayMesh { + pub identity: Arc, + pub peers: RwLock>, + pub gossip: Arc, + pub config: ClusterConfig, + pub incoming_tx: mpsc::Sender<(NodeId, MeshMessage)>, +} + +impl OverlayMesh { + pub fn new( + identity: Arc, + gossip: Arc, + config: ClusterConfig, + incoming_tx: mpsc::Sender<(NodeId, MeshMessage)>, + ) -> Self { + Self { + identity, + peers: RwLock::new(HashMap::new()), + gossip, + config, + incoming_tx, + } + } + + /// Start the WebSocket listener for incoming peer connections. + pub async fn start_listener(self: &Arc) -> anyhow::Result<()> { + let listen_addr = format!("{}:{}", self.config.bind_addr, self.config.mesh_port); + let listener = TcpListener::bind(&listen_addr).await?; + tracing::info!("Mesh overlay listening on {}", listen_addr); + + let mesh = self.clone(); + tokio::spawn(async move { + loop { + match listener.accept().await { + Ok((stream, addr)) => { + let mesh = mesh.clone(); + tokio::spawn(async move { + if let Err(e) = mesh.handle_incoming(stream, addr).await { + tracing::debug!("Incoming peer {} handshake failed: {}", addr, e); + } + }); + } + Err(e) => { + tracing::warn!("Mesh accept error: {}", e); + tokio::time::sleep(std::time::Duration::from_secs(1)).await; + } + } + } + }); + + Ok(()) + } + + /// Handle an incoming WebSocket connection (responder side of PQ handshake). + async fn handle_incoming(&self, stream: TcpStream, addr: SocketAddr) -> anyhow::Result<()> { + let ws = accept_async(stream).await?; + let (mut ws_tx, mut ws_rx) = ws.split(); + + // Step 1: Receive initiator's handshake (kem_pk + verify_pk + node_id) + let init_msg = ws_rx + .next() + .await + .ok_or_else(|| anyhow::anyhow!("No handshake received"))??; + + let init: HandshakeInit = serde_json::from_slice(&init_msg.into_data())?; + + // Dedup: if lower ID should be connector, reject + if init.node_id > self.identity.node_id { + // We have the lower ID -- we should be the connector, not the responder + // Drop this connection; we'll connect outbound instead + return Ok(()); + } + + // Step 2: Encapsulate shared secret with initiator's KEM public key + let (shared_secret, ciphertext) = crypto::encapsulate(&init.kem_pk)?; + + // Step 3: Sign the ciphertext + our node_id + let mut sign_data = Vec::new(); + sign_data.extend_from_slice(&ciphertext); + sign_data.extend_from_slice(&self.identity.node_id); + let signature = self.identity.sign(&sign_data); + + let reply = HandshakeReply { + node_id: self.identity.node_id, + ciphertext, + verify_pk: self.identity.verify_key_bytes(), + signature, + }; + + ws_tx + .send(WsMessage::Binary(serde_json::to_vec(&reply)?.into())) + .await?; + + // Step 4: Derive symmetric keys + let (send_cipher, recv_cipher) = + crypto::derive_keys(&shared_secret, &self.identity.node_id, &init.node_id); + + // Connection established -- run the encrypted message loop + self.run_peer_loop(init.node_id, addr, ws_tx, ws_rx, send_cipher, recv_cipher) + .await; + + Ok(()) + } + + /// Initiate an outbound connection to a discovered peer (initiator side). + pub async fn connect_to_peer(self: &Arc, peer_addr: SocketAddr, peer_node_id: &NodeId) -> anyhow::Result<()> { + // Dedup: lower NodeId is always the connector + if self.identity.node_id >= *peer_node_id { + // We have the higher ID -- wait for the other side to connect to us + return Ok(()); + } + + if self.is_connected(peer_node_id).await { + return Ok(()); + } + + let url = format!("ws://{}", peer_addr); + let (ws, _) = connect_async(&url).await?; + let (mut ws_tx, mut ws_rx) = ws.split(); + + // Step 1: Send our handshake (kem_pk + verify_pk + node_id) + let init = HandshakeInit { + node_id: self.identity.node_id, + kem_pk: self.identity.kem_pk_bytes(), + verify_pk: self.identity.verify_key_bytes(), + }; + + ws_tx + .send(WsMessage::Binary(serde_json::to_vec(&init)?.into())) + .await?; + + // Step 2: Receive responder's reply (ciphertext + verify_pk + signature) + let reply_msg = ws_rx + .next() + .await + .ok_or_else(|| anyhow::anyhow!("No handshake reply"))??; + + let reply: HandshakeReply = serde_json::from_slice(&reply_msg.into_data())?; + + // Step 3: Verify signature + let peer_verify_key = ed25519_dalek::VerifyingKey::from_bytes( + reply.verify_pk.as_slice().try_into() + .map_err(|_| anyhow::anyhow!("Invalid verify key length"))?, + )?; + + let mut sign_data = Vec::new(); + sign_data.extend_from_slice(&reply.ciphertext); + sign_data.extend_from_slice(&reply.node_id); + + if !MeshIdentity::verify(&peer_verify_key, &sign_data, &reply.signature) { + anyhow::bail!("Invalid handshake signature from peer"); + } + + // Step 4: Decapsulate shared secret + let shared_secret = self.identity.decapsulate(&reply.ciphertext)?; + + // Step 5: Derive symmetric keys (we are the initiator) + let (send_cipher, recv_cipher) = + crypto::derive_keys(&shared_secret, &self.identity.node_id, &reply.node_id); + + tracing::info!( + "PQ handshake complete with {} (ML-KEM-768 + AES-256-GCM)", + node_id_hex(&reply.node_id) + ); + + // Run encrypted message loop + let mesh = self.clone(); + let peer_id = reply.node_id; + tokio::spawn(async move { + mesh.run_peer_loop(peer_id, peer_addr, ws_tx, ws_rx, send_cipher, recv_cipher) + .await; + }); + + Ok(()) + } + + /// Run the encrypted message loop for a connected peer. + async fn run_peer_loop( + &self, + peer_id: NodeId, + addr: SocketAddr, + mut ws_tx: S, + mut ws_rx: R, + send_cipher: Aes256Gcm, + recv_cipher: Aes256Gcm, + ) where + S: futures::Sink + Unpin + Send + 'static, + R: futures::Stream> + Unpin + Send + 'static, + { + // Channel for outgoing messages (plaintext bytes, encrypted before sending) + let (out_tx, mut out_rx) = mpsc::channel::>(64); + + // Register peer + self.add_peer(peer_id, addr, out_tx).await; + + let incoming_tx = self.incoming_tx.clone(); + let send_cipher = Arc::new(send_cipher); + let recv_cipher = Arc::new(recv_cipher); + + // Sender task: encrypt and send + let send_c = send_cipher.clone(); + let sender = tokio::spawn(async move { + let mut nonce: u64 = 0; + while let Some(plaintext) = out_rx.recv().await { + nonce += 1; + let frame = crypto::encrypt(&send_c, nonce, &plaintext); + if ws_tx.send(WsMessage::Binary(frame.into())).await.is_err() { + break; + } + } + }); + + // Receiver task: decrypt and dispatch + let recv_c = recv_cipher.clone(); + let receiver = tokio::spawn(async move { + while let Some(Ok(msg)) = ws_rx.next().await { + let data = match msg { + WsMessage::Binary(b) => b.to_vec(), + WsMessage::Ping(_) | WsMessage::Pong(_) => continue, + WsMessage::Close(_) => break, + _ => continue, + }; + + match crypto::decrypt(&recv_c, &data) { + Ok(plaintext) => { + match serde_json::from_slice::(&plaintext) { + Ok(mesh_msg) => { + if incoming_tx.send((peer_id, mesh_msg)).await.is_err() { + break; + } + } + Err(e) => { + tracing::warn!("Invalid mesh message from {}: {}", node_id_hex(&peer_id), e); + } + } + } + Err(e) => { + tracing::warn!("Decrypt failed from {}: {}", node_id_hex(&peer_id), e); + break; // Crypto failure = drop connection + } + } + } + }); + + // Wait for either task to finish, then clean up + tokio::select! { + _ = sender => {}, + _ = receiver => {}, + } + + self.remove_peer(&peer_id).await; + } + + /// Check if we're connected to a peer. + pub async fn is_connected(&self, node_id: &NodeId) -> bool { + self.peers.read().await.contains_key(node_id) + } + + /// Number of active peer connections. + pub async fn peer_count(&self) -> usize { + self.peers.read().await.len() + } + + /// Send a message to a specific peer. + pub async fn send_to(&self, node_id: &NodeId, msg: &MeshMessage) -> anyhow::Result<()> { + let peers = self.peers.read().await; + let peer = peers + .get(node_id) + .ok_or_else(|| anyhow::anyhow!("Peer not connected"))?; + + let json = serde_json::to_vec(msg)?; + peer.send_tx + .send(json) + .await + .map_err(|_| anyhow::anyhow!("Peer channel closed")) + } + + /// Broadcast a message to all connected peers. + pub async fn broadcast(&self, msg: &MeshMessage) { + let json = match serde_json::to_vec(msg) { + Ok(j) => j, + Err(e) => { + tracing::error!("Failed to serialize mesh message: {}", e); + return; + } + }; + + let peers = self.peers.read().await; + for (id, peer) in peers.iter() { + if peer.send_tx.send(json.clone()).await.is_err() { + tracing::warn!("Failed to send to peer {}", node_id_hex(id)); + } + } + } + + /// Register a new peer connection. + pub async fn add_peer(&self, node_id: NodeId, addr: SocketAddr, send_tx: mpsc::Sender>) { + let mut peers = self.peers.write().await; + + if peers.len() >= self.config.max_peers && !peers.contains_key(&node_id) { + tracing::debug!( + "Max peers ({}) reached, rejecting {}", + self.config.max_peers, + node_id_hex(&node_id) + ); + return; + } + + tracing::info!("Peer connected: {} ({})", node_id_hex(&node_id), addr); + + peers.insert( + node_id, + PeerConnection { + node_id, + addr, + send_tx, + send_nonce: std::sync::atomic::AtomicU64::new(0), + }, + ); + } + + /// Remove a peer connection. + pub async fn remove_peer(&self, node_id: &NodeId) { + let mut peers = self.peers.write().await; + if peers.remove(node_id).is_some() { + tracing::info!("Peer disconnected: {}", node_id_hex(node_id)); + } + } + + /// Get list of connected peer IDs. + pub async fn connected_peer_ids(&self) -> Vec { + self.peers.read().await.keys().cloned().collect() + } +} diff --git a/src/cluster/router.rs b/src/cluster/router.rs new file mode 100644 index 00000000..1499fdd5 --- /dev/null +++ b/src/cluster/router.rs @@ -0,0 +1,173 @@ +//! Intelligent task routing across the mesh. +//! +//! Scores nodes based on capability match, load, VRAM, and hop distance. + +use super::gossip::GossipState; +use super::types::*; +use std::sync::Arc; + +/// Score a node for a given task. Returns None if hard requirements aren't met. +pub fn score_node(node: &NodeInfo, task: &TaskEnvelope) -> Option { + // Hard filters + if let Some(ref model) = task.required_model { + if node.capabilities.loaded_model.as_ref() != Some(model) { + return None; + } + } + + for tool in &task.required_tools { + if !node.capabilities.available_tools.contains(tool) { + return None; + } + } + + if let Some(min_vram) = task.min_vram_mb { + let free = node + .capabilities + .gpu + .as_ref() + .map(|g| g.vram_free_mb) + .unwrap_or(0); + if free < min_vram { + return None; + } + } + + // Soft scoring (higher is better) + let load_score = 1.0 - node.load.clamp(0.0, 1.0) as f64; + + let vram_score = node + .capabilities + .gpu + .as_ref() + .map(|g| { + if g.vram_total_mb > 0 { + g.vram_free_mb as f64 / g.vram_total_mb as f64 + } else { + 0.0 + } + }) + .unwrap_or(0.0); + + let memory_score = if node.capabilities.total_memory_mb > 0 { + node.capabilities.free_memory_mb as f64 / node.capabilities.total_memory_mb as f64 + } else { + 0.0 + }; + + let hop_score = if task.hop_count == 0 { + 1.0 + } else { + 1.0 / (task.hop_count as f64 + 1.0) + }; + + let model_bonus = if task.required_model.is_some() + && node.capabilities.loaded_model == task.required_model + { + 1.0 + } else { + 0.0 + }; + + Some( + load_score * 0.35 + + vram_score * 0.25 + + memory_score * 0.15 + + hop_score * 0.15 + + model_bonus * 0.10, + ) +} + +/// Select the best node for a task from the gossip membership. +pub async fn route_task( + gossip: &GossipState, + local_info: &NodeInfo, + task: &TaskEnvelope, +) -> Option { + let mut candidates: Vec<(NodeInfo, f64)> = Vec::new(); + + // Score local node + if let Some(score) = score_node(local_info, task) { + candidates.push((local_info.clone(), score)); + } + + // Score peers + for peer in gossip.alive_peers().await { + if let Some(score) = score_node(&peer, task) { + candidates.push((peer, score)); + } + } + + // Sort by score descending + candidates.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + + candidates.into_iter().next().map(|(node, _)| node) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::SocketAddr; + + fn make_node(load: f32, vram_free: u64, vram_total: u64, model: Option<&str>) -> NodeInfo { + NodeInfo { + id: [0; 16], + hostname: "test".into(), + mesh_addr: "127.0.0.1:9901".parse().unwrap(), + gateway_addr: None, + capabilities: NodeCapabilities { + gpu: Some(GpuInfo { + name: "RTX 3060".into(), + vram_total_mb: vram_total, + vram_free_mb: vram_free, + compute_capability: None, + }), + loaded_model: model.map(|s| s.to_string()), + available_tools: vec!["shell".into(), "read_file".into()], + free_memory_mb: 8000, + total_memory_mb: 16000, + }, + load, + version: "0.22.0".into(), + started_at: 0, + last_heartbeat: 0, + } + } + + fn make_task() -> TaskEnvelope { + TaskEnvelope { + task_id: "test".into(), + content: "test task".into(), + origin_node: [0; 16], + required_model: None, + required_tools: vec![], + min_vram_mb: None, + priority: 0, + hop_count: 0, + } + } + + #[test] + fn test_idle_node_preferred() { + let task = make_task(); + let idle = make_node(0.1, 4000, 6000, None); + let busy = make_node(0.9, 4000, 6000, None); + assert!(score_node(&idle, &task).unwrap() > score_node(&busy, &task).unwrap()); + } + + #[test] + fn test_model_hard_filter() { + let mut task = make_task(); + task.required_model = Some("llama3".into()); + let node = make_node(0.0, 4000, 6000, Some("qwen2.5")); + assert!(score_node(&node, &task).is_none()); + } + + #[test] + fn test_vram_hard_filter() { + let mut task = make_task(); + task.min_vram_mb = Some(8000); + let node = make_node(0.0, 4000, 6000, None); + assert!(score_node(&node, &task).is_none()); + } +} diff --git a/src/cluster/types.rs b/src/cluster/types.rs new file mode 100644 index 00000000..57d15071 --- /dev/null +++ b/src/cluster/types.rs @@ -0,0 +1,175 @@ +//! Core types for the OptimClaw mesh network. + +use serde::{Deserialize, Serialize}; +use std::net::SocketAddr; + +/// 16-byte node identifier, derived from blake3 hash of the node's public key. +pub type NodeId = [u8; 16]; + +/// Hex-encode a NodeId for display. +pub fn node_id_hex(id: &NodeId) -> String { + hex::encode(id) +} + +/// Parse a hex string back to NodeId. +pub fn node_id_from_hex(s: &str) -> Option { + let bytes = hex::decode(s).ok()?; + if bytes.len() != 16 { + return None; + } + let mut id = [0u8; 16]; + id.copy_from_slice(&bytes); + Some(id) +} + +/// Information about a node in the mesh. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct NodeInfo { + pub id: NodeId, + pub hostname: String, + pub mesh_addr: SocketAddr, + pub gateway_addr: Option, + pub capabilities: NodeCapabilities, + pub load: f32, + pub version: String, + pub started_at: i64, + pub last_heartbeat: i64, +} + +/// What a node can do. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct NodeCapabilities { + pub gpu: Option, + pub loaded_model: Option, + pub available_tools: Vec, + pub free_memory_mb: u64, + pub total_memory_mb: u64, +} + +/// GPU information. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GpuInfo { + pub name: String, + pub vram_total_mb: u64, + pub vram_free_mb: u64, + pub compute_capability: Option, +} + +/// Messages exchanged between mesh nodes (over encrypted WebSocket). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum MeshMessage { + // -- Gossip / Membership -- + Ping { + from: NodeId, + seq: u64, + piggyback: Vec, + }, + Ack { + from: NodeId, + seq: u64, + piggyback: Vec, + }, + IndirectPing { + origin: NodeId, + target: NodeId, + seq: u64, + }, + IndirectAck { + origin: NodeId, + target: NodeId, + seq: u64, + alive: bool, + }, + + // -- Peer exchange -- + PeerExchange { + nodes: Vec, + }, + + // -- Task routing -- + TaskRequest { + task_id: String, + envelope: TaskEnvelope, + }, + TaskAccept { + task_id: String, + node_id: NodeId, + }, + TaskReject { + task_id: String, + node_id: NodeId, + reason: String, + }, + TaskStream { + task_id: String, + chunk: TaskStreamChunk, + }, + TaskComplete { + task_id: String, + result: TaskResult, + }, + + // -- Lifecycle -- + Leaving { + id: NodeId, + }, +} + +/// Piggybacked membership state change. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MembershipUpdate { + pub node_id: NodeId, + pub incarnation: u64, + pub status: MemberStatus, +} + +/// Node status in the membership table. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum MemberStatus { + Alive, + Suspect, + Dead, + Left, +} + +/// A task to be routed across the mesh. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskEnvelope { + pub task_id: String, + pub content: String, + pub origin_node: NodeId, + pub required_model: Option, + pub required_tools: Vec, + pub min_vram_mb: Option, + pub priority: u8, + pub hop_count: u8, +} + +/// Streaming chunk from a task execution. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "kind")] +pub enum TaskStreamChunk { + Text { content: String }, + ToolStarted { name: String, summary: String }, + ToolCompleted { name: String, success: bool, output: String }, + Thinking, +} + +/// Final result of a task. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskResult { + pub success: bool, + pub response: String, + pub node_id: NodeId, + pub duration_ms: u64, +} + +/// UDP beacon packet (compact binary format). +pub const BEACON_MAGIC: &[u8; 6] = b"OMESH1"; +pub const BEACON_PORT: u16 = 9900; +pub const MESH_PORT: u16 = 9901; + +/// Beacon flags. +pub const FLAG_HAS_GPU: u8 = 0x01; +pub const FLAG_ACCEPTING_TASKS: u8 = 0x02; diff --git a/src/lib.rs b/src/lib.rs index dbdd2260..5d97591d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -43,6 +43,8 @@ pub mod app; pub mod boot_screen; pub mod bootstrap; pub mod channels; +#[cfg(feature = "cluster")] +pub mod cluster; pub mod cli; pub mod config; pub mod context;