xiaoski/PicoBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: stream Anthropic turns with signed replay state

Browse Source
This commit is contained in:
xiaoxixi 2026-07-17 17:11:26 +08:00
parent 6230ac8e36
commit c7ceb877a2
5 changed files with 522 additions and 213 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/agent/agent_loop.rs
View File

@ -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(),

596
src/providers/anthropic.rs
View File

@ -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<AnthropicMessage>,
max_tokens: u32,
temperature: Option<f32>,
stream: bool,
#[serde(skip_serializing_if = "Option::is_none")]
tools: Option<Vec<AnthropicTool>>,
#[serde(flatten)]
@ -157,6 +162,8 @@ fn convert_messages(messages: &[Message]) -> Vec<AnthropicMessage> {
"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<AnthropicMessage> {
.collect()
}
fn native_anthropic_content(message: &Message) -> Option<Vec<Value>> {
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<CacheControl>,
}
#[derive(Deserialize)]
struct AnthropicResponse {
id: Option<String>,
model: Option<String>,
#[serde(default)]
content: Vec<AnthropicContent>,
#[serde(default)]
usage: Option<AnthropicUsage>,
#[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<usize, Value>,
tool_json: HashMap<usize, String>,
usage: Usage,
finish_reason: Option<FinishReason>,
done_emitted: bool,
}
#[derive(Deserialize)]
struct AnthropicUsage {
#[serde(default)]
input_tokens: u32,
#[serde(default)]
output_tokens: u32,
#[serde(default)]
cache_read_input_tokens: Option<u32>,
#[serde(default)]
cache_creation_input_tokens: Option<u32>,
impl AnthropicSseDecoder {
fn push(&mut self, bytes: &[u8]) -> Result<Vec<ProviderChunk>, AnthropicStreamError> {
let frames = self.framer.push(bytes)?;
self.decode_frames(frames)
}
fn finish(&mut self) -> Result<Vec<ProviderChunk>, 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<String>,
) -> Result<Vec<ProviderChunk>, 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<ProviderChunk>) {
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<ProviderChunk>) {
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<ProviderChunk>) {
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<ProviderChunk>) {
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<ProviderChunk>) {
if self.done_emitted {
return;
}
let content = self.blocks.values().cloned().collect::<Vec<_>>();
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<usize> {
payload
.get("index")
.and_then(Value::as_u64)
.and_then(|value| usize::try_from(value).ok())
}
fn append_block_string(blocks: &mut BTreeMap<usize, Value>, 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<u32> {
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<ProviderChunk>,
reached_eof: bool,
}
async fn next_anthropic_chunk(
mut state: AnthropicHttpStream,
) -> Result<Option<(ProviderChunk, AnthropicHttpStream)>, 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::<serde_json::Value>(&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::<String>();
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)
));
}
}

66
src/providers/openai.rs
View File

@ -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<u8>,
}
impl SseFramer {
fn push(&mut self, bytes: &[u8]) -> Result<Vec<String>, OpenAIStreamError> {
self.buffer.extend_from_slice(bytes);
self.drain_frames(false)
}
fn finish(&mut self) -> Result<Vec<String>, OpenAIStreamError> {
self.drain_frames(true)
}
fn drain_frames(&mut self, finish: bool) -> Result<Vec<String>, 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::<Vec<_>>();
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<u8>) -> Result<Option<String>, 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::<Vec<_>>()
.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,

64
src/providers/stream.rs
View File

@ -13,6 +13,70 @@ pub type DynProviderError = Box<dyn Error + Send + Sync>;
pub type ProviderStreamItem = Result<ProviderChunk, DynProviderError>;
pub type ProviderStream = Pin<Box<dyn Stream<Item = ProviderStreamItem> + 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<u8>,
}
impl SseFramer {
pub(crate) fn push(&mut self, bytes: &[u8]) -> Result<Vec<String>, std::string::FromUtf8Error> {
self.buffer.extend_from_slice(bytes);
self.drain_frames(false)
}
pub(crate) fn finish(&mut self) -> Result<Vec<String>, std::string::FromUtf8Error> {
self.drain_frames(true)
}
fn drain_frames(&mut self, finish: bool) -> Result<Vec<String>, 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::<Vec<_>>();
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<u8>) -> Result<Option<String>, 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::<Vec<_>>()
.join("\n");
Ok((!data.is_empty()).then_some(data))
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FinishReason {

7
src/providers/traits.rs
View File

@ -10,6 +10,9 @@ pub struct Message {
pub content: Vec<ContentBlock>,
#[serde(skip_serializing_if = "Option::is_none")]
pub reasoning_content: Option<String>,
/// Opaque state replayed only by the provider that produced it.
#[serde(skip)]
pub provider_state: Option<crate::bus::ProviderReasoningState>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
#[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,