cortexai_crew/

streaming.rs

//! Unified Streaming for Workflows and Agents
//!
//! Provides token-level and event-level streaming across all workflow types:
//! - Graph node execution with LLM token streaming
//! - Handoff conversation streaming
//! - Workflow step streaming
//! - Real-time event callbacks
//!
//! ## Example
//!
//! ```rust,ignore
//! use cortexai_crew::streaming::{StreamingGraph, WorkflowEvent};
//!
//! // Create a streaming graph
//! let streaming = StreamingGraph::new(graph, backend);
//!
//! // Execute with event stream
//! let mut stream = streaming.execute_streaming(initial_state).await?;
//!
//! while let Some(event) = stream.next().await {
//!     match event? {
//!         WorkflowEvent::NodeStarted { node_id } => println!("Starting: {}", node_id),
//!         WorkflowEvent::TokenDelta { token } => print!("{}", token),
//!         WorkflowEvent::NodeCompleted { node_id, .. } => println!("\nDone: {}", node_id),
//!         WorkflowEvent::WorkflowCompleted { .. } => break,
//!         _ => {}
//!     }
//! }
//! ```

use crate::graph::{Graph, GraphState, GraphStatus, END};
use crate::handoff::{HandoffContext, HandoffResult, HandoffRouter, HandoffStatus};
use chrono::{DateTime, Utc};
use futures::Stream;
use cortexai_core::errors::CrewError;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::mpsc;

/// Unified workflow event for streaming
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum WorkflowEvent {
    // === Workflow-level events ===
    /// Workflow execution started
    WorkflowStarted {
        workflow_id: String,
        workflow_type: WorkflowType,
        timestamp: DateTime<Utc>,
    },

    /// Workflow execution completed
    WorkflowCompleted {
        workflow_id: String,
        status: StreamingStatus,
        duration_ms: u64,
        timestamp: DateTime<Utc>,
    },

    /// Workflow execution failed
    WorkflowFailed {
        workflow_id: String,
        error: String,
        timestamp: DateTime<Utc>,
    },

    // === Node-level events ===
    /// Node execution started
    NodeStarted {
        node_id: String,
        node_type: StreamingNodeType,
        timestamp: DateTime<Utc>,
    },

    /// Node execution completed
    NodeCompleted {
        node_id: String,
        duration_ms: u64,
        state_changes: Vec<String>,
        timestamp: DateTime<Utc>,
    },

    /// Node execution failed
    NodeFailed {
        node_id: String,
        error: String,
        timestamp: DateTime<Utc>,
    },

    // === Token-level events (LLM streaming) ===
    /// Token generated by LLM
    TokenDelta {
        token: String,
        node_id: Option<String>,
        agent_id: Option<String>,
    },

    /// Reasoning/thinking token (for models that expose this)
    ReasoningDelta {
        token: String,
        node_id: Option<String>,
    },

    /// Tool call started
    ToolCallStarted {
        tool_id: String,
        tool_name: String,
        node_id: Option<String>,
    },

    /// Tool call arguments streaming
    ToolCallDelta {
        tool_id: String,
        arguments_delta: String,
    },

    /// Tool call completed
    ToolCallCompleted {
        tool_id: String,
        result: Option<String>,
    },

    // === Handoff events ===
    /// Agent handoff occurred
    AgentHandoff {
        from_agent: String,
        to_agent: String,
        reason: String,
        timestamp: DateTime<Utc>,
    },

    /// Agent returned to caller
    AgentReturn {
        from_agent: String,
        to_agent: String,
        timestamp: DateTime<Utc>,
    },

    /// Message added to conversation
    MessageAdded {
        role: String,
        agent_id: Option<String>,
        content_preview: String,
        timestamp: DateTime<Utc>,
    },

    // === State events ===
    /// State checkpoint created
    CheckpointCreated {
        checkpoint_id: String,
        node_id: String,
        timestamp: DateTime<Utc>,
    },

    /// State was modified
    StateUpdated {
        keys_changed: Vec<String>,
        node_id: Option<String>,
    },

    // === Progress events ===
    /// Progress update
    Progress {
        current_step: usize,
        total_steps: Option<usize>,
        message: String,
    },

    /// Custom event for extensibility
    Custom {
        event_type: String,
        data: serde_json::Value,
    },
}

/// Type of workflow
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum WorkflowType {
    Graph,
    Handoff,
    Workflow,
    Subgraph,
}

