mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
feat: autonomous AI mesh network with post-quantum crypto
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
Generated
+186
-130
@@ -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]]
|
||||
|
||||
@@ -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"
|
||||
|
||||
+23
-2
@@ -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() {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user