zesdex/src/app/runtime/stream/mod.rs

pub mod turn;
pub mod tools;

use serde::{Deserialize, Serialize};
use serde_json::Value;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum StreamEvent {
    Token(String),
    Reasoning(String),
    ToolCallDelta {
        index: usize,
        id: Option<String>,
        name: Option<String>,
        arguments_delta: String,
    },
    Usage {
        prompt_tokens: u64,
        completion_tokens: u64,
        total_tokens: u64,
    },
    Done,
    Error(String),
}

pub struct SseParser {
    buffer: String,
    event_type: Option<String>,
    data_lines: Vec<String>,
}

impl SseParser {
    pub fn new() -> Self {
        SseParser {
            buffer: String::new(),
            event_type: None,
            data_lines: Vec::new(),
        }
    }

    pub fn feed(&mut self, chunk: &str) -> Vec<StreamEvent> {
        self.buffer.push_str(chunk);
        let mut events = Vec::new();
        while let Some(line_end) = self.buffer.find('\n') {
            let line = self.buffer[..line_end].trim_end_matches('\r').to_string();
            self.buffer = self.buffer[line_end + 1..].to_string();
            if line.is_empty() {
                if let Some(event) = self.flush_event() {
                    events.push(event);
                }
            } else if let Some(ty) = line.strip_prefix("event: ") {
                self.event_type = Some(ty.trim().to_string());
            } else if let Some(data) = line.strip_prefix("data: ") {
                self.data_lines.push(data.to_string());
            } else if line.starts_with("data:") {
                self.data_lines.push(String::new());
            }
        }
        events
    }

    fn flush_event(&mut self) -> Option<StreamEvent> {
        let data = self.data_lines.join("\n");
        self.data_lines.clear();
        let event_type = self.event_type.take().unwrap_or_default();
        if data.is_empty() || data == "[DONE]" {
            if data == "[DONE]" {
                return Some(StreamEvent::Done);
            }
            return None;
        }
        let value: Value = serde_json::from_str(&data).ok()?;
        if let Some(usage) = value.get("usage") {
            if !usage.is_null() {
                let prompt_tokens = usage.get("prompt_tokens").and_then(|v| v.as_u64()).unwrap_or(0);
                let completion_tokens = usage.get("completion_tokens").and_then(|v| v.as_u64()).unwrap_or(0);
                let total_tokens = usage.get("total_tokens").and_then(|v| v.as_u64())
                    .unwrap_or(prompt_tokens + completion_tokens);
                return Some(StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens });
            }
        }
        match event_type.as_str() {
            "message.stop" => Some(StreamEvent::Done),
            "message.start" => None,
            "message.delta" | "" => {
                let delta = value.get("delta").or_else(|| value.get("choices"))?;
                if let Some(choices) = delta.as_array() {
                    let choice = choices.first()?;
                    let delta = choice.get("delta")?;
                    if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
                        return Some(StreamEvent::Token(content.to_string()));
                    }
                    if let Some(reasoning) = delta.get("reasoning_content").and_then(|r| r.as_str()) {
                        return Some(StreamEvent::Reasoning(reasoning.to_string()));
                    }
                    if let Some(tool_calls) = delta.get("tool_calls").and_then(|tc| tc.as_array()) {
                        if let Some(tc) = tool_calls.first() {
                            let index = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
                            let id = tc.get("id").and_then(|i| i.as_str()).map(|s| s.to_string());
                            let name = tc.get("function")
                                .and_then(|f| f.get("name"))
                                .and_then(|n| n.as_str())
                                .map(|s| s.to_string());
                            let args_delta = tc.get("function")
                                .and_then(|f| f.get("arguments"))
                                .and_then(|a| a.as_str())
                                .unwrap_or("")
                                .to_string();
                            return Some(StreamEvent::ToolCallDelta {
                                index,
                                id,
                                name,
                                arguments_delta: args_delta,
                            });
                        }
                    }
                    let finish = choice.get("finish_reason");
                    if let Some(reason) = finish.and_then(|r| r.as_str()) {
                        if reason == "stop" || reason == "tool_calls" {
                            return Some(StreamEvent::Done);
                        }
                    }
                }
                if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
                    return Some(StreamEvent::Token(content.to_string()));
                }
                None
            }
            _ => None,
        }
    }

    /// Clears any partially-buffered SSE frame. Reserved for reconnect/retry flows that
    /// reuse a parser instance across requests rather than constructing a fresh one.
    #[allow(dead_code)]
    pub fn reset(&mut self) {
        self.buffer.clear();
        self.event_type = None;
        self.data_lines.clear();
    }
}

