gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 07:20:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: autonomous AI mesh network with post-quantum crypto

Browse Source 
Add peer-to-peer mesh cluster module (feature-gated behind `cluster`):

- UDP beacon discovery (port 9900) for zero-config LAN auto-discovery
- ML-KEM-768 (Kyber) post-quantum key exchange + AES-256-GCM encrypted
  WebSocket overlay mesh
- SWIM gossip protocol for membership and failure detection
- Intelligent task routing: scores nodes by load, VRAM, model match,
  hop distance
- Remote task execution via subprocess with streaming results
- REST API endpoints: /api/mesh/status, /api/mesh/nodes
- ed25519 signed beacons, persistent keypairs (~/.optimclaw/mesh_keys.json)
- Lazy tool loading (OPTIMCLAW_LAZY_TOOLS) to reduce system prompt size

Enable with: CLUSTER_ENABLED=1 optimclaw run

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
dingo
2026-03-29 13:10:35 +07:00
co-authored by Claude Opus 4.6
parent 6d9dbbb3b9
commit 93c27fa053
13 changed files with 2361 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+186 -130
View File
@@ -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]]
Cargo.toml
+8
View File
@@ -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"
src/agent/dispatcher.rs
+23 -2
View File
@@ -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() {
src/cluster/api.rs
+105
View File
@@ -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<Option<Arc<MeshNode>>> {
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<String>,
loaded_model: Option<String>,
tool_count: usize,
free_memory_mb: u64,
}
async fn mesh_status(
State(mesh): State<Option<Arc<MeshNode>>>,
) -> Json<MeshStatus> {
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<Option<Arc<MeshNode>>>,
) -> Json<Vec<MeshNodeInfo>> {
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)
}
src/cluster/beacon.rs
+136
View File
@@ -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<MeshIdentity>,
) -> anyhow::Result<mpsc::Receiver<DiscoveredPeer>> {
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<u8> {
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<DiscoveredPeer> {
// 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,
})
}
src/cluster/config.rs
+78
View File
@@ -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<T: std::str::FromStr>(key: &str, default: T) -> T {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
src/cluster/crypto.rs
+269
View File
@@ -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<u8>,
kem_sk: Vec<u8>,
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<Self> {
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<Self> {
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<u8> {
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<Vec<u8>> {
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<u8> {
self.kem_pk.as_bytes().to_vec()
}
/// Public verify key bytes for sharing.
pub fn verify_key_bytes(&self) -> Vec<u8> {
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<u8>, Vec<u8>)> {
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::<Sha256>::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<u8> {
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<Vec<u8>> {
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);
}
}
src/cluster/gossip.rs
+219
View File
@@ -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<Instant>,
}
/// The gossip membership table.
pub struct GossipState {
members: RwLock<HashMap<NodeId, MemberState>>,
local_id: NodeId,
local_incarnation: RwLock<u64>,
}
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<NodeInfo> {
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<NodeId>, Vec<NodeId>) {
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<NodeInfo> {
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::<usize>() % alive.len();
Some(alive[idx].info.clone())
}
/// Build piggybacked membership updates for gossip dissemination.
pub async fn membership_updates(&self) -> Vec<MembershipUpdate> {
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,
}
}
src/cluster/mod.rs
+598
View File
@@ -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<MeshIdentity>,
pub gossip: Arc<GossipState>,
pub overlay: Arc<OverlayMesh>,
pub event_tx: tokio::sync::broadcast::Sender<ClusterEvent>,
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<Arc<Self>> {
// 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<TaskResult> {
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<ClusterEvent> {
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<RwLock<HashMap<String, tokio::sync::oneshot::Sender<TaskResult>>>>;
/// 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<tokio::sync::broadcast::Sender<ClusterEvent>>,
) {
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::<Vec<&str>>()
.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<GpuInfo> {
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,
})
}
src/cluster/overlay.rs
+389
View File
@@ -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<Vec<u8>>,
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<u8>,
verify_pk: Vec<u8>,
}
#[derive(Serialize, Deserialize)]
struct HandshakeReply {
node_id: NodeId,
ciphertext: Vec<u8>,
verify_pk: Vec<u8>,
signature: Vec<u8>,
}
/// Manages the overlay mesh connections.
pub struct OverlayMesh {
pub identity: Arc<MeshIdentity>,
pub peers: RwLock<HashMap<NodeId, PeerConnection>>,
pub gossip: Arc<GossipState>,
pub config: ClusterConfig,
pub incoming_tx: mpsc::Sender<(NodeId, MeshMessage)>,
}
impl OverlayMesh {
pub fn new(
identity: Arc<MeshIdentity>,
gossip: Arc<GossipState>,
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<Self>) -> 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<Self>, 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<S, R>(
&self,
peer_id: NodeId,
addr: SocketAddr,
mut ws_tx: S,
mut ws_rx: R,
send_cipher: Aes256Gcm,
recv_cipher: Aes256Gcm,
) where
S: futures::Sink<WsMessage, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
R: futures::Stream<Item = Result<WsMessage, tokio_tungstenite::tungstenite::Error>> + Unpin + Send + 'static,
{
// Channel for outgoing messages (plaintext bytes, encrypted before sending)
let (out_tx, mut out_rx) = mpsc::channel::<Vec<u8>>(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::<MeshMessage>(&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<Vec<u8>>) {
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<NodeId> {
self.peers.read().await.keys().cloned().collect()
}
}
src/cluster/router.rs
+173
View File
@@ -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<f64> {
// 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<NodeInfo> {
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());
}
}
src/cluster/types.rs
+175
View File
@@ -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<NodeId> {
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<SocketAddr>,
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<GpuInfo>,
pub loaded_model: Option<String>,
pub available_tools: Vec<String>,
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<String>,
}
/// 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<MembershipUpdate>,
},
Ack {
from: NodeId,
seq: u64,
piggyback: Vec<MembershipUpdate>,
},
IndirectPing {
origin: NodeId,
target: NodeId,
seq: u64,
},
IndirectAck {
origin: NodeId,
target: NodeId,
seq: u64,
alive: bool,
},
// -- Peer exchange --
PeerExchange {
nodes: Vec<NodeInfo>,
},
// -- 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<String>,
pub required_tools: Vec<String>,
pub min_vram_mb: Option<u64>,
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;
src/lib.rs
+2
View File
@@ -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;