diff --git a/crates/goose/src/providers/bedrock.rs b/crates/goose/src/providers/bedrock.rs index cecc5cf87..c925b0633 100644 --- a/crates/goose/src/providers/bedrock.rs +++ b/crates/goose/src/providers/bedrock.rs @@ -31,7 +31,8 @@ use smithy_transport_reqwest::ReqwestHttpClient; use super::formats::bedrock::{ bedrock_anthropic_thinking_fields, bedrock_inference_config, from_bedrock_message, - from_bedrock_usage, to_bedrock_message_with_caching, to_bedrock_tool_config, + from_bedrock_usage, sanitize_json_unicode_tags, to_bedrock_message_with_caching, + to_bedrock_tool_config, }; pub(crate) const BEDROCK_PROVIDER_NAME: &str = "aws_bedrock"; @@ -793,9 +794,13 @@ fn process_stream_event( .with_arguments(object(serde_json::json!({})))) } else { match serde_json::from_str::(&input_json) { - Ok(parsed) => { - Ok(CallToolRequestParams::new(name).with_arguments(object(parsed))) - } + Ok(parsed) => sanitize_json_unicode_tags(parsed) + .map(|arguments| { + CallToolRequestParams::new(name).with_arguments(object(arguments)) + }) + .map_err(|error| { + ErrorData::new(ErrorCode::INVALID_PARAMS, error.to_string(), None) + }), Err(_) => Err(ErrorData::new( ErrorCode::INVALID_PARAMS, format!("Could not parse tool arguments: {}", input_json), @@ -1502,6 +1507,41 @@ mod tests { } } + #[test] + fn test_stream_tool_use_sanitizes_nested_arguments() { + let mut state = StreamBlockState::default(); + + process_stream_event( + tool_start_event(1, "tool-1", "lookup"), + &mut state, + TEST_MESSAGE_ID, + ); + process_stream_event( + tool_delta_event( + 1, + "{\"query\":\"visible\u{E0041}text\",\"nested\":[{\"cit\u{E0042}y\":\"東京🌍\u{E0043}\"}]}", + ), + &mut state, + TEST_MESSAGE_ID, + ); + + let (messages, _) = process_stream_event(stop_event(1), &mut state, TEST_MESSAGE_ID); + let MessageContent::ToolRequest(request) = &messages[0].content[0] else { + panic!("expected tool request"); + }; + let call = request + .tool_call + .as_ref() + .expect("expected valid tool call"); + assert_eq!( + call.arguments, + Some(object(serde_json::json!({ + "query": "visibletext", + "nested": [{"city": "東京🌍"}] + }))) + ); + } + #[test] fn test_stream_tool_use_invalid_json_yields_error_request() { let mut state = StreamBlockState::default(); diff --git a/crates/goose/src/providers/formats/bedrock.rs b/crates/goose/src/providers/formats/bedrock.rs index 806783baa..af8dda57d 100644 --- a/crates/goose/src/providers/formats/bedrock.rs +++ b/crates/goose/src/providers/formats/bedrock.rs @@ -21,7 +21,7 @@ use crate::providers::formats::anthropic::{ thinking_block_is_stale, thinking_budget_tokens, thinking_type_for_provider, ThinkingType, ANTHROPIC_PROVIDER_NAME, MIN_ANSWER_TOKENS, }; -use crate::utils::sanitize_unicode_tags; +use crate::utils::{sanitize_unicode_tags, strip_unicode_tags}; use goose_providers::conversation::token_usage::Usage; use goose_providers::model::ModelConfig; use once_cell::sync::Lazy; @@ -387,6 +387,7 @@ pub fn to_bedrock_tool(tool: &Tool) -> Result { if !input_schema.contains_key("type") { input_schema.insert("type".to_string(), Value::String("object".to_string())); } + let input_schema = sanitize_json_unicode_tags(Value::Object(input_schema))?; Ok(bedrock::Tool::ToolSpec( bedrock::ToolSpecification::builder() @@ -394,16 +395,40 @@ pub fn to_bedrock_tool(tool: &Tool) -> Result { .description( tool.description .as_ref() - .map(|d| d.to_string()) + .map(|d| strip_unicode_tags(d)) .unwrap_or_default(), ) .input_schema(bedrock::ToolInputSchema::Json(to_bedrock_json( - &Value::Object(input_schema), + &input_schema, ))) .build()?, )) } +pub(crate) fn sanitize_json_unicode_tags(value: Value) -> Result { + Ok(match value { + Value::String(text) => Value::String(strip_unicode_tags(&text)), + Value::Array(values) => Value::Array( + values + .into_iter() + .map(sanitize_json_unicode_tags) + .collect::>()?, + ), + Value::Object(values) => { + let mut sanitized = serde_json::Map::new(); + for (key, value) in values { + let key = strip_unicode_tags(&key); + if sanitized.contains_key(&key) { + bail!("JSON contains a duplicate key after Unicode tag sanitization"); + } + sanitized.insert(key, sanitize_json_unicode_tags(value)?); + } + Value::Object(sanitized) + } + value => value, + }) +} + fn args_to_value(args: Option>) -> Value { match args { Some(map) => Value::Object(map), @@ -493,11 +518,18 @@ pub fn from_bedrock_message(message: &bedrock::Message) -> Result { pub fn from_bedrock_content_block(block: &bedrock::ContentBlock) -> Result { Ok(match block { bedrock::ContentBlock::Text(text) => MessageContent::text(text), - bedrock::ContentBlock::ToolUse(tool_use) => MessageContent::tool_request( - tool_use.tool_use_id.to_string(), - Ok(CallToolRequestParams::new(tool_use.name.clone()) - .with_arguments(object(from_bedrock_json(&tool_use.input.clone())?))), - ), + bedrock::ContentBlock::ToolUse(tool_use) => { + let arguments = from_bedrock_json(&tool_use.input.clone()) + .and_then(sanitize_json_unicode_tags) + .map(|arguments| { + CallToolRequestParams::new(tool_use.name.clone()) + .with_arguments(object(arguments)) + }) + .map_err(|error| { + ErrorData::new(ErrorCode::INVALID_PARAMS, error.to_string(), None) + }); + MessageContent::tool_request(tool_use.tool_use_id.to_string(), arguments) + } bedrock::ContentBlock::ToolResult(tool_res) => MessageContent::tool_response( tool_res.tool_use_id.to_string(), if tool_res.content.is_empty() { @@ -827,6 +859,70 @@ mod tests { assert!(error_msg.contains("Failed to decode base64 image data")); } + #[test] + fn test_to_bedrock_tool_sanitizes_description_and_schema_metadata() -> Result<()> { + let tool = Tool::new( + "lookup", + "検索\u{E0041} ツール cafe\u{301}", + serde_json::Map::from_iter([ + ("type".to_string(), json!("object")), + ( + "properties".to_string(), + json!({ + "pro\u{E0042}mpt": { + "type": "string", + "description": "都市🌍\u{E0043}" + } + }), + ), + ]), + ); + + let bedrock_tool = to_bedrock_tool(&tool)?; + let spec = bedrock_tool + .as_tool_spec() + .expect("expected Bedrock tool specification"); + assert_eq!(spec.description(), Some("検索 ツール cafe\u{301}")); + + let schema = spec + .input_schema() + .expect("expected input schema") + .as_json() + .expect("expected JSON input schema"); + assert_eq!( + from_bedrock_json(schema)?, + json!({ + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "都市🌍" + } + } + }) + ); + + Ok(()) + } + + #[test] + fn test_to_bedrock_tool_rejects_schema_key_collision_after_sanitization() { + let mut properties = serde_json::Map::new(); + properties.insert("prompt".to_string(), json!({"type": "string"})); + properties.insert("pro\u{E0041}mpt".to_string(), json!({"type": "string"})); + let tool = Tool::new( + "lookup", + "Lookup", + serde_json::Map::from_iter([ + ("type".to_string(), json!("object")), + ("properties".to_string(), Value::Object(properties)), + ]), + ); + + let error = to_bedrock_tool(&tool).expect_err("sanitized keys must remain unique"); + assert!(error.to_string().contains("duplicate key")); + } + #[test] fn test_to_bedrock_message_content_image() -> Result<()> { let image = ImageContent::new(TEST_IMAGE_B64.to_string(), "image/png".to_string()); @@ -840,6 +936,61 @@ mod tests { Ok(()) } + #[test] + fn test_from_bedrock_tool_use_sanitizes_nested_arguments() -> Result<()> { + let original = bedrock::ContentBlock::ToolUse( + bedrock::ToolUseBlock::builder() + .tool_use_id("tool-1") + .name("lookup") + .input(to_bedrock_json(&json!({ + "query": "visible\u{E0041}text", + "nested": [{"cit\u{E0042}y": "東京🌍\u{E0043}"}], + "path": "cafe\u{301}.txt" + }))) + .build()?, + ); + + let MessageContent::ToolRequest(request) = from_bedrock_content_block(&original)? else { + panic!("expected tool request"); + }; + let call = request.tool_call.expect("expected valid tool call"); + assert_eq!( + call.arguments, + Some(object(json!({ + "query": "visibletext", + "nested": [{"city": "東京🌍"}], + "path": "cafe\u{301}.txt" + }))) + ); + + Ok(()) + } + + #[test] + fn test_from_bedrock_tool_use_collision_yields_error_request() -> Result<()> { + let original = bedrock::ContentBlock::ToolUse( + bedrock::ToolUseBlock::builder() + .tool_use_id("tool-1") + .name("lookup") + .input(to_bedrock_json(&json!({ + "prompt": "visible", + "pro\u{E0041}mpt": "hidden" + }))) + .build()?, + ); + + let MessageContent::ToolRequest(request) = from_bedrock_content_block(&original)? else { + panic!("expected tool request"); + }; + let error = request + .tool_call + .expect_err("sanitized key collision must produce an invalid tool call"); + assert_eq!(error.code, ErrorCode::INVALID_PARAMS); + assert!(error.message.contains("duplicate key")); + + Ok(()) + } + #[test] fn test_to_bedrock_tool_result_content_block_image() -> Result<()> { let content = ContentBlock::image(TEST_IMAGE_B64.to_string(), "image/png".to_string()); diff --git a/crates/goose/src/utils.rs b/crates/goose/src/utils.rs index dee7fc4ec..97dc2beb8 100644 --- a/crates/goose/src/utils.rs +++ b/crates/goose/src/utils.rs @@ -28,14 +28,16 @@ pub fn contains_unicode_tags(text: &str) -> bool { text.chars().any(is_in_unicode_tag_range) } +pub fn strip_unicode_tags(text: &str) -> String { + text.chars() + .filter(|&c| !is_in_unicode_tag_range(c)) + .collect() +} + /// Sanitize Unicode Tags Block characters from text pub fn sanitize_unicode_tags(text: &str) -> String { let normalized: String = text.nfc().collect(); - - normalized - .chars() - .filter(|&c| !is_in_unicode_tag_range(c)) - .collect() + strip_unicode_tags(&normalized) } /// Safely truncate a string at character boundaries, not byte boundaries @@ -134,6 +136,12 @@ mod tests { assert_eq!(cleaned, clean_text); } + #[test] + fn test_strip_unicode_tags_preserves_canonical_form() { + let decomposed = "cafe\u{301}\u{E0041}"; + assert_eq!(strip_unicode_tags(decomposed), "cafe\u{301}"); + } + #[test] fn test_sanitize_unicode_tags_empty_string() { let empty = "";