/// Status of streaming workflow completion
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum StreamingStatus {
    Success,
    Failed,
    MaxIterations,
    MaxHandoffs,
    Interrupted,
}

impl From<GraphStatus> for StreamingStatus {
    fn from(status: GraphStatus) -> Self {
        match status {
            GraphStatus::Success => StreamingStatus::Success,
            GraphStatus::Failed => StreamingStatus::Failed,
            GraphStatus::MaxIterations => StreamingStatus::MaxIterations,
            GraphStatus::Interrupted => StreamingStatus::Interrupted,
            GraphStatus::Paused => StreamingStatus::Interrupted,
        }
    }
}

impl From<HandoffStatus> for StreamingStatus {
    fn from(status: HandoffStatus) -> Self {
        match status {
            HandoffStatus::Completed => StreamingStatus::Success,
            HandoffStatus::MaxTurnsReached => StreamingStatus::MaxIterations,
            HandoffStatus::MaxHandoffsReached => StreamingStatus::MaxHandoffs,
            HandoffStatus::Interrupted => StreamingStatus::Interrupted,
        }
    }
}

/// Type of streaming node
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum StreamingNodeType {
    Regular,
    Subgraph,
    LLMCall,
    ToolExecution,
    Conditional,
    Human,
}

/// Stream of workflow events
pub type EventStream = Pin<Box<dyn Stream<Item = Result<WorkflowEvent, CrewError>> + Send>>;

/// Event callback function type
pub type EventCallback = Arc<dyn Fn(WorkflowEvent) + Send + Sync>;

/// Builder for event callbacks
pub struct EventCallbackBuilder {
    on_workflow_started: Option<EventCallback>,
    on_workflow_completed: Option<EventCallback>,
    on_node_started: Option<EventCallback>,
    on_node_completed: Option<EventCallback>,
    on_token: Option<EventCallback>,
    on_handoff: Option<EventCallback>,
    on_any: Option<EventCallback>,
}

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

impl EventCallbackBuilder {
    /// Create a new callback builder
    pub fn new() -> Self {
        Self {
            on_workflow_started: None,
            on_workflow_completed: None,
            on_node_started: None,
            on_node_completed: None,
            on_token: None,
            on_handoff: None,
            on_any: None,
        }
    }

    /// Set callback for workflow started
    pub fn on_workflow_started<F>(mut self, f: F) -> Self
    where
        F: Fn(WorkflowEvent) + Send + Sync + 'static,
    {
        self.on_workflow_started = Some(Arc::new(f));
        self
    }

    /// Set callback for workflow completed
    pub fn on_workflow_completed<F>(mut self, f: F) -> Self
    where
        F: Fn(WorkflowEvent) + Send + Sync + 'static,
    {
        self.on_workflow_completed = Some(Arc::new(f));
        self
    }

    /// Set callback for node started
    pub fn on_node_started<F>(mut self, f: F) -> Self
    where
        F: Fn(WorkflowEvent) + Send + Sync + 'static,
    {
        self.on_node_started = Some(Arc::new(f));
        self
    }

    /// Set callback for node completed
    pub fn on_node_completed<F>(mut self, f: F) -> Self
    where
        F: Fn(WorkflowEvent) + Send + Sync + 'static,
    {
        self.on_node_completed = Some(Arc::new(f));
        self
    }

    /// Set callback for tokens
    pub fn on_token<F>(mut self, f: F) -> Self
    where
        F: Fn(WorkflowEvent) + Send + Sync + 'static,
    {
        self.on_token = Some(Arc::new(f));
        self
    }

    /// Set callback for handoffs
    pub fn on_handoff<F>(mut self, f: F) -> Self
    where
        F: Fn(WorkflowEvent) + Send + Sync + 'static,
    {
        self.on_handoff = Some(Arc::new(f));
        self
    }

    /// Set callback for any event
    pub fn on_any<F>(mut self, f: F) -> Self
    where
        F: Fn(WorkflowEvent) + Send + Sync + 'static,
    {
        self.on_any = Some(Arc::new(f));
        self
    }

    /// Build the event handler
    pub fn build(self) -> EventHandler {
        EventHandler {
            on_workflow_started: self.on_workflow_started,
            on_workflow_completed: self.on_workflow_completed,
            on_node_started: self.on_node_started,
            on_node_completed: self.on_node_completed,
            on_token: self.on_token,
            on_handoff: self.on_handoff,
            on_any: self.on_any,
        }
    }
}

