@@ -2,10 +2,11 @@ use crate::conversation::token_usage::{CostSource, ProviderUsage};
22use crate :: conversation:: tool_result_serde;
33use crate :: mcp_utils:: extract_text_from_resource;
44use crate :: utils:: sanitize_unicode_tags;
5+ use base64:: Engine ;
56use chrono:: Utc ;
67use rmcp:: model:: {
78 CallToolRequestParams , CallToolResult , ContentBlock , ElicitationAction , ImageContent ,
8- JsonObject , PromptMessage , Role , TextContent ,
9+ JsonObject , PromptMessage , ResourceContents , Role , TextContent ,
910} ;
1011use serde:: { Deserialize , Deserializer , Serialize } ;
1112use std:: collections:: HashSet ;
@@ -53,18 +54,24 @@ where
5354 . map_err ( |e| Error :: custom ( format ! ( "Failed to deserialize MessageContent: {}" , e) ) ) ?;
5455
5556 for message_content in & mut content {
56- if let MessageContentBlock :: Text ( text_content) = message_content {
57- let original = & text_content. text ;
58- let sanitized = sanitize_unicode_tags ( original) ;
59- if * original != sanitized {
60- tracing:: info!(
61- original = %original,
62- sanitized = %sanitized,
63- removed_count = original. len( ) - sanitized. len( ) ,
64- "Unicode Tags sanitized during Message deserialization"
65- ) ;
66- text_content. text = sanitized;
57+ match message_content {
58+ MessageContentBlock :: Text ( text_content) => {
59+ let original = & text_content. text ;
60+ let sanitized = sanitize_unicode_tags ( original) ;
61+ if * original != sanitized {
62+ tracing:: info!(
63+ original = %original,
64+ sanitized = %sanitized,
65+ removed_count = original. len( ) - sanitized. len( ) ,
66+ "Unicode Tags sanitized during Message deserialization"
67+ ) ;
68+ text_content. text = sanitized;
69+ }
6770 }
71+ MessageContentBlock :: ToolResponse ( response) => {
72+ sanitize_tool_result_in_place ( & mut response. tool_result ) ;
73+ }
74+ _ => { }
6875 }
6976 }
7077
7683pub type ProviderMetadata = serde_json:: Map < String , serde_json:: Value > ;
7784pub type ToolResult < T > = Result < T , rmcp:: model:: ErrorData > ;
7885
86+ pub ( crate ) fn sanitize_tool_result_in_place ( tool_result : & mut ToolResult < CallToolResult > ) {
87+ match tool_result {
88+ Ok ( result) => {
89+ for content in & mut result. content {
90+ match content {
91+ ContentBlock :: Text ( text) => {
92+ text. text = sanitize_unicode_tags ( & text. text ) ;
93+ }
94+ ContentBlock :: Resource ( resource) => match & mut resource. resource {
95+ ResourceContents :: TextResourceContents { text, .. } => {
96+ * text = sanitize_unicode_tags ( text) ;
97+ }
98+ ResourceContents :: BlobResourceContents { blob, .. } => {
99+ let Ok ( bytes) =
100+ base64:: engine:: general_purpose:: STANDARD . decode ( blob. as_bytes ( ) )
101+ else {
102+ * blob = sanitize_unicode_tags ( blob) ;
103+ continue ;
104+ } ;
105+ let Ok ( text) = String :: from_utf8 ( bytes) else {
106+ continue ;
107+ } ;
108+ let sanitized = sanitize_unicode_tags ( & text) ;
109+ if text != sanitized {
110+ * blob = base64:: engine:: general_purpose:: STANDARD
111+ . encode ( sanitized. as_bytes ( ) ) ;
112+ }
113+ }
114+ _ => { }
115+ } ,
116+ _ => { }
117+ }
118+ }
119+ }
120+ Err ( error) => {
121+ error. message = sanitize_unicode_tags ( error. message . as_ref ( ) ) . into ( ) ;
122+ }
123+ }
124+ }
125+
126+ fn sanitize_tool_result ( mut tool_result : ToolResult < CallToolResult > ) -> ToolResult < CallToolResult > {
127+ sanitize_tool_result_in_place ( & mut tool_result) ;
128+ tool_result
129+ }
130+
79131#[ derive( Debug , Clone , PartialEq , Serialize , Deserialize ) ]
80132#[ serde( rename_all = "camelCase" ) ]
81133pub struct ToolRequest {
@@ -422,7 +474,7 @@ impl MessageContentBlock {
422474 pub fn tool_response < S : Into < String > > ( id : S , tool_result : ToolResult < CallToolResult > ) -> Self {
423475 MessageContentBlock :: ToolResponse ( ToolResponse {
424476 id : id. into ( ) ,
425- tool_result,
477+ tool_result : sanitize_tool_result ( tool_result ) ,
426478 metadata : None ,
427479 } )
428480 }
@@ -434,7 +486,7 @@ impl MessageContentBlock {
434486 ) -> Self {
435487 MessageContentBlock :: ToolResponse ( ToolResponse {
436488 id : id. into ( ) ,
437- tool_result,
489+ tool_result : sanitize_tool_result ( tool_result ) ,
438490 metadata : metadata. cloned ( ) ,
439491 } )
440492 }
@@ -1268,13 +1320,15 @@ pub struct TokenState {
12681320#[ cfg( test) ]
12691321mod tests {
12701322 use crate :: conversation:: message:: {
1271- ActionRequiredData , Message , MessageContentBlock , MessageMetadata ,
1323+ ActionRequiredData , Message , MessageContentBlock , MessageMetadata , ProviderMetadata ,
1324+ ToolResponse ,
12721325 } ;
1326+ use base64:: Engine ;
12731327 use rmcp:: model:: {
12741328 Annotations , CallToolResult , ElicitationAction , ErrorCode , ErrorData , ImageContent ,
1275- TextContent ,
1329+ ResourceContents , TextContent ,
12761330 } ;
1277- use rmcp:: model:: { CallToolRequestParams , ContentBlock , PromptMessage , ResourceContents , Role } ;
1331+ use rmcp:: model:: { CallToolRequestParams , ContentBlock , EmbeddedResource , PromptMessage , Role } ;
12781332 use rmcp:: object;
12791333 use serde_json:: Value ;
12801334
@@ -1292,6 +1346,251 @@ mod tests {
12921346 assert_eq ! ( message. as_concat_text( ) , clean_text) ;
12931347 }
12941348
1349+ #[ test]
1350+ fn test_tool_response_sanitizes_unicode_tags ( ) {
1351+ let content = MessageContentBlock :: tool_response (
1352+ "tool-1" ,
1353+ Ok ( CallToolResult :: success ( vec ! [ ContentBlock :: text(
1354+ "visible\u{E0041} \u{E0042} text" ,
1355+ ) ] ) ) ,
1356+ ) ;
1357+
1358+ let MessageContentBlock :: ToolResponse ( response) = content else {
1359+ panic ! ( "expected tool response" ) ;
1360+ } ;
1361+ let result = response. tool_result . unwrap ( ) ;
1362+ let ContentBlock :: Text ( text) = & result. content [ 0 ] else {
1363+ panic ! ( "expected text content" ) ;
1364+ } ;
1365+ assert_eq ! ( text. text, "visibletext" ) ;
1366+ }
1367+
1368+ #[ test]
1369+ fn test_tool_response_with_metadata_sanitizes_unicode_tags ( ) {
1370+ let mut metadata = ProviderMetadata :: new ( ) ;
1371+ metadata. insert ( "provider" . to_string ( ) , serde_json:: json!( "test" ) ) ;
1372+ let tagged = ContentBlock :: Text (
1373+ TextContent :: new ( "result\u{E0041} " )
1374+ . with_annotations ( Annotations :: default ( ) . with_audience ( vec ! [ Role :: Assistant ] ) ) ,
1375+ ) ;
1376+ let mut message = Message :: user ( ) ;
1377+
1378+ message. add_tool_response_with_metadata (
1379+ "tool-1" ,
1380+ Ok ( CallToolResult :: success ( vec ! [ tagged] ) ) ,
1381+ Some ( & metadata) ,
1382+ ) ;
1383+
1384+ let MessageContentBlock :: ToolResponse ( response) = & message. content [ 0 ] else {
1385+ panic ! ( "expected tool response" ) ;
1386+ } ;
1387+ assert_eq ! ( response. metadata. as_ref( ) , Some ( & metadata) ) ;
1388+ let result = response. tool_result . as_ref ( ) . unwrap ( ) ;
1389+ let text = & result. content [ 0 ] ;
1390+ let ContentBlock :: Text ( text) = text else {
1391+ panic ! ( "expected text content" ) ;
1392+ } ;
1393+ assert_eq ! (
1394+ text. annotations
1395+ . as_ref( )
1396+ . and_then( |value| value. audience. as_ref( ) ) ,
1397+ Some ( & vec![ Role :: Assistant ] )
1398+ ) ;
1399+ assert_eq ! ( text. text, "result" ) ;
1400+ }
1401+
1402+ #[ test]
1403+ fn test_tool_response_sanitizes_error_message ( ) {
1404+ let data = serde_json:: json!( { "retry" : false } ) ;
1405+ let content = MessageContentBlock :: tool_response (
1406+ "tool-1" ,
1407+ Err ( ErrorData :: new (
1408+ ErrorCode :: INTERNAL_ERROR ,
1409+ "error\u{E0041} text" ,
1410+ Some ( data. clone ( ) ) ,
1411+ ) ) ,
1412+ ) ;
1413+
1414+ let MessageContentBlock :: ToolResponse ( response) = content else {
1415+ panic ! ( "expected tool response" ) ;
1416+ } ;
1417+ let error = response. tool_result . unwrap_err ( ) ;
1418+ assert_eq ! ( error. message, "errortext" ) ;
1419+ assert_eq ! ( error. code, ErrorCode :: INTERNAL_ERROR ) ;
1420+ assert_eq ! ( error. data, Some ( data) ) ;
1421+ }
1422+
1423+ #[ test]
1424+ fn test_tool_response_sanitizes_text_resource ( ) {
1425+ let resource = ResourceContents :: TextResourceContents {
1426+ uri : "file:///result.txt" . to_string ( ) ,
1427+ mime_type : Some ( "text/plain" . to_string ( ) ) ,
1428+ text : "resource\u{E0041} text" . to_string ( ) ,
1429+ meta : None ,
1430+ } ;
1431+ let content = MessageContentBlock :: tool_response (
1432+ "tool-1" ,
1433+ Ok ( CallToolResult :: success ( vec ! [ ContentBlock :: Resource (
1434+ EmbeddedResource :: new( resource) ,
1435+ ) ] ) ) ,
1436+ ) ;
1437+
1438+ let MessageContentBlock :: ToolResponse ( response) = content else {
1439+ panic ! ( "expected tool response" ) ;
1440+ } ;
1441+ let result = response. tool_result . unwrap ( ) ;
1442+ let ContentBlock :: Resource ( resource) = & result. content [ 0 ] else {
1443+ panic ! ( "expected resource content" ) ;
1444+ } ;
1445+ let ResourceContents :: TextResourceContents {
1446+ uri,
1447+ mime_type,
1448+ text,
1449+ meta,
1450+ } = & resource. resource
1451+ else {
1452+ panic ! ( "expected text resource" ) ;
1453+ } ;
1454+ assert_eq ! ( uri, "file:///result.txt" ) ;
1455+ assert_eq ! ( mime_type. as_deref( ) , Some ( "text/plain" ) ) ;
1456+ assert_eq ! ( text, "resourcetext" ) ;
1457+ assert ! ( meta. is_none( ) ) ;
1458+ }
1459+
1460+ #[ test]
1461+ fn test_tool_response_sanitizes_utf8_blob_resource ( ) {
1462+ let blob =
1463+ base64:: engine:: general_purpose:: STANDARD . encode ( "resource\u{E0041} text" . as_bytes ( ) ) ;
1464+ let resource = ResourceContents :: BlobResourceContents {
1465+ uri : "file:///result.txt" . to_string ( ) ,
1466+ mime_type : Some ( "text/plain" . to_string ( ) ) ,
1467+ blob,
1468+ meta : None ,
1469+ } ;
1470+ let content = MessageContentBlock :: tool_response (
1471+ "tool-1" ,
1472+ Ok ( CallToolResult :: success ( vec ! [ ContentBlock :: Resource (
1473+ EmbeddedResource :: new( resource) ,
1474+ ) ] ) ) ,
1475+ ) ;
1476+
1477+ let MessageContentBlock :: ToolResponse ( response) = content else {
1478+ panic ! ( "expected tool response" ) ;
1479+ } ;
1480+ let result = response. tool_result . unwrap ( ) ;
1481+ let ContentBlock :: Resource ( resource) = & result. content [ 0 ] else {
1482+ panic ! ( "expected resource content" ) ;
1483+ } ;
1484+ let ResourceContents :: BlobResourceContents { blob, .. } = & resource. resource else {
1485+ panic ! ( "expected blob resource" ) ;
1486+ } ;
1487+ assert_eq ! (
1488+ base64:: engine:: general_purpose:: STANDARD
1489+ . decode( blob)
1490+ . unwrap( ) ,
1491+ b"resourcetext"
1492+ ) ;
1493+ }
1494+
1495+ #[ test]
1496+ fn test_tool_response_sanitizes_malformed_blob_resource ( ) {
1497+ let resource = ResourceContents :: BlobResourceContents {
1498+ uri : "file:///result.txt" . to_string ( ) ,
1499+ mime_type : Some ( "text/plain" . to_string ( ) ) ,
1500+ blob : "malformed\u{E0041} text" . to_string ( ) ,
1501+ meta : None ,
1502+ } ;
1503+ let content = MessageContentBlock :: tool_response (
1504+ "tool-1" ,
1505+ Ok ( CallToolResult :: success ( vec ! [ ContentBlock :: Resource (
1506+ EmbeddedResource :: new( resource) ,
1507+ ) ] ) ) ,
1508+ ) ;
1509+
1510+ let MessageContentBlock :: ToolResponse ( response) = content else {
1511+ panic ! ( "expected tool response" ) ;
1512+ } ;
1513+ let result = response. tool_result . unwrap ( ) ;
1514+ let ContentBlock :: Resource ( resource) = & result. content [ 0 ] else {
1515+ panic ! ( "expected resource content" ) ;
1516+ } ;
1517+ let ResourceContents :: BlobResourceContents { blob, .. } = & resource. resource else {
1518+ panic ! ( "expected blob resource" ) ;
1519+ } ;
1520+ assert_eq ! ( blob, "malformedtext" ) ;
1521+ }
1522+
1523+ #[ test]
1524+ fn test_deserialization_sanitizes_persisted_tool_response ( ) {
1525+ let message = Message :: new (
1526+ Role :: User ,
1527+ 1 ,
1528+ vec ! [ MessageContentBlock :: ToolResponse ( ToolResponse {
1529+ id: "tool-1" . to_string( ) ,
1530+ tool_result: Ok ( CallToolResult :: success( vec![ ContentBlock :: text(
1531+ "persisted\u{E0041} text" ,
1532+ ) ] ) ) ,
1533+ metadata: None ,
1534+ } ) ] ,
1535+ ) ;
1536+
1537+ let json = serde_json:: to_string ( & message) . unwrap ( ) ;
1538+ let deserialized: Message = serde_json:: from_str ( & json) . unwrap ( ) ;
1539+ let MessageContentBlock :: ToolResponse ( response) = & deserialized. content [ 0 ] else {
1540+ panic ! ( "expected tool response" ) ;
1541+ } ;
1542+ let result = response. tool_result . as_ref ( ) . unwrap ( ) ;
1543+ let ContentBlock :: Text ( text) = & result. content [ 0 ] else {
1544+ panic ! ( "expected text content" ) ;
1545+ } ;
1546+ assert_eq ! ( text. text, "persistedtext" ) ;
1547+ }
1548+
1549+ #[ test]
1550+ fn test_content_deserialization_sanitizes_persisted_tool_response ( ) {
1551+ let content = vec ! [ MessageContentBlock :: ToolResponse ( ToolResponse {
1552+ id: "tool-1" . to_string( ) ,
1553+ tool_result: Ok ( CallToolResult :: success( vec![ ContentBlock :: text(
1554+ "persisted\u{E0041} text" ,
1555+ ) ] ) ) ,
1556+ metadata: None ,
1557+ } ) ] ;
1558+
1559+ let json = serde_json:: to_string ( & content) . unwrap ( ) ;
1560+ let deserialized: Vec < MessageContentBlock > = serde_json:: from_str ( & json) . unwrap ( ) ;
1561+ let MessageContentBlock :: ToolResponse ( response) = & deserialized[ 0 ] else {
1562+ panic ! ( "expected tool response" ) ;
1563+ } ;
1564+ let result = response. tool_result . as_ref ( ) . unwrap ( ) ;
1565+ let ContentBlock :: Text ( text) = & result. content [ 0 ] else {
1566+ panic ! ( "expected text content" ) ;
1567+ } ;
1568+ assert_eq ! ( text. text, "persistedtext" ) ;
1569+ }
1570+
1571+ #[ test]
1572+ fn test_tool_response_sanitization_preserves_legitimate_content ( ) {
1573+ let text = ContentBlock :: Text (
1574+ TextContent :: new ( "世界 🌍 café" )
1575+ . with_annotations ( Annotations :: default ( ) . with_audience ( vec ! [ Role :: Assistant ] ) ) ,
1576+ ) ;
1577+ let image = ContentBlock :: Image (
1578+ ImageContent :: new ( "image-data" , "image/png" )
1579+ . with_annotations ( Annotations :: default ( ) . with_audience ( vec ! [ Role :: User ] ) ) ,
1580+ ) ;
1581+ let mut result = CallToolResult :: success ( vec ! [ text, image] ) ;
1582+ result. structured_content = Some ( serde_json:: json!( { "safe" : "世界" } ) ) ;
1583+ result. meta = Some ( rmcp:: model:: Meta ( object ! ( { "source" : "test" } ) ) ) ;
1584+ let expected = result. clone ( ) ;
1585+
1586+ let content = MessageContentBlock :: tool_response ( "tool-1" , Ok ( result) ) ;
1587+
1588+ let MessageContentBlock :: ToolResponse ( response) = content else {
1589+ panic ! ( "expected tool response" ) ;
1590+ } ;
1591+ assert_eq ! ( response. tool_result. unwrap( ) , expected) ;
1592+ }
1593+
12951594 #[ test]
12961595 fn test_message_serialization ( ) {
12971596 let message = Message :: assistant ( )
0 commit comments