codetether_agent/bus/

mod.rs

//! Central in-process agent message bus
//!
//! Provides a broadcast-based pub/sub bus that all local agents (sub-agents,
//! workers, the A2A server) can plug into.  Remote agents talk to the bus
//! through the gRPC / JSON-RPC bridge; local agents get zero-copy clones
//! via `BusHandle`.
//!
//! # Topic routing
//!
//! Every envelope carries a `topic` string that follows a hierarchical scheme:
//!
//! | Pattern | Semantics |
//! |---------|-----------|
//! | `agent.{id}` | Messages *to* a specific agent |
//! | `agent.{id}.events` | Events *from* a specific agent |
//! | `task.{id}` | All updates for a task |
//! | `swarm.{id}` | Swarm-level coordination |
//! | `broadcast` | Global announcements |
//! | `tools.{name}` | Tool-specific channels |

pub mod registry;
pub mod relay;
pub mod s3_sink;

use crate::a2a::types::{Artifact, Part, TaskState};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::broadcast;
use uuid::Uuid;

// ─── Envelope & Messages ─────────────────────────────────────────────────

/// Metadata wrapper for every message that travels through the bus.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BusEnvelope {
    /// Unique id for this envelope
    pub id: String,
    /// Hierarchical topic for routing (e.g. `agent.{id}`, `task.{id}`)
    pub topic: String,
    /// The agent that originated this message
    pub sender_id: String,
    /// Optional correlation id (links request → response)
    pub correlation_id: Option<String>,
    /// When the envelope was created
    pub timestamp: DateTime<Utc>,
    /// The payload
    pub message: BusMessage,
}

/// The set of messages the bus can carry.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum BusMessage {
    /// An agent has come online and is ready to receive work
    AgentReady {
        agent_id: String,
        capabilities: Vec<String>,
    },
    /// An agent is shutting down
    AgentShutdown { agent_id: String },
    /// Free-form message from one agent to another (mirrors A2A `Message`)
    AgentMessage {
        from: String,
        to: String,
        parts: Vec<Part>,
    },
    /// Task status changed
    TaskUpdate {
        task_id: String,
        state: TaskState,
        message: Option<String>,
    },
    /// A new artifact was produced for a task
    ArtifactUpdate { task_id: String, artifact: Artifact },
    /// A sub-agent published a shared result (replaces raw `ResultStore` publish)
    SharedResult {
        key: String,
        value: serde_json::Value,
        tags: Vec<String>,
    },
    /// Tool execution request (for shared tool dispatch)
    ToolRequest {
        request_id: String,
        agent_id: String,
        tool_name: String,
        arguments: serde_json::Value,
    },
    /// Tool execution response
    ToolResponse {
        request_id: String,
        agent_id: String,
        tool_name: String,
        result: String,
        success: bool,
    },
    /// Heartbeat (keep-alive / health signal)
    Heartbeat { agent_id: String, status: String },

    // ── Ralph loop messages ──────────────────────────────────────────
    /// An agent shares learnings from a Ralph iteration (insights,
    /// blockers, patterns discovered) so subsequent iterations or
    /// co-ordinating agents can build on them.
    RalphLearning {
        prd_id: String,
        story_id: String,
        iteration: usize,
        learnings: Vec<String>,
        context: serde_json::Value,
    },

    /// Context handoff between sequential Ralph stories.  The finishing
    /// story publishes what the *next* story should know.
    RalphHandoff {
        prd_id: String,
        from_story: String,
        to_story: String,
        context: serde_json::Value,
        progress_summary: String,
    },

    /// PRD-level progress update (aggregated across all stories).
    RalphProgress {
        prd_id: String,
        passed: usize,
        total: usize,
        iteration: usize,
        status: String,
    },

    // ── Full tool output ─────────────────────────────────────────────
    /// Full, untruncated tool output from an agent loop step.
    /// Published *before* the result is truncated for the LLM context
    /// window, so other agents and the bus log see the complete output.
    ToolOutputFull {
        agent_id: String,
        tool_name: String,
        output: String,
        success: bool,
        step: usize,
    },

    /// Internal reasoning / chain-of-thought from an agent step.
    /// Published so the training pipeline captures the model's
    /// thinking process alongside its actions.
    AgentThinking {
        agent_id: String,
        thinking: String,
        step: usize,
    },

    // ── Voice session messages ───────────────────────────────────────
    /// A voice session (LiveKit room) has been created.
    VoiceSessionStarted {
        room_name: String,
        agent_id: String,
        voice_id: String,
    },

    /// A transcript fragment from a voice session.
    VoiceTranscript {
        room_name: String,
        text: String,
        /// "user" or "agent"
        role: String,
        is_final: bool,
    },

    /// The voice agent's state changed (e.g. listening, thinking, speaking).
    VoiceAgentStateChanged { room_name: String, state: String },

    /// A voice session has ended.
    VoiceSessionEnded { room_name: String, reason: String },
}

