fix(kiro): 修复 Kiro thinking 流重写按字节截断导致的 UTF-8 边界 panic (#320)

- events.rs: 保留 <thinking>/</thinking> 尾部窗口时改为按 UTF-8 字符边界拆分,避免中文流式 chunk 在 String 切片时 panic
- tests.rs: 抽取 Kiro report_context 构造并补充普通文本/思维块内多字节内容回归测试
This commit is contained in:
Entropy.Xu
2026-04-23 22:21:27 +08:00
committed by GitHub
parent bb312ff0cf
commit ad58f4e852
2 changed files with 90 additions and 17 deletions

View File

@@ -11,6 +11,30 @@ use crate::GatewayError;
use super::super::AwsEventFrame;
use super::super::KiroClaudeStreamState;
fn floor_char_boundary(text: &str, index: usize) -> usize {
let mut boundary = index.min(text.len());
while boundary > 0 && !text.is_char_boundary(boundary) {
boundary -= 1;
}
boundary
}
fn split_preserving_trailing_bytes(
buffer: &str,
trailing_bytes: usize,
) -> Option<(String, String)> {
if buffer.len() <= trailing_bytes {
return None;
}
let split = floor_char_boundary(buffer, buffer.len() - trailing_bytes);
if split == 0 {
return None;
}
Some((buffer[..split].to_string(), buffer[split..].to_string()))
}
impl KiroClaudeStreamState {
pub(super) fn process_frame(&mut self, frame: AwsEventFrame) -> Result<Vec<u8>, GatewayError> {
let message_type = frame.headers.message_type().unwrap_or("event");
@@ -150,12 +174,12 @@ impl KiroClaudeStreamState {
}
let keep = "<thinking>".len();
if self.thinking_buffer.len() > keep {
let split = self.thinking_buffer.len() - keep;
let safe = self.thinking_buffer[..split].to_string();
if let Some((safe, remaining)) =
split_preserving_trailing_bytes(&self.thinking_buffer, keep)
{
if !safe.trim().is_empty() {
events.extend(self.emit_text_delta(&safe));
self.thinking_buffer = self.thinking_buffer[split..].to_string();
self.thinking_buffer = remaining;
}
}
break;
@@ -185,12 +209,12 @@ impl KiroClaudeStreamState {
}
let keep = "</thinking>".len();
if self.thinking_buffer.len() > keep {
let split = self.thinking_buffer.len() - keep;
let safe = self.thinking_buffer[..split].to_string();
if let Some((safe, remaining)) =
split_preserving_trailing_bytes(&self.thinking_buffer, keep)
{
if !safe.is_empty() {
events.extend(self.emit_thinking_delta(&safe));
self.thinking_buffer = self.thinking_buffer[split..].to_string();
self.thinking_buffer = remaining;
}
}
break;

View File

@@ -37,14 +37,26 @@ fn encode_frame(headers: Vec<u8>, payload: Vec<u8>) -> Vec<u8> {
out
}
#[test]
fn kiro_stream_rewriter_converts_text_events_to_claude_sse() {
let report_context = json!({
fn kiro_report_context(thinking_enabled: bool) -> Value {
let mut context = json!({
"provider_api_format": "claude:cli",
"client_api_format": "claude:cli",
"envelope_name": "kiro:generateAssistantResponse",
"mapped_model": "claude-sonnet-4.5"
});
if thinking_enabled {
context["original_request_body"] = json!({
"thinking": {
"type": "enabled"
}
});
}
context
}
#[test]
fn kiro_stream_rewriter_converts_text_events_to_claude_sse() {
let report_context = kiro_report_context(false);
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
let chunk = [
encode_event_frame(
@@ -76,12 +88,7 @@ fn kiro_stream_rewriter_converts_text_events_to_claude_sse() {
#[test]
fn kiro_stream_rewriter_converts_tool_use_to_claude_events() {
let report_context = json!({
"provider_api_format": "claude:cli",
"client_api_format": "claude:cli",
"envelope_name": "kiro:generateAssistantResponse",
"mapped_model": "claude-sonnet-4.5"
});
let report_context = kiro_report_context(false);
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
let chunk = [
encode_event_frame(
@@ -115,3 +122,45 @@ fn kiro_stream_rewriter_converts_tool_use_to_claude_events() {
assert!(text.contains("\"partial_json\":\"{\\\"city\\\":\\\"SF\\\"}\""));
assert!(text.contains("\"stop_reason\":\"tool_use\""));
}
#[test]
fn kiro_stream_rewriter_handles_multibyte_text_without_thinking_tag() {
let report_context = kiro_report_context(true);
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
let chunk = encode_event_frame(
"event",
Some("assistantResponseEvent"),
&json!({"content": "\n\n你好!有"}),
);
let first = rewriter
.push_chunk(&report_context, &chunk)
.expect("rewrite should succeed");
let rest = rewriter
.finish(&report_context)
.expect("finish should succeed");
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
assert!(text.contains("\"type\":\"text_delta\""));
assert!(text.contains("你好!有"));
}
#[test]
fn kiro_stream_rewriter_handles_multibyte_text_inside_thinking_block() {
let report_context = kiro_report_context(true);
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
let chunk = encode_event_frame(
"event",
Some("assistantResponseEvent"),
&json!({"content": "<thinking>\n\n你好!有"}),
);
let first = rewriter
.push_chunk(&report_context, &chunk)
.expect("rewrite should succeed");
let rest = rewriter
.finish(&report_context)
.expect("finish should succeed");
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
assert!(text.contains("\"type\":\"thinking_delta\""));
assert!(text.contains("你好!有"));
}