diff --git a/crates/goose-provider-types/src/conversation/message.rs b/crates/goose-provider-types/src/conversation/message.rs index ce6133a110b4..31614a925aa3 100644 --- a/crates/goose-provider-types/src/conversation/message.rs +++ b/crates/goose-provider-types/src/conversation/message.rs @@ -2,10 +2,11 @@ use crate::conversation::token_usage::{CostSource, ProviderUsage}; use crate::conversation::tool_result_serde; use crate::mcp_utils::extract_text_from_resource; use crate::utils::sanitize_unicode_tags; +use base64::Engine; use chrono::Utc; use rmcp::model::{ CallToolRequestParams, CallToolResult, ContentBlock, ElicitationAction, ImageContent, - JsonObject, PromptMessage, Role, TextContent, + JsonObject, PromptMessage, ResourceContents, Role, TextContent, }; use serde::{Deserialize, Deserializer, Serialize}; use std::collections::HashSet; @@ -53,18 +54,24 @@ where .map_err(|e| Error::custom(format!("Failed to deserialize MessageContent: {}", e)))?; for message_content in &mut content { - if let MessageContentBlock::Text(text_content) = message_content { - let original = &text_content.text; - let sanitized = sanitize_unicode_tags(original); - if *original != sanitized { - tracing::info!( - original = %original, - sanitized = %sanitized, - removed_count = original.len() - sanitized.len(), - "Unicode Tags sanitized during Message deserialization" - ); - text_content.text = sanitized; + match message_content { + MessageContentBlock::Text(text_content) => { + let original = &text_content.text; + let sanitized = sanitize_unicode_tags(original); + if *original != sanitized { + tracing::info!( + original = %original, + sanitized = %sanitized, + removed_count = original.len() - sanitized.len(), + "Unicode Tags sanitized during Message deserialization" + ); + text_content.text = sanitized; + } } + MessageContentBlock::ToolResponse(response) => { + sanitize_tool_result_in_place(&mut response.tool_result); + } + _ => {} } } @@ -76,6 +83,51 @@ where pub type ProviderMetadata = serde_json::Map; pub type ToolResult = Result; +pub(crate) fn sanitize_tool_result_in_place(tool_result: &mut ToolResult) { + match tool_result { + Ok(result) => { + for content in &mut result.content { + match content { + ContentBlock::Text(text) => { + text.text = sanitize_unicode_tags(&text.text); + } + ContentBlock::Resource(resource) => match &mut resource.resource { + ResourceContents::TextResourceContents { text, .. } => { + *text = sanitize_unicode_tags(text); + } + ResourceContents::BlobResourceContents { blob, .. } => { + let Ok(bytes) = + base64::engine::general_purpose::STANDARD.decode(blob.as_bytes()) + else { + *blob = sanitize_unicode_tags(blob); + continue; + }; + let Ok(text) = String::from_utf8(bytes) else { + continue; + }; + let sanitized = sanitize_unicode_tags(&text); + if text != sanitized { + *blob = base64::engine::general_purpose::STANDARD + .encode(sanitized.as_bytes()); + } + } + _ => {} + }, + _ => {} + } + } + } + Err(error) => { + error.message = sanitize_unicode_tags(error.message.as_ref()).into(); + } + } +} + +fn sanitize_tool_result(mut tool_result: ToolResult) -> ToolResult { + sanitize_tool_result_in_place(&mut tool_result); + tool_result +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct ToolRequest { @@ -386,7 +438,7 @@ impl MessageContentBlock { pub fn tool_response>(id: S, tool_result: ToolResult) -> Self { MessageContentBlock::ToolResponse(ToolResponse { id: id.into(), - tool_result, + tool_result: sanitize_tool_result(tool_result), metadata: None, }) } @@ -398,7 +450,7 @@ impl MessageContentBlock { ) -> Self { MessageContentBlock::ToolResponse(ToolResponse { id: id.into(), - tool_result, + tool_result: sanitize_tool_result(tool_result), metadata: metadata.cloned(), }) } @@ -1126,13 +1178,15 @@ pub struct TokenState { #[cfg(test)] mod tests { use crate::conversation::message::{ - ActionRequiredData, Message, MessageContentBlock, MessageMetadata, + ActionRequiredData, Message, MessageContentBlock, MessageMetadata, ProviderMetadata, + ToolResponse, }; + use base64::Engine; use rmcp::model::{ Annotations, CallToolResult, ElicitationAction, ErrorCode, ErrorData, ImageContent, - TextContent, + ResourceContents, TextContent, }; - use rmcp::model::{CallToolRequestParams, ContentBlock, PromptMessage, ResourceContents, Role}; + use rmcp::model::{CallToolRequestParams, ContentBlock, EmbeddedResource, PromptMessage, Role}; use rmcp::object; use serde_json::Value; @@ -1150,6 +1204,251 @@ mod tests { assert_eq!(message.as_concat_text(), clean_text); } + #[test] + fn test_tool_response_sanitizes_unicode_tags() { + let content = MessageContentBlock::tool_response( + "tool-1", + Ok(CallToolResult::success(vec![ContentBlock::text( + "visible\u{E0041}\u{E0042}text", + )])), + ); + + let MessageContentBlock::ToolResponse(response) = content else { + panic!("expected tool response"); + }; + let result = response.tool_result.unwrap(); + let ContentBlock::Text(text) = &result.content[0] else { + panic!("expected text content"); + }; + assert_eq!(text.text, "visibletext"); + } + + #[test] + fn test_tool_response_with_metadata_sanitizes_unicode_tags() { + let mut metadata = ProviderMetadata::new(); + metadata.insert("provider".to_string(), serde_json::json!("test")); + let tagged = ContentBlock::Text( + TextContent::new("result\u{E0041}") + .with_annotations(Annotations::default().with_audience(vec![Role::Assistant])), + ); + let mut message = Message::user(); + + message.add_tool_response_with_metadata( + "tool-1", + Ok(CallToolResult::success(vec![tagged])), + Some(&metadata), + ); + + let MessageContentBlock::ToolResponse(response) = &message.content[0] else { + panic!("expected tool response"); + }; + assert_eq!(response.metadata.as_ref(), Some(&metadata)); + let result = response.tool_result.as_ref().unwrap(); + let text = &result.content[0]; + let ContentBlock::Text(text) = text else { + panic!("expected text content"); + }; + assert_eq!( + text.annotations + .as_ref() + .and_then(|value| value.audience.as_ref()), + Some(&vec![Role::Assistant]) + ); + assert_eq!(text.text, "result"); + } + + #[test] + fn test_tool_response_sanitizes_error_message() { + let data = serde_json::json!({"retry": false}); + let content = MessageContentBlock::tool_response( + "tool-1", + Err(ErrorData::new( + ErrorCode::INTERNAL_ERROR, + "error\u{E0041}text", + Some(data.clone()), + )), + ); + + let MessageContentBlock::ToolResponse(response) = content else { + panic!("expected tool response"); + }; + let error = response.tool_result.unwrap_err(); + assert_eq!(error.message, "errortext"); + assert_eq!(error.code, ErrorCode::INTERNAL_ERROR); + assert_eq!(error.data, Some(data)); + } + + #[test] + fn test_tool_response_sanitizes_text_resource() { + let resource = ResourceContents::TextResourceContents { + uri: "file:///result.txt".to_string(), + mime_type: Some("text/plain".to_string()), + text: "resource\u{E0041}text".to_string(), + meta: None, + }; + let content = MessageContentBlock::tool_response( + "tool-1", + Ok(CallToolResult::success(vec![ContentBlock::Resource( + EmbeddedResource::new(resource), + )])), + ); + + let MessageContentBlock::ToolResponse(response) = content else { + panic!("expected tool response"); + }; + let result = response.tool_result.unwrap(); + let ContentBlock::Resource(resource) = &result.content[0] else { + panic!("expected resource content"); + }; + let ResourceContents::TextResourceContents { + uri, + mime_type, + text, + meta, + } = &resource.resource + else { + panic!("expected text resource"); + }; + assert_eq!(uri, "file:///result.txt"); + assert_eq!(mime_type.as_deref(), Some("text/plain")); + assert_eq!(text, "resourcetext"); + assert!(meta.is_none()); + } + + #[test] + fn test_tool_response_sanitizes_utf8_blob_resource() { + let blob = + base64::engine::general_purpose::STANDARD.encode("resource\u{E0041}text".as_bytes()); + let resource = ResourceContents::BlobResourceContents { + uri: "file:///result.txt".to_string(), + mime_type: Some("text/plain".to_string()), + blob, + meta: None, + }; + let content = MessageContentBlock::tool_response( + "tool-1", + Ok(CallToolResult::success(vec![ContentBlock::Resource( + EmbeddedResource::new(resource), + )])), + ); + + let MessageContentBlock::ToolResponse(response) = content else { + panic!("expected tool response"); + }; + let result = response.tool_result.unwrap(); + let ContentBlock::Resource(resource) = &result.content[0] else { + panic!("expected resource content"); + }; + let ResourceContents::BlobResourceContents { blob, .. } = &resource.resource else { + panic!("expected blob resource"); + }; + assert_eq!( + base64::engine::general_purpose::STANDARD + .decode(blob) + .unwrap(), + b"resourcetext" + ); + } + + #[test] + fn test_tool_response_sanitizes_malformed_blob_resource() { + let resource = ResourceContents::BlobResourceContents { + uri: "file:///result.txt".to_string(), + mime_type: Some("text/plain".to_string()), + blob: "malformed\u{E0041}text".to_string(), + meta: None, + }; + let content = MessageContentBlock::tool_response( + "tool-1", + Ok(CallToolResult::success(vec![ContentBlock::Resource( + EmbeddedResource::new(resource), + )])), + ); + + let MessageContentBlock::ToolResponse(response) = content else { + panic!("expected tool response"); + }; + let result = response.tool_result.unwrap(); + let ContentBlock::Resource(resource) = &result.content[0] else { + panic!("expected resource content"); + }; + let ResourceContents::BlobResourceContents { blob, .. } = &resource.resource else { + panic!("expected blob resource"); + }; + assert_eq!(blob, "malformedtext"); + } + + #[test] + fn test_deserialization_sanitizes_persisted_tool_response() { + let message = Message::new( + Role::User, + 1, + vec![MessageContentBlock::ToolResponse(ToolResponse { + id: "tool-1".to_string(), + tool_result: Ok(CallToolResult::success(vec![ContentBlock::text( + "persisted\u{E0041}text", + )])), + metadata: None, + })], + ); + + let json = serde_json::to_string(&message).unwrap(); + let deserialized: Message = serde_json::from_str(&json).unwrap(); + let MessageContentBlock::ToolResponse(response) = &deserialized.content[0] else { + panic!("expected tool response"); + }; + let result = response.tool_result.as_ref().unwrap(); + let ContentBlock::Text(text) = &result.content[0] else { + panic!("expected text content"); + }; + assert_eq!(text.text, "persistedtext"); + } + + #[test] + fn test_content_deserialization_sanitizes_persisted_tool_response() { + let content = vec![MessageContentBlock::ToolResponse(ToolResponse { + id: "tool-1".to_string(), + tool_result: Ok(CallToolResult::success(vec![ContentBlock::text( + "persisted\u{E0041}text", + )])), + metadata: None, + })]; + + let json = serde_json::to_string(&content).unwrap(); + let deserialized: Vec = serde_json::from_str(&json).unwrap(); + let MessageContentBlock::ToolResponse(response) = &deserialized[0] else { + panic!("expected tool response"); + }; + let result = response.tool_result.as_ref().unwrap(); + let ContentBlock::Text(text) = &result.content[0] else { + panic!("expected text content"); + }; + assert_eq!(text.text, "persistedtext"); + } + + #[test] + fn test_tool_response_sanitization_preserves_legitimate_content() { + let text = ContentBlock::Text( + TextContent::new("世界 🌍 café") + .with_annotations(Annotations::default().with_audience(vec![Role::Assistant])), + ); + let image = ContentBlock::Image( + ImageContent::new("image-data", "image/png") + .with_annotations(Annotations::default().with_audience(vec![Role::User])), + ); + let mut result = CallToolResult::success(vec![text, image]); + result.structured_content = Some(serde_json::json!({"safe": "世界"})); + result.meta = Some(rmcp::model::Meta(object!({"source": "test"}))); + let expected = result.clone(); + + let content = MessageContentBlock::tool_response("tool-1", Ok(result)); + + let MessageContentBlock::ToolResponse(response) = content else { + panic!("expected tool response"); + }; + assert_eq!(response.tool_result.unwrap(), expected); + } + #[test] fn test_message_serialization() { let message = Message::assistant() diff --git a/crates/goose-provider-types/src/conversation/tool_result_serde.rs b/crates/goose-provider-types/src/conversation/tool_result_serde.rs index 0f223a5db72e..da401ae14246 100644 --- a/crates/goose-provider-types/src/conversation/tool_result_serde.rs +++ b/crates/goose-provider-types/src/conversation/tool_result_serde.rs @@ -115,6 +115,7 @@ where pub mod call_tool_result { use super::*; + use crate::conversation::message::sanitize_tool_result_in_place; use rmcp::model::{CallToolResult, ContentBlock}; pub fn serialize( @@ -150,41 +151,44 @@ pub mod call_tool_result { let format = ResultFormat::deserialize(deserializer)?; - match format { + let mut result = match format { ResultFormat::SuccessWithCallToolResult { status, value } => { if status == "success" { - Ok(Ok(value)) + Ok(value) } else { - Err(serde::de::Error::custom(format!( + return Err(serde::de::Error::custom(format!( "Expected status 'success', got '{}'", status - ))) + ))); } } ResultFormat::SuccessWithContentVec { status, value } => { if status == "success" { - Ok(Ok(CallToolResult::success(value))) + Ok(CallToolResult::success(value)) } else { - Err(serde::de::Error::custom(format!( + return Err(serde::de::Error::custom(format!( "Expected status 'success', got '{}'", status - ))) + ))); } } ResultFormat::Error { status, error } => { if status == "error" { - Ok(Err(ErrorData { + Err(ErrorData { code: ErrorCode::INTERNAL_ERROR, message: Cow::from(error), data: None, - })) + }) } else { - Err(serde::de::Error::custom(format!( + return Err(serde::de::Error::custom(format!( "Expected status 'error', got '{}'", status - ))) + ))); } } - } + }; + + sanitize_tool_result_in_place(&mut result); + Ok(result) } } diff --git a/crates/goose/src/session/import_formats/claude_code.rs b/crates/goose/src/session/import_formats/claude_code.rs index 5622f701b050..1d2516ada58c 100644 --- a/crates/goose/src/session/import_formats/claude_code.rs +++ b/crates/goose/src/session/import_formats/claude_code.rs @@ -333,6 +333,57 @@ mod tests { assert!(content.iter().any(|c| c["type"] == "toolResponse")); } + #[test] + fn sanitizes_unicode_tags_in_tool_result() { + let jsonl = serde_json::json!({ + "type": "user", + "sessionId": "s", + "uuid": "u1", + "timestamp": "2026-01-01T00:00:00Z", + "cwd": "/tmp", + "message": { + "role": "user", + "content": [{ + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [{"type": "text", "text": "visible\u{E0041}世界"}] + }] + } + }) + .to_string(); + + let json = convert(&jsonl).unwrap(); + + assert!(json.contains("visible世界")); + assert!(!json.contains('\u{E0041}')); + } + + #[test] + fn sanitizes_unicode_tags_in_tool_result_error() { + let jsonl = serde_json::json!({ + "type": "user", + "sessionId": "s", + "uuid": "u1", + "timestamp": "2026-01-01T00:00:00Z", + "cwd": "/tmp", + "message": { + "role": "user", + "content": [{ + "type": "tool_result", + "tool_use_id": "toolu_1", + "is_error": true, + "content": "failed\u{E0041}café" + }] + } + }) + .to_string(); + + let json = convert(&jsonl).unwrap(); + + assert!(json.contains("failedcafé")); + assert!(!json.contains('\u{E0041}')); + } + #[test] fn emits_cache_token_breakdown() { let jsonl = r#"{"type":"user","sessionId":"s","uuid":"u1","timestamp":"2026-01-01T00:00:01Z","cwd":"/tmp","message":{"role":"user","content":"hi"}} diff --git a/crates/goose/src/session/import_formats/codex.rs b/crates/goose/src/session/import_formats/codex.rs index 7434293ba39a..4d886fb63d89 100644 --- a/crates/goose/src/session/import_formats/codex.rs +++ b/crates/goose/src/session/import_formats/codex.rs @@ -341,6 +341,34 @@ mod tests { .any(|c| c["type"] == "toolResponse")); } + #[test] + fn sanitizes_unicode_tags_in_function_call_output() { + let jsonl = [ + serde_json::json!({ + "timestamp": "2026-05-22T13:37:22Z", + "type": "session_meta", + "payload": {"id": "s", "cwd": "/w"} + }) + .to_string(), + serde_json::json!({ + "timestamp": "2026-05-22T13:37:23Z", + "type": "response_item", + "payload": { + "type": "function_call_output", + "call_id": "call_1", + "output": "visible\u{E0041}世界" + } + }) + .to_string(), + ] + .join("\n"); + + let json = convert(&jsonl).unwrap(); + + assert!(json.contains("visible世界")); + assert!(!json.contains('\u{E0041}')); + } + #[test] fn first_user_text_skips_context_blobs() { let jsonl = r#"{"timestamp":"2026-05-22T13:37:22Z","type":"session_meta","payload":{"id":"s","cwd":"/w"}}