rustytorch/crates/production/rtx-serving-api/src/streaming.rs

//! Streaming response generation with SSE and WebSocket support
//!
//! Provides comprehensive streaming capabilities including:
//! - Server-sent events (SSE) with proper flow control
//! - WebSocket streaming with bidirectional communication
//! - Chunked transfer encoding with backpressure handling
//! - Progressive generation with intermediate result streaming
//! - Stream multiplexing for concurrent requests
//! - Error handling and recovery in streaming scenarios

use anyhow::{Result, anyhow};
use axum::{
    extract::ws::{Message as WsMessage, WebSocket},
    response::sse::Event,
};
use bytes::Bytes;
use chrono::{DateTime, Utc};
use dashmap::DashMap;
use futures_util::{sink::SinkExt, stream::StreamExt};
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, sync::Arc, time::Duration};
use tokio::{sync::broadcast, time::interval};
use tokio_stream::Stream;
use uuid::Uuid;

/// Stream types supported by the system
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum StreamType {
    ServerSentEvents,
    WebSocket,
    ChunkedHttp,
}

/// Stream configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamConfig {
    pub stream_type: StreamType,
    pub buffer_size: usize,
    pub flush_interval: Duration,
    pub keep_alive_interval: Duration,
    pub max_concurrent_streams: usize,
    pub backpressure_threshold: usize,
    pub timeout: Duration,
    pub enable_compression: bool,
    pub chunk_size: usize,
}

impl Default for StreamConfig {
    fn default() -> Self {
        Self {
            stream_type: StreamType::ServerSentEvents,
            buffer_size: 1024,
            flush_interval: Duration::from_millis(100),
            keep_alive_interval: Duration::from_secs(30),
            max_concurrent_streams: 1000,
            backpressure_threshold: 10000,
            timeout: Duration::from_secs(300),
            enable_compression: true,
            chunk_size: 4096,
        }
    }
}

/// Streaming event types
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "data")]
pub enum StreamEvent {
    Start {
        stream_id: String,
        model: String,
        timestamp: DateTime<Utc>,
    },
    Token {
        token: String,
        token_id: u32,
        position: usize,
        probability: f32,
        cumulative_text: String,
    },
    PartialResponse {
        text: String,
        tokens_generated: usize,
        finish_reason: Option<String>,
    },
    Metadata {
        key: String,
        value: serde_json::Value,
    },
    Progress {
        current: usize,
        total: Option<usize>,
        percentage: Option<f32>,
    },
    Error {
        error: String,
        code: String,
        recoverable: bool,
    },
    Complete {
        final_text: String,
        total_tokens: usize,
        processing_time_ms: u64,
        finish_reason: String,
    },
    KeepAlive {
        timestamp: DateTime<Utc>,
    },
    Heartbeat,
}

impl StreamEvent {
    /// Convert to SSE event
    pub fn to_sse_event(&self) -> Result<Event> {
        let json_data = serde_json::to_string(self)?;

        let event_type = match self {
            Self::Start { .. } => "start",
            Self::Token { .. } => "token",
            Self::PartialResponse { .. } => "partial",
            Self::Metadata { .. } => "metadata",
            Self::Progress { .. } => "progress",
            Self::Error { .. } => "error",
            Self::Complete { .. } => "complete",
            Self::KeepAlive { .. } => "keep_alive",
            Self::Heartbeat => "heartbeat",
        };

        Ok(Event::default().event(event_type).data(json_data))
    }

    /// Convert to WebSocket message
    pub fn to_ws_message(&self) -> Result<WsMessage> {
        let json_data = serde_json::to_string(self)?;
        Ok(WsMessage::Text(json_data))
    }
}

/// Stream session state
#[derive(Debug, Clone)]
pub struct StreamSession {
    pub session_id: String,
    pub user_id: String,
    pub model_name: String,
    pub stream_type: StreamType,
    pub created_at: DateTime<Utc>,
    pub last_activity: DateTime<Utc>,
    pub tokens_streamed: usize,
    pub bytes_sent: u64,
    pub is_active: bool,
    pub error_count: u32,
    pub backpressure_events: u32,
}

