xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(ai): preserve UTF-8 in streamed responses

This commit is contained in:
t8y2 2026-06-29 01:25:43 +08:00
parent 7a5984a1fd
commit 2d2284a0ef
1 changed files with 61 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

114
crates/dbx-core/src/ai.rs
View File

@ -367,6 +367,17 @@ pub fn stream_data_payload(line: &str) -> Option<&str> {
None
}
fn drain_next_stream_line(buffer: &mut Vec<u8>) -> Result<Option<String>, String> {
let Some(pos) = buffer.iter().position(|byte| *byte == b'\n') else {
return Ok(None);
};
let mut line = buffer.drain(..=pos).collect::<Vec<u8>>();
if line.last() == Some(&b'\n') {
line.pop();
}
String::from_utf8(line).map(Some).map_err(|e| format!("AI stream returned invalid UTF-8: {e}"))
}
pub fn claude_stream_text(event: &serde_json::Value) -> Option<&str> {
if event["type"] == "content_block_delta" {
return event["delta"]["text"].as_str();
@ -870,15 +881,12 @@ async fn measure_first_stream_chunk(
is_claude: bool,
is_gemini: bool,
) -> Result<(u64, String), String> {
let mut buf = String::new();
let mut buf = Vec::new();
while let Some(chunk) = byte_stream.next().await {
let chunk = chunk.map_err(|e| format!("stream read error: {e}"))?;
buf.push_str(&String::from_utf8_lossy(&chunk));
while let Some(pos) = buf.find('\n') {
let line = buf[..pos].to_string();
buf = buf[pos + 1..].to_string();
buf.extend_from_slice(&chunk);
while let Some(line) = drain_next_stream_line(&mut buf)? {
let Some(data) = stream_data_payload(&line) else { continue };
if data == "[DONE]" {
// stream finished without content — not a real failure but rare
@ -1217,20 +1225,17 @@ async fn stream_claude(
}
let mut byte_stream = res.bytes_stream();
let mut buf = String::new();
let mut buf = Vec::new();
loop {
tokio::select! {
chunk = byte_stream.next() => {
let Some(chunk) = chunk else { break };
let chunk = chunk.map_err(|e| e.to_string())?;
buf.push_str(&String::from_utf8_lossy(&chunk));
buf.extend_from_slice(&chunk);
let mut finished = false;
while let Some(pos) = buf.find('\n') {
let line = buf[..pos].to_string();
buf = buf[pos + 1..].to_string();
while let Some(line) = drain_next_stream_line(&mut buf)? {
let Some(data) = stream_data_payload(&line) else { continue };
if data == "[DONE]" {
finished = true;
@ -1304,20 +1309,17 @@ async fn stream_openai(
}
let mut byte_stream = res.bytes_stream();
let mut buf = String::new();
let mut buf = Vec::new();
loop {
tokio::select! {
chunk = byte_stream.next() => {
let Some(chunk) = chunk else { break };
let chunk = chunk.map_err(|e| e.to_string())?;
buf.push_str(&String::from_utf8_lossy(&chunk));
buf.extend_from_slice(&chunk);
let mut finished = false;
while let Some(pos) = buf.find('\n') {
let line = buf[..pos].to_string();
buf = buf[pos + 1..].to_string();
while let Some(line) = drain_next_stream_line(&mut buf)? {
let Some(data) = stream_data_payload(&line) else { continue };
if data == "[DONE]" {
finished = true;
@ -1391,20 +1393,17 @@ async fn stream_responses_api(
}
let mut byte_stream = res.bytes_stream();
let mut buf = String::new();
let mut buf = Vec::new();
loop {
tokio::select! {
chunk = byte_stream.next() => {
let Some(chunk) = chunk else { break };
let chunk = chunk.map_err(|e| e.to_string())?;
buf.push_str(&String::from_utf8_lossy(&chunk));
buf.extend_from_slice(&chunk);
let mut finished = false;
while let Some(pos) = buf.find('\n') {
let line = buf[..pos].to_string();
buf = buf[pos + 1..].to_string();
while let Some(line) = drain_next_stream_line(&mut buf)? {
let Some(data) = stream_data_payload(&line) else { continue };
if data == "[DONE]" {
finished = true;
@ -1481,19 +1480,16 @@ async fn stream_gemini(
}
let mut byte_stream = res.bytes_stream();
let mut buf = String::new();
let mut buf = Vec::new();
loop {
tokio::select! {
chunk = byte_stream.next() => {
let Some(chunk) = chunk else { break };
let chunk = chunk.map_err(|e| e.to_string())?;
buf.push_str(&String::from_utf8_lossy(&chunk));
while let Some(pos) = buf.find('\n') {
let line = buf[..pos].to_string();
buf = buf[pos + 1..].to_string();
buf.extend_from_slice(&chunk);
while let Some(line) = drain_next_stream_line(&mut buf)? {
let Some(data) = stream_data_payload(&line) else { continue };
if let Ok(event) = serde_json::from_str::<serde_json::Value>(data) {
let text = gemini_text(&event);
@ -1679,7 +1675,7 @@ async fn stream_claude_with_tools(
}
let mut byte_stream = res.bytes_stream();
let mut buf = String::new();
let mut buf = Vec::new();
// Track the current content block index and type for tool_use blocks
let mut current_block_index: Option<u32> = None;
let mut current_block_type: Option<String> = None;
@ -1690,13 +1686,10 @@ async fn stream_claude_with_tools(
chunk = byte_stream.next() => {
let Some(chunk) = chunk else { break };
let chunk = chunk.map_err(|e| e.to_string())?;
buf.push_str(&String::from_utf8_lossy(&chunk));
buf.extend_from_slice(&chunk);
let mut finished = false;
while let Some(pos) = buf.find('\n') {
let line = buf[..pos].to_string();
buf = buf[pos + 1..].to_string();
while let Some(line) = drain_next_stream_line(&mut buf)? {
let Some(data) = stream_data_payload(&line) else { continue };
if data == "[DONE]" {
finished = true;
@ -1856,7 +1849,7 @@ async fn stream_openai_with_tools(
}
let mut byte_stream = res.bytes_stream();
let mut buf = String::new();
let mut buf = Vec::new();
let mut token_usage: Option<TokenUsage> = None;
loop {
@ -1864,13 +1857,10 @@ async fn stream_openai_with_tools(
chunk = byte_stream.next() => {
let Some(chunk) = chunk else { break };
let chunk = chunk.map_err(|e| e.to_string())?;
buf.push_str(&String::from_utf8_lossy(&chunk));
buf.extend_from_slice(&chunk);
let mut finished = false;
while let Some(pos) = buf.find('\n') {
let line = buf[..pos].to_string();
buf = buf[pos + 1..].to_string();
while let Some(line) = drain_next_stream_line(&mut buf)? {
let Some(data) = stream_data_payload(&line) else { continue };
if data == "[DONE]" {
finished = true;
@ -2014,7 +2004,7 @@ async fn stream_gemini_with_tools(
}
let mut byte_stream = res.bytes_stream();
let mut buf = String::new();
let mut buf = Vec::new();
let mut tool_call_idx: u32 = 0;
let mut token_usage: Option<TokenUsage> = None;
@ -2023,12 +2013,9 @@ async fn stream_gemini_with_tools(
chunk = byte_stream.next() => {
let Some(chunk) = chunk else { break };
let chunk = chunk.map_err(|e| e.to_string())?;
buf.push_str(&String::from_utf8_lossy(&chunk));
while let Some(pos) = buf.find('\n') {
let line = buf[..pos].to_string();
buf = buf[pos + 1..].to_string();
buf.extend_from_slice(&chunk);
while let Some(line) = drain_next_stream_line(&mut buf)? {
let Some(data) = stream_data_payload(&line) else { continue };
if let Ok(event) = serde_json::from_str::<serde_json::Value>(data) {
// Token usage (overwrite each chunk, keep last value)
@ -2190,13 +2177,34 @@ pub fn load_config(path: &Path) -> Result<Option<AiConfig>, String> {
mod tests {
use super::{
add_temperature_if_supported_for_config, build_ai_http_client, claude_headers, claude_system_prompt,
gemini_text, is_kimi_model, openai_response_text, openai_stream_reasoning, openai_stream_text,
parse_model_list_response, resolve_endpoint, resolve_model_list_endpoint, responses_max_output_tokens,
responses_text, supports_temperature, temperature_value, uses_anthropic_messages_api, validate_config,
AiApiStyle, AiAuthMethod, AiConfig, AiModelInfo, AiProvider, AiReasoningLevel, AUTHORIZATION,
CLAUDE_DEFAULT_SYSTEM, TEST_PROMPT,
drain_next_stream_line, gemini_text, is_kimi_model, openai_response_text, openai_stream_reasoning,
openai_stream_text, parse_model_list_response, resolve_endpoint, resolve_model_list_endpoint,
responses_max_output_tokens, responses_text, stream_data_payload, supports_temperature, temperature_value,
uses_anthropic_messages_api, validate_config, AiApiStyle, AiAuthMethod, AiConfig, AiModelInfo, AiProvider,
AiReasoningLevel, AUTHORIZATION, CLAUDE_DEFAULT_SYSTEM, TEST_PROMPT,
};
#[test]
fn stream_line_decoder_preserves_split_multibyte_utf8() {
let text = "\u{8bf4}\u{660e}";
let json = serde_json::json!({ "delta": text }).to_string();
let line = format!("data: {json}\n");
let bytes = line.as_bytes();
let split = bytes.iter().position(|byte| *byte >= 0x80).unwrap() + 1;
let mut buffer = Vec::new();
buffer.extend_from_slice(&bytes[..split]);
assert_eq!(drain_next_stream_line(&mut buffer).unwrap(), None);
buffer.extend_from_slice(&bytes[split..]);
let decoded = drain_next_stream_line(&mut buffer).unwrap().unwrap();
let payload = stream_data_payload(&decoded).unwrap();
let parsed: serde_json::Value = serde_json::from_str(payload).unwrap();
assert_eq!(parsed["delta"].as_str(), Some(text));
assert!(!decoded.contains('\u{fffd}'));
}
#[test]
fn ai_config_proxy_fields_default_for_legacy_config() {
let config: AiConfig = serde_json::from_value(serde_json::json!({