// ─── Constants ───────────────────────────────────────────────────────────

/// Default channel capacity (per bus instance)
const DEFAULT_BUS_CAPACITY: usize = 4096;

// ─── AgentBus ────────────────────────────────────────────────────────────

/// The central in-process message bus.
///
/// Internally this is a `tokio::sync::broadcast` channel so every subscriber
/// receives every message.  Filtering by topic is done on the consumer side
/// through `BusHandle::subscribe_topic` / `BusHandle::recv_filtered`.
pub struct AgentBus {
    tx: broadcast::Sender<BusEnvelope>,
    /// Registry of connected agents
    pub registry: Arc<registry::AgentRegistry>,
}

impl std::fmt::Debug for AgentBus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AgentBus")
            .field("subscribers", &self.tx.receiver_count())
            .finish()
    }
}

impl AgentBus {
    /// Create a new bus with default capacity.
    pub fn new() -> Self {
        Self::with_capacity(DEFAULT_BUS_CAPACITY)
    }

    /// Create a new bus with a specific channel capacity.
    pub fn with_capacity(capacity: usize) -> Self {
        let (tx, _) = broadcast::channel(capacity);
        Self {
            tx,
            registry: Arc::new(registry::AgentRegistry::new()),
        }
    }

    /// Wrap `self` in an `Arc` for sharing across tasks.
    pub fn into_arc(self) -> Arc<Self> {
        Arc::new(self)
    }

    /// Create a `BusHandle` scoped to a specific agent.
    pub fn handle(self: &Arc<Self>, agent_id: impl Into<String>) -> BusHandle {
        BusHandle {
            agent_id: agent_id.into(),
            bus: Arc::clone(self),
            rx: self.tx.subscribe(),
        }
    }

    /// Publish an envelope directly (low-level).
    pub fn publish(&self, envelope: BusEnvelope) -> usize {
        match &envelope.message {
            BusMessage::AgentReady {
                agent_id,
                capabilities,
            } => {
                self.registry.register_ready(agent_id, capabilities);
            }
            BusMessage::AgentShutdown { agent_id } => {
                self.registry.deregister(agent_id);
            }
            _ => {}
        }

        // Returns the number of receivers that got the message.
        // If there are no receivers the message is silently dropped.
        self.tx.send(envelope).unwrap_or(0)
    }

    /// Number of active receivers.
    pub fn receiver_count(&self) -> usize {
        self.tx.receiver_count()
    }
}

impl Default for AgentBus {
    fn default() -> Self {
        Self::new()
    }
}

// ─── BusHandle ───────────────────────────────────────────────────────────

/// A scoped handle that a single agent uses to send and receive messages.
///
/// The handle knows the agent's id and uses it as `sender_id` on outgoing
/// envelopes.  It also provides filtered receive helpers.
pub struct BusHandle {
    agent_id: String,
    bus: Arc<AgentBus>,
    rx: broadcast::Receiver<BusEnvelope>,
}