impl StreamSession {
    /// Create new stream session
    #[must_use]
    pub fn new(user_id: String, model_name: String, stream_type: StreamType) -> Self {
        let now = Utc::now();
        Self {
            session_id: Uuid::new_v4().to_string(),
            user_id,
            model_name,
            stream_type,
            created_at: now,
            last_activity: now,
            tokens_streamed: 0,
            bytes_sent: 0,
            is_active: true,
            error_count: 0,
            backpressure_events: 0,
        }
    }

    /// Update activity timestamp
    pub fn update_activity(&mut self) {
        self.last_activity = Utc::now();
    }

    /// Check if session has timed out
    #[must_use]
    pub fn is_timed_out(&self, timeout: Duration) -> bool {
        let elapsed = Utc::now().signed_duration_since(self.last_activity);
        elapsed.to_std().unwrap_or(Duration::ZERO) > timeout
    }
}

/// Stream multiplexer for handling multiple concurrent streams
pub struct StreamMultiplexer {
    sessions: Arc<DashMap<String, StreamSession>>,
    config: StreamConfig,
    broadcaster: broadcast::Sender<StreamEvent>,
}

impl StreamMultiplexer {
    /// Create new stream multiplexer
    #[must_use]
    pub fn new(config: StreamConfig) -> Self {
        let (broadcaster, _) = broadcast::channel(10000);

        Self {
            sessions: Arc::new(DashMap::new()),
            config,
            broadcaster,
        }
    }

    /// Create new stream session
    pub fn create_session(
        &self,
        user_id: String,
        model_name: String,
        stream_type: StreamType,
    ) -> Result<String> {
        if self.sessions.len() >= self.config.max_concurrent_streams {
            return Err(anyhow!("Maximum concurrent streams exceeded"));
        }

        let session = StreamSession::new(user_id, model_name, stream_type);
        let session_id = session.session_id.clone();

        self.sessions.insert(session_id.clone(), session);

        Ok(session_id)
    }

    /// Send event to specific session
    pub async fn send_to_session(&self, session_id: &str, event: StreamEvent) -> Result<()> {
        if let Some(mut session) = self.sessions.get_mut(session_id) {
            session.update_activity();
            session.tokens_streamed += 1;

            // Calculate bytes (approximate)
            let bytes = serde_json::to_string(&event)?.len() as u64;
            session.bytes_sent += bytes;

            // Send via broadcaster (receivers will filter by session_id).
            // A `SendError` here just means there are currently no active
            // subscribers listening on the broadcast channel, which is a
            // normal condition (not a failure of the session update itself),
            // so it is intentionally ignored rather than propagated.
            let tagged_event = StreamEvent::Metadata {
                key: "session_id".to_string(),
                value: serde_json::json!(session_id),
            };

            let _ = self.broadcaster.send(tagged_event);
            let _ = self.broadcaster.send(event);

            Ok(())
        } else {
            Err(anyhow!("Session not found: {session_id}"))
        }
    }

    /// Broadcast event to all sessions
    pub async fn broadcast(&self, event: StreamEvent) -> Result<()> {
        if self.broadcaster.send(event).is_err() {
            return Err(anyhow!("Failed to broadcast event"));
        }
        Ok(())
    }

    /// Close session
    #[must_use]
    pub fn close_session(&self, session_id: &str) -> Option<StreamSession> {
        self.sessions.remove(session_id).map(|(_, mut session)| {
            session.is_active = false;
            session
        })
    }

    /// Get session info
    #[must_use]
    pub fn get_session(&self, session_id: &str) -> Option<StreamSession> {
        self.sessions.get(session_id).map(|s| s.clone())
    }

    /// Get all active sessions
    #[must_use]
    pub fn get_active_sessions(&self) -> Vec<StreamSession> {
        self.sessions
            .iter()
            .filter(|entry| entry.value().is_active)
            .map(|entry| entry.value().clone())
            .collect()
    }