/// Handler for workflow events
pub struct EventHandler {
    on_workflow_started: Option<EventCallback>,
    on_workflow_completed: Option<EventCallback>,
    on_node_started: Option<EventCallback>,
    on_node_completed: Option<EventCallback>,
    on_token: Option<EventCallback>,
    on_handoff: Option<EventCallback>,
    on_any: Option<EventCallback>,
}

impl EventHandler {
    /// Handle an event
    pub fn handle(&self, event: &WorkflowEvent) {
        // Call specific handler
        match event {
            WorkflowEvent::WorkflowStarted { .. } => {
                if let Some(cb) = &self.on_workflow_started {
                    cb(event.clone());
                }
            }
            WorkflowEvent::WorkflowCompleted { .. } | WorkflowEvent::WorkflowFailed { .. } => {
                if let Some(cb) = &self.on_workflow_completed {
                    cb(event.clone());
                }
            }
            WorkflowEvent::NodeStarted { .. } => {
                if let Some(cb) = &self.on_node_started {
                    cb(event.clone());
                }
            }
            WorkflowEvent::NodeCompleted { .. } | WorkflowEvent::NodeFailed { .. } => {
                if let Some(cb) = &self.on_node_completed {
                    cb(event.clone());
                }
            }
            WorkflowEvent::TokenDelta { .. } | WorkflowEvent::ReasoningDelta { .. } => {
                if let Some(cb) = &self.on_token {
                    cb(event.clone());
                }
            }
            WorkflowEvent::AgentHandoff { .. } | WorkflowEvent::AgentReturn { .. } => {
                if let Some(cb) = &self.on_handoff {
                    cb(event.clone());
                }
            }
            _ => {}
        }

        // Call any handler
        if let Some(cb) = &self.on_any {
            cb(event.clone());
        }
    }
}

/// Event emitter for sending events to multiple receivers
pub struct EventEmitter {
    senders: Vec<mpsc::UnboundedSender<WorkflowEvent>>,
    handler: Option<EventHandler>,
}

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

impl EventEmitter {
    /// Create a new emitter
    pub fn new() -> Self {
        Self {
            senders: Vec::new(),
            handler: None,
        }
    }

    /// Create with an event handler
    pub fn with_handler(handler: EventHandler) -> Self {
        Self {
            senders: Vec::new(),
            handler: Some(handler),
        }
    }

    /// Subscribe to events
    pub fn subscribe(&mut self) -> mpsc::UnboundedReceiver<WorkflowEvent> {
        let (tx, rx) = mpsc::unbounded_channel();
        self.senders.push(tx);
        rx
    }

    /// Emit an event
    pub fn emit(&self, event: WorkflowEvent) {
        // Send to all subscribers
        for sender in &self.senders {
            let _ = sender.send(event.clone());
        }

        // Call handler
        if let Some(handler) = &self.handler {
            handler.handle(&event);
        }
    }

    /// Emit workflow started event
    pub fn workflow_started(&self, workflow_id: &str, workflow_type: WorkflowType) {
        self.emit(WorkflowEvent::WorkflowStarted {
            workflow_id: workflow_id.to_string(),
            workflow_type,
            timestamp: Utc::now(),
        });
    }

    /// Emit workflow completed event
    pub fn workflow_completed(&self, workflow_id: &str, status: StreamingStatus, duration_ms: u64) {
        self.emit(WorkflowEvent::WorkflowCompleted {
            workflow_id: workflow_id.to_string(),
            status,
            duration_ms,
            timestamp: Utc::now(),
        });
    }

    /// Emit workflow failed event
    pub fn workflow_failed(&self, workflow_id: &str, error: &str) {
        self.emit(WorkflowEvent::WorkflowFailed {
            workflow_id: workflow_id.to_string(),
            error: error.to_string(),
            timestamp: Utc::now(),
        });
    }

    /// Emit node started event
    pub fn node_started(&self, node_id: &str, node_type: StreamingNodeType) {
        self.emit(WorkflowEvent::NodeStarted {
            node_id: node_id.to_string(),
            node_type,
            timestamp: Utc::now(),
        });
    }