impl BusHandle {
    /// The agent id this handle belongs to.
    pub fn agent_id(&self) -> &str {
        &self.agent_id
    }

    /// Send a message on the given topic.
    pub fn send(&self, topic: impl Into<String>, message: BusMessage) -> usize {
        self.send_with_correlation(topic, message, None)
    }

    /// Send a message with a correlation id.
    pub fn send_with_correlation(
        &self,
        topic: impl Into<String>,
        message: BusMessage,
        correlation_id: Option<String>,
    ) -> usize {
        let envelope = BusEnvelope {
            id: Uuid::new_v4().to_string(),
            topic: topic.into(),
            sender_id: self.agent_id.clone(),
            correlation_id,
            timestamp: Utc::now(),
            message,
        };
        self.bus.publish(envelope)
    }

    /// Announce this agent as ready.
    pub fn announce_ready(&self, capabilities: Vec<String>) -> usize {
        self.send(
            "broadcast",
            BusMessage::AgentReady {
                agent_id: self.agent_id.clone(),
                capabilities,
            },
        )
    }

    /// Announce this agent is shutting down.
    pub fn announce_shutdown(&self) -> usize {
        self.send(
            "broadcast",
            BusMessage::AgentShutdown {
                agent_id: self.agent_id.clone(),
            },
        )
    }

    /// Announce a task status update.
    pub fn send_task_update(
        &self,
        task_id: &str,
        state: TaskState,
        message: Option<String>,
    ) -> usize {
        self.send(
            format!("task.{task_id}"),
            BusMessage::TaskUpdate {
                task_id: task_id.to_string(),
                state,
                message,
            },
        )
    }

    /// Announce an artifact update.
    pub fn send_artifact_update(&self, task_id: &str, artifact: Artifact) -> usize {
        self.send(
            format!("task.{task_id}"),
            BusMessage::ArtifactUpdate {
                task_id: task_id.to_string(),
                artifact,
            },
        )
    }

    /// Send a direct message to another agent.
    pub fn send_to_agent(&self, to: &str, parts: Vec<Part>) -> usize {
        self.send(
            format!("agent.{to}"),
            BusMessage::AgentMessage {
                from: self.agent_id.clone(),
                to: to.to_string(),
                parts,
            },
        )
    }

    /// Publish a shared result (visible to the entire swarm).
    pub fn publish_shared_result(
        &self,
        key: impl Into<String>,
        value: serde_json::Value,
        tags: Vec<String>,
    ) -> usize {
        let key = key.into();
        self.send(
            format!("results.{}", &key),
            BusMessage::SharedResult { key, value, tags },
        )
    }

    // ── Ralph helpers ────────────────────────────────────────────────

    /// Publish learnings from a Ralph iteration so other agents / future
    /// iterations can build on them.
    pub fn publish_ralph_learning(
        &self,
        prd_id: &str,
        story_id: &str,
        iteration: usize,
        learnings: Vec<String>,
        context: serde_json::Value,
    ) -> usize {
        self.send(
            format!("ralph.{prd_id}"),
            BusMessage::RalphLearning {
                prd_id: prd_id.to_string(),
                story_id: story_id.to_string(),
                iteration,
                learnings,
                context,
            },
        )
    }

    /// Publish a context handoff between sequential Ralph stories.
    pub fn publish_ralph_handoff(
        &self,
        prd_id: &str,
        from_story: &str,
        to_story: &str,
        context: serde_json::Value,
        progress_summary: &str,
    ) -> usize {
        self.send(
            format!("ralph.{prd_id}"),
            BusMessage::RalphHandoff {
                prd_id: prd_id.to_string(),
                from_story: from_story.to_string(),
                to_story: to_story.to_string(),
                context,
                progress_summary: progress_summary.to_string(),
            },
        )
    }

