quvyta-cli 0.1.0-alpha.5

//! A streamed reply in the Anthropic Messages shape, put together as it arrives.
//!
//! The server sends server-sent events: `event:` and `data:` lines, a blank line between events.
//! Only the `data` matters, since each JSON object names its own type. Text and thinking are
//! handed on piece by piece as they come, so the screen can show them growing; a tool call is
//! handed on once its name is known, and its input, which arrives as fragments of JSON, is read
//! when its block ends.

use std::io;

use serde_json::{Map, Value, json};

use super::ProviderError;

/// Something that arrived while a reply streams in.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Delta {
    /// More of the answer's text.
    Text(String),
    /// More of the model's thinking, for providers that show it.
    Thinking(String),
    /// The model began a tool call with this name.
    ToolCall(String),
}

/// Why the model stopped.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Stop {
    /// It finished its turn.
    EndTurn,
    /// It asked for tools and waits for their results.
    ToolUse,
    /// It ran out of the tokens it was allowed to write.
    MaxTokens,
    /// Anything else the server named, kept as it was said.
    Other(String),
}

/// A whole reply.
#[derive(Debug, Clone, PartialEq)]
pub struct Reply {
    /// The content blocks in the Messages shape, ready to go back into the conversation as the
    /// assistant's message: `text`, `thinking` and `tool_use` blocks.
    pub content: Vec<Value>,
    /// Why the model stopped.
    pub stop: Stop,
    /// The input tokens the server counted, when it said.
    pub input_tokens: Option<u64>,
    /// The output tokens the server counted, when it said.
    pub output_tokens: Option<u64>,
}

/// One content block while it is being written.
#[derive(Debug)]
struct Block {
    value: Value,
    /// The index the server gave the block, which its pieces name.
    index: Option<u64>,
    /// The fragments of a tool call's input, joined when the block ends.
    input: String,
    /// The text fields the pieces grow (`text`, `thinking`, `signature`), kept as strings while
    /// they grow and put into the block at the end: adding a piece to a JSON value copies the
    /// whole text every time, which turns a long reply into a slow one.
    growing: Vec<(&'static str, String)>,
}

impl Block {
    fn new(value: Value, index: Option<u64>) -> Self {
        Self { value, index, input: String::new(), growing: Vec::new() }
    }