    /// Cleanup timed out sessions
    pub async fn cleanup_timed_out_sessions(&self) -> usize {
        let timeout = self.config.timeout;
        let mut removed = 0;

        let timed_out_sessions: Vec<String> = self
            .sessions
            .iter()
            .filter(|entry| entry.value().is_timed_out(timeout))
            .map(|entry| entry.key().clone())
            .collect();

        for session_id in timed_out_sessions {
            if self.sessions.remove(&session_id).is_some() {
                removed += 1;

                // Send timeout event
                let _ = self
                    .send_to_session(
                        &session_id,
                        StreamEvent::Error {
                            error: "Session timed out".to_string(),
                            code: "TIMEOUT".to_string(),
                            recoverable: false,
                        },
                    )
                    .await;
            }
        }

        removed
    }

    /// Subscribe to stream events
    #[must_use]
    pub fn subscribe(&self) -> broadcast::Receiver<StreamEvent> {
        self.broadcaster.subscribe()
    }
}

/// SSE stream handler
pub struct SseStreamHandler {
    multiplexer: Arc<StreamMultiplexer>,
    session_id: String,
}

impl SseStreamHandler {
    /// Create new SSE stream handler
    #[must_use]
    pub fn new(multiplexer: Arc<StreamMultiplexer>, session_id: String) -> Self {
        Self {
            multiplexer,
            session_id,
        }
    }

    /// Create SSE stream
    pub async fn create_stream(self) -> impl Stream<Item = Result<Event, axum::Error>> + use<> {
        let mut receiver = self.multiplexer.subscribe();
        let session_id = self.session_id.clone();
        let multiplexer = self.multiplexer.clone();

        // Send initial start event
        let start_event = StreamEvent::Start {
            stream_id: session_id.clone(),
            model: "unknown".to_string(), // Would be set from request
            timestamp: Utc::now(),
        };

        let _ = multiplexer.send_to_session(&session_id, start_event).await;

        async_stream::stream! {
            // Send keep-alive events periodically
            let mut keep_alive_interval = interval(multiplexer.config.keep_alive_interval);

            loop {
                tokio::select! {
                    event_result = receiver.recv() => {
                        match event_result {
                            Ok(event) => {
                                match event.to_sse_event() {
                                    Ok(sse_event) => yield Ok(sse_event),
                                    Err(e) => yield Err(axum::Error::new(format!("Event conversion error: {e}"))),
                                }
                            }
                            Err(broadcast::error::RecvError::Lagged(_)) => {
                                // Skip lagged messages
                                continue;
                            }
                            Err(broadcast::error::RecvError::Closed) => {
                                break;
                            }
                        }
                    }

                    _ = keep_alive_interval.tick() => {
                        let keep_alive = StreamEvent::KeepAlive {
                            timestamp: Utc::now(),
                        };

                        match keep_alive.to_sse_event() {
                            Ok(sse_event) => yield Ok(sse_event),
                            Err(e) => yield Err(axum::Error::new(format!("Keep-alive error: {e}"))),
                        }
                    }
                }
            }
        }
    }
}

/// WebSocket stream handler
pub struct WebSocketStreamHandler {
    multiplexer: Arc<StreamMultiplexer>,
}

impl WebSocketStreamHandler {
    /// Create new WebSocket stream handler
    #[must_use]
    pub fn new(multiplexer: Arc<StreamMultiplexer>) -> Self {
        Self { multiplexer }
    }