    /// Publish PRD-level progress.
    pub fn publish_ralph_progress(
        &self,
        prd_id: &str,
        passed: usize,
        total: usize,
        iteration: usize,
        status: &str,
    ) -> usize {
        self.send(
            format!("ralph.{prd_id}"),
            BusMessage::RalphProgress {
                prd_id: prd_id.to_string(),
                passed,
                total,
                iteration,
                status: status.to_string(),
            },
        )
    }

    /// Drain all accumulated Ralph learnings for a PRD (non-blocking).
    pub fn drain_ralph_learnings(&mut self, prd_id: &str) -> Vec<BusEnvelope> {
        let prefix = format!("ralph.{prd_id}");
        let mut out = Vec::new();
        while let Some(env) = self.try_recv() {
            if env.topic.starts_with(&prefix) {
                if matches!(
                    &env.message,
                    BusMessage::RalphLearning { .. } | BusMessage::RalphHandoff { .. }
                ) {
                    out.push(env);
                }
            }
        }
        out
    }

    // ── Voice helpers ────────────────────────────────────────────────

    /// Announce that a voice session has started.
    pub fn send_voice_session_started(&self, room_name: &str, voice_id: &str) -> usize {
        self.send(
            format!("voice.{room_name}"),
            BusMessage::VoiceSessionStarted {
                room_name: room_name.to_string(),
                agent_id: self.agent_id.clone(),
                voice_id: voice_id.to_string(),
            },
        )
    }

    /// Publish a transcript fragment from a voice session.
    pub fn send_voice_transcript(
        &self,
        room_name: &str,
        text: &str,
        role: &str,
        is_final: bool,
    ) -> usize {
        self.send(
            format!("voice.{room_name}"),
            BusMessage::VoiceTranscript {
                room_name: room_name.to_string(),
                text: text.to_string(),
                role: role.to_string(),
                is_final,
            },
        )
    }

    /// Announce a voice agent state change.
    pub fn send_voice_agent_state(&self, room_name: &str, state: &str) -> usize {
        self.send(
            format!("voice.{room_name}"),
            BusMessage::VoiceAgentStateChanged {
                room_name: room_name.to_string(),
                state: state.to_string(),
            },
        )
    }

    /// Announce that a voice session has ended.
    pub fn send_voice_session_ended(&self, room_name: &str, reason: &str) -> usize {
        self.send(
            format!("voice.{room_name}"),
            BusMessage::VoiceSessionEnded {
                room_name: room_name.to_string(),
                reason: reason.to_string(),
            },
        )
    }

    /// Receive the next envelope (blocks until available).
    pub async fn recv(&mut self) -> Option<BusEnvelope> {
        loop {
            match self.rx.recv().await {
                Ok(env) => return Some(env),
                Err(broadcast::error::RecvError::Lagged(n)) => {
                    tracing::warn!(
                        agent_id = %self.agent_id,
                        skipped = n,
                        "Bus handle lagged, skipping messages"
                    );
                    continue;
                }
                Err(broadcast::error::RecvError::Closed) => return None,
            }
        }
    }

    /// Receive the next envelope whose topic starts with the given prefix.
    pub async fn recv_topic(&mut self, prefix: &str) -> Option<BusEnvelope> {
        loop {
            match self.recv().await {
                Some(env) if env.topic.starts_with(prefix) => return Some(env),
                Some(_) => continue, // wrong topic, skip
                None => return None,
            }
        }
    }

    /// Receive the next envelope addressed to this agent.
    pub async fn recv_mine(&mut self) -> Option<BusEnvelope> {
        let prefix = format!("agent.{}", self.agent_id);
        self.recv_topic(&prefix).await
    }

    /// Try to receive without blocking. Returns `None` if no message is
    /// queued.
    pub fn try_recv(&mut self) -> Option<BusEnvelope> {
        loop {
            match self.rx.try_recv() {
                Ok(env) => return Some(env),
                Err(broadcast::error::TryRecvError::Lagged(n)) => {
                    tracing::warn!(
                        agent_id = %self.agent_id,
                        skipped = n,
                        "Bus handle lagged (try_recv), skipping"
                    );
                    continue;
                }
                Err(broadcast::error::TryRecvError::Empty)
                | Err(broadcast::error::TryRecvError::Closed) => return None,
            }
        }
    }