    /// Adds `piece` to the text field `field`.
    fn grow(&mut self, field: &'static str, piece: &str) {
        if let Some((_, text)) = self.growing.iter_mut().find(|(name, _)| *name == field) {
            text.push_str(piece);
            return;
        }
        let mut text = self.value.get(field).and_then(Value::as_str).unwrap_or_default().to_owned();
        text.push_str(piece);
        self.growing.push((field, text));
    }
}

/// Reads the event stream in `lines` to its end and puts the reply together, handing each piece
/// to `on_delta` as it comes. `cancelled` is asked between events; when it says yes the reply
/// read so far is returned as it is.
///
/// # Errors
///
/// When the server sends an error event, when the stream breaks, or when it ends without saying
/// the message is complete.
pub fn read(
    lines: impl IntoIterator<Item = io::Result<String>>,
    on_delta: &mut dyn FnMut(Delta),
    cancelled: &dyn Fn() -> bool,
) -> Result<Reply, ProviderError> {
    let mut blocks: Vec<Block> = Vec::new();
    let mut reply = Reply { content: Vec::new(), stop: Stop::EndTurn, input_tokens: None, output_tokens: None };
    let mut data = String::new();
    let mut complete = false;
    for line in lines {
        let line = line.map_err(|error| ProviderError::Unreadable { what: error.to_string() })?;
        if let Some(rest) = line.strip_prefix("data:") {
            data.push_str(rest.strip_prefix(' ').unwrap_or(rest));
            continue;
        }
        if !line.is_empty() || data.is_empty() {
            continue;
        }
        let event = std::mem::take(&mut data);
        if event == "[DONE]" {
            complete = true;
            break;
        }
        let event: Value = serde_json::from_str(&event)
            .map_err(|error| ProviderError::Unreadable { what: format!("{error}: {event}") })?;
        if apply(&event, &mut blocks, &mut reply, on_delta)? {
            complete = true;
            break;
        }
        if cancelled() {
            complete = true;
            break;
        }
    }
    // Lines that stop coming because the reply was stopped end it as a stop does between events.
    if !complete && cancelled() {
        complete = true;
    }
    if !complete {
        return Err(ProviderError::Unreadable { what: "the reply ended before the message was complete".to_owned() });
    }
    reply.content = blocks.into_iter().map(finish).collect();
    Ok(reply)
}

/// Applies one event and says whether it was the last.
fn apply(
    event: &Value,
    blocks: &mut Vec<Block>,
    reply: &mut Reply,
    on_delta: &mut dyn FnMut(Delta),
) -> Result<bool, ProviderError> {
    let kind = event.get("type").and_then(Value::as_str).unwrap_or_default();
    match kind {
        "message_start" => {
            let usage = event.pointer("/message/usage");
            reply.input_tokens = usage.and_then(|usage| usage.get("input_tokens")).and_then(Value::as_u64);
        }
        "content_block_start" => {
            let block = event.get("content_block").cloned().unwrap_or_else(|| json!({"type": "text", "text": ""}));
            if block.get("type").and_then(Value::as_str) == Some("tool_use") {
                on_delta(Delta::ToolCall(block.get("name").and_then(Value::as_str).unwrap_or_default().to_owned()));
            }
            blocks.push(Block::new(block, event.get("index").and_then(Value::as_u64)));
        }
        "content_block_delta" => {
            // A piece goes to the block its index names; servers that write two blocks at once
            // (thinking and text, for some) interleave their pieces. The latest block of that
            // index wins, and a piece that names none goes to the block being written.
            let index = event.get("index").and_then(Value::as_u64);
            let at = index
                .and_then(|index| blocks.iter().rposition(|block| block.index == Some(index)))
                .or_else(|| blocks.len().checked_sub(1));
            let Some(block) = at.and_then(|at| blocks.get_mut(at)) else {
                return Err(ProviderError::Unreadable { what: "a delta arrived before any block".to_owned() });
            };
            let delta = event.get("delta").cloned().unwrap_or_default();
            let text = |field: &str| delta.get(field).and_then(Value::as_str).unwrap_or_default().to_owned();
            match delta.get("type").and_then(Value::as_str).unwrap_or_default() {
                "text_delta" => {
                    let piece = text("text");
                    block.grow("text", &piece);
                    on_delta(Delta::Text(piece));
                }
                "thinking_delta" => {
                    let piece = text("thinking");
                    block.grow("thinking", &piece);
                    on_delta(Delta::Thinking(piece));
                }
                "signature_delta" => block.grow("signature", &text("signature")),
                "input_json_delta" => block.input.push_str(&text("partial_json")),
                _ => {}
            }
        }
        "message_delta" => {
            if let Some(stop) = event.pointer("/delta/stop_reason").and_then(Value::as_str) {
                reply.stop = match stop {
                    "end_turn" | "stop_sequence" => Stop::EndTurn,
                    "tool_use" => Stop::ToolUse,
                    "max_tokens" => Stop::MaxTokens,
                    other => Stop::Other(other.to_owned()),
                };
            }
            if let Some(tokens) = event.pointer("/usage/output_tokens").and_then(Value::as_u64) {
                reply.output_tokens = Some(tokens);
            }
            // The counts here are the whole reply's. Some servers (OpenRouter among them) say 0
            // input tokens when the reply starts and the real count only here; a 0 never
            // replaces a count the start already gave.
            if let Some(tokens) = event.pointer("/usage/input_tokens").and_then(Value::as_u64)
                && (tokens > 0 || reply.input_tokens.is_none())
            {
                reply.input_tokens = Some(tokens);
            }
        }
        "message_stop" => return Ok(true),
        "error" => {
            let said = event.pointer("/error/message").and_then(Value::as_str).unwrap_or("error").to_owned();
            return Err(ProviderError::Stopped { said });
        }
        _ => {}
    }
    Ok(false)
}

/// A block as it goes back into the conversation: a tool call gets its input read from the
/// fragments. The API takes only an object as a call's input and refuses every request after one
/// that is not, so an input that is not a JSON object (cut in half, or never JSON) is kept whole
/// under [`UNPARSED`] in an object of its own: the model still sees what it sent.
fn finish(block: Block) -> Value {
    let Block { mut value, input, growing, .. } = block;
    if let Some(object) = value.as_object_mut() {
        for (field, text) in growing {
            object.insert(field.to_owned(), Value::String(text));
        }
    }
    if value.get("type").and_then(Value::as_str) == Some("tool_use")
        && let Some(object) = value.as_object_mut()
    {
        let parsed = if input.trim().is_empty() {
            Value::Object(Map::new())
        } else {
            match serde_json::from_str(&input) {
                Ok(Value::Object(fields)) => Value::Object(fields),
                _ => json!({ UNPARSED: input }),
            }
        };
        object.insert("input".to_owned(), parsed);
    }
    value
}

/// The field a call's input is kept under when it is not a JSON object.
pub const UNPARSED: &str = "unparsed_input";

#[cfg(test)]
mod tests {
    use std::io::BufRead as _;

