a2a_llm/
tool_call.rs

1//! Incremental tool-call assembly for streaming LLM responses.
2//!
3//! Providers stream a tool call as a sequence of
4//! [`ToolCallChunk`](super::LlmStreamEvent::ToolCallChunk)s (an id, an optional
5//! name, and a fragment of the JSON arguments) followed by a finalized
6//! [`ToolCall`]. [`ToolCallAccumulator`] folds those chunks — keyed by call id,
7//! so interleaved calls stay separate — into running [`PartialToolCall`]s that a
8//! UI can render live, and reconciles the authoritative final call.
9
10use super::ToolCall;
11
12/// A tool call assembled so far from streamed chunks.
13#[derive(Debug, Clone, Default, PartialEq, Eq)]
14pub struct PartialToolCall {
15    /// Provider-assigned call id (stable across this call's chunks).
16    pub id: String,
17    /// Function name, once a chunk has carried it.
18    pub name: Option<String>,
19    /// JSON arguments accumulated so far (may be partial/unparseable mid-stream).
20    pub arguments: String,
21    /// Set once a finalized [`ToolCall`] has reconciled this entry.
22    pub complete: bool,
23}
24
25impl PartialToolCall {
26    fn to_tool_call(&self) -> ToolCall {
27        ToolCall {
28            id: self.id.clone(),
29            name: self.name.clone().unwrap_or_default(),
30            arguments: self.arguments.clone(),
31        }
32    }
33}
34
35/// Folds streamed [`ToolCallChunk`](super::LlmStreamEvent::ToolCallChunk)s into
36/// complete [`ToolCall`]s, preserving first-seen order.
37#[derive(Debug, Default)]
38pub struct ToolCallAccumulator {
39    calls: Vec<PartialToolCall>,
40}
41
42impl ToolCallAccumulator {
43    /// Create an empty accumulator.
44    pub fn new() -> Self {
45        Self::default()
46    }
47
48    fn index_of(&mut self, id: &str) -> usize {
49        if let Some(i) = self.calls.iter().position(|c| c.id == id) {
50            return i;
51        }
52        self.calls.push(PartialToolCall {
53            id: id.to_string(),
54            ..Default::default()
55        });
56        self.calls.len() - 1
57    }
58
59    /// Apply one streamed chunk, returning the running partial for this id. A
60    /// non-empty `name` overrides; `args_delta` is appended.
61    pub fn push(&mut self, id: &str, name: Option<&str>, args_delta: &str) -> &PartialToolCall {
62        let idx = self.index_of(id);
63        let call = &mut self.calls[idx];
64        if let Some(n) = name
65            && !n.is_empty()
66        {
67            call.name = Some(n.to_string());
68        }
69        call.arguments.push_str(args_delta);
70        &self.calls[idx]
71    }
72
73    /// Reconcile a finalized [`ToolCall`]: its name and arguments are
74    /// authoritative and replace whatever was accumulated, marking the entry
75    /// complete.
76    pub fn finalize(&mut self, call: ToolCall) {
77        let idx = self.index_of(&call.id);
78        let entry = &mut self.calls[idx];
79        entry.name = Some(call.name);
80        entry.arguments = call.arguments;
81        entry.complete = true;
82    }
83
84    /// The running partial for `id`, if any.
85    pub fn partial(&self, id: &str) -> Option<&PartialToolCall> {
86        self.calls.iter().find(|c| c.id == id)
87    }
88
89    /// Calls reconciled by [`finalize`](Self::finalize), as concrete
90    /// [`ToolCall`]s, without clearing state.
91    pub fn completed(&self) -> Vec<ToolCall> {
92        self.calls
93            .iter()
94            .filter(|c| c.complete)
95            .map(PartialToolCall::to_tool_call)
96            .collect()
97    }
98
99    /// Drain every accumulated call as a [`ToolCall`], clearing the accumulator.
100    pub fn drain_completed(&mut self) -> Vec<ToolCall> {
101        let out = self
102            .calls
103            .iter()
104            .map(PartialToolCall::to_tool_call)
105            .collect();
106        self.calls.clear();
107        out
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn folds_interleaved_calls_by_id() {
117        let mut acc = ToolCallAccumulator::new();
118        acc.push("a", Some("add"), "{\"x\":");
119        acc.push("b", Some("mul"), "{\"y\":");
120        acc.push("a", None, "1}");
121        acc.push("b", None, "2}");
122
123        assert_eq!(acc.partial("a").unwrap().name.as_deref(), Some("add"));
124        assert_eq!(acc.partial("a").unwrap().arguments, "{\"x\":1}");
125        assert_eq!(acc.partial("b").unwrap().arguments, "{\"y\":2}");
126    }
127
128    #[test]
129    fn finalize_is_authoritative_and_marks_complete() {
130        let mut acc = ToolCallAccumulator::new();
131        acc.push("a", Some("add"), "{\"x\":1"); // truncated mid-stream
132        assert!(acc.completed().is_empty());
133
134        acc.finalize(ToolCall {
135            id: "a".to_string(),
136            name: "add".to_string(),
137            arguments: "{\"x\":1,\"y\":2}".to_string(),
138        });
139
140        let done = acc.completed();
141        assert_eq!(done.len(), 1);
142        assert_eq!(done[0].arguments, "{\"x\":1,\"y\":2}");
143    }
144
145    #[test]
146    fn drain_empties() {
147        let mut acc = ToolCallAccumulator::new();
148        acc.push("a", Some("add"), "{}");
149        assert_eq!(acc.drain_completed().len(), 1);
150        assert!(acc.drain_completed().is_empty());
151    }
152}
a2a_llm/tool_call.rs

a2a_llm/
tool_call.rs
