use crate::message::{Message, MessagePart};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::fmt;
#[cfg(feature = "streaming")]
pub mod handler;
pub mod rate_limit;
#[cfg(feature = "streaming")]
pub use handler::{DetectedRateLimit, RateLimitConfig, RateLimitKind};
pub use rate_limit::{RateLimiter, TokenBucket};
#[derive(Debug)]
#[non_exhaustive]
pub enum StreamError {
InvalidToolInputJson(serde_json::Error, String),
}
impl fmt::Display for StreamError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
StreamError::InvalidToolInputJson(err, raw) => {
write!(f, "invalid tool input JSON: {err} (raw_len={})", raw.len())
}
}
}
}
impl std::error::Error for StreamError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
StreamError::InvalidToolInputJson(err, _) => Some(err),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum StreamEvent {
MessageStart(MessageStart),
PartStart(PartStart),
IndexedDelta(IndexedDelta),
PartStop {
#[serde(default)]
index: Option<usize>,
},
MessageDelta(MessageDelta),
MessageStop,
Ping,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageStart {
pub message: MessageMetadata,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageMetadata {
pub id: String,
pub role: String,
pub model: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PartStart {
pub index: usize,
pub part: Option<MessagePart>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexedDelta {
pub index: usize,
pub delta: DeltaPart,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
#[non_exhaustive]
pub enum DeltaPart {
#[serde(rename = "text_delta")]
Text {
text: String,
},
#[serde(rename = "tool_call_delta")]
ToolCall {
partial_json: Value,
},
#[serde(rename = "input_json_delta")]
InputJson {
partial_json: String,
},
#[serde(rename = "thinking_delta")]
Thinking {
text: String,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum StreamStopReason {
ToolCall,
MaxTokens,
StopSequence,
EndTurn,
}
impl StreamStopReason {
#[must_use]
pub fn from_api_str(s: &str) -> Option<Self> {
match s {
"tool_call" | "tool_use" => Some(Self::ToolCall),
"max_tokens" => Some(Self::MaxTokens),
"stop_sequence" => Some(Self::StopSequence),
"end_turn" => Some(Self::EndTurn),
_ => None,
}
}
#[must_use]
pub fn to_api_str(self) -> &'static str {
match self {
Self::ToolCall => "tool_call",
Self::MaxTokens => "max_tokens",
Self::StopSequence => "stop_sequence",
Self::EndTurn => "end_turn",
}
}
#[must_use]
pub fn should_continue_tool_loop(self) -> bool {
matches!(self, Self::ToolCall)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageDelta {
pub delta: MessageDeltaPayload,
pub usage: Option<Usage>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageDeltaPayload {
pub stop_reason: Option<String>,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq)]
pub struct Usage {
pub input_tokens: u32,
pub output_tokens: u32,
}
impl Usage {
#[must_use]
pub fn new(input_tokens: u32, output_tokens: u32) -> Self {
Self {
input_tokens,
output_tokens,
}
}
#[must_use]
pub fn total_tokens(self) -> u32 {
self.input_tokens.saturating_add(self.output_tokens)
}
}
#[derive(Debug, Default)]
pub struct StreamAccumulator {
completed: Vec<MessagePart>,
open: Vec<OpenPart>,
model: Option<String>,
usage: Option<Usage>,
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
enum OpenPartKind {
#[default]
Text,
Tool,
Thinking,
}
#[derive(Debug, Default)]
struct OpenPart {
index: usize,
kind: OpenPartKind,
text: String,
thinking: String,
tool_id: String,
tool_name: String,
tool_input: String,
}
impl StreamAccumulator {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn process(&mut self, event: &StreamEvent) -> Result<(), StreamError> {
match event {
StreamEvent::MessageStart(msg_start) => {
self.model = Some(msg_start.message.model.clone());
Ok(())
}
StreamEvent::PartStart(part_start) => {
let kind = match &part_start.part {
Some(MessagePart::ToolCall { .. }) => OpenPartKind::Tool,
Some(_) => OpenPartKind::Text,
None => OpenPartKind::Thinking,
};
let mut slot = OpenPart {
index: part_start.index,
kind,
..Default::default()
};
if let Some(MessagePart::ToolCall { id, name, .. }) = &part_start.part {
slot.tool_id.clone_from(id);
slot.tool_name.clone_from(name);
}
self.open.push(slot);
Ok(())
}
StreamEvent::IndexedDelta(delta) => {
let kind = match &delta.delta {
DeltaPart::InputJson { .. } | DeltaPart::ToolCall { .. } => OpenPartKind::Tool,
DeltaPart::Text { .. } => OpenPartKind::Text,
DeltaPart::Thinking { .. } => OpenPartKind::Thinking,
};
let Some(slot) = self
.open
.iter_mut()
.find(|s| s.kind == kind && s.index == delta.index)
else {
return Ok(());
};
match &delta.delta {
DeltaPart::Text { text } => {
slot.text.push_str(text);
}
DeltaPart::InputJson { partial_json } => {
slot.tool_input.push_str(partial_json);
}
DeltaPart::ToolCall { partial_json } => {
if let Some(s) = partial_json.as_str() {
slot.tool_input.push_str(s);
}
}
DeltaPart::Thinking { text } => {
slot.thinking.push_str(text);
}
}
Ok(())
}
StreamEvent::PartStop { index } => {
let pos = match index {
Some(i) => self.open.iter().position(|s| s.index == *i),
None => (!self.open.is_empty()).then_some(0),
};
let Some(pos) = pos else {
return Ok(());
};
let slot = self.open.remove(pos);
let flushed = match slot.kind {
OpenPartKind::Text if !slot.text.is_empty() => {
Some(MessagePart::text(slot.text))
}
OpenPartKind::Tool if !slot.tool_name.is_empty() => {
let input: Value = if slot.tool_input.is_empty() {
Value::Object(serde_json::Map::new())
} else {
serde_json::from_str(&slot.tool_input).map_err(|e| {
StreamError::InvalidToolInputJson(e, slot.tool_input.clone())
})?
};
Some(MessagePart::tool_call(slot.tool_id, slot.tool_name, input))
}
_ => None,
};
if let Some(part) = flushed {
self.completed.push(part);
}
Ok(())
}
StreamEvent::MessageDelta(delta) => {
self.usage = delta.usage;
Ok(())
}
StreamEvent::MessageStop | StreamEvent::Ping => Ok(()),
}
}
#[must_use]
pub fn peek_parts(&self) -> &[MessagePart] {
&self.completed
}
#[must_use]
pub fn build(self) -> Message {
Message {
role: crate::message::Role::Assistant,
parts: self.completed,
}
}
#[must_use]
pub fn usage(&self) -> Option<&Usage> {
self.usage.as_ref()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_stream_stop_reason_from_api_str() {
assert_eq!(
StreamStopReason::from_api_str("tool_call"),
Some(StreamStopReason::ToolCall)
);
assert_eq!(
StreamStopReason::from_api_str("tool_use"),
Some(StreamStopReason::ToolCall)
);
assert_eq!(
StreamStopReason::from_api_str("max_tokens"),
Some(StreamStopReason::MaxTokens)
);
assert_eq!(
StreamStopReason::from_api_str("end_turn"),
Some(StreamStopReason::EndTurn)
);
assert_eq!(
StreamStopReason::from_api_str("stop_sequence"),
Some(StreamStopReason::StopSequence)
);
assert_eq!(StreamStopReason::from_api_str("unknown"), None);
}
#[test]
fn test_stream_stop_reason_to_api_str() {
assert_eq!(StreamStopReason::ToolCall.to_api_str(), "tool_call");
assert_eq!(StreamStopReason::MaxTokens.to_api_str(), "max_tokens");
assert_eq!(StreamStopReason::EndTurn.to_api_str(), "end_turn");
assert_eq!(StreamStopReason::StopSequence.to_api_str(), "stop_sequence");
}
#[test]
fn test_stream_stop_reason_should_continue() {
assert!(StreamStopReason::ToolCall.should_continue_tool_loop());
assert!(!StreamStopReason::EndTurn.should_continue_tool_loop());
assert!(!StreamStopReason::MaxTokens.should_continue_tool_loop());
}
#[test]
fn test_usage() {
let usage = Usage::new(100, 50);
assert_eq!(usage.input_tokens, 100);
assert_eq!(usage.output_tokens, 50);
assert_eq!(usage.total_tokens(), 150);
}
#[test]
fn test_usage_default() {
let usage = Usage::default();
assert_eq!(usage.total_tokens(), 0);
}
#[test]
fn test_accumulator_text_message() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::MessageStart(MessageStart {
message: MessageMetadata {
id: "msg_1".to_string(),
role: "assistant".to_string(),
model: "test-model".to_string(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(crate::message::MessagePart::text("")),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::Text {
text: "Hello".to_string(),
},
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::Text {
text: " world".to_string(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
acc.process(&StreamEvent::MessageDelta(MessageDelta {
delta: MessageDeltaPayload {
stop_reason: Some("end_turn".to_string()),
},
usage: Some(Usage::new(10, 5)),
}))
.unwrap();
acc.process(&StreamEvent::MessageStop).unwrap();
let msg = acc.build();
assert_eq!(msg.role, crate::message::Role::Assistant);
assert_eq!(msg.parts.len(), 1);
assert_eq!(msg.parts[0].as_text(), Some("Hello world"));
}
#[test]
fn test_accumulator_tool_call() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::tool_call(
"tool_1",
"read_file",
Value::Object(serde_json::Map::new()),
)),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::InputJson {
partial_json: r#"{"path":"/tmp/test"}"#.to_string(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
let msg = acc.build();
assert_eq!(msg.parts.len(), 1);
assert!(msg.parts[0].is_tool_call());
}
#[test]
fn test_accumulator_interleaved_tool_calls() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::tool_call(
"call_a",
"echo",
Value::Object(serde_json::Map::new()),
)),
}))
.unwrap();
acc.process(&StreamEvent::PartStart(PartStart {
index: 1,
part: Some(MessagePart::tool_call(
"call_b",
"search",
Value::Object(serde_json::Map::new()),
)),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::InputJson {
partial_json: r#"{"msg":"#.to_string(),
},
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 1,
delta: DeltaPart::InputJson {
partial_json: r#"{"q":"#.to_string(),
},
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::InputJson {
partial_json: r#""a"}"#.to_string(),
},
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 1,
delta: DeltaPart::InputJson {
partial_json: r#""b"}"#.to_string(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
let msg = acc.build();
assert_eq!(msg.parts.len(), 2);
match &msg.parts[0] {
MessagePart::ToolCall { name, input, .. } => {
assert_eq!(name, "echo");
assert_eq!(input, &serde_json::json!({"msg": "a"}));
}
other => panic!("expected ToolCall, got {other:?}"),
}
match &msg.parts[1] {
MessagePart::ToolCall { name, input, .. } => {
assert_eq!(name, "search");
assert_eq!(input, &serde_json::json!({"q": "b"}));
}
other => panic!("expected ToolCall, got {other:?}"),
}
}
#[test]
fn test_accumulator_empty() {
let acc = StreamAccumulator::new();
let msg = acc.build();
assert_eq!(msg.parts.len(), 0);
}
#[test]
fn test_accumulator_usage() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::MessageDelta(MessageDelta {
delta: MessageDeltaPayload { stop_reason: None },
usage: Some(Usage::new(100, 50)),
}))
.unwrap();
assert_eq!(acc.usage().unwrap().total_tokens(), 150);
}
#[test]
fn test_stream_event_variants() {
let _ = StreamEvent::Ping;
let _ = StreamEvent::MessageStop;
let _ = StreamEvent::PartStop { index: None };
}
#[test]
fn test_accumulator_invalid_tool_json() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::tool_call(
"tool_1",
"bad_tool",
Value::Object(serde_json::Map::new()),
)),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::InputJson {
partial_json: "not valid json{".to_string(),
},
}))
.unwrap();
let result = acc.process(&StreamEvent::PartStop { index: None });
assert!(result.is_err());
let err = result.unwrap_err();
match &err {
StreamError::InvalidToolInputJson(_, raw) => {
assert_eq!(raw, "not valid json{");
}
}
}
#[test]
fn test_accumulator_tool_call_empty_input() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::tool_call(
"tool_1",
"no_args",
Value::Object(serde_json::Map::new()),
)),
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
let msg = acc.build();
assert_eq!(msg.parts.len(), 1);
assert!(msg.parts[0].is_tool_call());
}
#[test]
fn test_accumulator_ignores_delta_with_mismatched_index() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::text("")),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 1,
delta: DeltaPart::Text {
text: "ignored".into(),
},
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::Text {
text: "hello".into(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
let msg = acc.build();
assert_eq!(msg.parts.len(), 1);
assert_eq!(msg.parts[0].as_text(), Some("hello"));
}
#[test]
fn test_accumulator_ignores_input_json_with_mismatched_index() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::text("")),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 5,
delta: DeltaPart::InputJson {
partial_json: "{\"bad\":true}".into(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
let msg = acc.build();
assert!(msg.parts.is_empty());
}
#[test]
fn test_accumulator_delta_tool_call_string_value() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::tool_call("id1", "search", Value::Null)),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::ToolCall {
partial_json: Value::String("{\"q\":\"rust\"}".into()),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
let msg = acc.build();
assert_eq!(msg.parts.len(), 1);
if let MessagePart::ToolCall { input, .. } = &msg.parts[0] {
assert_eq!(input["q"], "rust");
} else {
panic!("expected ToolCall");
}
}
#[test]
fn test_accumulator_delta_tool_call_non_string_ignored() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::tool_call("id1", "search", Value::Null)),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::ToolCall {
partial_json: Value::Number(42.into()),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
let msg = acc.build();
assert_eq!(msg.parts.len(), 1);
if let MessagePart::ToolCall { input, .. } = &msg.parts[0] {
assert!(input.is_object());
}
}
#[test]
fn test_accumulator_ping_no_op() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::Ping).unwrap();
assert!(acc.usage().is_none());
let msg = acc.build();
assert!(msg.parts.is_empty());
}
#[test]
fn test_accumulator_message_stop_no_op() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::text("")),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::Text { text: "hi".into() },
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
acc.process(&StreamEvent::MessageStop).unwrap();
let msg = acc.build();
assert_eq!(msg.parts.len(), 1);
assert_eq!(msg.parts[0].as_text(), Some("hi"));
}
#[test]
fn test_accumulator_multiple_text_parts() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::text("")),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::Text {
text: "hello".into(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
acc.process(&StreamEvent::PartStart(PartStart {
index: 1,
part: Some(MessagePart::text("")),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 1,
delta: DeltaPart::Text {
text: "world".into(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
let msg = acc.build();
assert_eq!(msg.parts.len(), 2);
assert_eq!(msg.parts[0].as_text(), Some("hello"));
assert_eq!(msg.parts[1].as_text(), Some("world"));
}
#[test]
fn test_accumulator_message_delta_overwrites_usage() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::MessageDelta(MessageDelta {
delta: MessageDeltaPayload {
stop_reason: Some("end_turn".into()),
},
usage: Some(Usage::new(100, 50)),
}))
.unwrap();
assert_eq!(acc.usage().unwrap().input_tokens, 100);
assert_eq!(acc.usage().unwrap().output_tokens, 50);
acc.process(&StreamEvent::MessageDelta(MessageDelta {
delta: MessageDeltaPayload {
stop_reason: Some("max_tokens".into()),
},
usage: Some(Usage::new(200, 75)),
}))
.unwrap();
assert_eq!(acc.usage().unwrap().input_tokens, 200);
assert_eq!(acc.usage().unwrap().output_tokens, 75);
}
#[test]
fn accumulator_drops_thinking_not_into_text() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 1,
part: None,
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 1,
delta: DeltaPart::Thinking {
text: "reasoning here".into(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
let msg = acc.build();
let text = msg.parts.iter().find_map(|p| match p {
MessagePart::Text { text } => Some(text.as_str()),
_ => None,
});
assert!(
!text.unwrap_or("").contains("reasoning here"),
"reasoning must not leak into the message text: {text:?}"
);
}
#[test]
fn input_json_never_enters_text_slot() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::text("")),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::Text {
text: "answer".into(),
},
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::InputJson {
partial_json: r#"{"msg":"hi"}"#.into(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: Some(0) })
.unwrap();
let msg = acc.build();
assert_eq!(
msg.parts.len(),
1,
"tool-argument deltas must not open or flush extra parts"
);
match &msg.parts[0] {
MessagePart::Text { text } => assert_eq!(text, "answer"),
other => panic!("expected Text, got {other:?}"),
}
}
#[test]
fn addressed_part_stop_closes_named_slot() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::tool_call("call_a", "echo", Value::Null)),
}))
.unwrap();
acc.process(&StreamEvent::PartStart(PartStart {
index: 1,
part: Some(MessagePart::tool_call("call_b", "search", Value::Null)),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 1,
delta: DeltaPart::InputJson {
partial_json: r#"{"q":"rust"}"#.into(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: Some(1) })
.unwrap();
let msg = acc.build();
assert_eq!(
msg.parts.len(),
1,
"the named slot flushes; the still-open slot does not"
);
match &msg.parts[0] {
MessagePart::ToolCall { id, input, .. } => {
assert_eq!(id, "call_b");
assert_eq!(input, &serde_json::json!({"q": "rust"}));
}
other => panic!("expected ToolCall, got {other:?}"),
}
}
#[test]
fn part_stop_without_index_keeps_fifo() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 3,
part: Some(MessagePart::tool_call("call_a", "echo", Value::Null)),
}))
.unwrap();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::text("")),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::Text { text: "hi".into() },
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
acc.process(&StreamEvent::PartStop { index: None }).unwrap();
let msg = acc.build();
assert_eq!(
msg.parts.len(),
2,
"both slots flush in open order, not index order"
);
match &msg.parts[0] {
MessagePart::ToolCall { id, .. } => assert_eq!(
id, "call_a",
"the first-opened slot closes first regardless of its index"
),
other => panic!("expected ToolCall, got {other:?}"),
}
match &msg.parts[1] {
MessagePart::Text { text } => assert_eq!(text, "hi"),
other => panic!("expected Text, got {other:?}"),
}
}
#[test]
fn part_stop_deserializes_without_index_for_back_compat() {
let legacy: StreamEvent = serde_json::from_str(r#"{"type":"part_stop"}"#)
.expect("events serialized before the index existed must still load");
assert!(
matches!(legacy, StreamEvent::PartStop { index: None }),
"a missing index deserializes as the legacy FIFO close"
);
let addressed: StreamEvent = serde_json::from_str(r#"{"type":"part_stop","index":2}"#)
.expect("the addressed form must load");
assert!(matches!(
addressed,
StreamEvent::PartStop { index: Some(2) }
));
let round: StreamEvent =
serde_json::from_str(&serde_json::to_string(&addressed).unwrap()).unwrap();
assert!(
matches!(round, StreamEvent::PartStop { index: Some(2) }),
"the addressed form must survive a serialize/deserialize round trip"
);
}
#[test]
fn addressed_part_stop_for_unknown_index_is_a_noop() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::text("")),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::Text { text: "hi".into() },
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: Some(7) })
.unwrap();
let msg = acc.build();
assert!(
msg.parts.is_empty(),
"a stop naming no open slot closes nothing; the still-open slot drops at build"
);
}
#[test]
fn addressed_stop_closes_oldest_slot_at_shared_index() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 1,
part: None,
}))
.unwrap();
acc.process(&StreamEvent::PartStart(PartStart {
index: 1,
part: Some(MessagePart::tool_call("call_b", "search", Value::Null)),
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: Some(1) })
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 1,
delta: DeltaPart::InputJson {
partial_json: r#"{"q":"rust"}"#.into(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: Some(1) })
.unwrap();
let msg = acc.build();
assert_eq!(
msg.parts.len(),
1,
"the first close hits the older thinking slot; the tool slot stays open for its deltas"
);
match &msg.parts[0] {
MessagePart::ToolCall { input, .. } => assert_eq!(
input,
&serde_json::json!({"q": "rust"}),
"arguments arriving after the shared-index close must still land in the tool slot"
),
other => panic!("expected ToolCall, got {other:?}"),
}
}
#[test]
fn text_delta_never_enters_tool_slot() {
let mut acc = StreamAccumulator::new();
acc.process(&StreamEvent::PartStart(PartStart {
index: 0,
part: Some(MessagePart::tool_call("call_1", "echo", Value::Null)),
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::Text {
text: "stray narration".into(),
},
}))
.unwrap();
acc.process(&StreamEvent::IndexedDelta(IndexedDelta {
index: 0,
delta: DeltaPart::InputJson {
partial_json: r#"{"msg":"hi"}"#.into(),
},
}))
.unwrap();
acc.process(&StreamEvent::PartStop { index: Some(0) })
.unwrap();
let msg = acc.build();
assert_eq!(
msg.parts.len(),
1,
"a text delta must not open or flush an extra part in a tool lane"
);
match &msg.parts[0] {
MessagePart::ToolCall { input, .. } => assert_eq!(
input,
&serde_json::json!({"msg": "hi"}),
"the tool arguments must be untouched by the stray text delta"
),
other => panic!("expected ToolCall, got {other:?}"),
}
}
#[test]
fn deltapart_thinking_serde_roundtrip() {
let delta = DeltaPart::Thinking { text: "hmm".into() };
let json = serde_json::to_string(&delta).unwrap();
assert_eq!(json, r#"{"type":"thinking_delta","text":"hmm"}"#);
let parsed: DeltaPart = serde_json::from_str(&json).unwrap();
match &parsed {
DeltaPart::Thinking { text } => assert_eq!(text, "hmm"),
other => panic!("expected Thinking, got {other:?}"),
}
}
#[test]
fn deltapart_thinking_empty_text_roundtrip() {
let delta = DeltaPart::Thinking {
text: String::new(),
};
let json = serde_json::to_string(&delta).unwrap();
let parsed: DeltaPart = serde_json::from_str(&json).unwrap();
match &parsed {
DeltaPart::Thinking { text } => assert_eq!(text, ""),
other => panic!("expected Thinking with empty text, got {other:?}"),
}
}
#[test]
fn deltapart_thinking_match_compiles() {
let delta = DeltaPart::Thinking { text: "x".into() };
let result = match &delta {
DeltaPart::Text { text } => format!("text:{text}"),
DeltaPart::Thinking { text } => format!("thinking:{text}"),
DeltaPart::ToolCall { .. } | DeltaPart::InputJson { .. } => "other".into(),
};
assert_eq!(result, "thinking:x");
}
}