    /// Emit node completed event
    pub fn node_completed(&self, node_id: &str, duration_ms: u64, state_changes: Vec<String>) {
        self.emit(WorkflowEvent::NodeCompleted {
            node_id: node_id.to_string(),
            duration_ms,
            state_changes,
            timestamp: Utc::now(),
        });
    }

    /// Emit token delta event
    pub fn token_delta(&self, token: &str, node_id: Option<&str>, agent_id: Option<&str>) {
        self.emit(WorkflowEvent::TokenDelta {
            token: token.to_string(),
            node_id: node_id.map(|s| s.to_string()),
            agent_id: agent_id.map(|s| s.to_string()),
        });
    }

    /// Emit handoff event
    pub fn agent_handoff(&self, from: &str, to: &str, reason: &str) {
        self.emit(WorkflowEvent::AgentHandoff {
            from_agent: from.to_string(),
            to_agent: to.to_string(),
            reason: reason.to_string(),
            timestamp: Utc::now(),
        });
    }
}

/// Streaming wrapper for Graph execution
pub struct StreamingGraph {
    graph: Arc<Graph>,
    emitter: Arc<EventEmitter>,
}

impl StreamingGraph {
    /// Create a new streaming graph
    pub fn new(graph: Graph) -> Self {
        Self {
            graph: Arc::new(graph),
            emitter: Arc::new(EventEmitter::new()),
        }
    }

    /// Create with event handler
    pub fn with_handler(graph: Graph, handler: EventHandler) -> Self {
        Self {
            graph: Arc::new(graph),
            emitter: Arc::new(EventEmitter::with_handler(handler)),
        }
    }

    /// Get the emitter for subscribing to events
    pub fn emitter(&self) -> &EventEmitter {
        &self.emitter
    }

    /// Execute with streaming events
    pub async fn execute_streaming(
        &self,
        initial_state: GraphState,
    ) -> Result<(GraphState, GraphStatus), CrewError> {
        let workflow_id = uuid::Uuid::new_v4().to_string();
        let start_time = std::time::Instant::now();

        self.emitter
            .workflow_started(&workflow_id, WorkflowType::Graph);

        let mut state = initial_state;
        state.metadata.started_at = Some(Utc::now());
        state.metadata.iterations = 0;

        let mut current_node = self.graph.entry_node.clone();

        let result = loop {
            // Check max iterations
            if state.metadata.iterations >= self.graph.config.max_iterations {
                break (state, GraphStatus::MaxIterations);
            }

            // Check for END node
            if current_node == END {
                break (state, GraphStatus::Success);
            }

            // Get node
            let node = match self.graph.nodes.get(&current_node) {
                Some(n) => n,
                None => {
                    self.emitter.workflow_failed(
                        &workflow_id,
                        &format!("Node not found: {}", current_node),
                    );
                    return Err(CrewError::TaskNotFound(format!(
                        "Node not found: {}",
                        current_node
                    )));
                }
            };

            // Emit node started
            self.emitter
                .node_started(&current_node, StreamingNodeType::Regular);
            let node_start = std::time::Instant::now();

            // Track state keys before execution
            let keys_before: Vec<String> = state
                .data
                .as_object()
                .map(|o| o.keys().cloned().collect())
                .unwrap_or_default();

            // Execute node
            state.metadata.visited_nodes.push(current_node.clone());
            state.metadata.iterations += 1;

            state = match node.executor.call(state).await {
                Ok(s) => s,
                Err(e) => {
                    self.emitter.emit(WorkflowEvent::NodeFailed {
                        node_id: current_node.clone(),
                        error: e.to_string(),
                        timestamp: Utc::now(),
                    });
                    self.emitter.workflow_failed(&workflow_id, &e.to_string());
                    return Err(e);
                }
            };

            // Detect state changes
            let keys_after: Vec<String> = state
                .data
                .as_object()
                .map(|o| o.keys().cloned().collect())
                .unwrap_or_default();

            let state_changes: Vec<String> = keys_after
                .iter()
                .filter(|k| !keys_before.contains(k))
                .cloned()
                .collect();

            // Emit node completed
            self.emitter.node_completed(
                &current_node,
                node_start.elapsed().as_millis() as u64,
                state_changes,
            );

            // Find next node
            current_node = self.graph.find_next_node(&current_node, &state)?;
        };

        let duration_ms = start_time.elapsed().as_millis() as u64;
        self.emitter
            .workflow_completed(&workflow_id, result.1.into(), duration_ms);

        Ok(result)
    }
}