    /// Handle WebSocket connection
    pub async fn handle_connection(
        &self,
        socket: WebSocket,
        user_id: String,
        model_name: String,
    ) -> Result<()> {
        let session_id = self.multiplexer.create_session(
            user_id.clone(),
            model_name.clone(),
            StreamType::WebSocket,
        )?;

        let (mut sender, mut receiver) = socket.split();
        let multiplexer = self.multiplexer.clone();
        let session_id_clone = session_id.clone();

        // Subscribe to stream events
        let mut event_receiver = self.multiplexer.subscribe();

        // Spawn task to handle incoming WebSocket messages
        let incoming_task = tokio::spawn(async move {
            while let Some(msg) = receiver.next().await {
                match msg {
                    Ok(WsMessage::Text(text)) => {
                        // Handle incoming text message (e.g., configuration updates)
                        if let Err(e) =
                            Self::handle_incoming_message(&multiplexer, &session_id, &text).await
                        {
                            tracing::error!("Error handling incoming message: {}", e);
                        }
                    }
                    Ok(WsMessage::Close(_)) => {
                        tracing::info!("WebSocket connection closed by client");
                        break;
                    }
                    Ok(WsMessage::Ping(_data)) => {
                        // Respond to ping with pong
                        if let Err(e) = multiplexer
                            .send_to_session(&session_id, StreamEvent::Heartbeat)
                            .await
                        {
                            tracing::error!("Error sending heartbeat: {}", e);
                        }
                    }
                    Err(e) => {
                        tracing::error!("WebSocket error: {}", e);
                        break;
                    }
                    _ => {
                        // Ignore other message types
                    }
                }
            }
        });

        // Spawn task to handle outgoing WebSocket messages
        let outgoing_task = tokio::spawn(async move {
            while let Ok(event) = event_receiver.recv().await {
                match event.to_ws_message() {
                    Ok(ws_msg) => {
                        if let Err(e) = sender.send(ws_msg).await {
                            tracing::error!("Error sending WebSocket message: {}", e);
                            break;
                        }
                    }
                    Err(e) => {
                        tracing::error!("Error converting event to WebSocket message: {}", e);
                    }
                }
            }
        });

        // Wait for either task to complete
        tokio::select! {
            _ = incoming_task => {},
            _ = outgoing_task => {},
        }

        // Cleanup session
        self.multiplexer.close_session(&session_id_clone);

        Ok(())
    }

    /// Handle incoming WebSocket message
    async fn handle_incoming_message(
        multiplexer: &StreamMultiplexer,
        session_id: &str,
        message: &str,
    ) -> Result<()> {
        // Parse incoming message as JSON
        let parsed: serde_json::Value = serde_json::from_str(message)?;

        if let Some(msg_type) = parsed.get("type").and_then(|t| t.as_str()) {
            match msg_type {
                "ping" => {
                    // Respond with pong
                    multiplexer
                        .send_to_session(session_id, StreamEvent::Heartbeat)
                        .await?;
                }
                "config" => {
                    // Handle configuration update
                    if let Some(config) = parsed.get("data") {
                        multiplexer
                            .send_to_session(
                                session_id,
                                StreamEvent::Metadata {
                                    key: "config_updated".to_string(),
                                    value: config.clone(),
                                },
                            )
                            .await?;
                    }
                }
                "cancel" => {
                    // Handle cancellation request
                    multiplexer
                        .send_to_session(
                            session_id,
                            StreamEvent::Error {
                                error: "Generation cancelled by user".to_string(),
                                code: "USER_CANCELLED".to_string(),
                                recoverable: false,
                            },
                        )
                        .await?;
                }
                _ => {
                    tracing::warn!("Unknown message type: {}", msg_type);
                }
            }
        }

        Ok(())
    }
}

/// Chunked HTTP stream handler
pub struct ChunkedStreamHandler {
    multiplexer: Arc<StreamMultiplexer>,
}

impl ChunkedStreamHandler {
    /// Create new chunked stream handler
    #[must_use]
    pub fn new(multiplexer: Arc<StreamMultiplexer>) -> Self {
        Self { multiplexer }
    }

    /// Create chunked response stream
    pub async fn create_stream(
        &self,
        user_id: String,
        model_name: String,
    ) -> Result<impl Stream<Item = Result<Bytes, std::io::Error>> + use<>> {
        let _session_id =
            self.multiplexer
                .create_session(user_id, model_name, StreamType::ChunkedHttp)?;

        let mut receiver = self.multiplexer.subscribe();
        let multiplexer = self.multiplexer.clone();

        Ok(async_stream::stream! {
            let mut buffer = Vec::new();
            let chunk_size = multiplexer.config.chunk_size;

            while let Ok(event) = receiver.recv().await {
                match serde_json::to_vec(&event) {
                    Ok(mut event_bytes) => {
                        event_bytes.push(b'\n'); // Add newline delimiter
                        buffer.extend(event_bytes);

                        // Flush buffer when it reaches chunk size
                        if buffer.len() >= chunk_size {
                            let chunk = buffer.clone();
                            buffer.clear();
                            yield Ok(Bytes::from(chunk));
                        }
                    }
                    Err(e) => {
                        let error_msg = format!("Serialization error: {e}\n");
                        yield Ok(Bytes::from(error_msg.into_bytes()));
                    }
                }
            }

            // Flush remaining buffer
            if !buffer.is_empty() {
                yield Ok(Bytes::from(buffer));
            }
        })
    }
}

