qcli/provider/
reply.rs

1//! A streamed reply in the Anthropic Messages shape, put together as it arrives.
2//!
3//! The server sends server-sent events: `event:` and `data:` lines, a blank line between events.
4//! Only the `data` matters, since each JSON object names its own type. Text and thinking are
5//! handed on piece by piece as they come, so the screen can show them growing; a tool call is
6//! handed on once its name is known, and its input, which arrives as fragments of JSON, is read
7//! when its block ends.
8
9use std::io;
10
11use serde_json::{Map, Value, json};
12
13use super::ProviderError;
14
15/// Something that arrived while a reply streams in.
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub enum Delta {
18    /// More of the answer's text.
19    Text(String),
20    /// More of the model's thinking, for providers that show it.
21    Thinking(String),
22    /// The model began a tool call with this name.
23    ToolCall(String),
24}
25
26/// Why the model stopped.
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub enum Stop {
29    /// It finished its turn.
30    EndTurn,
31    /// It asked for tools and waits for their results.
32    ToolUse,
33    /// It ran out of the tokens it was allowed to write.
34    MaxTokens,
35    /// Anything else the server named, kept as it was said.
36    Other(String),
37}
38
39/// A whole reply.
40#[derive(Debug, Clone, PartialEq)]
41pub struct Reply {
42    /// The content blocks in the Messages shape, ready to go back into the conversation as the
43    /// assistant's message: `text`, `thinking` and `tool_use` blocks.
44    pub content: Vec<Value>,
45    /// Why the model stopped.
46    pub stop: Stop,
47    /// The input tokens the server counted, when it said.
48    pub input_tokens: Option<u64>,
49    /// The output tokens the server counted, when it said.
50    pub output_tokens: Option<u64>,
51}
52
53/// One content block while it is being written.
54#[derive(Debug)]
55struct Block {
56    value: Value,
57    /// The index the server gave the block, which its pieces name.
58    index: Option<u64>,
59    /// The fragments of a tool call's input, joined when the block ends.
60    input: String,
61    /// The text fields the pieces grow (`text`, `thinking`, `signature`), kept as strings while
62    /// they grow and put into the block at the end: adding a piece to a JSON value copies the
63    /// whole text every time, which turns a long reply into a slow one.
64    growing: Vec<(&'static str, String)>,
65}
66
67impl Block {
68    fn new(value: Value, index: Option<u64>) -> Self {
69        Self { value, index, input: String::new(), growing: Vec::new() }
70    }
71
72    /// Adds `piece` to the text field `field`.
73    fn grow(&mut self, field: &'static str, piece: &str) {
74        if let Some((_, text)) = self.growing.iter_mut().find(|(name, _)| *name == field) {
75            text.push_str(piece);
76            return;
77        }
78        let mut text = self.value.get(field).and_then(Value::as_str).unwrap_or_default().to_owned();
79        text.push_str(piece);
80        self.growing.push((field, text));
81    }
82}
83
84/// Reads the event stream in `lines` to its end and puts the reply together, handing each piece
85/// to `on_delta` as it comes. `cancelled` is asked between events; when it says yes the reply
86/// read so far is returned as it is.
87///
88/// # Errors
89///
90/// When the server sends an error event, when the stream breaks, or when it ends without saying
91/// the message is complete.
92pub fn read(
93    lines: impl IntoIterator<Item = io::Result<String>>,
94    on_delta: &mut dyn FnMut(Delta),
95    cancelled: &dyn Fn() -> bool,
96) -> Result<Reply, ProviderError> {
97    let mut blocks: Vec<Block> = Vec::new();
98    let mut reply = Reply { content: Vec::new(), stop: Stop::EndTurn, input_tokens: None, output_tokens: None };
99    let mut data = String::new();
100    let mut complete = false;
101    for line in lines {
102        let line = line.map_err(|error| ProviderError::Unreadable { what: error.to_string() })?;
103        if let Some(rest) = line.strip_prefix("data:") {
104            data.push_str(rest.strip_prefix(' ').unwrap_or(rest));
105            continue;
106        }
107        if !line.is_empty() || data.is_empty() {
108            continue;
109        }
110        let event = std::mem::take(&mut data);
111        if event == "[DONE]" {
112            complete = true;
113            break;
114        }
115        let event: Value = serde_json::from_str(&event)
116            .map_err(|error| ProviderError::Unreadable { what: format!("{error}: {event}") })?;
117        if apply(&event, &mut blocks, &mut reply, on_delta)? {
118            complete = true;
119            break;
120        }
121        if cancelled() {
122            complete = true;
123            break;
124        }
125    }
126    // Lines that stop coming because the reply was stopped end it as a stop does between events.
127    if !complete && cancelled() {
128        complete = true;
129    }
130    if !complete {
131        return Err(ProviderError::Unreadable { what: "the reply ended before the message was complete".to_owned() });
132    }
133    reply.content = blocks.into_iter().map(finish).collect();
134    Ok(reply)
135}
136
137/// Applies one event and says whether it was the last.
138fn apply(
139    event: &Value,
140    blocks: &mut Vec<Block>,
141    reply: &mut Reply,
142    on_delta: &mut dyn FnMut(Delta),
143) -> Result<bool, ProviderError> {
144    let kind = event.get("type").and_then(Value::as_str).unwrap_or_default();
145    match kind {
146        "message_start" => {
147            let usage = event.pointer("/message/usage");
148            reply.input_tokens = usage.and_then(|usage| usage.get("input_tokens")).and_then(Value::as_u64);
149        }
150        "content_block_start" => {
151            let block = event.get("content_block").cloned().unwrap_or_else(|| json!({"type": "text", "text": ""}));
152            if block.get("type").and_then(Value::as_str) == Some("tool_use") {
153                on_delta(Delta::ToolCall(block.get("name").and_then(Value::as_str).unwrap_or_default().to_owned()));
154            }
155            blocks.push(Block::new(block, event.get("index").and_then(Value::as_u64)));
156        }
157        "content_block_delta" => {
158            // A piece goes to the block its index names; servers that write two blocks at once
159            // (thinking and text, for some) interleave their pieces. The latest block of that
160            // index wins, and a piece that names none goes to the block being written.
161            let index = event.get("index").and_then(Value::as_u64);
162            let at = index
163                .and_then(|index| blocks.iter().rposition(|block| block.index == Some(index)))
164                .or_else(|| blocks.len().checked_sub(1));
165            let Some(block) = at.and_then(|at| blocks.get_mut(at)) else {
166                return Err(ProviderError::Unreadable { what: "a delta arrived before any block".to_owned() });
167            };
168            let delta = event.get("delta").cloned().unwrap_or_default();
169            let text = |field: &str| delta.get(field).and_then(Value::as_str).unwrap_or_default().to_owned();
170            match delta.get("type").and_then(Value::as_str).unwrap_or_default() {
171                "text_delta" => {
172                    let piece = text("text");
173                    block.grow("text", &piece);
174                    on_delta(Delta::Text(piece));
175                }
176                "thinking_delta" => {
177                    let piece = text("thinking");
178                    block.grow("thinking", &piece);
179                    on_delta(Delta::Thinking(piece));
180                }
181                "signature_delta" => block.grow("signature", &text("signature")),
182                "input_json_delta" => block.input.push_str(&text("partial_json")),
183                _ => {}
184            }
185        }
186        "message_delta" => {
187            if let Some(stop) = event.pointer("/delta/stop_reason").and_then(Value::as_str) {
188                reply.stop = match stop {
189                    "end_turn" | "stop_sequence" => Stop::EndTurn,
190                    "tool_use" => Stop::ToolUse,
191                    "max_tokens" => Stop::MaxTokens,
192                    other => Stop::Other(other.to_owned()),
193                };
194            }
195            if let Some(tokens) = event.pointer("/usage/output_tokens").and_then(Value::as_u64) {
196                reply.output_tokens = Some(tokens);
197            }
198            // The counts here are the whole reply's. Some servers (OpenRouter among them) say 0
199            // input tokens when the reply starts and the real count only here; a 0 never
200            // replaces a count the start already gave.
201            if let Some(tokens) = event.pointer("/usage/input_tokens").and_then(Value::as_u64)
202                && (tokens > 0 || reply.input_tokens.is_none())
203            {
204                reply.input_tokens = Some(tokens);
205            }
206        }
207        "message_stop" => return Ok(true),
208        "error" => {
209            let said = event.pointer("/error/message").and_then(Value::as_str).unwrap_or("error").to_owned();
210            return Err(ProviderError::Stopped { said });
211        }
212        _ => {}
213    }
214    Ok(false)
215}
216
217/// A block as it goes back into the conversation: a tool call gets its input read from the
218/// fragments. The API takes only an object as a call's input and refuses every request after one
219/// that is not, so an input that is not a JSON object (cut in half, or never JSON) is kept whole
220/// under [`UNPARSED`] in an object of its own: the model still sees what it sent.
221fn finish(block: Block) -> Value {
222    let Block { mut value, input, growing, .. } = block;
223    if let Some(object) = value.as_object_mut() {
224        for (field, text) in growing {
225            object.insert(field.to_owned(), Value::String(text));
226        }
227    }
228    if value.get("type").and_then(Value::as_str) == Some("tool_use")
229        && let Some(object) = value.as_object_mut()
230    {
231        let parsed = if input.trim().is_empty() {
232            Value::Object(Map::new())
233        } else {
234            match serde_json::from_str(&input) {
235                Ok(Value::Object(fields)) => Value::Object(fields),
236                _ => json!({ UNPARSED: input }),
237            }
238        };
239        object.insert("input".to_owned(), parsed);
240    }
241    value
242}
243
244/// The field a call's input is kept under when it is not a JSON object.
245pub const UNPARSED: &str = "unparsed_input";
246
247#[cfg(test)]
248mod tests {
249    use std::io::BufRead as _;
250
251    use super::*;
252
253    /// The events of a reply, written the way servers send them.
254    fn stream(events: &[Value]) -> String {
255        events
256            .iter()
257            .map(|event| format!("event: {}\ndata: {event}\n\n", event["type"].as_str().unwrap_or("")))
258            .collect()
259    }
260
261    fn collect(text: &str) -> (Result<Reply, ProviderError>, Vec<Delta>) {
262        let mut deltas = Vec::new();
263        let reply = read(text.as_bytes().lines(), &mut |delta| deltas.push(delta), &|| false);
264        (reply, deltas)
265    }
266
267    #[test]
268    fn text_and_thinking_arrive_in_pieces_and_end_up_whole() {
269        let text = stream(&[
270            json!({"type": "message_start", "message": {"usage": {"input_tokens": 14}}}),
271            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": "", "signature": ""}}),
272            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "Short "}}),
273            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "answer."}}),
274            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": "sig"}}),
275            json!({"type": "content_block_stop", "index": 0}),
276            json!({"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}),
277            json!({"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "Hey"}}),
278            json!({"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": " there"}}),
279            json!({"type": "content_block_stop", "index": 1}),
280            json!({"type": "ping"}),
281            json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 9}}),
282            json!({"type": "message_stop"}),
283        ]);
284        let (reply, deltas) = collect(&text);
285        let reply = reply.expect("a reply");
286        assert_eq!(
287            deltas,
288            [
289                Delta::Thinking("Short ".into()),
290                Delta::Thinking("answer.".into()),
291                Delta::Text("Hey".into()),
292                Delta::Text(" there".into())
293            ]
294        );
295        assert_eq!(
296            reply.content,
297            [
298                json!({"type": "thinking", "thinking": "Short answer.", "signature": "sig"}),
299                json!({"type": "text", "text": "Hey there"})
300            ]
301        );
302        assert_eq!(reply.stop, Stop::EndTurn);
303        assert_eq!((reply.input_tokens, reply.output_tokens), (Some(14), Some(9)));
304    }
305
306    #[test]
307    fn the_input_count_said_only_at_the_end_is_kept() {
308        let text = stream(&[
309            json!({"type": "message_start", "message": {"usage": {"input_tokens": 0}}}),
310            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
311            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi"}}),
312            json!({"type": "content_block_stop", "index": 0}),
313            json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"input_tokens": 2817, "output_tokens": 41}}),
314            json!({"type": "message_stop"}),
315        ]);
316        let (reply, _) = collect(&text);
317        let reply = reply.expect("a reply");
318        assert_eq!((reply.input_tokens, reply.output_tokens), (Some(2817), Some(41)));
319    }
320
321    #[test]
322    fn a_zero_at_the_end_does_not_undo_the_starts_count() {
323        let text = stream(&[
324            json!({"type": "message_start", "message": {"usage": {"input_tokens": 1534}}}),
325            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
326            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi"}}),
327            json!({"type": "content_block_stop", "index": 0}),
328            json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"input_tokens": 0, "output_tokens": 7}}),
329            json!({"type": "message_stop"}),
330        ]);
331        let (reply, _) = collect(&text);
332        let reply = reply.expect("a reply");
333        assert_eq!(reply.input_tokens, Some(1534));
334    }
335
336    #[test]
337    fn a_tool_call_is_announced_by_name_and_its_input_read_from_fragments() {
338        let text = stream(&[
339            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "t1", "name": "read_file", "input": {}}}),
340            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"path\": \"src/"}}),
341            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "main.rs\"}"}}),
342            json!({"type": "content_block_stop", "index": 0}),
343            json!({"type": "message_delta", "delta": {"stop_reason": "tool_use"}}),
344            json!({"type": "message_stop"}),
345        ]);
346        let (reply, deltas) = collect(&text);
347        let reply = reply.expect("a reply");
348        assert_eq!(deltas, [Delta::ToolCall("read_file".into())]);
349        assert_eq!(
350            reply.content,
351            [json!({"type": "tool_use", "id": "t1", "name": "read_file", "input": {"path": "src/main.rs"}})]
352        );
353        assert_eq!(reply.stop, Stop::ToolUse);
354    }
355
356    #[test]
357    fn an_error_event_and_a_cut_stream_are_errors() {
358        let error = stream(&[json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}})]);
359        assert_eq!(collect(&error).0, Err(ProviderError::Stopped { said: "Overloaded".into() }));
360        let cut = stream(&[
361            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
362        ]);
363        assert!(matches!(collect(&cut).0, Err(ProviderError::Unreadable { .. })), "no message_stop, no reply");
364    }
365
366    #[test]
367    fn a_cancelled_reply_keeps_what_had_arrived() {
368        let text = stream(&[
369            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
370            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Half"}}),
371            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " more"}}),
372        ]);
373        let mut seen = 0;
374        let reply = read(text.as_bytes().lines(), &mut |_| seen += 1, &|| true).expect("what arrived");
375        assert_eq!(reply.content, [json!({"type": "text", "text": ""})], "stopped after the first event");
376        assert_eq!(seen, 0);
377    }
378
379    /// Reads `text` as a stream that stopped early: the error it has to give, and every piece it
380    /// showed while reading.
381    fn cut(text: &str) -> (ProviderError, Vec<Delta>) {
382        let mut deltas = Vec::new();
383        let error = read(text.as_bytes().lines(), &mut |delta| deltas.push(delta), &|| false)
384            .expect_err("an unfinished reply is not a reply");
385        let ProviderError::Unreadable { what } = &error else { panic!("a cut stream, got {error:?}") };
386        assert!(what.contains("ended before the message was complete"), "{what}");
387        (error, deltas)
388    }
389
390    #[test]
391    fn a_stream_that_stops_in_the_middle_of_a_json_line_is_an_error_and_not_a_reply() {
392        // The connection goes in the middle of a line, so the last line has no end of line.
393        cut("event: content_block_start\ndata: {\"type\": \"content_bl");
394        // An event whose blank line never came is not an event yet either.
395        cut("event: content_block_start\ndata: {\"type\": \"message_stop\"}");
396        // What came whole before the cut is not handed on as a reply, and no piece is shown.
397        let whole = stream(&[json!({"type": "message_start", "message": {"usage": {"input_tokens": 3}}})]);
398        let (_, deltas) = cut(&format!("{whole}data: {{\"type\": \"mess"));
399        assert!(deltas.is_empty(), "no piece of an unfinished reply is shown");
400    }
401
402    #[test]
403    fn pieces_of_two_blocks_written_at_once_go_each_to_its_own_block() {
404        let text = stream(&[
405            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}),
406            json!({"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}),
407            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "weigh "}}),
408            json!({"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "Hello"}}),
409            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "it"}}),
410            json!({"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": " there"}}),
411            json!({"type": "message_stop"}),
412        ]);
413        let (reply, _) = collect(&text);
414        assert_eq!(
415            reply.expect("a reply").content,
416            [json!({"type": "thinking", "thinking": "weigh it"}), json!({"type": "text", "text": "Hello there"})],
417            "each piece is in the block its index names"
418        );
419    }
420
421    #[test]
422    fn a_piece_that_names_no_block_goes_to_the_block_being_written() {
423        let text = stream(&[
424            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
425            json!({"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": "second"}}),
426            json!({"type": "content_block_delta", "delta": {"type": "text_delta", "text": " on"}}),
427            json!({"type": "content_block_delta", "index": 7, "delta": {"type": "text_delta", "text": " and on"}}),
428            json!({"type": "message_stop"}),
429        ]);
430        let (reply, _) = collect(&text);
431        assert_eq!(
432            reply.expect("a reply").content,
433            [json!({"type": "text", "text": ""}), json!({"type": "text", "text": "second on and on"})]
434        );
435    }
436
437    #[test]
438    fn one_index_with_two_starts_is_two_blocks_and_neither_is_lost() {
439        let text = stream(&[
440            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "one", "name": "list_dir", "input": {}}}),
441            json!({"type": "content_block_stop", "index": 0}),
442            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "two", "name": "read_file", "input": {}}}),
443            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"path\": \"a\"}"}}),
444            json!({"type": "message_stop"}),
445        ]);
446        let (reply, deltas) = collect(&text);
447        let reply = reply.expect("a reply");
448        assert_eq!(deltas, [Delta::ToolCall("list_dir".into()), Delta::ToolCall("read_file".into())]);
449        assert_eq!(
450            reply.content,
451            [
452                json!({"type": "tool_use", "id": "one", "name": "list_dir", "input": {}}),
453                json!({"type": "tool_use", "id": "two", "name": "read_file", "input": {"path": "a"}})
454            ]
455        );
456    }
457
458    #[test]
459    fn a_thinking_block_with_no_signature_comes_through_as_it_arrived() {
460        let text = stream(&[
461            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}),
462            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "weighing it"}}),
463            json!({"type": "message_stop"}),
464        ]);
465        let (reply, _) = collect(&text);
466        assert_eq!(reply.expect("a reply").content, [json!({"type": "thinking", "thinking": "weighing it"})]);
467    }
468
469    #[test]
470    fn a_stop_reason_that_is_none_of_the_three_is_kept_as_it_was_said() {
471        for (said, wanted) in [
472            ("refusal", Stop::Other("refusal".to_owned())),
473            ("pause_turn", Stop::Other("pause_turn".to_owned())),
474            ("", Stop::Other(String::new())),
475        ] {
476            let text = stream(&[
477                json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
478                json!({"type": "message_delta", "delta": {"stop_reason": said}}),
479                json!({"type": "message_stop"}),
480            ]);
481            assert_eq!(collect(&text).0.expect("a reply").stop, wanted, "{said:?}");
482        }
483    }
484
485    #[test]
486    fn a_tool_input_of_a_megabyte_comes_through_as_it_was_sent() {
487        let path = "p".repeat(1024 * 1024);
488        let text = stream(&[
489            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "big", "name": "read_file", "input": {}}}),
490            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"path\": \""}}),
491            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": path}}),
492            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "\"}"}}),
493            json!({"type": "message_stop"}),
494        ]);
495        let (reply, _) = collect(&text);
496        let reply = reply.expect("a reply");
497        assert_eq!(
498            reply.content[0]["input"]["path"].as_str().map(str::len),
499            Some(1024 * 1024),
500            "the whole path is in the input"
501        );
502    }
503
504    #[test]
505    fn a_tool_call_with_no_name_is_announced_empty_and_the_input_is_read_as_it_came() {
506        let text = stream(&[
507            json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "x"}}),
508            json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"path\": \"a\"}"}}),
509            json!({"type": "message_stop"}),
510        ]);
511        let (reply, deltas) = collect(&text);
512        assert_eq!(deltas, [Delta::ToolCall("".into())], "the name is empty, not missing");
513        assert_eq!(reply.expect("a reply").content[0]["input"]["path"], "a");
514    }
515
516    #[test]
517    fn a_delta_that_arrives_before_any_block_is_an_error_that_says_so() {
518        let text =
519            stream(&[json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}})]);
520        assert_eq!(
521            collect(&text).0,
522            Err(ProviderError::Unreadable { what: "a delta arrived before any block".to_owned() })
523        );
524    }
525}
qcli/provider/reply.rs

qcli/provider/
reply.rs
