mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 13:10:21 +08:00
fix(ai): sanitize Claude Read pages passthrough
This commit is contained in:
@@ -8,6 +8,44 @@ fn utf8(bytes: Vec<u8>) -> String {
|
||||
String::from_utf8(bytes).expect("utf8 should decode")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_claude_local_stream_rewriter_sanitizes_read_input_json_delta() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
"needs_conversion": false,
|
||||
});
|
||||
let mut rewriter =
|
||||
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
|
||||
let mut output = rewriter
|
||||
.push_chunk(
|
||||
b"event: content_block_start\n\
|
||||
data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_read_1\",\"name\":\"Read\",\"input\":{}}}\n\n",
|
||||
)
|
||||
.expect("start should be accepted");
|
||||
output.extend(
|
||||
rewriter
|
||||
.push_chunk(
|
||||
b"event: content_block_delta\n\
|
||||
data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"file_path\\\":\\\"/tmp/a.txt\\\",\\\"pages\\\":\\\"\\\"}\"}}\n\n",
|
||||
)
|
||||
.expect("delta should be accepted"),
|
||||
);
|
||||
output.extend(
|
||||
rewriter
|
||||
.push_chunk(
|
||||
b"event: content_block_stop\n\
|
||||
data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
|
||||
)
|
||||
.expect("stop should flush sanitized delta"),
|
||||
);
|
||||
|
||||
let output_text = utf8(output);
|
||||
assert!(output_text.contains("\"name\":\"Read\""));
|
||||
assert!(output_text.contains("\\\"file_path\\\":\\\"/tmp/a.txt\\\""));
|
||||
assert!(!output_text.contains("\\\"pages\\\":\\\"\\\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn standard_sync_bridge_converts_openai_chat_sync_json_to_openai_chat_sse() {
|
||||
let outcome = maybe_bridge_standard_sync_json_to_stream(
|
||||
|
||||
@@ -46,6 +46,44 @@ pub fn remove_empty_pages_from_tool_arguments(tool_name: &str, arguments: &str)
|
||||
serde_json::to_string(&value).unwrap_or_else(|_| arguments.to_string())
|
||||
}
|
||||
|
||||
pub fn remove_empty_pages_from_tool_input_value(tool_name: &str, input: &Value) -> Value {
|
||||
if tool_name != "Read" || input.get("pages").and_then(Value::as_str) != Some("") {
|
||||
return input.clone();
|
||||
}
|
||||
let Some(object) = input.as_object() else {
|
||||
return input.clone();
|
||||
};
|
||||
let mut object = object.clone();
|
||||
object.remove("pages");
|
||||
Value::Object(object)
|
||||
}
|
||||
|
||||
pub fn sanitize_claude_read_tool_inputs(value: &mut Value) -> bool {
|
||||
let Some(content) = value.get_mut("content").and_then(Value::as_array_mut) else {
|
||||
return false;
|
||||
};
|
||||
let mut changed = false;
|
||||
for block in content {
|
||||
let Some(block_object) = block.as_object_mut() else {
|
||||
continue;
|
||||
};
|
||||
if block_object.get("type").and_then(Value::as_str) != Some("tool_use")
|
||||
|| block_object.get("name").and_then(Value::as_str) != Some("Read")
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let Some(input) = block_object.get("input") else {
|
||||
continue;
|
||||
};
|
||||
let sanitized = remove_empty_pages_from_tool_input_value("Read", input);
|
||||
if sanitized != *input {
|
||||
block_object.insert("input".to_string(), sanitized);
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
changed
|
||||
}
|
||||
|
||||
pub fn prepare_local_success_response_parts(
|
||||
headers: &BTreeMap<String, String>,
|
||||
body_json: &Value,
|
||||
@@ -134,7 +172,8 @@ mod tests {
|
||||
build_generated_tool_call_id, build_local_success_background_report,
|
||||
build_local_success_conversion_background_report, canonicalize_tool_arguments,
|
||||
prepare_local_success_response_parts, prepare_local_success_response_parts_owned,
|
||||
remove_empty_pages_from_tool_arguments, LocalSyncReportParts,
|
||||
remove_empty_pages_from_tool_arguments, sanitize_claude_read_tool_inputs,
|
||||
LocalSyncReportParts,
|
||||
};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
@@ -175,6 +214,58 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitizes_claude_read_tool_inputs_only() {
|
||||
let mut value = serde_json::json!({
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"name": "Read",
|
||||
"input": {
|
||||
"file_path": "/tmp/a.txt",
|
||||
"limit": 20,
|
||||
"pages": ""
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"name": "Search",
|
||||
"input": {
|
||||
"query": "",
|
||||
"pages": ""
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"name": "Read",
|
||||
"input": {
|
||||
"pages": "1-2"
|
||||
}
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
assert!(sanitize_claude_read_tool_inputs(&mut value));
|
||||
assert_eq!(
|
||||
value["content"][0]["input"],
|
||||
serde_json::json!({
|
||||
"file_path": "/tmp/a.txt",
|
||||
"limit": 20,
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
value["content"][1]["input"],
|
||||
serde_json::json!({
|
||||
"query": "",
|
||||
"pages": "",
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
value["content"][2]["input"],
|
||||
serde_json::json!({"pages": "1-2"})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_local_success_response_parts_normalizes_headers() {
|
||||
let headers = BTreeMap::from([
|
||||
|
||||
@@ -1,7 +1,13 @@
|
||||
use serde_json::Value;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::formats::openai::image::stream::{OpenAiImageChatStreamState, OpenAiImageStreamState};
|
||||
use crate::formats::shared::model_directives::model_directive_display_model_from_report_context;
|
||||
use crate::formats::shared::response::{
|
||||
remove_empty_pages_from_tool_arguments, remove_empty_pages_from_tool_input_value,
|
||||
};
|
||||
use crate::formats::shared::sse::encode_json_sse;
|
||||
use crate::formats::shared::stream_core::StreamingStandardFormatMatrix;
|
||||
use crate::formats::shared::AiSurfaceFinalizeError;
|
||||
use crate::provider_compat::kiro_stream::KiroToClaudeCliStreamState;
|
||||
@@ -16,6 +22,7 @@ pub enum FinalizeStreamRewriteMode {
|
||||
ModelDirectiveDisplay,
|
||||
OpenAiImage,
|
||||
OpenAiImageToOpenAiChat,
|
||||
ClaudeReadToolSanitize,
|
||||
Standard,
|
||||
KiroToClaudeCli,
|
||||
KiroToClaudeCliThenStandard,
|
||||
@@ -72,6 +79,9 @@ pub fn resolve_finalize_stream_rewrite_mode(
|
||||
// Parsing→rebuilding only adds overhead and may lose information
|
||||
// (encrypted_content, original item IDs, etc.).
|
||||
if is_same_format_family(provider_api_format.as_str(), client_api_format.as_str()) {
|
||||
if provider_api_format == "claude:messages" && client_api_format == "claude:messages" {
|
||||
return Some(FinalizeStreamRewriteMode::ClaudeReadToolSanitize);
|
||||
}
|
||||
return model_directive_display_model_from_report_context(report_context)
|
||||
.map(|_| FinalizeStreamRewriteMode::ModelDirectiveDisplay);
|
||||
}
|
||||
@@ -96,9 +106,16 @@ pub fn resolve_finalize_stream_rewrite_mode(
|
||||
provider_api_format.as_str(),
|
||||
)
|
||||
{
|
||||
if provider_api_format == "claude:messages" {
|
||||
return Some(FinalizeStreamRewriteMode::ClaudeReadToolSanitize);
|
||||
}
|
||||
return Some(FinalizeStreamRewriteMode::ModelDirectiveDisplay);
|
||||
}
|
||||
|
||||
if provider_api_format == "claude:messages" && client_api_format == "claude:messages" {
|
||||
return Some(FinalizeStreamRewriteMode::ClaudeReadToolSanitize);
|
||||
}
|
||||
|
||||
(provider_api_format == client_api_format
|
||||
&& provider_adaptation_should_unwrap_stream_envelope(
|
||||
envelope_name.as_str(),
|
||||
@@ -112,6 +129,7 @@ enum AiSurfaceStreamRewriteState {
|
||||
ModelDirectiveDisplay,
|
||||
OpenAiImage(Box<OpenAiImageStreamState>),
|
||||
OpenAiImageToOpenAiChat(Box<OpenAiImageChatStreamState>),
|
||||
ClaudeReadToolSanitize(Box<ClaudeReadToolStreamSanitizer>),
|
||||
Standard(Box<StreamingStandardFormatMatrix>),
|
||||
KiroToClaudeCli(Box<KiroToClaudeCliStreamState>),
|
||||
KiroToClaudeCliThenStandard {
|
||||
@@ -143,6 +161,11 @@ pub fn maybe_build_ai_surface_stream_rewriter<'a>(
|
||||
Box::<OpenAiImageChatStreamState>::default(),
|
||||
)
|
||||
}
|
||||
FinalizeStreamRewriteMode::ClaudeReadToolSanitize => {
|
||||
AiSurfaceStreamRewriteState::ClaudeReadToolSanitize(
|
||||
Box::<ClaudeReadToolStreamSanitizer>::default(),
|
||||
)
|
||||
}
|
||||
FinalizeStreamRewriteMode::Standard => {
|
||||
AiSurfaceStreamRewriteState::Standard(Box::<StreamingStandardFormatMatrix>::default())
|
||||
}
|
||||
@@ -173,6 +196,9 @@ impl AiSurfaceStreamRewriter<'_> {
|
||||
AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(state) => {
|
||||
state.push_chunk(self.report_context, chunk)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::ClaudeReadToolSanitize(state) => {
|
||||
state.push_chunk(self.report_context, chunk)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
|
||||
state.push_chunk(self.report_context, chunk)
|
||||
}
|
||||
@@ -200,6 +226,9 @@ impl AiSurfaceStreamRewriter<'_> {
|
||||
AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(state) => {
|
||||
state.finish(self.report_context)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::ClaudeReadToolSanitize(state) => {
|
||||
state.finish(self.report_context)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
|
||||
state.finish(self.report_context)
|
||||
}
|
||||
@@ -246,12 +275,272 @@ impl AiSurfaceStreamRewriter<'_> {
|
||||
}
|
||||
AiSurfaceStreamRewriteState::OpenAiImage(_)
|
||||
| AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(_)
|
||||
| AiSurfaceStreamRewriteState::ClaudeReadToolSanitize(_)
|
||||
| AiSurfaceStreamRewriteState::KiroToClaudeCli(_)
|
||||
| AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { .. } => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ClaudeReadToolBlockState {
|
||||
name: String,
|
||||
buffered_input_json: String,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ClaudeReadToolStreamSanitizer {
|
||||
buffered: Vec<u8>,
|
||||
blocks: BTreeMap<usize, ClaudeReadToolBlockState>,
|
||||
}
|
||||
|
||||
impl ClaudeReadToolStreamSanitizer {
|
||||
fn push_chunk(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
chunk: &[u8],
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
let mut output = Vec::new();
|
||||
while let Some(record) = drain_next_sse_record(&mut self.buffered) {
|
||||
output.extend(self.transform_record(report_context, record)?);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.buffered.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let record = std::mem::take(&mut self.buffered);
|
||||
self.transform_record(report_context, record)
|
||||
}
|
||||
|
||||
fn transform_record(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
record: Vec<u8>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let Some((event, mut payload)) = parse_sse_record_json(&record) else {
|
||||
return rewrite_model_directive_stream_record(report_context, record);
|
||||
};
|
||||
let event_type = payload
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(event.as_deref().unwrap_or_default())
|
||||
.to_string();
|
||||
let mut output = match event_type.as_str() {
|
||||
"content_block_start" => {
|
||||
self.transform_content_block_start(event.as_deref(), payload, record)?
|
||||
}
|
||||
"content_block_delta" => self.transform_content_block_delta(payload, record)?,
|
||||
"content_block_stop" => self.transform_content_block_stop(payload, record)?,
|
||||
_ => {
|
||||
if !rewrite_stream_payload_model_from_context(report_context, &mut payload) {
|
||||
return Ok(record);
|
||||
}
|
||||
encode_json_sse(event.as_deref(), &payload)?
|
||||
}
|
||||
};
|
||||
if model_directive_display_model_from_report_context(report_context).is_some()
|
||||
&& !matches!(event_type.as_str(), "message_start" | "message_delta")
|
||||
{
|
||||
output = rewrite_model_directive_stream_record(report_context, output)?;
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn transform_content_block_start(
|
||||
&mut self,
|
||||
event: Option<&str>,
|
||||
mut payload: Value,
|
||||
original_record: Vec<u8>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let index = payload
|
||||
.get("index")
|
||||
.and_then(Value::as_u64)
|
||||
.map(|value| value as usize)
|
||||
.unwrap_or(0);
|
||||
let Some(block) = payload
|
||||
.get_mut("content_block")
|
||||
.and_then(Value::as_object_mut)
|
||||
else {
|
||||
return Ok(original_record);
|
||||
};
|
||||
let block_type = block
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
if block_type != "tool_use" {
|
||||
return Ok(original_record);
|
||||
}
|
||||
let name = block
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
self.blocks.insert(
|
||||
index,
|
||||
ClaudeReadToolBlockState {
|
||||
name: name.clone(),
|
||||
buffered_input_json: String::new(),
|
||||
},
|
||||
);
|
||||
if sanitize_claude_tool_input_object(block, &name) {
|
||||
encode_json_sse(event, &payload)
|
||||
} else {
|
||||
Ok(original_record)
|
||||
}
|
||||
}
|
||||
|
||||
fn transform_content_block_delta(
|
||||
&mut self,
|
||||
payload: Value,
|
||||
original_record: Vec<u8>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let index = payload
|
||||
.get("index")
|
||||
.and_then(Value::as_u64)
|
||||
.map(|value| value as usize)
|
||||
.unwrap_or(0);
|
||||
let delta_type = payload
|
||||
.get("delta")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|delta| delta.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let partial_json = payload
|
||||
.get("delta")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|delta| delta.get("partial_json"))
|
||||
.and_then(Value::as_str);
|
||||
if delta_type != "input_json_delta" {
|
||||
return Ok(original_record);
|
||||
}
|
||||
let Some(state) = self.blocks.get_mut(&index) else {
|
||||
return Ok(original_record);
|
||||
};
|
||||
if state.name != "Read" {
|
||||
return Ok(original_record);
|
||||
}
|
||||
if let Some(partial_json) = partial_json {
|
||||
state.buffered_input_json.push_str(partial_json);
|
||||
}
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
fn transform_content_block_stop(
|
||||
&mut self,
|
||||
payload: Value,
|
||||
original_record: Vec<u8>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let index = payload
|
||||
.get("index")
|
||||
.and_then(Value::as_u64)
|
||||
.map(|value| value as usize)
|
||||
.unwrap_or(0);
|
||||
let Some(state) = self.blocks.remove(&index) else {
|
||||
return Ok(original_record);
|
||||
};
|
||||
let mut output = Vec::new();
|
||||
if state.name == "Read" && !state.buffered_input_json.is_empty() {
|
||||
let partial_json =
|
||||
remove_empty_pages_from_tool_arguments("Read", &state.buffered_input_json);
|
||||
if !partial_json.is_empty() {
|
||||
output.extend(encode_json_sse(
|
||||
Some("content_block_delta"),
|
||||
&json!({
|
||||
"type": "content_block_delta",
|
||||
"index": index,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": partial_json,
|
||||
}
|
||||
}),
|
||||
)?);
|
||||
}
|
||||
}
|
||||
if output.is_empty() {
|
||||
output = original_record;
|
||||
} else {
|
||||
output.extend(original_record);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
}
|
||||
|
||||
fn sanitize_claude_tool_input_object(block: &mut Map<String, Value>, name: &str) -> bool {
|
||||
let Some(input) = block.get("input") else {
|
||||
return false;
|
||||
};
|
||||
let sanitized = remove_empty_pages_from_tool_input_value(name, input);
|
||||
if sanitized == *input {
|
||||
return false;
|
||||
}
|
||||
block.insert("input".to_string(), sanitized);
|
||||
true
|
||||
}
|
||||
|
||||
fn drain_next_sse_record(buffer: &mut Vec<u8>) -> Option<Vec<u8>> {
|
||||
let mut line_start = 0usize;
|
||||
let mut index = 0usize;
|
||||
while index < buffer.len() {
|
||||
if buffer[index] != b'\n' {
|
||||
index += 1;
|
||||
continue;
|
||||
}
|
||||
let line_end = index + 1;
|
||||
let line = &buffer[line_start..line_end];
|
||||
let line_without_newline = line
|
||||
.strip_suffix(b"\n")
|
||||
.unwrap_or(line)
|
||||
.strip_suffix(b"\r")
|
||||
.unwrap_or_else(|| line.strip_suffix(b"\n").unwrap_or(line));
|
||||
if line_without_newline.is_empty() {
|
||||
return Some(buffer.drain(..line_end).collect());
|
||||
}
|
||||
line_start = line_end;
|
||||
index = line_end;
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn parse_sse_record_json(record: &[u8]) -> Option<(Option<String>, Value)> {
|
||||
let text = std::str::from_utf8(record).ok()?;
|
||||
let mut event = None;
|
||||
let mut data = String::new();
|
||||
for line in text.lines() {
|
||||
let line = line.strip_suffix('\r').unwrap_or(line);
|
||||
if let Some(value) = line.strip_prefix("event:") {
|
||||
event = Some(value.trim().to_string());
|
||||
} else if let Some(value) = line.strip_prefix("data:") {
|
||||
if !data.is_empty() {
|
||||
data.push('\n');
|
||||
}
|
||||
data.push_str(value.trim_start());
|
||||
}
|
||||
}
|
||||
if data.trim().is_empty() || data.trim() == "[DONE]" {
|
||||
return None;
|
||||
}
|
||||
let value = serde_json::from_str::<Value>(data.trim()).ok()?;
|
||||
Some((event, value))
|
||||
}
|
||||
|
||||
fn rewrite_model_directive_stream_record(
|
||||
report_context: &Value,
|
||||
record: Vec<u8>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let mut output = Vec::new();
|
||||
for line in record.split_inclusive(|byte| *byte == b'\n') {
|
||||
output.extend(rewrite_model_directive_stream_line(
|
||||
report_context,
|
||||
line.to_vec(),
|
||||
)?);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn rewrite_model_directive_stream_line(
|
||||
report_context: &Value,
|
||||
line: Vec<u8>,
|
||||
@@ -309,6 +598,14 @@ fn rewrite_stream_payload_model(value: &mut Value, display_model: &str) -> bool
|
||||
changed
|
||||
}
|
||||
|
||||
fn rewrite_stream_payload_model_from_context(report_context: &Value, value: &mut Value) -> bool {
|
||||
let Some(display_model) = model_directive_display_model_from_report_context(report_context)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
rewrite_stream_payload_model(value, &display_model)
|
||||
}
|
||||
|
||||
fn transform_standard_bytes(
|
||||
standard: &mut StreamingStandardFormatMatrix,
|
||||
report_context: &Value,
|
||||
@@ -647,14 +944,106 @@ data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"thinki
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_claude_without_display_model_passes_through_verbatim() {
|
||||
// Claude→Claude without display model: no rewriter needed at all.
|
||||
fn same_format_claude_uses_read_tool_sanitizer_without_display_model() {
|
||||
// Claude→Claude needs a narrow sanitizer for Claude Code Read input.
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
"needs_conversion": true,
|
||||
});
|
||||
assert!(maybe_build_ai_surface_stream_rewriter(Some(&report_context)).is_none());
|
||||
assert_eq!(
|
||||
resolve_finalize_stream_rewrite_mode(&report_context),
|
||||
Some(FinalizeStreamRewriteMode::ClaudeReadToolSanitize)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_claude_stream_sanitizes_read_start_input() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
"needs_conversion": false,
|
||||
});
|
||||
let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context))
|
||||
.expect("same-format claude sanitizer should exist");
|
||||
let output = rewriter
|
||||
.push_chunk(
|
||||
b"event: content_block_start\n\
|
||||
data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_read_1\",\"name\":\"Read\",\"input\":{\"file_path\":\"/tmp/a.txt\",\"limit\":20,\"pages\":\"\"}}}\n\n",
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let output = String::from_utf8(output).expect("output should be utf8");
|
||||
|
||||
assert!(output.contains("\"name\":\"Read\""));
|
||||
assert!(output.contains("\"file_path\":\"/tmp/a.txt\""));
|
||||
assert!(!output.contains("\"pages\":\"\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_claude_stream_sanitizes_read_input_json_delta() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
"needs_conversion": false,
|
||||
});
|
||||
let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context))
|
||||
.expect("same-format claude sanitizer should exist");
|
||||
let mut output = rewriter
|
||||
.push_chunk(
|
||||
b"event: content_block_start\n\
|
||||
data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_read_1\",\"name\":\"Read\",\"input\":{}}}\n\n",
|
||||
)
|
||||
.expect("start should rewrite");
|
||||
output.extend(
|
||||
rewriter
|
||||
.push_chunk(
|
||||
b"event: content_block_delta\n\
|
||||
data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"file_path\\\":\\\"/tmp/a.txt\\\",\"}}\n\n\
|
||||
event: content_block_delta\n\
|
||||
data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"\\\"limit\\\":20,\\\"pages\\\":\\\"\\\"}\"}}\n\n",
|
||||
)
|
||||
.expect("deltas should buffer"),
|
||||
);
|
||||
let buffered_output = String::from_utf8(output.clone()).expect("output should be utf8");
|
||||
assert!(!buffered_output.contains("input_json_delta"));
|
||||
|
||||
output.extend(
|
||||
rewriter
|
||||
.push_chunk(
|
||||
b"event: content_block_stop\n\
|
||||
data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
|
||||
)
|
||||
.expect("stop should flush sanitized delta"),
|
||||
);
|
||||
let output = String::from_utf8(output).expect("output should be utf8");
|
||||
|
||||
assert!(output.contains("event: content_block_delta"));
|
||||
assert!(output.contains("\\\"limit\\\":20"));
|
||||
assert!(!output.contains("\\\"pages\\\":\\\"\\\""));
|
||||
assert!(output.contains("event: content_block_stop"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_claude_stream_preserves_other_tool_empty_pages() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
"needs_conversion": false,
|
||||
});
|
||||
let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context))
|
||||
.expect("same-format claude sanitizer should exist");
|
||||
let output = rewriter
|
||||
.push_chunk(
|
||||
b"event: content_block_start\n\
|
||||
data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_search_1\",\"name\":\"Search\",\"input\":{}}}\n\n\
|
||||
event: content_block_delta\n\
|
||||
data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"query\\\":\\\"\\\",\\\"pages\\\":\\\"\\\"}\"}}\n\n",
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let output = String::from_utf8(output).expect("output should be utf8");
|
||||
|
||||
assert!(output.contains("\"name\":\"Search\""));
|
||||
assert!(output.contains("\\\"pages\\\":\\\"\\\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -21,7 +21,10 @@ use serde_json::{json, Map, Value};
|
||||
use super::AiSurfaceFinalizeError;
|
||||
use crate::formats::gemini::generate_content::stream::GeminiProviderState;
|
||||
use crate::formats::shared::model_directives::model_directive_display_model_from_report_context;
|
||||
use crate::formats::shared::response::remove_empty_pages_from_tool_arguments;
|
||||
use crate::formats::shared::response::{
|
||||
remove_empty_pages_from_tool_arguments, remove_empty_pages_from_tool_input_value,
|
||||
sanitize_claude_read_tool_inputs,
|
||||
};
|
||||
use crate::formats::shared::stream_core::common::{
|
||||
content_part_from_openai_image_generation_item, map_openai_finish_reason_to_gemini,
|
||||
parse_json_arguments_value, CanonicalContentPart, CanonicalStreamEvent, CanonicalUsage,
|
||||
@@ -480,8 +483,13 @@ fn maybe_build_standard_same_format_sync_body(
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut body_json = body_json.clone();
|
||||
if expected_api_format == "claude:messages" {
|
||||
sanitize_claude_read_tool_inputs(&mut body_json);
|
||||
}
|
||||
|
||||
Some(client_body_with_report_context_model(
|
||||
body_json.clone(),
|
||||
body_json,
|
||||
report_context,
|
||||
&client_api_format,
|
||||
))
|
||||
@@ -2530,13 +2538,20 @@ pub fn aggregate_claude_stream_sync_response(body: &[u8]) -> Option<Value> {
|
||||
}
|
||||
}
|
||||
"tool_use" => {
|
||||
let tool_name = block
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
if let Some(input) = block.get("input") {
|
||||
let sanitized = remove_empty_pages_from_tool_input_value(&tool_name, input);
|
||||
if sanitized != *input {
|
||||
block.insert("input".to_string(), sanitized);
|
||||
}
|
||||
}
|
||||
if !state.partial_json.is_empty() {
|
||||
let tool_name = block
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let arguments =
|
||||
remove_empty_pages_from_tool_arguments(tool_name, &state.partial_json);
|
||||
remove_empty_pages_from_tool_arguments(&tool_name, &state.partial_json);
|
||||
let input = serde_json::from_str::<Value>(&arguments)
|
||||
.unwrap_or(Value::String(arguments));
|
||||
block.insert("input".to_string(), input);
|
||||
@@ -3119,6 +3134,31 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aggregates_claude_stream_removes_empty_pages_from_start_tool_input() {
|
||||
let body = concat!(
|
||||
"event: message_start\n",
|
||||
"data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_123\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-4-5\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null}}\n\n",
|
||||
"event: content_block_start\n",
|
||||
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"toolu_read\",\"name\":\"Read\",\"input\":{\"file_path\":\"/tmp/a.txt\",\"limit\":20,\"pages\":\"\"}}}\n\n",
|
||||
"event: content_block_stop\n",
|
||||
"data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
|
||||
"event: message_stop\n",
|
||||
"data: {\"type\":\"message_stop\"}\n\n",
|
||||
);
|
||||
|
||||
let aggregated =
|
||||
aggregate_claude_stream_sync_response(body.as_bytes()).expect("body should aggregate");
|
||||
|
||||
assert_eq!(
|
||||
aggregated["content"][0]["input"],
|
||||
json!({
|
||||
"file_path": "/tmp/a.txt",
|
||||
"limit": 20,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aggregates_claude_stream_preserves_empty_pages_for_non_read_tool_input() {
|
||||
let body = concat!(
|
||||
@@ -3418,6 +3458,67 @@ mod tests {
|
||||
assert_eq!(body_json, provider_body_json);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_claude_sync_body_sanitizes_read_tool_input() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
"needs_conversion": false,
|
||||
});
|
||||
let provider_body_json = json!({
|
||||
"id": "msg_read",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_read",
|
||||
"name": "Read",
|
||||
"input": {
|
||||
"file_path": "/tmp/a.txt",
|
||||
"limit": 20,
|
||||
"pages": ""
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_search",
|
||||
"name": "Search",
|
||||
"input": {
|
||||
"query": "",
|
||||
"pages": ""
|
||||
}
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
let body_json = maybe_build_standard_same_format_sync_body_from_normalized_payload(
|
||||
"claude_chat_sync_finalize",
|
||||
200,
|
||||
Some(&report_context),
|
||||
Some(&provider_body_json),
|
||||
None,
|
||||
)
|
||||
.expect("same-format sync body should succeed")
|
||||
.expect("body should exist");
|
||||
|
||||
assert_eq!(
|
||||
body_json["content"][0]["input"],
|
||||
json!({
|
||||
"file_path": "/tmp/a.txt",
|
||||
"limit": 20,
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
body_json["content"][1]["input"],
|
||||
json!({
|
||||
"query": "",
|
||||
"pages": "",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_format_sync_response_restores_model_directive_display_model() {
|
||||
let report_context = json!({
|
||||
|
||||
@@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::formats::openai::shared::map_thinking_budget_to_openai_reasoning_effort;
|
||||
use crate::formats::shared::response::remove_empty_pages_from_tool_input_value;
|
||||
|
||||
pub use crate::protocol::stream::{CanonicalStreamEvent, CanonicalStreamFrame};
|
||||
|
||||
@@ -3825,7 +3826,7 @@ pub(crate) fn canonical_block_to_claude(
|
||||
input,
|
||||
extensions,
|
||||
} => {
|
||||
let input = claude_tool_use_input_without_empty_read_pages(name, input);
|
||||
let input = remove_empty_pages_from_tool_input_value(name, input);
|
||||
let mut out = Map::new();
|
||||
out.insert("type".to_string(), Value::String("tool_use".to_string()));
|
||||
out.insert(
|
||||
@@ -3868,18 +3869,6 @@ pub(crate) fn canonical_block_to_claude(
|
||||
}
|
||||
}
|
||||
|
||||
fn claude_tool_use_input_without_empty_read_pages(name: &str, input: &Value) -> Value {
|
||||
if name != "Read" || input.get("pages").and_then(Value::as_str) != Some("") {
|
||||
return input.clone();
|
||||
}
|
||||
let Some(object) = input.as_object() else {
|
||||
return input.clone();
|
||||
};
|
||||
let mut object = object.clone();
|
||||
object.remove("pages");
|
||||
Value::Object(object)
|
||||
}
|
||||
|
||||
fn canonical_tool_result_content_to_claude(
|
||||
output: Option<&Value>,
|
||||
content_text: Option<&str>,
|
||||
|
||||
Reference in New Issue
Block a user