Vapora/crates/vapora-agents/src/messages.rs

// vapora-agents: NATS message protocol for inter-agent communication
// Phase 2: Message types for agent coordination

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

/// Agent message envelope for NATS pub/sub
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AgentMessage {
    TaskAssigned(TaskAssignment),
    TaskStarted(TaskStarted),
    TaskProgress(TaskProgress),
    TaskCompleted(TaskCompleted),
    TaskFailed(TaskFailed),
    Heartbeat(Heartbeat),
    AgentRegistered(AgentRegistered),
    AgentStopped(AgentStopped),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskAssignment {
    pub id: String,
    pub agent_id: String,
    pub required_role: String,
    pub title: String,
    pub description: String,
    pub context: String,
    pub priority: u32,
    pub deadline: Option<DateTime<Utc>>,
    pub assigned_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskStarted {
    pub task_id: String,
    pub agent_id: String,
    pub started_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskProgress {
    pub task_id: String,
    pub agent_id: String,
    pub progress_percent: u32,
    pub current_step: String,
    pub estimated_completion: Option<DateTime<Utc>>,
    pub updated_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskCompleted {
    pub task_id: String,
    pub agent_id: String,
    pub result: String,
    pub artifacts: Vec<String>,
    pub tokens_used: u64,
    pub duration_ms: u64,
    pub completed_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskFailed {
    pub task_id: String,
    pub agent_id: String,
    pub error: String,
    pub retry_count: u32,
    pub can_retry: bool,
    pub failed_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Heartbeat {
    pub agent_id: String,
    pub status: String,
    pub load: f64,
    pub active_tasks: u32,
    pub total_tasks_completed: u64,
    pub uptime_seconds: u64,
    pub timestamp: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentRegistered {
    pub agent_id: String,
    pub role: String,
    pub version: String,
    pub capabilities: Vec<String>,
    pub registered_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentStopped {
    pub agent_id: String,
    pub role: String,
    pub reason: String,
    pub stopped_at: DateTime<Utc>,
}

impl AgentMessage {
    /// Serialize message to JSON bytes for NATS
    pub fn to_bytes(&self) -> Result<Vec<u8>, serde_json::Error> {
        serde_json::to_vec(self)
    }

    /// Deserialize message from JSON bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, serde_json::Error> {
        serde_json::from_slice(bytes)
    }

    /// Get message type as string
    pub fn message_type(&self) -> &str {
        match self {
            AgentMessage::TaskAssigned(_) => "task_assigned",
            AgentMessage::TaskStarted(_) => "task_started",
            AgentMessage::TaskProgress(_) => "task_progress",
            AgentMessage::TaskCompleted(_) => "task_completed",
            AgentMessage::TaskFailed(_) => "task_failed",
            AgentMessage::Heartbeat(_) => "heartbeat",
            AgentMessage::AgentRegistered(_) => "agent_registered",
            AgentMessage::AgentStopped(_) => "agent_stopped",
        }
    }
}

/// NATS subjects for agent communication
pub mod subjects {
    pub const TASKS_ASSIGNED: &str = "vapora.tasks.assigned";
    pub const TASKS_STARTED: &str = "vapora.tasks.started";
    pub const TASKS_PROGRESS: &str = "vapora.tasks.progress";
    pub const TASKS_COMPLETED: &str = "vapora.tasks.completed";
    pub const TASKS_FAILED: &str = "vapora.tasks.failed";
    pub const AGENT_HEARTBEAT: &str = "vapora.agent.heartbeat";
    pub const AGENT_REGISTERED: &str = "vapora.agent.registered";
    pub const AGENT_STOPPED: &str = "vapora.agent.stopped";

    /// Get subject for a specific agent role
    pub fn agent_role_subject(role: &str) -> String {
        format!("vapora.agent.role.{}", role)
    }

    /// Get subject for a specific task
    pub fn task_subject(task_id: &str) -> String {
        format!("vapora.task.{}", task_id)
    }
}

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

    #[test]
    fn test_message_serialization() {
        let msg = AgentMessage::TaskAssigned(TaskAssignment {
            id: "task-123".to_string(),
            agent_id: "agent-001".to_string(),
            required_role: "developer".to_string(),
            title: "Test task".to_string(),
            description: "Test description".to_string(),
            context: "{}".to_string(),
            priority: 80,
            deadline: None,
            assigned_at: Utc::now(),
        });

        let bytes = msg.to_bytes().unwrap();
        let deserialized = AgentMessage::from_bytes(&bytes).unwrap();

        assert_eq!(msg.message_type(), deserialized.message_type());
    }

    #[test]
    fn test_heartbeat_message() {
        let heartbeat = Heartbeat {
            agent_id: "agent-001".to_string(),
            status: "active".to_string(),
            load: 0.5,
            active_tasks: 2,
            total_tasks_completed: 100,
            uptime_seconds: 3600,
            timestamp: Utc::now(),
        };

        let msg = AgentMessage::Heartbeat(heartbeat);
        assert_eq!(msg.message_type(), "heartbeat");
    }

    #[test]
    fn test_subject_generation() {
        assert_eq!(subjects::agent_role_subject("developer"), "vapora.agent.role.developer");
        assert_eq!(subjects::task_subject("task-123"), "vapora.task.task-123");
    }
}