/// Stream manager that coordinates all streaming functionality
pub struct StreamManager {
    multiplexer: Arc<StreamMultiplexer>,
    sse_handler: SseStreamHandler,
    ws_handler: WebSocketStreamHandler,
    chunked_handler: ChunkedStreamHandler,
}

impl StreamManager {
    /// Create new stream manager
    #[must_use]
    pub fn new(config: StreamConfig) -> Self {
        let multiplexer = Arc::new(StreamMultiplexer::new(config));

        // Create a dummy session ID for SSE handler (will be overridden)
        let dummy_session_id = Uuid::new_v4().to_string();
        let sse_handler = SseStreamHandler::new(multiplexer.clone(), dummy_session_id);
        let ws_handler = WebSocketStreamHandler::new(multiplexer.clone());
        let chunked_handler = ChunkedStreamHandler::new(multiplexer.clone());

        Self {
            multiplexer,
            sse_handler,
            ws_handler,
            chunked_handler,
        }
    }

    /// Create SSE stream for user
    pub async fn create_sse_stream(
        &self,
        user_id: String,
        model_name: String,
    ) -> Result<impl Stream<Item = Result<Event, axum::Error>> + use<>> {
        let session_id =
            self.multiplexer
                .create_session(user_id, model_name, StreamType::ServerSentEvents)?;

        let handler = SseStreamHandler::new(self.multiplexer.clone(), session_id);
        Ok(handler.create_stream().await)
    }

    /// Handle WebSocket connection
    pub async fn handle_websocket(
        &self,
        socket: WebSocket,
        user_id: String,
        model_name: String,
    ) -> Result<()> {
        self.ws_handler
            .handle_connection(socket, user_id, model_name)
            .await
    }

    /// Create chunked HTTP stream
    pub async fn create_chunked_stream(
        &self,
        user_id: String,
        model_name: String,
    ) -> Result<impl Stream<Item = Result<Bytes, std::io::Error>> + use<>> {
        self.chunked_handler
            .create_stream(user_id, model_name)
            .await
    }

    /// Send event to specific session
    pub async fn send_event(&self, session_id: &str, event: StreamEvent) -> Result<()> {
        self.multiplexer.send_to_session(session_id, event).await
    }

    /// Broadcast event to all sessions
    pub async fn broadcast_event(&self, event: StreamEvent) -> Result<()> {
        self.multiplexer.broadcast(event).await
    }

    /// Get stream statistics
    #[must_use]
    pub fn get_stream_stats(&self) -> StreamStats {
        let sessions = self.multiplexer.get_active_sessions();
        let total_sessions = sessions.len();
        let total_tokens_streamed: usize = sessions.iter().map(|s| s.tokens_streamed).sum();
        let total_bytes_sent: u64 = sessions.iter().map(|s| s.bytes_sent).sum();

        let mut by_type = HashMap::new();
        for session in &sessions {
            *by_type.entry(session.stream_type).or_insert(0) += 1;
        }

        StreamStats {
            total_active_sessions: total_sessions,
            sessions_by_type: by_type,
            total_tokens_streamed,
            total_bytes_sent,
            average_tokens_per_session: if total_sessions > 0 {
                total_tokens_streamed as f64 / total_sessions as f64
            } else {
                0.0
            },
        }
    }

    /// Start cleanup task
    #[must_use]
    pub fn start_cleanup_task(&self) -> tokio::task::JoinHandle<()> {
        let multiplexer = self.multiplexer.clone();

        tokio::spawn(async move {
            let mut interval = interval(Duration::from_secs(60)); // Cleanup every minute

            loop {
                interval.tick().await;
                let removed = multiplexer.cleanup_timed_out_sessions().await;
                if removed > 0 {
                    tracing::info!("Cleaned up {} timed out stream sessions", removed);
                }
            }
        })
    }
}

