ag_ui/client/
chunks.rs

1//! Normalizing `*_CHUNK` events into explicit start/content/end triples.
2//!
3//! A producer that cannot bracket its output — most provider adapters, because
4//! the upstream API does not tell them a message has ended until the next one
5//! begins — sends `TEXT_MESSAGE_CHUNK`, `TOOL_CALL_CHUNK` and
6//! `REASONING_MESSAGE_CHUNK` instead. Those events fold start, content and end
7//! into one, and they carry their id and name **only on the first chunk**:
8//!
9//! ```text
10//! TEXT_MESSAGE_CHUNK { messageId: "msg-1", delta: "Hel" }
11//! TEXT_MESSAGE_CHUNK { delta: "lo" }
12//! TEXT_MESSAGE_CHUNK { messageId: "msg-2", delta: "Bye" }   <- msg-1 just ended
13//! ```
14//!
15//! So the id has to be remembered, and the end of one stream is only knowable
16//! from the start of the next — or from the end of the run. That bookkeeping is
17//! this module.
18//!
19//! ```
20//! use ag_ui::client::chunks::normalize_all;
21//! use ag_ui::{Event, EventType, MessageId};
22//!
23//! let events = normalize_all([
24//!     Event::text_message_chunk(Some(MessageId::new("msg-1")), Some("Hel".into())),
25//!     Event::text_message_chunk(None, Some("lo".into())),
26//! ])?;
27//!
28//! let types: Vec<EventType> = events.iter().map(Event::event_type).collect();
29//! assert_eq!(types, [
30//!     EventType::TextMessageStart,
31//!     EventType::TextMessageContent,
32//!     EventType::TextMessageContent,
33//!     EventType::TextMessageEnd,
34//! ]);
35//! # Ok::<(), ag_ui::client::Error>(())
36//! ```
37//!
38//! # Subagents
39//!
40//! "The previous chunk" is only meaningful *per subagent* once several stream
41//! at once, so the shorthand resolves within the sending subagent's own
42//! stream: one stream may be open per owner, and a chunk that names no id
43//! continues the stream of the subagent it is attributed to. A chunk that
44//! carries neither an id nor a `subagentRunId` continues the parent's open
45//! stream when there is one, and otherwise the sole open stream; when several
46//! subagents' streams could all claim it there is nothing to resolve it
47//! against, and it is rejected rather than guessed at. When streaming
48//! concurrently, attribute every chunk — or repeat the id.
49
50use crate::{
51    Event, MessageId, ReasoningMessageChunkEvent, SubagentRunId, TextMessageChunkEvent,
52    TextMessageStartEvent, ToolCallChunkEvent, ToolCallId, ToolCallStartEvent,
53};
54
55use crate::client::error::{Error, Result};
56
57/// Which family of stream is open.
58#[derive(Clone, Copy, Debug, PartialEq, Eq)]
59enum Kind {
60    Text,
61    Tool,
62    Reasoning,
63}
64
65impl Kind {
66    const fn chunk_name(self) -> &'static str {
67        match self {
68            Self::Text => "TEXT_MESSAGE_CHUNK",
69            Self::Tool => "TOOL_CALL_CHUNK",
70            Self::Reasoning => "REASONING_MESSAGE_CHUNK",
71        }
72    }
73
74    const fn id_name(self) -> &'static str {
75        match self {
76            Self::Text | Self::Reasoning => "messageId",
77            Self::Tool => "toolCallId",
78        }
79    }
80
81    const fn noun(self) -> &'static str {
82        match self {
83            Self::Text => "message",
84            Self::Tool => "call",
85            Self::Reasoning => "reasoning message",
86        }
87    }
88}
89
90/// A stream in flight. At most one per owner: a new stream from the same
91/// owner is what ends the previous one.
92#[derive(Clone, Debug)]
93struct Open {
94    kind: Kind,
95    /// The message id or tool call id, as a plain string.
96    id: String,
97    /// Whether this normalizer emitted the opening event and therefore owes the
98    /// closing one. Streams the producer opened explicitly close themselves.
99    owed: bool,
100    /// The subagent the stream belongs to; `None` for the parent agent.
101    owner: Option<SubagentRunId>,
102}
103
104/// Expands chunk events into the explicit events the rest of the protocol is
105/// written in.
106///
107/// Feed every event of a run through it, in order, and apply what comes out.
108/// Events that are not chunks pass through untouched — but they are still
109/// *observed*, so that an explicitly opened message can absorb a following
110/// bare chunk, and so an explicit end is not duplicated.
111#[derive(Clone, Debug, Default)]
112pub struct ChunkNormalizer {
113    open: Vec<Open>,
114}
115
116impl ChunkNormalizer {
117    /// A normalizer with nothing open.
118    pub fn new() -> Self {
119        Self::default()
120    }
121
122    /// Whether any stream is currently open.
123    pub fn is_open(&self) -> bool {
124        !self.open.is_empty()
125    }
126
127    /// Expands one event, appending the result to `out`.
128    ///
129    /// Appends between zero and three events: the synthesized end of the
130    /// previous stream, the synthesized start of this one, and the content.
131    ///
132    /// # Errors
133    ///
134    /// A chunk that names no stream and has none to continue — none open, or
135    /// several subagents' open and no attribution to pick one — or a tool call
136    /// chunk that opens a call without naming the tool, is a protocol
137    /// violation: there is no id to attach the payload to, and guessing one
138    /// would put text in the wrong message.
139    pub fn normalize(&mut self, event: Event, out: &mut Vec<Event>) -> Result<()> {
140        match event {
141            Event::TextMessageChunk(chunk) => self.text_chunk(chunk, out),
142            Event::ToolCallChunk(chunk) => self.tool_chunk(chunk, out),
143            Event::ReasoningMessageChunk(chunk) => self.reasoning_chunk(chunk, out),
144            other => {
145                self.observe(&other, out);
146                out.push(other);
147                Ok(())
148            }
149        }
150    }
151
152    /// Closes whatever the normalizer opened but never got to end.
153    ///
154    /// Call this when the transport stream ends. A chunk stream that never
155    /// terminates is the normal case, not an error: the last message of a run
156    /// has nothing after it to imply its end.
157    pub fn finish(&mut self, out: &mut Vec<Event>) {
158        self.close_all(out);
159    }
160
161    // ---- chunk families -------------------------------------------------
162
163    /// The stream a chunk belongs to: the one it names, or the open one of its
164    /// kind that its attribution resolves to.
165    fn stream_id(
166        &self,
167        kind: Kind,
168        named: Option<&str>,
169        tag: &Option<SubagentRunId>,
170    ) -> Result<String> {
171        if let Some(id) = named {
172            return Ok(id.to_owned());
173        }
174        let missing = |why: String| {
175            Error::protocol(format!(
176                "{} carries no {} and {why}",
177                kind.chunk_name(),
178                kind.id_name()
179            ))
180        };
181        if tag.is_some() {
182            return self.current(kind, tag).ok_or_else(|| {
183                missing(format!(
184                    "subagent {:?} has no {} open",
185                    tag.as_ref().map(|id| id.as_str()).unwrap_or_default(),
186                    kind.noun()
187                ))
188            });
189        }
190        if let Some(id) = self.current(kind, &None) {
191            return Ok(id);
192        }
193        let mut of_kind = self.open.iter().filter(|open| open.kind == kind);
194        match (of_kind.next(), of_kind.next()) {
195            (Some(only), None) => Ok(only.id.clone()),
196            (None, _) => Err(missing(format!("no {} is open", kind.noun()))),
197            (Some(_), Some(_)) => Err(missing(format!(
198                "several subagents have a {} open; attribute the chunk",
199                kind.noun()
200            ))),
201        }
202    }
203
204    /// Closes the owner's open stream, emits `start`, and takes on the debt
205    /// for the matching end event.
206    fn begin(
207        &mut self,
208        kind: Kind,
209        id: &str,
210        owner: Option<SubagentRunId>,
211        start: Event,
212        out: &mut Vec<Event>,
213    ) {
214        self.close_owner(&owner, out);
215        out.push(start);
216        self.open.push(Open {
217            kind,
218            id: id.to_owned(),
219            owed: true,
220            owner,
221        });
222    }
223
224    fn text_chunk(&mut self, chunk: TextMessageChunkEvent, out: &mut Vec<Event>) -> Result<()> {
225        let id = MessageId::new(self.stream_id(
226            Kind::Text,
227            chunk.message_id.as_ref().map(MessageId::as_str),
228            &chunk.subagent_run_id,
229        )?);
230
231        let owner = self.owner_of(Kind::Text, id.as_str(), &chunk.subagent_run_id);
232        if self.current(Kind::Text, &owner).as_deref() != Some(id.as_str()) {
233            let mut start = TextMessageStartEvent::new(id.clone(), chunk.role.unwrap_or_default());
234            start.name = chunk.name;
235            start.base = chunk.base.clone();
236            start.subagent_run_id = owner.clone();
237            self.begin(Kind::Text, id.as_str(), owner.clone(), start.into(), out);
238        }
239
240        if let Some(delta) = chunk.delta {
241            let mut content = crate::TextMessageContentEvent::new(id, delta);
242            content.base = chunk.base;
243            content.subagent_run_id = owner;
244            out.push(content.into());
245        }
246        Ok(())
247    }
248
249    fn tool_chunk(&mut self, chunk: ToolCallChunkEvent, out: &mut Vec<Event>) -> Result<()> {
250        let id = ToolCallId::new(self.stream_id(
251            Kind::Tool,
252            chunk.tool_call_id.as_ref().map(ToolCallId::as_str),
253            &chunk.subagent_run_id,
254        )?);
255
256        let owner = self.owner_of(Kind::Tool, id.as_str(), &chunk.subagent_run_id);
257        if self.current(Kind::Tool, &owner).as_deref() != Some(id.as_str()) {
258            // Checked before anything is emitted: a call with no name cannot be
259            // opened, and the previous stream should not be closed on the way
260            // to finding that out.
261            let Some(name) = chunk.tool_call_name else {
262                return Err(Error::protocol(format!(
263                    "TOOL_CALL_CHUNK opens tool call {id:?} without a toolCallName"
264                )));
265            };
266            let mut start = ToolCallStartEvent::new(id.clone(), name);
267            start.parent_message_id = chunk.parent_message_id;
268            start.base = chunk.base.clone();
269            start.subagent_run_id = owner.clone();
270            self.begin(Kind::Tool, id.as_str(), owner.clone(), start.into(), out);
271        }
272
273        if let Some(delta) = chunk.delta {
274            let mut args = crate::ToolCallArgsEvent::new(id, delta);
275            args.base = chunk.base;
276            args.subagent_run_id = owner;
277            out.push(args.into());
278        }
279        Ok(())
280    }
281
282    fn reasoning_chunk(
283        &mut self,
284        chunk: ReasoningMessageChunkEvent,
285        out: &mut Vec<Event>,
286    ) -> Result<()> {
287        let id = MessageId::new(self.stream_id(
288            Kind::Reasoning,
289            chunk.message_id.as_ref().map(MessageId::as_str),
290            &chunk.subagent_run_id,
291        )?);
292
293        let owner = self.owner_of(Kind::Reasoning, id.as_str(), &chunk.subagent_run_id);
294        if self.current(Kind::Reasoning, &owner).as_deref() != Some(id.as_str()) {
295            let mut start = crate::ReasoningMessageStartEvent::new(id.clone());
296            start.base = chunk.base.clone();
297            start.subagent_run_id = owner.clone();
298            self.begin(
299                Kind::Reasoning,
300                id.as_str(),
301                owner.clone(),
302                start.into(),
303                out,
304            );
305        }
306
307        if let Some(delta) = chunk.delta {
308            let mut content = crate::ReasoningMessageContentEvent::new(id, delta);
309            content.base = chunk.base;
310            content.subagent_run_id = owner;
311            out.push(content.into());
312        }
313        Ok(())
314    }
315
316    // ---- explicit events ------------------------------------------------
317
318    /// Tracks what an explicit (non-chunk) event does to the open streams.
319    ///
320    /// Only events that belong to a stream, the one that answers a tool call,
321    /// and the two that end a run touch them: a `STATE_DELTA` between two chunks
322    /// of one message must not split that message in half.
323    fn observe(&mut self, event: &Event, out: &mut Vec<Event>) {
324        match event {
325            Event::TextMessageStart(e) => {
326                self.open_explicit(Kind::Text, e.message_id.as_str(), &e.subagent_run_id, out);
327            }
328            Event::TextMessageContent(e) => {
329                self.open_explicit(Kind::Text, e.message_id.as_str(), &e.subagent_run_id, out);
330            }
331            Event::TextMessageEnd(e) => self.close_explicit(Kind::Text, e.message_id.as_str()),
332
333            Event::ToolCallStart(e) => {
334                self.open_explicit(Kind::Tool, e.tool_call_id.as_str(), &e.subagent_run_id, out);
335            }
336            Event::ToolCallArgs(e) => {
337                self.open_explicit(Kind::Tool, e.tool_call_id.as_str(), &e.subagent_run_id, out);
338            }
339            Event::ToolCallEnd(e) => self.close_explicit(Kind::Tool, e.tool_call_id.as_str()),
340            // A result answers a call, so the call is over — and the protocol
341            // puts `TOOL_CALL_END` before it. A chunk-streamed call has no end
342            // of its own, so without this the result overtakes the terminator
343            // this normalizer still owes. It also ends whatever else its owner
344            // had streaming: a result cannot interleave with an open message
345            // either, and the party answering has clearly moved on.
346            Event::ToolCallResult(e) => {
347                self.close_id(Kind::Tool, e.tool_call_id.as_str(), out);
348                self.close_owner(&e.subagent_run_id, out);
349            }
350
351            Event::ReasoningMessageStart(e) => {
352                self.open_explicit(
353                    Kind::Reasoning,
354                    e.message_id.as_str(),
355                    &e.subagent_run_id,
356                    out,
357                );
358            }
359            Event::ReasoningMessageContent(e) => {
360                self.open_explicit(
361                    Kind::Reasoning,
362                    e.message_id.as_str(),
363                    &e.subagent_run_id,
364                    out,
365                );
366            }
367            Event::ReasoningMessageEnd(e) => {
368                self.close_explicit(Kind::Reasoning, e.message_id.as_str());
369            }
370            // A reasoning block closing implies its message has closed.
371            Event::ReasoningEnd(e) => self.close_id(Kind::Reasoning, e.message_id.as_str(), out),
372            Event::RunFinished(_) | Event::RunError(_) => self.close_all(out),
373            _ => {}
374        }
375    }
376
377    /// An explicit event for a stream: closes the owner's other open stream,
378    /// then takes ownership of this one. Nothing is owed — a stream the
379    /// producer opened, the producer ends.
380    ///
381    /// An explicit event for the stream *this* normalizer opened leaves the
382    /// debt in place: the end is still owed until the producer sends one.
383    fn open_explicit(
384        &mut self,
385        kind: Kind,
386        id: &str,
387        tag: &Option<SubagentRunId>,
388        out: &mut Vec<Event>,
389    ) {
390        let owner = self.owner_of(kind, id, tag);
391        if self.current(kind, &owner).as_deref() != Some(id) {
392            self.close_owner(&owner, out);
393            self.open.push(Open {
394                kind,
395                id: id.to_owned(),
396                owed: false,
397                owner,
398            });
399        }
400    }
401
402    /// The producer ended a stream itself: forget it without emitting.
403    fn close_explicit(&mut self, kind: Kind, id: &str) {
404        self.open
405            .retain(|open| !(open.kind == kind && open.id == id));
406    }
407
408    // ---- bookkeeping ----------------------------------------------------
409
410    /// The owner an event's stream belongs to: the event's tag when it has
411    /// one, otherwise the owner of the stream already open under that id,
412    /// otherwise the parent agent. An untagged continuation of a subagent's
413    /// chunk stream is legal on the wire and must not be read as the parent's.
414    fn owner_of(&self, kind: Kind, id: &str, tag: &Option<SubagentRunId>) -> Option<SubagentRunId> {
415        if tag.is_some() {
416            return tag.clone();
417        }
418        self.open
419            .iter()
420            .find(|open| open.kind == kind && open.id == id)
421            .and_then(|open| open.owner.clone())
422    }
423
424    /// The id of `owner`'s open stream, if it is of this kind.
425    fn current(&self, kind: Kind, owner: &Option<SubagentRunId>) -> Option<String> {
426        self.open
427            .iter()
428            .find(|open| &open.owner == owner)
429            .filter(|open| open.kind == kind)
430            .map(|open| open.id.clone())
431    }
432
433    /// Emits the end event for `owner`'s open stream, if this normalizer owes
434    /// one.
435    fn close_owner(&mut self, owner: &Option<SubagentRunId>, out: &mut Vec<Event>) {
436        if let Some(index) = self.open.iter().position(|open| &open.owner == owner) {
437            let open = self.open.remove(index);
438            Self::settle(open, out);
439        }
440    }
441
442    /// Closes the stream with this id, whoever owns it.
443    fn close_id(&mut self, kind: Kind, id: &str, out: &mut Vec<Event>) {
444        if let Some(index) = self
445            .open
446            .iter()
447            .position(|open| open.kind == kind && open.id == id)
448        {
449            let open = self.open.remove(index);
450            Self::settle(open, out);
451        }
452    }
453
454    fn close_all(&mut self, out: &mut Vec<Event>) {
455        for open in std::mem::take(&mut self.open) {
456            Self::settle(open, out);
457        }
458    }
459
460    /// Emits the end this normalizer owes for a stream, carrying the stream's
461    /// attribution.
462    fn settle(open: Open, out: &mut Vec<Event>) {
463        if !open.owed {
464            return;
465        }
466        let mut end = match open.kind {
467            Kind::Text => Event::text_message_end(MessageId::new(open.id)),
468            Kind::Tool => Event::tool_call_end(ToolCallId::new(open.id)),
469            Kind::Reasoning => Event::reasoning_message_end(MessageId::new(open.id)),
470        };
471        if let Some(owner) = open.owner {
472            end.set_subagent_run_id(owner);
473        }
474        out.push(end);
475    }
476}
477
478/// Normalizes a whole run in one call, closing anything left open at the end.
479///
480/// The streaming form is [`ChunkNormalizer`]; this is the convenience for
481/// tests, recorded streams, and anything else that has all the events already.
482pub fn normalize_all(events: impl IntoIterator<Item = Event>) -> Result<Vec<Event>> {
483    let mut normalizer = ChunkNormalizer::new();
484    let mut out = Vec::new();
485    for event in events {
486        normalizer.normalize(event, &mut out)?;
487    }
488    normalizer.finish(&mut out);
489    Ok(out)
490}
ag_ui/client/chunks.rs

ag_ui/client/
chunks.rs