    use super::*;

    /// The events of a reply, written the way servers send them.
    fn stream(events: &[Value]) -> String {
        events
            .iter()
            .map(|event| format!("event: {}\ndata: {event}\n\n", event["type"].as_str().unwrap_or("")))
            .collect()
    }

    fn collect(text: &str) -> (Result<Reply, ProviderError>, Vec<Delta>) {
        let mut deltas = Vec::new();
        let reply = read(text.as_bytes().lines(), &mut |delta| deltas.push(delta), &|| false);
        (reply, deltas)
    }

    #[test]
    fn text_and_thinking_arrive_in_pieces_and_end_up_whole() {
        let text = stream(&[
            json!({"type": "message_start", "message": {"usage": {"input_tokens": 14}}}),
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": "", "signature": ""}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "Short "}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "answer."}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": "sig"}}),
            json!({"type": "content_block_stop", "index": 0}),
            json!({"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}),
            json!({"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "Hey"}}),
            json!({"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": " there"}}),
            json!({"type": "content_block_stop", "index": 1}),
            json!({"type": "ping"}),
            json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 9}}),
            json!({"type": "message_stop"}),
        ]);
        let (reply, deltas) = collect(&text);
        let reply = reply.expect("a reply");
        assert_eq!(
            deltas,
            [
                Delta::Thinking("Short ".into()),
                Delta::Thinking("answer.".into()),
                Delta::Text("Hey".into()),
                Delta::Text(" there".into())
            ]
        );
        assert_eq!(
            reply.content,
            [
                json!({"type": "thinking", "thinking": "Short answer.", "signature": "sig"}),
                json!({"type": "text", "text": "Hey there"})
            ]
        );
        assert_eq!(reply.stop, Stop::EndTurn);
        assert_eq!((reply.input_tokens, reply.output_tokens), (Some(14), Some(9)));
    }

    #[test]
    fn the_input_count_said_only_at_the_end_is_kept() {
        let text = stream(&[
            json!({"type": "message_start", "message": {"usage": {"input_tokens": 0}}}),
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi"}}),
            json!({"type": "content_block_stop", "index": 0}),
            json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"input_tokens": 2817, "output_tokens": 41}}),
            json!({"type": "message_stop"}),
        ]);
        let (reply, _) = collect(&text);
        let reply = reply.expect("a reply");
        assert_eq!((reply.input_tokens, reply.output_tokens), (Some(2817), Some(41)));
    }

    #[test]
    fn a_zero_at_the_end_does_not_undo_the_starts_count() {
        let text = stream(&[
            json!({"type": "message_start", "message": {"usage": {"input_tokens": 1534}}}),
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi"}}),
            json!({"type": "content_block_stop", "index": 0}),
            json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"input_tokens": 0, "output_tokens": 7}}),
            json!({"type": "message_stop"}),
        ]);
        let (reply, _) = collect(&text);
        let reply = reply.expect("a reply");
        assert_eq!(reply.input_tokens, Some(1534));
    }

    #[test]
    fn a_tool_call_is_announced_by_name_and_its_input_read_from_fragments() {
        let text = stream(&[
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "t1", "name": "read_file", "input": {}}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"path\": \"src/"}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "main.rs\"}"}}),
            json!({"type": "content_block_stop", "index": 0}),
            json!({"type": "message_delta", "delta": {"stop_reason": "tool_use"}}),
            json!({"type": "message_stop"}),
        ]);
        let (reply, deltas) = collect(&text);
        let reply = reply.expect("a reply");
        assert_eq!(deltas, [Delta::ToolCall("read_file".into())]);
        assert_eq!(
            reply.content,
            [json!({"type": "tool_use", "id": "t1", "name": "read_file", "input": {"path": "src/main.rs"}})]
        );
        assert_eq!(reply.stop, Stop::ToolUse);
    }

    #[test]
    fn an_error_event_and_a_cut_stream_are_errors() {
        let error = stream(&[json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}})]);
        assert_eq!(collect(&error).0, Err(ProviderError::Stopped { said: "Overloaded".into() }));
        let cut = stream(&[
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
        ]);
        assert!(matches!(collect(&cut).0, Err(ProviderError::Unreadable { .. })), "no message_stop, no reply");
    }

    #[test]
    fn a_cancelled_reply_keeps_what_had_arrived() {
        let text = stream(&[
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Half"}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " more"}}),
        ]);
        let mut seen = 0;
        let reply = read(text.as_bytes().lines(), &mut |_| seen += 1, &|| true).expect("what arrived");
        assert_eq!(reply.content, [json!({"type": "text", "text": ""})], "stopped after the first event");
        assert_eq!(seen, 0);
    }

    /// Reads `text` as a stream that stopped early: the error it has to give, and every piece it
    /// showed while reading.
    fn cut(text: &str) -> (ProviderError, Vec<Delta>) {
        let mut deltas = Vec::new();
        let error = read(text.as_bytes().lines(), &mut |delta| deltas.push(delta), &|| false)
            .expect_err("an unfinished reply is not a reply");
        let ProviderError::Unreadable { what } = &error else { panic!("a cut stream, got {error:?}") };
        assert!(what.contains("ended before the message was complete"), "{what}");
        (error, deltas)
    }

    #[test]
    fn a_stream_that_stops_in_the_middle_of_a_json_line_is_an_error_and_not_a_reply() {
        // The connection goes in the middle of a line, so the last line has no end of line.
        cut("event: content_block_start\ndata: {\"type\": \"content_bl");
        // An event whose blank line never came is not an event yet either.
        cut("event: content_block_start\ndata: {\"type\": \"message_stop\"}");
        // What came whole before the cut is not handed on as a reply, and no piece is shown.
        let whole = stream(&[json!({"type": "message_start", "message": {"usage": {"input_tokens": 3}}})]);
        let (_, deltas) = cut(&format!("{whole}data: {{\"type\": \"mess"));
        assert!(deltas.is_empty(), "no piece of an unfinished reply is shown");
    }

    #[test]
    fn pieces_of_two_blocks_written_at_once_go_each_to_its_own_block() {
        let text = stream(&[
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}),
            json!({"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "weigh "}}),
            json!({"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "Hello"}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "it"}}),
            json!({"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": " there"}}),
            json!({"type": "message_stop"}),
        ]);
        let (reply, _) = collect(&text);
        assert_eq!(
            reply.expect("a reply").content,
            [json!({"type": "thinking", "thinking": "weigh it"}), json!({"type": "text", "text": "Hello there"})],
            "each piece is in the block its index names"
        );
    }

    #[test]
    fn a_piece_that_names_no_block_goes_to_the_block_being_written() {
        let text = stream(&[
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
            json!({"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": "second"}}),
            json!({"type": "content_block_delta", "delta": {"type": "text_delta", "text": " on"}}),
            json!({"type": "content_block_delta", "index": 7, "delta": {"type": "text_delta", "text": " and on"}}),
            json!({"type": "message_stop"}),
        ]);
        let (reply, _) = collect(&text);
        assert_eq!(
            reply.expect("a reply").content,
            [json!({"type": "text", "text": ""}), json!({"type": "text", "text": "second on and on"})]
        );
    }

    #[test]
    fn one_index_with_two_starts_is_two_blocks_and_neither_is_lost() {
        let text = stream(&[
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "one", "name": "list_dir", "input": {}}}),
            json!({"type": "content_block_stop", "index": 0}),
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "two", "name": "read_file", "input": {}}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"path\": \"a\"}"}}),
            json!({"type": "message_stop"}),
        ]);
        let (reply, deltas) = collect(&text);
        let reply = reply.expect("a reply");
        assert_eq!(deltas, [Delta::ToolCall("list_dir".into()), Delta::ToolCall("read_file".into())]);
        assert_eq!(
            reply.content,
            [
                json!({"type": "tool_use", "id": "one", "name": "list_dir", "input": {}}),
                json!({"type": "tool_use", "id": "two", "name": "read_file", "input": {"path": "a"}})
            ]
        );
    }

    #[test]
    fn a_thinking_block_with_no_signature_comes_through_as_it_arrived() {
        let text = stream(&[
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "weighing it"}}),
            json!({"type": "message_stop"}),
        ]);
        let (reply, _) = collect(&text);
        assert_eq!(reply.expect("a reply").content, [json!({"type": "thinking", "thinking": "weighing it"})]);
    }

    #[test]
    fn a_stop_reason_that_is_none_of_the_three_is_kept_as_it_was_said() {
        for (said, wanted) in [
            ("refusal", Stop::Other("refusal".to_owned())),
            ("pause_turn", Stop::Other("pause_turn".to_owned())),
            ("", Stop::Other(String::new())),
        ] {
            let text = stream(&[
                json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
                json!({"type": "message_delta", "delta": {"stop_reason": said}}),
                json!({"type": "message_stop"}),
            ]);
            assert_eq!(collect(&text).0.expect("a reply").stop, wanted, "{said:?}");
        }
    }

    #[test]
    fn a_tool_input_of_a_megabyte_comes_through_as_it_was_sent() {
        let path = "p".repeat(1024 * 1024);
        let text = stream(&[
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "big", "name": "read_file", "input": {}}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"path\": \""}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": path}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "\"}"}}),
            json!({"type": "message_stop"}),
        ]);
        let (reply, _) = collect(&text);
        let reply = reply.expect("a reply");
        assert_eq!(
            reply.content[0]["input"]["path"].as_str().map(str::len),
            Some(1024 * 1024),
            "the whole path is in the input"
        );
    }

    #[test]
    fn a_tool_call_with_no_name_is_announced_empty_and_the_input_is_read_as_it_came() {
        let text = stream(&[
            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "x"}}),
            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"path\": \"a\"}"}}),
            json!({"type": "message_stop"}),
        ]);
        let (reply, deltas) = collect(&text);
        assert_eq!(deltas, [Delta::ToolCall("".into())], "the name is empty, not missing");
        assert_eq!(reply.expect("a reply").content[0]["input"]["path"], "a");
    }

    #[test]
    fn a_delta_that_arrives_before_any_block_is_an_error_that_says_so() {
        let text =
            stream(&[json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}})]);
        assert_eq!(
            collect(&text).0,
            Err(ProviderError::Unreadable { what: "a delta arrived before any block".to_owned() })
        );
    }
}