refactor(proxy): extract take_sse_block helper with CRLF delimiter support

Replace inline `buffer.find("\n\n")` SSE splitting logic across streaming,
streaming_responses, response_handler, and response_processor with a shared
`take_sse_block` function that handles both `\n\n` and `\r\n\r\n` delimiters.
This commit is contained in:
YoVinchen
2026-04-06 11:09:23 +08:00
parent 602c5717b2
commit abbf43ea6f
5 changed files with 49 additions and 21 deletions
+2 -5
View File
@@ -2,7 +2,7 @@
//!
//! 实现 OpenAI SSE → Anthropic SSE 格式转换
use crate::proxy::sse::strip_sse_field;
use crate::proxy::sse::{strip_sse_field, take_sse_block};
use bytes::Bytes;
use futures::stream::{Stream, StreamExt};
use serde::{Deserialize, Serialize};
@@ -110,10 +110,7 @@ pub fn create_anthropic_sse_stream<E: std::error::Error + Send + 'static>(
let text = String::from_utf8_lossy(&bytes);
buffer.push_str(&text);
while let Some(pos) = buffer.find("\n\n") {
let line = buffer[..pos].to_string();
buffer = buffer[pos + 2..].to_string();
while let Some(line) = take_sse_block(&mut buffer) {
if line.trim().is_empty() {
continue;
}
@@ -9,7 +9,7 @@
//! 与 Chat Completions 的 delta chunk 模型完全不同,需要独立的状态机处理。
use super::transform_responses::{build_anthropic_usage_from_responses, map_responses_stop_reason};
use crate::proxy::sse::strip_sse_field;
use crate::proxy::sse::{strip_sse_field, take_sse_block};
use bytes::Bytes;
use futures::stream::{Stream, StreamExt};
use serde_json::{json, Value};
@@ -122,10 +122,7 @@ pub fn create_anthropic_sse_stream_from_responses<E: std::error::Error + Send +
buffer.push_str(&text);
// SSE 事件由 \n\n 分隔
while let Some(pos) = buffer.find("\n\n") {
let block = buffer[..pos].to_string();
buffer = buffer[pos + 2..].to_string();
while let Some(block) = take_sse_block(&mut buffer) {
if block.trim().is_empty() {
continue;
}
+2 -5
View File
@@ -5,7 +5,7 @@
use super::session::ProxySession;
use super::usage::parser::TokenUsage;
use super::ProxyError;
use crate::proxy::sse::strip_sse_field;
use crate::proxy::sse::{strip_sse_field, take_sse_block};
use bytes::Bytes;
use futures::stream::{Stream, StreamExt};
use serde_json::Value;
@@ -86,10 +86,7 @@ impl StreamHandler {
buffer.push_str(&text);
// 提取完整事件
while let Some(pos) = buffer.find("\n\n") {
let event_text = buffer[..pos].to_string();
buffer = buffer[pos + 2..].to_string();
while let Some(event_text) = take_sse_block(&mut buffer) {
for line in event_text.lines() {
if let Some(data) = strip_sse_field(line, "data") {
if data.trim() != "[DONE]" {
+2 -5
View File
@@ -7,7 +7,7 @@ use super::{
handler_context::{RequestContext, StreamingTimeoutConfig},
hyper_client::ProxyResponse,
server::ProxyState,
sse::strip_sse_field,
sse::{strip_sse_field, take_sse_block},
usage::parser::TokenUsage,
ProxyError,
};
@@ -623,10 +623,7 @@ pub fn create_logged_passthrough_stream(
buffer.push_str(&text);
// 尝试解析并记录完整的 SSE 事件
while let Some(pos) = buffer.find("\n\n") {
let event_text = buffer[..pos].to_string();
buffer = buffer[pos + 2..].to_string();
while let Some(event_text) = take_sse_block(&mut buffer) {
if !event_text.trim().is_empty() {
// 提取 data 部分并尝试解析为 JSON
for line in event_text.lines() {
+41 -1
View File
@@ -4,9 +4,27 @@ pub(crate) fn strip_sse_field<'a>(line: &'a str, field: &str) -> Option<&'a str>
.or_else(|| line.strip_prefix(&format!("{field}:")))
}
#[inline]
pub(crate) fn take_sse_block(buffer: &mut String) -> Option<String> {
let mut best: Option<(usize, usize)> = None;
for (delimiter, len) in [("\r\n\r\n", 4usize), ("\n\n", 2usize)] {
if let Some(pos) = buffer.find(delimiter) {
if best.is_none_or(|(best_pos, _)| pos < best_pos) {
best = Some((pos, len));
}
}
}
let (pos, len) = best?;
let block = buffer[..pos].to_string();
buffer.drain(..pos + len);
Some(block)
}
#[cfg(test)]
mod tests {
use super::strip_sse_field;
use super::{strip_sse_field, take_sse_block};
#[test]
fn strip_sse_field_accepts_optional_space() {
@@ -28,4 +46,26 @@ mod tests {
);
assert_eq!(strip_sse_field("id:1", "data"), None);
}
#[test]
fn take_sse_block_supports_lf_delimiters() {
let mut buffer = "data: {\"ok\":true}\n\nrest".to_string();
assert_eq!(
take_sse_block(&mut buffer),
Some("data: {\"ok\":true}".to_string())
);
assert_eq!(buffer, "rest");
}
#[test]
fn take_sse_block_supports_crlf_delimiters() {
let mut buffer = "data: {\"ok\":true}\r\n\r\nrest".to_string();
assert_eq!(
take_sse_block(&mut buffer),
Some("data: {\"ok\":true}".to_string())
);
assert_eq!(buffer, "rest");
}
}