myco_model/
accumulator.rs

1//! Assemble validated message parts without owning a stream, retry policy, or event sink.
2
3use super::{
4    Content, ContentDelta, ContentStart, GenerateError, GenerateOutput, MessagePart, TokenUsage,
5    ToolUse, ToolUseDelta, ToolUseStart, TurnEndReason,
6};
7
8#[derive(Default)]
9pub struct MessageAccumulator {
10    started: bool,
11    content: Vec<Option<Content>>,
12    tool_uses: Vec<Option<IncompleteToolUse>>,
13    turn_end_reason: Option<TurnEndReason>,
14    usage: Option<TokenUsage>,
15}
16
17impl MessageAccumulator {
18    pub fn push(&mut self, part: &MessagePart) -> Result<(), GenerateError> {
19        if !self.started {
20            return self.start(part);
21        }
22        match part {
23            MessagePart::MessageStart => return Err(malformed("unexpected MessageStart")),
24            MessagePart::ContentStart(start) => {
25                let (index, block) = start_block(start.clone());
26                ensure_slot(&mut self.content, index, block);
27            }
28            MessagePart::ContentDelta(delta) => apply_content_delta(&mut self.content, delta)?,
29            MessagePart::ToolUseStart(ToolUseStart {
30                index,
31                name,
32                thought_signature,
33            }) => {
34                ensure_slot(
35                    &mut self.tool_uses,
36                    *index,
37                    IncompleteToolUse {
38                        name: name.clone(),
39                        input_json: String::new(),
40                        thought_signature: thought_signature.clone(),
41                    },
42                );
43            }
44            MessagePart::ToolUseDelta(delta) => self.append_tool_input(delta)?,
45            MessagePart::TurnEndReason(reason) => self.turn_end_reason = Some(reason.clone()),
46            MessagePart::Usage(usage) => {
47                self.usage = Some(self.usage.map_or(*usage, |prev| prev.merge(*usage)))
48            }
49        }
50        Ok(())
51    }
52
53    fn start(&mut self, part: &MessagePart) -> Result<(), GenerateError> {
54        if !matches!(part, MessagePart::MessageStart) {
55            return Err(malformed(
56                "first item is not MessageStart. Did you accidentally drain the stream already?",
57            ));
58        }
59        self.started = true;
60        Ok(())
61    }
62
63    fn append_tool_input(&mut self, delta: &ToolUseDelta) -> Result<(), GenerateError> {
64        let index = delta.index;
65        let tool = self
66            .tool_uses
67            .get_mut(index)
68            .and_then(Option::as_mut)
69            .ok_or_else(|| malformed(format!("tool use delta index {index} is out of bounds")))?;
70        tool.input_json.push_str(&delta.input_json_delta);
71        Ok(())
72    }
73
74    pub fn finish(self) -> Result<GenerateOutput, GenerateError> {
75        if !self.started {
76            return Err(malformed(
77                "empty stream. Did you accidentally drain the stream already?",
78            ));
79        }
80        let content = filled_slots(self.content, "content block")?;
81        let tool_uses = filled_slots(self.tool_uses, "tool use")?
82            .into_iter()
83            .map(IncompleteToolUse::finish)
84            .collect::<Result<Vec<_>, _>>()?;
85        let turn_end_reason = self
86            .turn_end_reason
87            .ok_or_else(|| malformed("no turn end reason provided"))?;
88        if turn_end_reason == TurnEndReason::ToolUse && tool_uses.is_empty() {
89            return Err(malformed(
90                "turn ended in tool_use but streamed zero tool uses",
91            ));
92        }
93        Ok(GenerateOutput {
94            content,
95            tool_uses,
96            turn_end_reason,
97            usage: self.usage,
98        })
99    }
100}
101
102struct IncompleteToolUse {
103    name: String,
104    input_json: String,
105    thought_signature: Option<String>,
106}
107
108impl IncompleteToolUse {
109    fn finish(self) -> Result<ToolUse, GenerateError> {
110        let json = if self.input_json.is_empty() {
111            "{}"
112        } else {
113            &self.input_json
114        };
115        let input = serde_json::from_str(json)
116            .map_err(|error| malformed(format!("tool use input JSON is invalid: {error}")))?;
117        Ok(ToolUse {
118            name: self.name,
119            input,
120            thought_signature: self.thought_signature,
121        })
122    }
123}
124
125fn malformed(message: impl std::fmt::Display) -> GenerateError {
126    GenerateError::MalformedResponseError(format!("Malformed stream: {message}"))
127}
128
129fn filled_slots<T>(slots: Vec<Option<T>>, name: &str) -> Result<Vec<T>, GenerateError> {
130    slots
131        .into_iter()
132        .enumerate()
133        .map(|(index, slot)| {
134            slot.ok_or_else(|| malformed(format!("missing {name} at index {index}")))
135        })
136        .collect()
137}
138
139fn ensure_slot<T>(slots: &mut Vec<Option<T>>, index: usize, value: T) {
140    while slots.len() <= index {
141        slots.push(None);
142    }
143    slots[index] = Some(value);
144}
145
146/// The empty [`Content`] block a [`ContentStart`] opens, with its index.
147fn start_block(start: ContentStart) -> (usize, Content) {
148    match start {
149        ContentStart::Text { index } => (
150            index,
151            Content::Text {
152                text: String::new(),
153            },
154        ),
155        ContentStart::Image { index } => (
156            index,
157            Content::Image {
158                source: String::new(),
159            },
160        ),
161        ContentStart::Thinking {
162            index,
163            signature,
164            redacted,
165        } => (
166            index,
167            Content::Thinking {
168                text: String::new(),
169                signature,
170                redacted,
171            },
172        ),
173    }
174}
175
176/// Append a [`ContentDelta`] to its opened block; the slot must exist and be
177/// the matching kind (redacted thinking swallows its deltas).
178fn apply_content_delta(
179    content: &mut [Option<Content>],
180    delta: &ContentDelta,
181) -> Result<(), GenerateError> {
182    let index = match delta {
183        ContentDelta::Text { index, .. }
184        | ContentDelta::Image { index, .. }
185        | ContentDelta::Thinking { index, .. }
186        | ContentDelta::Signature { index, .. } => *index,
187    };
188    match (content.get_mut(index).and_then(Option::as_mut), delta) {
189        (Some(Content::Text { text }), ContentDelta::Text { delta, .. }) => text.push_str(delta),
190        (Some(Content::Image { source }), ContentDelta::Image { delta, .. }) => {
191            source.push_str(delta);
192        }
193        (Some(Content::Thinking { text, redacted, .. }), ContentDelta::Thinking { delta, .. }) => {
194            if !*redacted {
195                text.push_str(delta);
196            }
197        }
198        (
199            Some(Content::Thinking {
200                signature,
201                redacted,
202                ..
203            }),
204            ContentDelta::Signature { delta, .. },
205        ) => {
206            if !*redacted {
207                signature.get_or_insert_with(String::new).push_str(delta);
208            }
209        }
210        _ => {
211            return Err(GenerateError::MalformedResponseError(format!(
212                "Malformed stream: content delta at index {index}: out of bounds or wrong kind"
213            )));
214        }
215    }
216    Ok(())
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    fn assemble(parts: &[MessagePart]) -> Result<GenerateOutput, GenerateError> {
224        let mut accumulator = MessageAccumulator::default();
225        for part in parts {
226            accumulator.push(part)?;
227        }
228        accumulator.finish()
229    }
230
231    #[test]
232    fn tool_input_fragments_form_one_json_value() {
233        let output = assemble(&[
234            MessagePart::MessageStart,
235            MessagePart::ToolUseStart(ToolUseStart {
236                index: 0,
237                name: "bash".into(),
238                thought_signature: None,
239            }),
240            MessagePart::ToolUseDelta(ToolUseDelta {
241                index: 0,
242                input_json_delta: "{\"command\":".into(),
243            }),
244            MessagePart::ToolUseDelta(ToolUseDelta {
245                index: 0,
246                input_json_delta: "\"pwd\"}".into(),
247            }),
248            MessagePart::TurnEndReason(TurnEndReason::ToolUse),
249        ])
250        .unwrap();
251        assert_eq!(output.tool_uses.len(), 1);
252        assert_eq!(output.tool_uses[0].name, "bash");
253        assert_eq!(
254            output.tool_uses[0].input,
255            serde_json::json!({"command": "pwd"})
256        );
257    }
258
259    #[test]
260    fn incomplete_and_malformed_messages_are_rejected() {
261        let cases = [
262            vec![],
263            vec![MessagePart::TurnEndReason(TurnEndReason::EndTurn)],
264            vec![MessagePart::MessageStart],
265            vec![MessagePart::MessageStart, MessagePart::MessageStart],
266            vec![
267                MessagePart::MessageStart,
268                MessagePart::ContentDelta(ContentDelta::Text {
269                    index: 0,
270                    delta: "unopened".into(),
271                }),
272            ],
273            vec![
274                MessagePart::MessageStart,
275                MessagePart::ContentStart(ContentStart::Text { index: 1 }),
276                MessagePart::TurnEndReason(TurnEndReason::EndTurn),
277            ],
278            vec![
279                MessagePart::MessageStart,
280                MessagePart::ToolUseDelta(ToolUseDelta {
281                    index: 0,
282                    input_json_delta: "{}".into(),
283                }),
284            ],
285            vec![
286                MessagePart::MessageStart,
287                MessagePart::ToolUseStart(ToolUseStart {
288                    index: 0,
289                    name: "bash".into(),
290                    thought_signature: None,
291                }),
292                MessagePart::ToolUseDelta(ToolUseDelta {
293                    index: 0,
294                    input_json_delta: "{".into(),
295                }),
296                MessagePart::TurnEndReason(TurnEndReason::ToolUse),
297            ],
298        ];
299        for parts in cases {
300            assert!(
301                matches!(
302                    assemble(&parts),
303                    Err(GenerateError::MalformedResponseError(_))
304                ),
305                "{parts:?}"
306            );
307        }
308    }
309
310    #[test]
311    fn tool_use_requires_a_call_but_empty_end_turn_is_valid() {
312        let output = assemble(&[
313            MessagePart::MessageStart,
314            MessagePart::TurnEndReason(TurnEndReason::EndTurn),
315        ])
316        .unwrap();
317        assert!(output.content.is_empty());
318        assert!(output.tool_uses.is_empty());
319        for content in [false, true] {
320            let mut parts = vec![MessagePart::MessageStart];
321            if content {
322                parts.extend([
323                    MessagePart::ContentStart(ContentStart::Text { index: 0 }),
324                    MessagePart::ContentDelta(ContentDelta::Text {
325                        index: 0,
326                        delta: "unfinished draft".into(),
327                    }),
328                ]);
329            }
330            parts.push(MessagePart::TurnEndReason(TurnEndReason::ToolUse));
331            assert!(matches!(
332                assemble(&parts),
333                Err(GenerateError::MalformedResponseError(_))
334            ));
335        }
336    }
337}
myco_model/accumulator.rs

myco_model/
accumulator.rs