/// Streaming wrapper for Handoff execution
pub struct StreamingHandoff {
    router: Arc<HandoffRouter>,
    emitter: Arc<EventEmitter>,
}

impl StreamingHandoff {
    /// Create a new streaming handoff router
    pub fn new(router: HandoffRouter) -> Self {
        Self {
            router: Arc::new(router),
            emitter: Arc::new(EventEmitter::new()),
        }
    }

    /// Create with event handler
    pub fn with_handler(router: HandoffRouter, handler: EventHandler) -> Self {
        Self {
            router: Arc::new(router),
            emitter: Arc::new(EventEmitter::with_handler(handler)),
        }
    }

    /// Get the emitter
    pub fn emitter(&self) -> &EventEmitter {
        &self.emitter
    }

    /// Execute with streaming events
    pub async fn execute_streaming(
        &self,
        context: HandoffContext,
    ) -> Result<HandoffResult, CrewError> {
        let workflow_id = context.conversation_id.clone();
        let start_time = std::time::Instant::now();

        self.emitter
            .workflow_started(&workflow_id, WorkflowType::Handoff);

        // Execute the handoff router
        let result = self.router.run(context).await?;

        let duration_ms = start_time.elapsed().as_millis() as u64;

        // Emit events for the handoff history
        for record in &result.context.handoff_history {
            if record.is_return {
                self.emitter.emit(WorkflowEvent::AgentReturn {
                    from_agent: record.from_agent.clone(),
                    to_agent: record.to_agent.clone(),
                    timestamp: record.timestamp,
                });
            } else {
                self.emitter.emit(WorkflowEvent::AgentHandoff {
                    from_agent: record.from_agent.clone(),
                    to_agent: record.to_agent.clone(),
                    reason: record.reason.clone(),
                    timestamp: record.timestamp,
                });
            }
        }

        self.emitter
            .workflow_completed(&workflow_id, result.status.into(), duration_ms);

        Ok(result)
    }
}

/// Collector for aggregating streamed tokens
pub struct TokenCollector {
    tokens: Vec<String>,
    node_tokens: HashMap<String, Vec<String>>,
    current_node: Option<String>,
}

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

impl TokenCollector {
    /// Create a new collector
    pub fn new() -> Self {
        Self {
            tokens: Vec::new(),
            node_tokens: HashMap::new(),
            current_node: None,
        }
    }

    /// Set current node context
    pub fn set_node(&mut self, node_id: Option<String>) {
        self.current_node = node_id;
    }

    /// Add a token
    pub fn add_token(&mut self, token: &str) {
        self.tokens.push(token.to_string());

        if let Some(node_id) = &self.current_node {
            self.node_tokens
                .entry(node_id.clone())
                .or_default()
                .push(token.to_string());
        }
    }

    /// Get all collected tokens as a string
    pub fn collected(&self) -> String {
        self.tokens.join("")
    }

    /// Get tokens for a specific node
    pub fn tokens_for_node(&self, node_id: &str) -> String {
        self.node_tokens
            .get(node_id)
            .map(|t| t.join(""))
            .unwrap_or_default()
    }

    /// Get all tokens
    pub fn all_tokens(&self) -> &[String] {
        &self.tokens
    }

    /// Clear the collector
    pub fn clear(&mut self) {
        self.tokens.clear();
        self.node_tokens.clear();
        self.current_node = None;
    }
}

