locode_provider/
assemble.rs

1//! Streaming tool-call assembly: accumulate partial-JSON args, parse at stop.
2//!
3//! Every studied harness consumes an SSE stream and **accumulates a tool call's
4//! arguments as a raw string** across `input_json_delta` fragments, because a single
5//! delta is not valid JSON — then parses once when the block closes. This helper
6//! encodes exactly that rule (raw string per content-block index, parse at
7//! [`ToolCallAssembler::finish`]). v0 is non-streaming, but the helper is built and
8//! tested standalone so the streaming wire drops straight in later.
9
10use std::collections::BTreeMap;
11
12use locode_protocol::ContentBlock;
13use serde_json::Value;
14use thiserror::Error;
15
16/// Accumulates streamed tool-call fragments by content-block index.
17#[derive(Debug, Default)]
18pub struct ToolCallAssembler {
19    // Keyed by content-block index so `finish` yields calls in stream order.
20    partials: BTreeMap<usize, Partial>,
21}
22
23#[derive(Debug)]
24struct Partial {
25    id: String,
26    name: String,
27    json: String,
28}
29
30/// A failure assembling streamed tool-call arguments.
31#[derive(Debug, Error)]
32pub enum AssembleError {
33    /// A fragment arrived for an index with no preceding [`ToolCallAssembler::begin`].
34    #[error("no tool call started at content-block index {0}")]
35    MissingStart(usize),
36    /// The accumulated argument buffer was not valid JSON.
37    #[error("invalid tool-call arguments JSON at index {index}: {source}")]
38    InvalidJson {
39        /// The content-block index whose buffer failed to parse.
40        index: usize,
41        /// The underlying JSON parse error.
42        source: serde_json::Error,
43    },
44}
45
46impl ToolCallAssembler {
47    /// A fresh, empty assembler.
48    #[must_use]
49    pub fn new() -> Self {
50        Self::default()
51    }
52
53    /// Begin a tool call at `index` (from the stream's `content_block_start`).
54    pub fn begin(&mut self, index: usize, id: impl Into<String>, name: impl Into<String>) {
55        self.partials.insert(
56            index,
57            Partial {
58                id: id.into(),
59                name: name.into(),
60                json: String::new(),
61            },
62        );
63    }
64
65    /// Append a raw partial-JSON fragment for `index`. **Never parsed here.**
66    ///
67    /// # Errors
68    /// [`AssembleError::MissingStart`] if no [`ToolCallAssembler::begin`] was seen for
69    /// `index`.
70    pub fn push_json(&mut self, index: usize, fragment: &str) -> Result<(), AssembleError> {
71        match self.partials.get_mut(&index) {
72            Some(partial) => {
73                partial.json.push_str(fragment);
74                Ok(())
75            }
76            None => Err(AssembleError::MissingStart(index)),
77        }
78    }
79
80    /// Finish: parse each accumulated buffer once, in index order, into
81    /// [`ContentBlock::ToolUse`] blocks.
82    ///
83    /// An empty/whitespace buffer becomes `{}` — Anthropic sends empty input for a
84    /// no-argument tool call.
85    ///
86    /// # Errors
87    /// [`AssembleError::InvalidJson`] if any accumulated buffer is not valid JSON.
88    pub fn finish(self) -> Result<Vec<ContentBlock>, AssembleError> {
89        let mut blocks = Vec::with_capacity(self.partials.len());
90        for (index, partial) in self.partials {
91            let input: Value = if partial.json.trim().is_empty() {
92                Value::Object(serde_json::Map::new())
93            } else {
94                serde_json::from_str(&partial.json)
95                    .map_err(|source| AssembleError::InvalidJson { index, source })?
96            };
97            blocks.push(ContentBlock::ToolUse {
98                id: partial.id,
99                name: partial.name,
100                input,
101            });
102        }
103        Ok(blocks)
104    }
105}
locode_provider/assemble.rs

locode_provider/
assemble.rs