/// Fallback parser for providers that send bare JSON chunks instead of SSE-framed
/// `data: ...` lines. Not used by the `SseParser` streaming path (which handles
/// standard SSE framing directly), kept for providers/tests that feed raw chunks.
#[allow(dead_code)]
pub fn parse_stream_chunk(data: &str) -> Option<StreamEvent> {
    let value: Value = serde_json::from_str(data).ok()?;
    if value == Value::Null {
        return None;
    }
    let choices = value.get("choices")?.as_array()?;
    let choice = choices.first()?;
    let delta = choice.get("delta")?;
    if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
        return Some(StreamEvent::Token(content.to_string()));
    }
    if let Some(reasoning) = delta.get("reasoning_content").and_then(|r| r.as_str()) {
        return Some(StreamEvent::Reasoning(reasoning.to_string()));
    }
    if let Some(finish) = choice.get("finish_reason").and_then(|r| r.as_str()) {
        if finish == "stop" || finish == "tool_calls" {
            return Some(StreamEvent::Done);
        }
    }
    if let Some(tool_calls) = delta.get("tool_calls").and_then(|tc| tc.as_array()) {
        if let Some(tc) = tool_calls.first() {
            let index = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
            let id = tc.get("id").and_then(|i| i.as_str()).map(|s| s.to_string());
            let name = tc.get("function")
                .and_then(|f| f.get("name"))
                .and_then(|n| n.as_str())
                .map(|s| s.to_string());
            let args = tc.get("function")
                .and_then(|f| f.get("arguments"))
                .and_then(|a| a.as_str())
                .unwrap_or("")
                .to_string();
            return Some(StreamEvent::ToolCallDelta {
                index,
                id,
                name,
                arguments_delta: args,
            });
        }
    }
    None
}

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

    #[test]
    fn feed_parses_single_token_chunk() {
        let mut p = SseParser::new();
        let events = p.feed("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n");
        assert_eq!(events.len(), 1);
        match &events[0] {
            StreamEvent::Token(t) => assert_eq!(t, "hello"),
            other => panic!("expected Token, got {:?}", other),
        }
    }

    #[test]
    fn feed_handles_chunk_split_mid_line() {
        let mut p = SseParser::new();
        let e1 = p.feed("data: {\"choices\":[{\"delta\":{\"content\":\"partial");
        assert!(e1.is_empty(), "no event until the line and blank separator complete");
        let e2 = p.feed("\"}}]}\n\n");
        assert_eq!(e2.len(), 1);
        match &e2[0] {
            StreamEvent::Token(t) => assert_eq!(t, "partial"),
            other => panic!("expected Token, got {:?}", other),
        }
    }

    #[test]
    fn feed_emits_done_on_done_sentinel() {
        let mut p = SseParser::new();
        let events = p.feed("data: [DONE]\n\n");
        assert_eq!(events.len(), 1);
        assert!(matches!(events[0], StreamEvent::Done));
    }

    #[test]
    fn feed_emits_done_on_finish_reason_stop() {
        let mut p = SseParser::new();
        let events = p.feed(
            "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
        );
        assert_eq!(events.len(), 1);
        assert!(matches!(events[0], StreamEvent::Done));
    }

    #[test]
    fn feed_parses_tool_call_delta() {
        let mut p = SseParser::new();
        let events = p.feed(
            "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"bash\",\"arguments\":\"{\\\"cmd\\\"\"}}]}}]}\n\n",
        );
        assert_eq!(events.len(), 1);
        match &events[0] {
            StreamEvent::ToolCallDelta { index, id, name, arguments_delta } => {
                assert_eq!(*index, 0);
                assert_eq!(id.as_deref(), Some("call_1"));
                assert_eq!(name.as_deref(), Some("bash"));
                assert_eq!(arguments_delta, "{\"cmd\"");
            }
            other => panic!("expected ToolCallDelta, got {:?}", other),
        }
    }

    #[test]
    fn feed_parses_usage_chunk() {
        let mut p = SseParser::new();
        let events = p.feed(
            "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}\n\n",
        );
        assert_eq!(events.len(), 1);
        match &events[0] {
            StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens } => {
                assert_eq!(*prompt_tokens, 10);
                assert_eq!(*completion_tokens, 5);
                assert_eq!(*total_tokens, 15);
            }
            other => panic!("expected Usage, got {:?}", other),
        }
    }

    #[test]
    fn feed_ignores_empty_data_lines() {
        let mut p = SseParser::new();
        let events = p.feed(": comment\n\n");
        assert!(events.is_empty());
    }

    #[test]
    fn feed_multiple_events_across_one_chunk() {
        let mut p = SseParser::new();
        let chunk = "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n";
        let events = p.feed(chunk);
        assert_eq!(events.len(), 2);
        match (&events[0], &events[1]) {
            (StreamEvent::Token(a), StreamEvent::Token(b)) => {
                assert_eq!(a, "a");
                assert_eq!(b, "b");
            }
            other => panic!("expected two Tokens, got {:?}", other),
        }
    }
}