    /// Access the agent registry.
    pub fn registry(&self) -> &Arc<registry::AgentRegistry> {
        &self.bus.registry
    }

    /// Get mutable access to the underlying receiver.
    /// This allows consuming the receiver in stream operations.
    pub fn into_receiver(self) -> broadcast::Receiver<BusEnvelope> {
        self.rx
    }
}

// ─── Tests ───────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_bus_send_recv() {
        let bus = AgentBus::new().into_arc();
        let mut handle_a = bus.handle("agent-a");
        let mut handle_b = bus.handle("agent-b");

        handle_a.send_to_agent(
            "agent-b",
            vec![Part::Text {
                text: "hello".into(),
            }],
        );

        // Both handles receive the broadcast
        let env = handle_b.recv().await.unwrap();
        assert_eq!(env.topic, "agent.agent-b");
        match &env.message {
            BusMessage::AgentMessage { from, to, .. } => {
                assert_eq!(from, "agent-a");
                assert_eq!(to, "agent-b");
            }
            other => panic!("unexpected message: {other:?}"),
        }

        // handle_a also receives it (broadcast semantics)
        let env_a = handle_a.try_recv().unwrap();
        assert_eq!(env_a.topic, "agent.agent-b");
    }

    #[tokio::test]
    async fn test_bus_task_update() {
        let bus = AgentBus::new().into_arc();
        let handle = bus.handle("worker-1");

        let h2 = bus.handle("observer");
        // need mutable for recv
        let mut h2 = h2;

        handle.send_task_update("task-42", TaskState::Working, Some("processing".into()));

        let env = h2.recv().await.unwrap();
        assert_eq!(env.topic, "task.task-42");
        match &env.message {
            BusMessage::TaskUpdate { task_id, state, .. } => {
                assert_eq!(task_id, "task-42");
                assert_eq!(*state, TaskState::Working);
            }
            other => panic!("unexpected: {other:?}"),
        }
    }

    #[tokio::test]
    async fn test_bus_no_receivers() {
        let bus = AgentBus::new().into_arc();
        // No handles created — send should succeed with 0 receivers
        let env = BusEnvelope {
            id: "test".into(),
            topic: "broadcast".into(),
            sender_id: "nobody".into(),
            correlation_id: None,
            timestamp: Utc::now(),
            message: BusMessage::Heartbeat {
                agent_id: "nobody".into(),
                status: "ok".into(),
            },
        };
        let count = bus.publish(env);
        assert_eq!(count, 0);
    }

    #[tokio::test]
    async fn test_recv_topic_filter() {
        let bus = AgentBus::new().into_arc();
        let handle = bus.handle("agent-x");
        let mut listener = bus.handle("listener");

        // Send to two different topics
        handle.send(
            "task.1",
            BusMessage::TaskUpdate {
                task_id: "1".into(),
                state: TaskState::Working,
                message: None,
            },
        );
        handle.send(
            "task.2",
            BusMessage::TaskUpdate {
                task_id: "2".into(),
                state: TaskState::Completed,
                message: None,
            },
        );

        // recv_topic("task.2") should skip the first and return the second
        let env = listener.recv_topic("task.2").await.unwrap();
        match &env.message {
            BusMessage::TaskUpdate { task_id, .. } => assert_eq!(task_id, "2"),
            other => panic!("unexpected: {other:?}"),
        }
    }

    #[tokio::test]
    async fn test_ready_shutdown_syncs_registry() {
        let bus = AgentBus::new().into_arc();
        let handle = bus.handle("planner-1");

        handle.announce_ready(vec!["plan".to_string(), "review".to_string()]);
        assert!(bus.registry.get("planner-1").is_some());

        handle.announce_shutdown();
        assert!(bus.registry.get("planner-1").is_none());
    }
}