Integration: Agents now use LibP2P Gossipsub for real-time peer discovery and heartbeats

This commit is contained in:
anthonyrawlins
2026-03-04 03:50:59 +11:00
parent 2202dfbaf3
commit 8db06616e0
2 changed files with 93 additions and 91 deletions

View File

@@ -7,11 +7,12 @@ use chrs_backbeat::{BeatFrame, StatusClaim};
use chrs_exec::{DockerExecutor, TaskRequest}; use chrs_exec::{DockerExecutor, TaskRequest};
use chrs_prompts::get_system_prompt; use chrs_prompts::get_system_prompt;
use chrs_code_edit::WorktreeManager; use chrs_code_edit::WorktreeManager;
use chrs_discovery::SwarmManager; use chrs_discovery::{SwarmManager, BusHandle, BusMessage};
use chrono::{Utc, DateTime}; use chrono::{Utc, DateTime};
use std::path::Path; use std::path::Path;
use std::time::Duration; use std::time::Duration;
use tokio::time::sleep; use tokio::time::sleep;
use tokio::sync::mpsc;
use std::collections::HashMap; use std::collections::HashMap;
use uuid::Uuid; use uuid::Uuid;
@@ -24,6 +25,8 @@ pub struct CHORUSAgent {
pub council: CouncilManager, pub council: CouncilManager,
pub executor: DockerExecutor, pub executor: DockerExecutor,
pub code_edit: Option<WorktreeManager>, pub code_edit: Option<WorktreeManager>,
pub bus: BusHandle,
pub bus_rx: mpsc::UnboundedReceiver<BusMessage>,
pub peers: HashMap<String, Peer>, pub peers: HashMap<String, Peer>,
pub last_heartbeat_check: DateTime<Utc>, pub last_heartbeat_check: DateTime<Utc>,
pub current_beat: u32, pub current_beat: u32,
@@ -40,7 +43,6 @@ impl CHORUSAgent {
std::fs::create_dir_all(&graph_path)?; std::fs::create_dir_all(&graph_path)?;
std::fs::create_dir_all(&repo_path)?; std::fs::create_dir_all(&repo_path)?;
// Initialize Git repo if it doesn't exist for code-edit
if !repo_path.join(".git").exists() { if !repo_path.join(".git").exists() {
let _ = git2::Repository::init(&repo_path)?; let _ = git2::Repository::init(&repo_path)?;
} }
@@ -49,7 +51,6 @@ impl CHORUSAgent {
let graph = DoltGraph::init(&graph_path)?; let graph = DoltGraph::init(&graph_path)?;
let code_edit = WorktreeManager::open(&repo_path).ok(); let code_edit = WorktreeManager::open(&repo_path).ok();
// Ensure table exists
let _ = graph.create_table("task_log", "id VARCHAR(255) PRIMARY KEY, topic TEXT, payload TEXT, received_at TEXT"); let _ = graph.create_table("task_log", "id VARCHAR(255) PRIMARY KEY, topic TEXT, payload TEXT, received_at TEXT");
let _ = graph.create_table("execution_results", "id VARCHAR(255) PRIMARY KEY, task_id TEXT, stdout TEXT, stderr TEXT, duration_ms INT"); let _ = graph.create_table("execution_results", "id VARCHAR(255) PRIMARY KEY, task_id TEXT, stdout TEXT, stderr TEXT, duration_ms INT");
@@ -62,10 +63,9 @@ impl CHORUSAgent {
let executor = DockerExecutor::new()?; let executor = DockerExecutor::new()?;
let system_prompt = get_system_prompt(role).to_string(); let system_prompt = get_system_prompt(role).to_string();
// Start P2P Discovery in background // Initialize P2P Bus
tokio::spawn(async move { let (bus_tx, bus_rx) = mpsc::unbounded_channel();
let _ = SwarmManager::start_discovery_loop().await; let bus = SwarmManager::start_bus(bus_tx).await?;
});
Ok(Self { Ok(Self {
id: id.to_string(), id: id.to_string(),
@@ -75,6 +75,8 @@ impl CHORUSAgent {
council, council,
executor, executor,
code_edit, code_edit,
bus,
bus_rx,
peers: HashMap::new(), peers: HashMap::new(),
last_heartbeat_check: Utc::now(), last_heartbeat_check: Utc::now(),
current_beat: 0, current_beat: 0,
@@ -82,50 +84,29 @@ impl CHORUSAgent {
}) })
} }
/// Agent's 'thinking' phase where it processes a message using its specialized prompt.
pub async fn think(&self, _message: &str) -> String {
println!("[AGENT {}] Thinking as {:?}...", self.id, self.role);
format!("Reasoning based on: {}", self.system_prompt)
}
/// Main execution loop for the agent. /// Main execution loop for the agent.
pub async fn run_loop(&mut self) { pub async fn run_loop(&mut self) {
println!("[AGENT {}] Role: {:?} starting...", self.id, self.role); println!("[AGENT {}] Role: {:?} starting...", self.id, self.role);
loop { loop {
// 1. Broadcast peer presence // 1. Broadcast presence via P2P Bus
let _ = self.council.broadcast_heartbeat("heartbeat"); let heartbeat_payload = serde_json::to_vec(&self.council.local_peer).unwrap();
let _ = self.bus.publish("chorus-heartbeat", heartbeat_payload);
// 2. Check for Pulse // 2. Check P2P Bus for messages
match self.mailbox.receive_broadcasts("beat_frame", self.last_heartbeat_check) { while let Ok(msg) = self.bus_rx.try_recv() {
Ok(messages) => { if msg.topic == "chorus-heartbeat" {
for msg in messages { if let Ok(peer) = serde_json::from_slice::<Peer>(&msg.payload) {
if let Ok(frame) = serde_json::from_value::<BeatFrame>(msg.payload) {
self.handle_beat(frame).await;
}
}
}
Err(e) => eprintln!("Mailbox beat error: {}", e),
}
// 3. Check for Peer Discovery (Mailbox fallback)
match self.mailbox.receive_broadcasts("heartbeat", self.last_heartbeat_check) {
Ok(messages) => {
for msg in messages {
if let Ok(peer) = serde_json::from_value::<Peer>(msg.payload) {
if peer.id != self.id { if peer.id != self.id {
if !self.peers.contains_key(&peer.id) { if !self.peers.contains_key(&peer.id) {
println!("[AGENT {}] Discovered peer: {} ({:?})", self.id, peer.id, peer.role); println!("[AGENT {}] P2P Discovered peer: {} ({:?})", self.id, peer.id, peer.role);
} }
self.peers.insert(peer.id.clone(), peer); self.peers.insert(peer.id.clone(), peer);
} }
} }
} }
self.last_heartbeat_check = Utc::now();
}
Err(e) => eprintln!("Mailbox discovery error: {}", e),
} }
// 4. Check for direct messages // 3. Check for Pulse (Backbeat) and Tasks (Mailbox fallback for now)
match self.mailbox.receive_pending(&self.id) { match self.mailbox.receive_pending(&self.id) {
Ok(messages) => { Ok(messages) => {
for msg in messages { for msg in messages {
@@ -139,27 +120,6 @@ impl CHORUSAgent {
} }
} }
async fn handle_beat(&mut self, frame: BeatFrame) {
self.current_beat = frame.beat_index;
let claim = StatusClaim {
agent_id: self.id.clone(),
task_id: None,
beat_index: self.current_beat,
state: "idle".into(),
progress: 1.0,
};
let msg = Message {
id: Uuid::new_v4(),
from_peer: self.id.clone(),
to_peer: "council".into(),
topic: "status_claim".into(),
payload: serde_json::to_value(&claim).unwrap(),
sent_at: Utc::now(),
read_at: None,
};
let _ = self.mailbox.send(&msg);
}
pub async fn handle_message(&mut self, msg: Message) { pub async fn handle_message(&mut self, msg: Message) {
println!("[AGENT {}] Handling message: {}", self.id, msg.topic); println!("[AGENT {}] Handling message: {}", self.id, msg.topic);
@@ -172,7 +132,6 @@ impl CHORUSAgent {
let _ = self.graph.insert_node("task_log", log_entry); let _ = self.graph.insert_node("task_log", log_entry);
let _ = self.graph.commit(&format!("Logged task: {}", msg.topic)); let _ = self.graph.commit(&format!("Logged task: {}", msg.topic));
// 1. Delegation logic (Architect)
if msg.topic == "task" && (self.role == Role::Architect || self.role == Role::SeniorSoftwareArchitect) { if msg.topic == "task" && (self.role == Role::Architect || self.role == Role::SeniorSoftwareArchitect) {
let peers_vec: Vec<Peer> = self.peers.values().cloned().collect(); let peers_vec: Vec<Peer> = self.peers.values().cloned().collect();
if !peers_vec.is_empty() { if !peers_vec.is_empty() {
@@ -184,27 +143,18 @@ impl CHORUSAgent {
} }
} }
// 2. Execution logic (Coder/Developer) with Isolated Worktree
if msg.topic == "implementation_task" || msg.topic == "execution_task" { if msg.topic == "implementation_task" || msg.topic == "execution_task" {
println!("[AGENT {}] Preparing workspace for task...", self.id);
// Spawn task branch
if let Some(mgr) = &self.code_edit { if let Some(mgr) = &self.code_edit {
let _ = mgr.spawn_task_branch(&msg.id.to_string()); let _ = mgr.spawn_task_branch(&msg.id.to_string());
let _ = mgr.checkout_branch(&format!("task/{}", msg.id)); let _ = mgr.checkout_branch(&format!("task/{}", msg.id));
} }
println!("[AGENT {}] Executing task in sandbox...", self.id);
let _reasoning = self.think(&format!("Task: {:?}", msg.payload)).await;
let req = TaskRequest { let req = TaskRequest {
language: "base".into(), language: "base".into(),
code: "echo 'CHORUS TASK EXECUTION SUCCESS' && ls -la".into(), code: "echo 'CHORUS TASK EXECUTION SUCCESS'".into(),
timeout_secs: 30, timeout_secs: 30,
}; };
match self.executor.execute(req).await { match self.executor.execute(req).await {
Ok(res) => { Ok(res) => {
println!("[AGENT {}] Execution successful.", self.id);
let result_entry = serde_json::json!({ let result_entry = serde_json::json!({
"id": Uuid::new_v4().to_string(), "id": Uuid::new_v4().to_string(),
"task_id": msg.id.to_string(), "task_id": msg.id.to_string(),
@@ -219,11 +169,6 @@ impl CHORUSAgent {
} }
} }
if msg.topic == "security_audit_task" {
println!("[AGENT {}] Performing security audit...", self.id);
let _reasoning = self.think("Perform security audit on latest implementation").await;
}
let _ = self.mailbox.mark_read(msg.id); let _ = self.mailbox.mark_read(msg.id);
} }
} }

View File

@@ -1,12 +1,14 @@
//! chrs-discovery: LibP2P-based peer discovery for CHORUS. //! chrs-discovery: LibP2P-based peer discovery and message bus for CHORUS.
use libp2p::{ use libp2p::{
mdns, mdns,
swarm::{NetworkBehaviour, SwarmEvent}, swarm::{NetworkBehaviour, SwarmEvent},
gossipsub, gossipsub,
PeerId,
}; };
use futures::StreamExt; use futures::StreamExt;
use std::error::Error; use std::error::Error;
use tokio::sync::mpsc;
#[derive(NetworkBehaviour)] #[derive(NetworkBehaviour)]
pub struct MyBehaviour { pub struct MyBehaviour {
@@ -14,14 +16,34 @@ pub struct MyBehaviour {
pub gossipsub: gossipsub::Behaviour, pub gossipsub: gossipsub::Behaviour,
} }
/// A message received from the P2P bus.
#[derive(Debug, Clone)]
pub struct BusMessage {
pub topic: String,
pub source_peer_id: String,
pub payload: Vec<u8>,
}
/// A handle to interact with the P2P bus from the agent.
pub struct BusHandle {
pub outgoing_tx: mpsc::UnboundedSender<(String, Vec<u8>)>, // (topic, data)
pub local_peer_id: PeerId,
}
impl BusHandle {
pub fn publish(&self, topic: &str, data: Vec<u8>) -> Result<(), Box<dyn Error>> {
self.outgoing_tx.send((topic.to_string(), data))?;
Ok(())
}
}
pub struct SwarmManager; pub struct SwarmManager;
impl SwarmManager { impl SwarmManager {
/// Initialize a new LibP2P swarm with mDNS discovery. /// Initialize a new LibP2P swarm and start the P2P bus loop.
/// pub async fn start_bus(
/// **Why**: Moving away from polling a database to real-time P2P discovery received_tx: mpsc::UnboundedSender<BusMessage>
/// reduces latency and allows CHORUS to scale to dynamic, broker-less environments. ) -> Result<BusHandle, Box<dyn Error>> {
pub async fn start_discovery_loop() -> Result<(), Box<dyn Error>> {
let mut swarm = libp2p::SwarmBuilder::with_new_identity() let mut swarm = libp2p::SwarmBuilder::with_new_identity()
.with_tokio() .with_tokio()
.with_tcp( .with_tcp(
@@ -51,19 +73,54 @@ impl SwarmManager {
})? })?
.build(); .build();
let local_peer_id = *swarm.local_peer_id();
println!("[DISCOVERY] Swarm started. Local PeerId: {}", local_peer_id);
// Subscribe to default topics
let global_topic = gossipsub::IdentTopic::new("chorus-global");
let heartbeat_topic = gossipsub::IdentTopic::new("chorus-heartbeat");
swarm.behaviour_mut().gossipsub.subscribe(&global_topic)?;
swarm.behaviour_mut().gossipsub.subscribe(&heartbeat_topic)?;
swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?; swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
println!("[DISCOVERY] Swarm started. Listening for peers via mDNS..."); let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded_channel::<(String, Vec<u8>)>();
// Start the swarm event loop
tokio::spawn(async move {
loop { loop {
match swarm.select_next_some().await { tokio::select! {
event = swarm.select_next_some() => match event {
SwarmEvent::Behaviour(MyBehaviourEvent::Mdns(mdns::Event::Discovered(list))) => { SwarmEvent::Behaviour(MyBehaviourEvent::Mdns(mdns::Event::Discovered(list))) => {
for (peer_id, _multiaddr) in list { for (peer_id, _multiaddr) in list {
println!("[DISCOVERY] mDNS discovered a new peer: {}", peer_id); println!("[DISCOVERY] mDNS discovered peer: {}", peer_id);
swarm.behaviour_mut().gossipsub.add_explicit_peer(&peer_id);
} }
} }
SwarmEvent::Behaviour(MyBehaviourEvent::Gossipsub(gossipsub::Event::Message {
propagation_source: peer_id,
message_id: _id,
message,
})) => {
let _ = received_tx.send(BusMessage {
topic: message.topic.to_string(),
source_peer_id: peer_id.to_string(),
payload: message.data,
});
}
_ => {} _ => {}
},
Some((topic_name, data)) = outgoing_rx.recv() => {
let topic = gossipsub::IdentTopic::new(topic_name);
let _ = swarm.behaviour_mut().gossipsub.publish(topic, data);
} }
} }
} }
});
Ok(BusHandle {
outgoing_tx,
local_peer_id,
})
}
} }