From c7ceb877a26886a140455e1fbeec7973f75a34a6 Mon Sep 17 00:00:00 2001 From: xiaoxixi Date: Fri, 17 Jul 2026 17:11:26 +0800 Subject: [PATCH] feat: stream Anthropic turns with signed replay state --- src/agent/agent_loop.rs | 2 + src/providers/anthropic.rs | 596 +++++++++++++++++++++++++++---------- src/providers/openai.rs | 66 +--- src/providers/stream.rs | 64 ++++ src/providers/traits.rs | 7 + 5 files changed, 522 insertions(+), 213 deletions(-) diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 1d8a494..b3032d4 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -571,6 +571,7 @@ impl AgentLoop { role: m.role.clone(), content, reasoning_content: m.reasoning_content.clone(), + provider_state: m.provider_state.clone(), tool_call_id: m.tool_call_id.clone(), name: m.tool_name.clone(), tool_calls: m.tool_calls.clone(), @@ -1357,6 +1358,7 @@ mod tests { role: chat_message.role.clone(), content, reasoning_content: None, + provider_state: None, tool_call_id: chat_message.tool_call_id.clone(), name: chat_message.tool_name.clone(), tool_calls: chat_message.tool_calls.clone(), diff --git a/src/providers/anthropic.rs b/src/providers/anthropic.rs index 131c979..5ff2429 100644 --- a/src/providers/anthropic.rs +++ b/src/providers/anthropic.rs @@ -1,15 +1,19 @@ use async_trait::async_trait; +use futures_util::stream; use reqwest::Client; -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; +use serde::Serialize; +use serde_json::Value; +use std::collections::{BTreeMap, HashMap, VecDeque}; use std::time::Duration; +use thiserror::Error; +use super::stream::SseFramer; use super::traits::Usage; use super::{ - ChatCompletionRequest, ChatCompletionResponse, DynProviderError, LLMProvider, Message, - ProviderStream, Tool, ToolCall, provider_stream_from_response, + ChatCompletionRequest, DynProviderError, FinishReason, LLMProvider, Message, ProviderChunk, + ProviderStream, Tool, }; -use crate::bus::message::ContentBlock; +use crate::bus::{ProviderReasoningState, message::ContentBlock}; use crate::storage::Storage; use std::sync::Arc; @@ -130,6 +134,7 @@ struct AnthropicRequest { messages: Vec, max_tokens: u32, temperature: Option, + stream: bool, #[serde(skip_serializing_if = "Option::is_none")] tools: Option>, #[serde(flatten)] @@ -157,6 +162,8 @@ fn convert_messages(messages: &[Message]) -> Vec { "tool_use_id": tool_call_id, "content": convert_content_blocks(&message.content, false), })] + } else if let Some(native) = native_anthropic_content(message) { + native } else { let mut blocks = convert_content_blocks(&message.content, message.role == "system"); if let Some(tool_calls) = message @@ -180,6 +187,26 @@ fn convert_messages(messages: &[Message]) -> Vec { .collect() } +fn native_anthropic_content(message: &Message) -> Option> { + if message.role != "assistant" { + return None; + } + let state = message.provider_state.as_ref()?; + if state.provider != "anthropic" { + return None; + } + let blocks = state.payload.get("content")?.as_array()?; + blocks + .iter() + .all(|block| { + block + .as_object() + .and_then(|value| value.get("type")) + .is_some() + }) + .then(|| blocks.clone()) +} + #[derive(Serialize)] struct AnthropicTool { name: String, @@ -189,48 +216,290 @@ struct AnthropicTool { cache_control: Option, } -#[derive(Deserialize)] -struct AnthropicResponse { - id: Option, - model: Option, - #[serde(default)] - content: Vec, - #[serde(default)] - usage: Option, +#[derive(Debug, Error)] +enum AnthropicStreamError { + #[error("invalid UTF-8 in Anthropic SSE event: {0}")] + Utf8(#[from] std::string::FromUtf8Error), + #[error("invalid Anthropic SSE payload: {0}")] + Json(#[from] serde_json::Error), + #[error("Anthropic stream error: {0}")] + Api(String), + #[error("Anthropic stream ended without message_stop")] + MissingFinish, } -#[derive(Deserialize)] -#[serde(tag = "type", rename_all = "snake_case")] -enum AnthropicContent { - Text { - #[serde(alias = "content")] - text: String, - }, - Thinking { - #[serde(alias = "content")] - thinking: String, - }, - #[serde(rename = "tool_use")] - ToolUse { - id: String, - name: String, - #[serde(alias = "arguments")] - input: serde_json::Value, - }, - #[serde(other)] - Unknown, +#[derive(Default)] +struct AnthropicSseDecoder { + framer: SseFramer, + blocks: BTreeMap, + tool_json: HashMap, + usage: Usage, + finish_reason: Option, + done_emitted: bool, } -#[derive(Deserialize)] -struct AnthropicUsage { - #[serde(default)] - input_tokens: u32, - #[serde(default)] - output_tokens: u32, - #[serde(default)] - cache_read_input_tokens: Option, - #[serde(default)] - cache_creation_input_tokens: Option, +impl AnthropicSseDecoder { + fn push(&mut self, bytes: &[u8]) -> Result, AnthropicStreamError> { + let frames = self.framer.push(bytes)?; + self.decode_frames(frames) + } + + fn finish(&mut self) -> Result, AnthropicStreamError> { + let frames = self.framer.finish()?; + let chunks = self.decode_frames(frames)?; + if !self.done_emitted { + return Err(AnthropicStreamError::MissingFinish); + } + Ok(chunks) + } + + fn decode_frames( + &mut self, + frames: Vec, + ) -> Result, AnthropicStreamError> { + let mut chunks = Vec::new(); + for data in frames { + let payload: Value = serde_json::from_str(&data)?; + match payload.get("type").and_then(Value::as_str) { + Some("message_start") => self.message_start(&payload, &mut chunks), + Some("content_block_start") => self.block_start(&payload, &mut chunks), + Some("content_block_delta") => self.block_delta(&payload, &mut chunks), + Some("content_block_stop") => self.block_stop(&payload), + Some("message_delta") => self.message_delta(&payload, &mut chunks), + Some("message_stop") => self.message_stop(&mut chunks), + Some("error") => { + let message = payload + .pointer("/error/message") + .and_then(Value::as_str) + .unwrap_or("unknown streaming error"); + return Err(AnthropicStreamError::Api(message.to_string())); + } + Some("ping") | None | Some(_) => {} + } + } + Ok(chunks) + } + + fn message_start(&mut self, payload: &Value, chunks: &mut Vec) { + let message = payload.get("message").unwrap_or(&Value::Null); + let id = message + .get("id") + .and_then(Value::as_str) + .unwrap_or_default(); + let model = message + .get("model") + .and_then(Value::as_str) + .unwrap_or_default(); + chunks.push(ProviderChunk::Metadata { + id: id.to_string(), + model: model.to_string(), + }); + if let Some(usage) = message.get("usage") { + update_anthropic_usage(&mut self.usage, usage); + chunks.push(ProviderChunk::Usage(self.usage.clone())); + } + } + + fn block_start(&mut self, payload: &Value, chunks: &mut Vec) { + let Some(index) = event_index(payload) else { + return; + }; + let Some(block) = payload.get("content_block") else { + return; + }; + self.blocks.insert(index, block.clone()); + match block.get("type").and_then(Value::as_str) { + Some("text") => { + if let Some(text) = block.get("text").and_then(Value::as_str) + && !text.is_empty() + { + chunks.push(ProviderChunk::Text(text.to_string())); + } + } + Some("thinking") => { + if let Some(thinking) = block.get("thinking").and_then(Value::as_str) + && !thinking.is_empty() + { + chunks.push(ProviderChunk::Reasoning(thinking.to_string())); + } + } + Some("tool_use") => { + chunks.push(ProviderChunk::ToolCallStart { + index, + id: block.get("id").and_then(Value::as_str).map(str::to_string), + name: block + .get("name") + .and_then(Value::as_str) + .map(str::to_string), + }); + let input = block.get("input").cloned().unwrap_or(Value::Null); + if !input.is_null() && input != serde_json::json!({}) { + chunks.push(ProviderChunk::ToolCallArguments { + index, + delta: input.to_string(), + }); + } + } + _ => {} + } + } + + fn block_delta(&mut self, payload: &Value, chunks: &mut Vec) { + let Some(index) = event_index(payload) else { + return; + }; + let Some(delta) = payload.get("delta") else { + return; + }; + match delta.get("type").and_then(Value::as_str) { + Some("text_delta") => { + if let Some(text) = delta.get("text").and_then(Value::as_str) { + append_block_string(&mut self.blocks, index, "text", text); + if !text.is_empty() { + chunks.push(ProviderChunk::Text(text.to_string())); + } + } + } + Some("thinking_delta") => { + if let Some(thinking) = delta.get("thinking").and_then(Value::as_str) { + append_block_string(&mut self.blocks, index, "thinking", thinking); + if !thinking.is_empty() { + chunks.push(ProviderChunk::Reasoning(thinking.to_string())); + } + } + } + Some("signature_delta") => { + if let Some(signature) = delta.get("signature").and_then(Value::as_str) { + append_block_string(&mut self.blocks, index, "signature", signature); + } + } + Some("input_json_delta") => { + if let Some(partial) = delta.get("partial_json").and_then(Value::as_str) { + self.tool_json.entry(index).or_default().push_str(partial); + if !partial.is_empty() { + chunks.push(ProviderChunk::ToolCallArguments { + index, + delta: partial.to_string(), + }); + } + } + } + _ => {} + } + } + + fn block_stop(&mut self, payload: &Value) { + let Some(index) = event_index(payload) else { + return; + }; + let Some(json) = self.tool_json.remove(&index) else { + return; + }; + let input = serde_json::from_str(&json).unwrap_or(Value::Null); + if let Some(block) = self.blocks.get_mut(&index) + && let Some(object) = block.as_object_mut() + { + object.insert("input".to_string(), input); + } + } + + fn message_delta(&mut self, payload: &Value, chunks: &mut Vec) { + if let Some(reason) = payload + .pointer("/delta/stop_reason") + .and_then(Value::as_str) + { + self.finish_reason = Some(FinishReason::from_provider(reason)); + } + if let Some(usage) = payload.get("usage") { + update_anthropic_usage(&mut self.usage, usage); + chunks.push(ProviderChunk::Usage(self.usage.clone())); + } + } + + fn message_stop(&mut self, chunks: &mut Vec) { + if self.done_emitted { + return; + } + let content = self.blocks.values().cloned().collect::>(); + chunks.push(ProviderChunk::ProviderState(ProviderReasoningState { + provider: "anthropic".to_string(), + payload: serde_json::json!({ "version": 1, "content": content }), + })); + chunks.push(ProviderChunk::Done( + self.finish_reason.clone().unwrap_or(FinishReason::Stop), + )); + self.done_emitted = true; + } +} + +fn event_index(payload: &Value) -> Option { + payload + .get("index") + .and_then(Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) +} + +fn append_block_string(blocks: &mut BTreeMap, index: usize, key: &str, delta: &str) { + let Some(object) = blocks.get_mut(&index).and_then(Value::as_object_mut) else { + return; + }; + let value = object + .entry(key.to_string()) + .or_insert_with(|| Value::String(String::new())); + if let Some(current) = value.as_str() { + *value = Value::String(format!("{current}{delta}")); + } +} + +fn update_anthropic_usage(usage: &mut Usage, value: &Value) { + if let Some(input) = json_u32(value, "input_tokens") { + usage.prompt_tokens = input; + } + if let Some(output) = json_u32(value, "output_tokens") { + usage.completion_tokens = output; + } + if let Some(cache_read) = json_u32(value, "cache_read_input_tokens") { + usage.cached_tokens = Some(cache_read); + usage.cache_read_input_tokens = Some(cache_read); + } + if let Some(cache_creation) = json_u32(value, "cache_creation_input_tokens") { + usage.cache_creation_input_tokens = Some(cache_creation); + } + usage.total_tokens = usage.prompt_tokens.saturating_add(usage.completion_tokens); +} + +fn json_u32(value: &Value, key: &str) -> Option { + value + .get(key) + .and_then(Value::as_u64) + .and_then(|number| u32::try_from(number).ok()) +} + +struct AnthropicHttpStream { + response: reqwest::Response, + decoder: AnthropicSseDecoder, + pending: VecDeque, + reached_eof: bool, +} + +async fn next_anthropic_chunk( + mut state: AnthropicHttpStream, +) -> Result, DynProviderError> { + loop { + if let Some(chunk) = state.pending.pop_front() { + return Ok(Some((chunk, state))); + } + if state.reached_eof { + return Ok(None); + } + match state.response.chunk().await? { + Some(bytes) => state.pending.extend(state.decoder.push(&bytes)?), + None => { + state.pending.extend(state.decoder.finish()?); + state.reached_eof = true; + } + } + } } #[async_trait] @@ -260,6 +529,7 @@ impl LLMProvider for AnthropicProvider { messages: convert_messages(&request.messages), max_tokens, temperature: request.temperature.or(self.temperature), + stream: true, tools, extra: self.model_extra.clone(), }; @@ -292,10 +562,8 @@ impl LLMProvider for AnthropicProvider { })?; let status = resp.status(); - let body_text = resp.text().await?; - tracing::debug!(status = %status, resp_body = %body_text, "LLM response"); - if !status.is_success() { + let body_text = resp.text().await?; let error_msg = serde_json::from_str::(&body_text) .ok() .and_then(|v| { @@ -327,111 +595,16 @@ impl LLMProvider for AnthropicProvider { } return Err(format!("API error ({}): {}", status.as_u16(), error_msg).into()); } - - let anthropic_resp: AnthropicResponse = match serde_json::from_str(&body_text) { - Ok(response) => response, - Err(e) => { - let err_msg = format!("decode error: {} | body: {}", e, &body_text); - if let Some(ref storage) = self.storage { - let dur = start.elapsed().as_millis() as u64; - if let Err(error) = storage - .append_llm_call( - &self.name, - &self.model_id, - &req_body_str, - Some(&body_text), - Some(&err_msg), - dur, - ) - .await - { - tracing::warn!("failed to persist LLM call (decode error): {}", error); - } - } - return Err(err_msg.into()); - } - }; - - let mut content = String::new(); - let mut reasoning = None; - let mut tool_calls = Vec::new(); - - for c in &anthropic_resp.content { - match c { - AnthropicContent::Text { text } => { - if !text.is_empty() { - if !content.is_empty() { - content.push('\n'); - } - content.push_str(text); - } - } - AnthropicContent::Thinking { thinking } => { - reasoning = Some(thinking.clone()); - } - AnthropicContent::Unknown => {} - AnthropicContent::ToolUse { id, name, input } => { - tool_calls.push(ToolCall { - id: id.clone(), - name: name.clone(), - arguments: input.clone(), - }); - } - } - } - - let response = ChatCompletionResponse { - id: anthropic_resp.id.unwrap_or_default(), - model: anthropic_resp.model.unwrap_or_default(), - content, - reasoning_content: reasoning, - provider_state: None, - tool_calls, - usage: Usage { - prompt_tokens: anthropic_resp - .usage - .as_ref() - .map(|u| u.input_tokens) - .unwrap_or(0), - completion_tokens: anthropic_resp - .usage - .as_ref() - .map(|u| u.output_tokens) - .unwrap_or(0), - total_tokens: anthropic_resp - .usage - .as_ref() - .map(|u| u.input_tokens + u.output_tokens) - .unwrap_or(0), - cached_tokens: anthropic_resp - .usage - .as_ref() - .and_then(|u| u.cache_read_input_tokens), - cache_read_input_tokens: anthropic_resp - .usage - .as_ref() - .and_then(|u| u.cache_read_input_tokens), - cache_creation_input_tokens: anthropic_resp - .usage - .as_ref() - .and_then(|u| u.cache_creation_input_tokens), + tracing::debug!(status = %status, "Anthropic streaming response started"); + Ok(Box::pin(stream::try_unfold( + AnthropicHttpStream { + response: resp, + decoder: AnthropicSseDecoder::default(), + pending: VecDeque::new(), + reached_eof: false, }, - }; - - if let Some(ref storage) = self.storage { - let _ = storage - .append_llm_call( - &self.name, - &self.model_id, - &req_body_str, - Some(&body_text), - None, - start.elapsed().as_millis() as u64, - ) - .await; - } - - Ok(provider_stream_from_response(response)) + next_anthropic_chunk, + ))) } fn ptype(&self) -> &str { @@ -450,6 +623,7 @@ impl LLMProvider for AnthropicProvider { #[cfg(test)] mod tests { use super::*; + use crate::providers::ProviderResponseAccumulator; use serde_json::json; #[test] @@ -491,6 +665,7 @@ mod tests { ContentBlock::image_url("data:image/png;base64,AAAA"), ], reasoning_content: None, + provider_state: None, tool_call_id: Some("call_1".to_string()), name: Some("file_read".to_string()), tool_calls: None, @@ -506,4 +681,125 @@ mod tests { assert_eq!(result["content"][1]["source"]["media_type"], "image/png"); assert_eq!(result["content"][1]["source"]["data"], "AAAA"); } + + #[test] + fn native_stream_decodes_thinking_signature_tools_usage_and_replay_state() { + let events = [ + json!({"type":"message_start","message":{"id":"msg_1","model":"claude-test","usage":{"input_tokens":11,"output_tokens":0,"cache_read_input_tokens":3}}}), + json!({"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""}}), + json!({"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"check "}}), + json!({"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"facts"}}), + json!({"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sig=="}}), + json!({"type":"content_block_stop","index":0}), + json!({"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"tool_1","name":"lookup","input":{}}}), + json!({"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"q\":"}}), + json!({"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"\"rust\"}"}}), + json!({"type":"content_block_stop","index":1}), + json!({"type":"content_block_start","index":2,"content_block":{"type":"text","text":""}}), + json!({"type":"content_block_delta","index":2,"delta":{"type":"text_delta","text":"answer"}}), + json!({"type":"content_block_stop","index":2}), + json!({"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":7}}), + json!({"type":"message_stop"}), + ]; + let wire = events + .iter() + .map(|event| format!("event: ignored\ndata: {event}\n\n")) + .collect::(); + let mut decoder = AnthropicSseDecoder::default(); + let mut accumulator = ProviderResponseAccumulator::default(); + for bytes in wire.as_bytes().chunks(7) { + for chunk in decoder.push(bytes).unwrap() { + accumulator.push(chunk); + } + } + for chunk in decoder.finish().unwrap() { + accumulator.push(chunk); + } + let response = accumulator.finish(); + + assert_eq!(response.id, "msg_1"); + assert_eq!(response.model, "claude-test"); + assert_eq!(response.reasoning_content.as_deref(), Some("check facts")); + assert_eq!(response.content, "answer"); + assert_eq!(response.tool_calls.len(), 1); + assert_eq!(response.tool_calls[0].id, "tool_1"); + assert_eq!(response.tool_calls[0].arguments, json!({"q":"rust"})); + assert_eq!(response.usage.prompt_tokens, 11); + assert_eq!(response.usage.completion_tokens, 7); + assert_eq!(response.usage.total_tokens, 18); + assert_eq!(response.usage.cache_read_input_tokens, Some(3)); + + let state = response.provider_state.unwrap(); + assert_eq!(state.provider, "anthropic"); + assert_eq!(state.payload["content"][0]["thinking"], "check facts"); + assert_eq!(state.payload["content"][0]["signature"], "sig=="); + assert_eq!(state.payload["content"][1]["input"], json!({"q":"rust"})); + assert_eq!(state.payload["content"][2]["text"], "answer"); + } + + #[test] + fn matching_provider_state_replays_native_blocks_without_generic_duplicates() { + let native = json!([ + {"type":"thinking","thinking":"signed thought","signature":"sig=="}, + {"type":"tool_use","id":"tool_1","name":"lookup","input":{"q":"rust"}} + ]); + let message = Message { + role: "assistant".into(), + content: vec![ContentBlock::text("generic text must not be appended")], + reasoning_content: Some("display copy".into()), + provider_state: Some(ProviderReasoningState { + provider: "anthropic".into(), + payload: json!({"version":1,"content":native}), + }), + tool_call_id: None, + name: None, + tool_calls: Some(vec![crate::providers::ToolCall { + id: "duplicate".into(), + name: "duplicate".into(), + arguments: json!({}), + }]), + }; + + let converted = convert_messages(&[message]); + + assert_eq!(converted[0].content.len(), 2); + assert_eq!(converted[0].content[0]["signature"], "sig=="); + assert_eq!(converted[0].content[1]["id"], "tool_1"); + } + + #[test] + fn foreign_provider_state_is_not_replayed_to_anthropic() { + let message = Message { + role: "assistant".into(), + content: vec![ContentBlock::text("answer")], + reasoning_content: Some("unsigned display reasoning".into()), + provider_state: Some(ProviderReasoningState { + provider: "openai".into(), + payload: json!({"private":"state"}), + }), + tool_call_id: None, + name: None, + tool_calls: None, + }; + + let converted = convert_messages(&[message]); + + assert_eq!( + converted[0].content, + vec![json!({"type":"text","text":"answer"})] + ); + } + + #[test] + fn stream_requires_message_stop() { + let mut decoder = AnthropicSseDecoder::default(); + decoder + .push(b"data: {\"type\":\"message_start\",\"message\":{}}\n\n") + .unwrap(); + + assert!(matches!( + decoder.finish(), + Err(AnthropicStreamError::MissingFinish) + )); + } } diff --git a/src/providers/openai.rs b/src/providers/openai.rs index c898844..c550211 100644 --- a/src/providers/openai.rs +++ b/src/providers/openai.rs @@ -6,6 +6,7 @@ use std::collections::{HashMap, VecDeque}; use std::time::Duration; use thiserror::Error; +use super::stream::SseFramer; use super::{ ChatCompletionRequest, DynProviderError, FinishReason, LLMProvider, Message, ProviderChunk, ProviderStream, Usage, @@ -230,69 +231,6 @@ enum OpenAIStreamError { MissingFinish, } -#[derive(Default)] -struct SseFramer { - buffer: Vec, -} - -impl SseFramer { - fn push(&mut self, bytes: &[u8]) -> Result, OpenAIStreamError> { - self.buffer.extend_from_slice(bytes); - self.drain_frames(false) - } - - fn finish(&mut self) -> Result, OpenAIStreamError> { - self.drain_frames(true) - } - - fn drain_frames(&mut self, finish: bool) -> Result, OpenAIStreamError> { - let mut frames = Vec::new(); - loop { - let delimiter = find_sse_delimiter(&self.buffer); - let Some((position, delimiter_len)) = delimiter else { - break; - }; - let frame = self.buffer.drain(..position).collect::>(); - self.buffer.drain(..delimiter_len); - if let Some(data) = sse_data(frame)? { - frames.push(data); - } - } - if finish && !self.buffer.is_empty() { - let frame = std::mem::take(&mut self.buffer); - if let Some(data) = sse_data(frame)? { - frames.push(data); - } - } - Ok(frames) - } -} - -fn find_sse_delimiter(buffer: &[u8]) -> Option<(usize, usize)> { - let lf = buffer.windows(2).position(|window| window == b"\n\n"); - let crlf = buffer.windows(4).position(|window| window == b"\r\n\r\n"); - match (lf, crlf) { - (Some(left), Some(right)) if left <= right => Some((left, 2)), - (Some(_), Some(right)) => Some((right, 4)), - (Some(position), None) => Some((position, 2)), - (None, Some(position)) => Some((position, 4)), - (None, None) => None, - } -} - -fn sse_data(frame: Vec) -> Result, OpenAIStreamError> { - let frame = String::from_utf8(frame)?; - let data = frame - .lines() - .filter_map(|line| { - line.strip_prefix("data:") - .map(|value| value.strip_prefix(' ').unwrap_or(value)) - }) - .collect::>() - .join("\n"); - Ok((!data.is_empty()).then_some(data)) -} - #[derive(Clone, Copy, PartialEq, Eq)] enum InlineMode { Text, @@ -744,6 +682,7 @@ mod tests { role: "assistant".to_string(), content: vec![ContentBlock::text("calling tool")], reasoning_content: None, + provider_state: None, tool_call_id: None, name: None, tool_calls: Some(vec![ToolCall { @@ -786,6 +725,7 @@ mod tests { ContentBlock::image_url("data:image/png;base64,AAAA"), ], reasoning_content: None, + provider_state: None, tool_call_id: Some("call_2".to_string()), name: Some("file_read".to_string()), tool_calls: None, diff --git a/src/providers/stream.rs b/src/providers/stream.rs index 7236eca..2013053 100644 --- a/src/providers/stream.rs +++ b/src/providers/stream.rs @@ -13,6 +13,70 @@ pub type DynProviderError = Box; pub type ProviderStreamItem = Result; pub type ProviderStream = Pin + Send>>; +/// Incremental framing shared by SSE-based providers. It accepts arbitrary +/// byte/UTF-8 boundaries and returns only joined `data:` payloads. +#[derive(Default)] +pub(crate) struct SseFramer { + buffer: Vec, +} + +impl SseFramer { + pub(crate) fn push(&mut self, bytes: &[u8]) -> Result, std::string::FromUtf8Error> { + self.buffer.extend_from_slice(bytes); + self.drain_frames(false) + } + + pub(crate) fn finish(&mut self) -> Result, std::string::FromUtf8Error> { + self.drain_frames(true) + } + + fn drain_frames(&mut self, finish: bool) -> Result, std::string::FromUtf8Error> { + let mut frames = Vec::new(); + loop { + let Some((position, delimiter_len)) = find_sse_delimiter(&self.buffer) else { + break; + }; + let frame = self.buffer.drain(..position).collect::>(); + self.buffer.drain(..delimiter_len); + if let Some(data) = sse_data(frame)? { + frames.push(data); + } + } + if finish && !self.buffer.is_empty() { + let frame = std::mem::take(&mut self.buffer); + if let Some(data) = sse_data(frame)? { + frames.push(data); + } + } + Ok(frames) + } +} + +fn find_sse_delimiter(buffer: &[u8]) -> Option<(usize, usize)> { + let lf = buffer.windows(2).position(|window| window == b"\n\n"); + let crlf = buffer.windows(4).position(|window| window == b"\r\n\r\n"); + match (lf, crlf) { + (Some(left), Some(right)) if left <= right => Some((left, 2)), + (Some(_), Some(right)) => Some((right, 4)), + (Some(position), None) => Some((position, 2)), + (None, Some(position)) => Some((position, 4)), + (None, None) => None, + } +} + +fn sse_data(frame: Vec) -> Result, std::string::FromUtf8Error> { + let frame = String::from_utf8(frame)?; + let data = frame + .lines() + .filter_map(|line| { + line.strip_prefix("data:") + .map(|value| value.strip_prefix(' ').unwrap_or(value)) + }) + .collect::>() + .join("\n"); + Ok((!data.is_empty()).then_some(data)) +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum FinishReason { diff --git a/src/providers/traits.rs b/src/providers/traits.rs index 55d82d0..5258ce1 100644 --- a/src/providers/traits.rs +++ b/src/providers/traits.rs @@ -10,6 +10,9 @@ pub struct Message { pub content: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub reasoning_content: Option, + /// Opaque state replayed only by the provider that produced it. + #[serde(skip)] + pub provider_state: Option, #[serde(skip_serializing_if = "Option::is_none")] pub tool_call_id: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -24,6 +27,7 @@ impl Message { role: "user".to_string(), content: vec![ContentBlock::text(content)], reasoning_content: None, + provider_state: None, tool_call_id: None, name: None, tool_calls: None, @@ -35,6 +39,7 @@ impl Message { role: "assistant".to_string(), content: vec![ContentBlock::text(content)], reasoning_content: None, + provider_state: None, tool_call_id: None, name: None, tool_calls: None, @@ -46,6 +51,7 @@ impl Message { role: "system".to_string(), content: vec![ContentBlock::text(content)], reasoning_content: None, + provider_state: None, tool_call_id: None, name: None, tool_calls: None, @@ -61,6 +67,7 @@ impl Message { role: "tool".to_string(), content: vec![ContentBlock::text(content)], reasoning_content: None, + provider_state: None, tool_call_id: Some(tool_call_id.into()), name: Some(tool_name.into()), tool_calls: None,