Vapora/crates/vapora-backend/src/api/websocket.rs

// vapora-backend: WebSocket handler for real-time workflow updates
// Phase 3: Stream workflow progress to connected clients

use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::broadcast;
use tracing::{debug, error};

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct WorkflowUpdate {
    pub workflow_id: String,
    pub status: String,
    pub progress: u32,
    pub message: String,
    pub timestamp: chrono::DateTime<chrono::Utc>,
}

impl WorkflowUpdate {
    pub fn new(workflow_id: String, status: String, progress: u32, message: String) -> Self {
        Self {
            workflow_id,
            status,
            progress,
            message,
            timestamp: chrono::Utc::now(),
        }
    }
}

/// Broadcaster for workflow updates
pub struct WorkflowBroadcaster {
    tx: broadcast::Sender<WorkflowUpdate>,
}

impl WorkflowBroadcaster {
    pub fn new() -> Self {
        let (tx, _) = broadcast::channel(100);
        Self { tx }
    }

    /// Send workflow update to all subscribers
    pub fn send_update(&self, update: WorkflowUpdate) {
        debug!(
            "Broadcasting update for workflow {}: {} ({}%)",
            update.workflow_id, update.message, update.progress
        );

        if let Err(e) = self.tx.send(update) {
            error!("Failed to broadcast update: {}", e);
        }
    }

    /// Subscribe to workflow updates
    pub fn subscribe(&self) -> broadcast::Receiver<WorkflowUpdate> {
        self.tx.subscribe()
    }

    /// Get subscriber count
    pub fn subscriber_count(&self) -> usize {
        self.tx.receiver_count()
    }
}

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

impl Clone for WorkflowBroadcaster {
    fn clone(&self) -> Self {
        Self {
            tx: self.tx.clone(),
        }
    }
}

// Note: WebSocket support requires ws feature in axum
// For Phase 4, we focus on the broadcaster infrastructure
// WebSocket handlers would be added when the ws feature is enabled

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

    #[test]
    fn test_broadcaster_creation() {
        let broadcaster = WorkflowBroadcaster::new();
        assert_eq!(broadcaster.subscriber_count(), 0);
    }

    #[test]
    fn test_subscribe() {
        let broadcaster = WorkflowBroadcaster::new();
        let _rx = broadcaster.subscribe();
        assert_eq!(broadcaster.subscriber_count(), 1);
    }

    #[tokio::test]
    async fn test_send_update() {
        let broadcaster = WorkflowBroadcaster::new();
        let mut rx = broadcaster.subscribe();

        let update = WorkflowUpdate::new(
            "wf-1".to_string(),
            "in_progress".to_string(),
            50,
            "Step 1 completed".to_string(),
        );

        broadcaster.send_update(update.clone());

        let received = rx.recv().await.unwrap();
        assert_eq!(received.workflow_id, "wf-1");
        assert_eq!(received.progress, 50);
    }

    #[tokio::test]
    async fn test_multiple_subscribers() {
        let broadcaster = WorkflowBroadcaster::new();
        let mut rx1 = broadcaster.subscribe();
        let mut rx2 = broadcaster.subscribe();

        let update = WorkflowUpdate::new(
            "wf-1".to_string(),
            "completed".to_string(),
            100,
            "All steps completed".to_string(),
        );

        broadcaster.send_update(update);

        let received1 = rx1.recv().await.unwrap();
        let received2 = rx2.recv().await.unwrap();

        assert_eq!(received1.workflow_id, received2.workflow_id);
        assert_eq!(received1.progress, 100);
        assert_eq!(received2.progress, 100);
    }

    #[test]
    fn test_update_serialization() {
        let update = WorkflowUpdate::new(
            "wf-1".to_string(),
            "running".to_string(),
            75,
            "Almost done".to_string(),
        );

        let json = serde_json::to_string(&update).unwrap();
        let deserialized: WorkflowUpdate = serde_json::from_str(&json).unwrap();

        assert_eq!(deserialized.workflow_id, "wf-1");
        assert_eq!(deserialized.progress, 75);
    }
}