use crate::ids::ToolCallId;
use crate::message::{ContentBlock, FinishReason, Usage};
use crate::provider::ProviderEvent;
#[derive(Clone, Debug, PartialEq)]
pub enum AssemblyError {
DoubleStart,
EventBeforeStart,
EventAfterCompleted,
InvalidBlockReference { block: u32 },
BlockKindMismatch { block: u32 },
InvalidToolCallArguments { block: u32, raw: String },
TruncatedWithToolCalls,
EndedWithoutCompletion,
}
#[derive(Clone, Debug, PartialEq)]
pub enum Assembled {
Complete {
blocks: Vec<ContentBlock>,
usage: Option<Usage>,
finish_reason: FinishReason,
},
Truncated {
blocks: Vec<ContentBlock>,
usage: Option<Usage>,
},
Invalid {
error: AssemblyError,
},
}
#[derive(Debug, Clone, Default, PartialEq)]
enum OpenKind {
#[default]
None,
Text,
Reasoning,
ToolCall {
id: String,
name: String,
json: String,
},
}
#[derive(Debug)]
enum Slot {
Current,
New,
}
#[derive(Debug, Default)]
pub struct ResponseAssembler {
started: bool,
finish_reason: Option<FinishReason>,
blocks: Vec<ContentBlock>,
open_index: Option<u32>,
open_kind: OpenKind,
usage: Option<Usage>,
terminal_error: Option<AssemblyError>,
}
impl ResponseAssembler {
pub fn new() -> Self {
Self::default()
}
pub fn push(&mut self, event: ProviderEvent) -> Result<(), AssemblyError> {
if let Err(error) = self.push_inner(event) {
self.terminal_error = Some(error.clone());
Err(error)
} else {
Ok(())
}
}
fn push_inner(&mut self, event: ProviderEvent) -> Result<(), AssemblyError> {
if self.finish_reason.is_some() {
return Err(AssemblyError::EventAfterCompleted);
}
match event {
ProviderEvent::ResponseStarted => {
if self.started {
return Err(AssemblyError::DoubleStart);
}
self.started = true;
Ok(())
}
ProviderEvent::TextDelta { block, text } => {
self.require_started()?;
match self.slot(block)? {
Slot::Current => match &self.open_kind {
OpenKind::Text => {
if let Some(ContentBlock::Text { text: t }) = self.blocks.last_mut() {
t.push_str(&text);
}
Ok(())
}
_ => Err(AssemblyError::BlockKindMismatch { block }),
},
Slot::New => {
self.close_open_block()?;
self.open_index = Some(block);
self.open_kind = OpenKind::Text;
self.blocks.push(ContentBlock::Text { text });
Ok(())
}
}
}
ProviderEvent::ReasoningDelta { block, text } => {
self.require_started()?;
match self.slot(block)? {
Slot::Current => match &mut self.open_kind {
OpenKind::Reasoning => {
if let Some(ContentBlock::Reasoning { text: t, .. }) =
self.blocks.last_mut()
{
t.push_str(&text);
}
Ok(())
}
_ => Err(AssemblyError::BlockKindMismatch { block }),
},
Slot::New => {
self.close_open_block()?;
self.open_index = Some(block);
self.open_kind = OpenKind::Reasoning;
self.blocks.push(ContentBlock::Reasoning {
text,
provider_state: None,
});
Ok(())
}
}
}
ProviderEvent::ToolCallStarted { block, id, name } => {
self.require_started()?;
match self.slot(block)? {
Slot::Current => Err(AssemblyError::BlockKindMismatch { block }),
Slot::New => {
self.close_open_block()?;
self.open_index = Some(block);
self.open_kind = OpenKind::ToolCall {
id,
name,
json: String::new(),
};
Ok(())
}
}
}
ProviderEvent::ToolCallArgumentsDelta { block, json } => {
self.require_started()?;
match self.slot(block)? {
Slot::Current => match &mut self.open_kind {
OpenKind::ToolCall { json: buf, .. } => {
buf.push_str(&json);
Ok(())
}
_ => Err(AssemblyError::BlockKindMismatch { block }),
},
Slot::New => Err(AssemblyError::InvalidBlockReference { block }),
}
}
ProviderEvent::UsageUpdated(usage) => {
self.require_started()?;
self.usage = Some(usage);
Ok(())
}
ProviderEvent::ResponseCompleted { finish_reason } => {
self.require_started()?;
if let Err(error) = self.close_open_block() {
self.terminal_error = Some(error.clone());
return Err(error);
}
self.finish_reason = Some(finish_reason);
Ok(())
}
}
}
pub fn finalize(self) -> Assembled {
if let Some(error) = self.terminal_error {
return Assembled::Invalid { error };
}
let Some(finish_reason) = self.finish_reason else {
return Assembled::Invalid {
error: AssemblyError::EndedWithoutCompletion,
};
};
let has_tool_call = self
.blocks
.iter()
.any(|b| matches!(b, ContentBlock::ToolCall { .. }));
match finish_reason {
FinishReason::Length if has_tool_call => Assembled::Invalid {
error: AssemblyError::TruncatedWithToolCalls,
},
FinishReason::Length => Assembled::Truncated {
blocks: self.blocks,
usage: self.usage,
},
reason => Assembled::Complete {
blocks: self.blocks,
usage: self.usage,
finish_reason: reason,
},
}
}
fn require_started(&self) -> Result<(), AssemblyError> {
if self.started {
Ok(())
} else {
Err(AssemblyError::EventBeforeStart)
}
}
fn slot(&self, block: u32) -> Result<Slot, AssemblyError> {
match self.open_index {
Some(current) if current == block => Ok(Slot::Current),
Some(current) if current > block => Err(AssemblyError::InvalidBlockReference { block }),
_ => Ok(Slot::New),
}
}
fn close_open_block(&mut self) -> Result<(), AssemblyError> {
let kind = std::mem::replace(&mut self.open_kind, OpenKind::None);
let index = self.open_index;
self.open_index = None;
if let OpenKind::ToolCall { id, name, json } = kind {
let arguments = if json.trim().is_empty() {
serde_json::Value::Object(Default::default())
} else {
serde_json::from_str(&json).map_err(|_| {
AssemblyError::InvalidToolCallArguments {
block: index.expect("open_index set while ToolCall open"),
raw: json,
}
})?
};
self.blocks.push(ContentBlock::ToolCall {
id: ToolCallId::from(id),
name,
arguments,
});
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ids::ToolCallId;
use crate::message::{ContentBlock, FinishReason, Usage};
fn assembler_with(events: Vec<ProviderEvent>) -> ResponseAssembler {
let mut a = ResponseAssembler::new();
for e in events {
a.push(e).unwrap();
}
a
}
#[test]
fn assembles_text_reasoning_and_tool_call_in_order() {
let a = assembler_with(vec![
ProviderEvent::ResponseStarted,
ProviderEvent::TextDelta {
block: 0,
text: "ä½ å¥½".into(),
},
ProviderEvent::ReasoningDelta {
block: 1,
text: "think".into(),
},
ProviderEvent::ToolCallStarted {
block: 2,
id: "c1".into(),
name: "read".into(),
},
ProviderEvent::ToolCallArgumentsDelta {
block: 2,
json: "{\"path\":".into(),
},
ProviderEvent::ToolCallArgumentsDelta {
block: 2,
json: "\"/a\"}".into(),
},
ProviderEvent::UsageUpdated(Usage {
input_tokens: 10,
output_tokens: 5,
cached_input_tokens: None,
}),
ProviderEvent::ResponseCompleted {
finish_reason: FinishReason::Stop,
},
]);
match a.finalize() {
Assembled::Complete {
blocks,
usage,
finish_reason,
} => {
assert_eq!(finish_reason, FinishReason::Stop);
assert_eq!(usage.unwrap().input_tokens, 10);
assert_eq!(blocks.len(), 3);
assert_eq!(
blocks[0],
ContentBlock::Text {
text: "ä½ å¥½".into()
}
);
match &blocks[2] {
ContentBlock::ToolCall {
id,
name,
arguments,
} => {
assert_eq!(id, &ToolCallId::from("c1"));
assert_eq!(name, "read");
assert_eq!(arguments, &serde_json::json!({"path": "/a"}));
}
other => panic!("unexpected block: {other:?}"),
}
}
other => panic!("expected complete, got {other:?}"),
}
}
#[test]
fn delta_before_start_is_rejected() {
let mut a = ResponseAssembler::new();
let err = a
.push(ProviderEvent::TextDelta {
block: 0,
text: "x".into(),
})
.unwrap_err();
assert!(matches!(err, AssemblyError::EventBeforeStart));
}
#[test]
fn double_start_is_rejected() {
let mut a = ResponseAssembler::new();
a.push(ProviderEvent::ResponseStarted).unwrap();
assert!(matches!(
a.push(ProviderEvent::ResponseStarted).unwrap_err(),
AssemblyError::DoubleStart
));
}
#[test]
fn event_after_completed_is_rejected() {
let mut a = assembler_with(vec![
ProviderEvent::ResponseStarted,
ProviderEvent::ResponseCompleted {
finish_reason: FinishReason::Stop,
},
]);
assert!(matches!(
a.push(ProviderEvent::TextDelta {
block: 0,
text: "x".into()
})
.unwrap_err(),
AssemblyError::EventAfterCompleted
));
}
#[test]
fn delta_for_lower_block_index_is_rejected() {
let mut a = assembler_with(vec![
ProviderEvent::ResponseStarted,
ProviderEvent::TextDelta {
block: 1,
text: "a".into(),
},
]);
assert!(matches!(
a.push(ProviderEvent::TextDelta {
block: 0,
text: "b".into()
})
.unwrap_err(),
AssemblyError::InvalidBlockReference { block: 0 }
));
}
#[test]
fn same_block_with_different_kind_is_rejected() {
let mut a = assembler_with(vec![
ProviderEvent::ResponseStarted,
ProviderEvent::TextDelta {
block: 0,
text: "a".into(),
},
]);
assert!(matches!(
a.push(ProviderEvent::ReasoningDelta {
block: 0,
text: "b".into()
})
.unwrap_err(),
AssemblyError::BlockKindMismatch { block: 0 }
));
}
#[test]
fn invalid_tool_call_json_is_rejected() {
let mut a = ResponseAssembler::new();
for e in [
ProviderEvent::ResponseStarted,
ProviderEvent::ToolCallStarted {
block: 0,
id: "c1".into(),
name: "t".into(),
},
ProviderEvent::ToolCallArgumentsDelta {
block: 0,
json: "{\"broken".into(),
},
ProviderEvent::ResponseCompleted {
finish_reason: FinishReason::Stop,
},
] {
let _ = a.push(e);
}
match a.finalize() {
Assembled::Invalid { error } => {
assert!(matches!(
error,
AssemblyError::InvalidToolCallArguments { block: 0, .. }
))
}
other => panic!("expected invalid, got {other:?}"),
}
}
#[test]
fn tool_call_without_argument_deltas_uses_empty_object() {
let a = assembler_with(vec![
ProviderEvent::ResponseStarted,
ProviderEvent::ToolCallStarted {
block: 0,
id: "c1".into(),
name: "t".into(),
},
ProviderEvent::ResponseCompleted {
finish_reason: FinishReason::Stop,
},
]);
match a.finalize() {
Assembled::Complete { blocks, .. } => match &blocks[0] {
ContentBlock::ToolCall { arguments, .. } => {
assert_eq!(arguments, &serde_json::json!({}))
}
other => panic!("unexpected block: {other:?}"),
},
other => panic!("expected complete, got {other:?}"),
}
}
#[test]
fn length_with_closed_text_blocks_is_truncated() {
let a = assembler_with(vec![
ProviderEvent::ResponseStarted,
ProviderEvent::TextDelta {
block: 0,
text: "partial answer".into(),
},
ProviderEvent::ResponseCompleted {
finish_reason: FinishReason::Length,
},
]);
assert!(matches!(a.finalize(), Assembled::Truncated { .. }));
}
#[test]
fn length_with_tool_call_blocks_is_invalid() {
let a = assembler_with(vec![
ProviderEvent::ResponseStarted,
ProviderEvent::ToolCallStarted {
block: 0,
id: "c1".into(),
name: "t".into(),
},
ProviderEvent::ToolCallArgumentsDelta {
block: 0,
json: "{}".into(),
},
ProviderEvent::ResponseCompleted {
finish_reason: FinishReason::Length,
},
]);
match a.finalize() {
Assembled::Invalid { error } => {
assert!(matches!(error, AssemblyError::TruncatedWithToolCalls))
}
other => panic!("expected invalid, got {other:?}"),
}
}
#[test]
fn stream_end_without_completion_is_invalid() {
let a = assembler_with(vec![ProviderEvent::ResponseStarted]);
assert!(matches!(
a.finalize(),
Assembled::Invalid {
error: AssemblyError::EndedWithoutCompletion
}
));
}
#[test]
fn later_usage_wins() {
let a = assembler_with(vec![
ProviderEvent::ResponseStarted,
ProviderEvent::UsageUpdated(Usage {
input_tokens: 1,
output_tokens: 1,
cached_input_tokens: None,
}),
ProviderEvent::UsageUpdated(Usage {
input_tokens: 9,
output_tokens: 2,
cached_input_tokens: Some(3),
}),
ProviderEvent::ResponseCompleted {
finish_reason: FinishReason::Stop,
},
]);
match a.finalize() {
Assembled::Complete { usage: Some(u), .. } => {
assert_eq!(u.input_tokens, 9);
assert_eq!(u.cached_input_tokens, Some(3));
}
other => panic!("expected complete, got {other:?}"),
}
}
#[test]
fn error_poisons_finalize_even_if_stream_completes() {
let mut a = ResponseAssembler::new();
a.push(ProviderEvent::ResponseStarted).unwrap();
let _ = a.push(ProviderEvent::TextDelta {
block: 0,
text: "a".into(),
});
let _ = a.push(ProviderEvent::ReasoningDelta {
block: 0,
text: "bad kind".into(),
});
let _ = a.push(ProviderEvent::ToolCallStarted {
block: 1,
id: "c1".into(),
name: "t".into(),
});
let _ = a.push(ProviderEvent::ResponseCompleted {
finish_reason: FinishReason::Stop,
});
match a.finalize() {
Assembled::Invalid { error } => {
assert!(matches!(
error,
AssemblyError::BlockKindMismatch { block: 0 }
))
}
other => panic!("expected invalid, got {other:?}"),
}
}
#[test]
fn block_index_gaps_are_accepted() {
let a = assembler_with(vec![
ProviderEvent::ResponseStarted,
ProviderEvent::TextDelta {
block: 0,
text: "a".into(),
},
ProviderEvent::TextDelta {
block: 5,
text: "b".into(),
},
ProviderEvent::ResponseCompleted {
finish_reason: FinishReason::Stop,
},
]);
match a.finalize() {
Assembled::Complete { blocks, .. } => {
assert_eq!(blocks.len(), 2);
assert_eq!(blocks[0], ContentBlock::Text { text: "a".into() });
assert_eq!(blocks[1], ContentBlock::Text { text: "b".into() });
}
other => panic!("expected complete, got {other:?}"),
}
}
#[test]
fn two_tool_calls_assemble_in_stream_order() {
let a = assembler_with(vec![
ProviderEvent::ResponseStarted,
ProviderEvent::ToolCallStarted {
block: 0,
id: "c1".into(),
name: "a".into(),
},
ProviderEvent::ToolCallArgumentsDelta {
block: 0,
json: "{\"x\":1}".into(),
},
ProviderEvent::ToolCallStarted {
block: 1,
id: "c2".into(),
name: "b".into(),
},
ProviderEvent::ToolCallArgumentsDelta {
block: 1,
json: "{}".into(),
},
ProviderEvent::ResponseCompleted {
finish_reason: FinishReason::Stop,
},
]);
match a.finalize() {
Assembled::Complete { blocks, .. } => {
assert_eq!(blocks.len(), 2);
assert!(
matches!(&blocks[0], ContentBlock::ToolCall { id, .. } if id == &ToolCallId::from("c1"))
);
assert!(
matches!(&blocks[1], ContentBlock::ToolCall { id, arguments, .. }
if id == &ToolCallId::from("c2") && arguments == &serde_json::json!({}))
);
}
other => panic!("expected complete, got {other:?}"),
}
}
}