/// Statistics about a streaming execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamingStats {
    /// Total tokens streamed
    pub total_tokens: usize,
    /// Tokens per node
    pub tokens_per_node: HashMap<String, usize>,
    /// Total duration in milliseconds
    pub duration_ms: u64,
    /// Number of nodes executed
    pub nodes_executed: usize,
    /// Number of handoffs
    pub handoffs: usize,
    /// Events emitted
    pub events_emitted: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::graph::GraphBuilder;
    use std::sync::atomic::{AtomicUsize, Ordering};

    #[tokio::test]
    async fn test_event_emitter() {
        let mut emitter = EventEmitter::new();
        let mut rx = emitter.subscribe();

        emitter.workflow_started("test-1", WorkflowType::Graph);

        let event = rx.recv().await.unwrap();
        match event {
            WorkflowEvent::WorkflowStarted { workflow_id, .. } => {
                assert_eq!(workflow_id, "test-1");
            }
            _ => panic!("Wrong event type"),
        }
    }

    #[tokio::test]
    async fn test_event_callback_builder() {
        let started_count = Arc::new(AtomicUsize::new(0));
        let completed_count = Arc::new(AtomicUsize::new(0));

        let started_clone = started_count.clone();
        let completed_clone = completed_count.clone();

        let handler = EventCallbackBuilder::new()
            .on_workflow_started(move |_| {
                started_clone.fetch_add(1, Ordering::SeqCst);
            })
            .on_workflow_completed(move |_| {
                completed_clone.fetch_add(1, Ordering::SeqCst);
            })
            .build();

        let emitter = EventEmitter::with_handler(handler);

        emitter.workflow_started("test", WorkflowType::Graph);
        emitter.workflow_completed("test", StreamingStatus::Success, 100);

        assert_eq!(started_count.load(Ordering::SeqCst), 1);
        assert_eq!(completed_count.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_streaming_graph() {
        let graph = GraphBuilder::new("test")
            .add_node("step1", |mut state: GraphState| async move {
                state.set("value", 1);
                Ok(state)
            })
            .add_node("step2", |mut state: GraphState| async move {
                let v: i32 = state.get("value").unwrap_or(0);
                state.set("value", v + 1);
                Ok(state)
            })
            .add_edge("step1", "step2")
            .add_edge("step2", END)
            .set_entry("step1")
            .build()
            .unwrap();

        let events = Arc::new(std::sync::Mutex::new(Vec::new()));
        let events_clone = events.clone();

        let handler = EventCallbackBuilder::new()
            .on_any(move |event| {
                events_clone.lock().unwrap().push(event);
            })
            .build();

        let streaming = StreamingGraph::with_handler(graph, handler);
        let (state, status) = streaming
            .execute_streaming(GraphState::new())
            .await
            .unwrap();

        assert_eq!(status, GraphStatus::Success);
        assert_eq!(state.get::<i32>("value"), Some(2));

        let captured = events.lock().unwrap();
        assert!(captured.len() >= 5); // started + 2*(node_started + node_completed) + completed
    }

    #[tokio::test]
    async fn test_token_collector() {
        let mut collector = TokenCollector::new();

        collector.set_node(Some("node1".to_string()));
        collector.add_token("Hello");
        collector.add_token(" ");
        collector.add_token("World");

        collector.set_node(Some("node2".to_string()));
        collector.add_token("!");

        assert_eq!(collector.collected(), "Hello World!");
        assert_eq!(collector.tokens_for_node("node1"), "Hello World");
        assert_eq!(collector.tokens_for_node("node2"), "!");
    }

    #[tokio::test]
    async fn test_workflow_event_serialization() {
        let event = WorkflowEvent::TokenDelta {
            token: "test".to_string(),
            node_id: Some("node1".to_string()),
            agent_id: None,
        };

        let json = serde_json::to_string(&event).unwrap();
        assert!(json.contains("TokenDelta"));
        assert!(json.contains("test"));

        let parsed: WorkflowEvent = serde_json::from_str(&json).unwrap();
        match parsed {
            WorkflowEvent::TokenDelta { token, .. } => assert_eq!(token, "test"),
            _ => panic!("Wrong type"),
        }
    }

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

        emitter.node_started("node1", StreamingNodeType::Regular);

        // Both should receive
        let e1 = rx1.recv().await.unwrap();
        let e2 = rx2.recv().await.unwrap();

        match (e1, e2) {
            (
                WorkflowEvent::NodeStarted { node_id: n1, .. },
                WorkflowEvent::NodeStarted { node_id: n2, .. },
            ) => {
                assert_eq!(n1, "node1");
                assert_eq!(n2, "node1");
            }
            _ => panic!("Wrong events"),
        }
    }

    #[test]
    fn test_streaming_status_conversions() {
        assert_eq!(
            StreamingStatus::from(GraphStatus::Success),
            StreamingStatus::Success
        );
        assert_eq!(
            StreamingStatus::from(GraphStatus::Failed),
            StreamingStatus::Failed
        );
        assert_eq!(
            StreamingStatus::from(HandoffStatus::Completed),
            StreamingStatus::Success
        );
        assert_eq!(
            StreamingStatus::from(HandoffStatus::MaxHandoffsReached),
            StreamingStatus::MaxHandoffs
        );
    }
}