/// Stream statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamStats {
    pub total_active_sessions: usize,
    pub sessions_by_type: HashMap<StreamType, usize>,
    pub total_tokens_streamed: usize,
    pub total_bytes_sent: u64,
    pub average_tokens_per_session: f64,
}

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

    #[test]
    fn test_stream_config_default() {
        let config = StreamConfig::default();
        assert_eq!(config.stream_type, StreamType::ServerSentEvents);
        assert_eq!(config.buffer_size, 1024);
        assert!(config.timeout > Duration::ZERO);
    }

    #[test]
    fn test_stream_session_creation() {
        let session = StreamSession::new(
            "user123".to_string(),
            "gpt-4".to_string(),
            StreamType::WebSocket,
        );

        assert_eq!(session.user_id, "user123");
        assert_eq!(session.model_name, "gpt-4");
        assert_eq!(session.stream_type, StreamType::WebSocket);
        assert!(session.is_active);
        assert_eq!(session.tokens_streamed, 0);
    }

    #[test]
    fn test_stream_event_conversion() {
        let event = StreamEvent::Token {
            token: "hello".to_string(),
            token_id: 123,
            position: 5,
            probability: 0.8,
            cumulative_text: "hello world".to_string(),
        };

        // Test SSE conversion - verify it succeeds
        let sse_event = event.to_sse_event();
        assert!(sse_event.is_ok(), "SSE conversion should succeed");

        // Test WebSocket conversion
        let ws_message = event.to_ws_message().unwrap();
        if let WsMessage::Text(text) = ws_message {
            assert!(text.contains("hello"));
        }
    }

    #[tokio::test]
    async fn test_stream_multiplexer() {
        let config = StreamConfig::default();
        let multiplexer = StreamMultiplexer::new(config);

        // Create session
        let session_id = multiplexer
            .create_session(
                "user1".to_string(),
                "gpt-4".to_string(),
                StreamType::ServerSentEvents,
            )
            .unwrap();

        // Send event to session
        let event = StreamEvent::Token {
            token: "test".to_string(),
            token_id: 1,
            position: 0,
            probability: 1.0,
            cumulative_text: "test".to_string(),
        };

        let result = multiplexer.send_to_session(&session_id, event).await;
        assert!(result.is_ok());

        // Check session was updated
        let session = multiplexer.get_session(&session_id).unwrap();
        assert_eq!(session.tokens_streamed, 1);
        assert!(session.bytes_sent > 0);
    }

    #[tokio::test]
    async fn test_session_timeout() {
        let mut config = StreamConfig::default();
        config.timeout = Duration::from_millis(100); // Very short timeout

        let multiplexer = StreamMultiplexer::new(config);

        let session_id = multiplexer
            .create_session(
                "user1".to_string(),
                "gpt-4".to_string(),
                StreamType::ServerSentEvents,
            )
            .unwrap();

        // Wait for timeout
        sleep(Duration::from_millis(150)).await;

        // Cleanup should remove the session
        let removed = multiplexer.cleanup_timed_out_sessions().await;
        assert_eq!(removed, 1);

        // Session should no longer exist
        assert!(multiplexer.get_session(&session_id).is_none());
    }

    #[test]
    fn test_stream_manager_creation() {
        let config = StreamConfig::default();
        let manager = StreamManager::new(config);

        let stats = manager.get_stream_stats();
        assert_eq!(stats.total_active_sessions, 0);
        assert_eq!(stats.total_tokens_streamed, 0);
    }

    #[test]
    fn test_session_activity_update() {
        let mut session = StreamSession::new(
            "user1".to_string(),
            "gpt-4".to_string(),
            StreamType::ServerSentEvents,
        );

        let initial_activity = session.last_activity;

        // Small delay to ensure timestamp difference
        std::thread::sleep(Duration::from_millis(1));

        session.update_activity();
        assert!(session.last_activity > initial_activity);
    }

    #[test]
    fn test_backpressure_handling() {
        let session = StreamSession::new(
            "user1".to_string(),
            "gpt-4".to_string(),
            StreamType::ServerSentEvents,
        );

        // Test that new sessions start with zero backpressure events
        assert_eq!(session.backpressure_events, 0);
    }
}