a3s_code_core/mcp/
result.rs

1//! Loss-aware projection of MCP tool results into A3S Code tool output.
2
3use crate::llm::Attachment;
4use crate::mcp::protocol::{CallToolResult, ResourceContent, ToolContent};
5use crate::tools::{ToolContext, ToolOutput};
6use anyhow::{anyhow, Context, Result};
7use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
8use serde_json::{json, Value};
9use std::path::{Path, PathBuf};
10use tokio::io::AsyncWriteExt;
11
12const MAX_CONTENT_ITEMS: usize = 64;
13const MAX_ARTIFACT_BYTES: usize = 16 * 1024 * 1024;
14const MAX_TOTAL_ARTIFACT_BYTES: usize = 32 * 1024 * 1024;
15const MAX_STRUCTURED_CONTENT_BYTES: usize = 4 * 1024 * 1024;
16const MAX_PROTOCOL_META_BYTES: usize = 1024 * 1024;
17const MAX_CACHED_ARTIFACTS_PER_TOOL: usize = 64;
18const MAX_CACHED_ARTIFACT_BYTES_PER_TOOL: u64 = 128 * 1024 * 1024;
19
20/// Convert an MCP result to model-visible text without discarding structured
21/// content, decoded images, embedded resources, or protocol metadata.
22pub(crate) async fn project_tool_result(
23    tool_name: &str,
24    result: &CallToolResult,
25    context: &ToolContext,
26) -> Result<ToolOutput> {
27    if result.content.len() > MAX_CONTENT_ITEMS {
28        return Err(anyhow!(
29            "MCP tool returned {} content items; the limit is {}",
30            result.content.len(),
31            MAX_CONTENT_ITEMS
32        ));
33    }
34    validate_json_size(
35        "structuredContent",
36        result.structured_content.as_ref(),
37        MAX_STRUCTURED_CONTENT_BYTES,
38    )?;
39    validate_json_size("_meta", result.meta.as_ref(), MAX_PROTOCOL_META_BYTES)?;
40
41    let mut text_parts = Vec::new();
42    let mut images = Vec::new();
43    let mut artifacts = Vec::new();
44    let mut content = Vec::with_capacity(result.content.len());
45    let mut decoded_bytes = 0usize;
46
47    for item in &result.content {
48        match item {
49            ToolContent::Text { text } => {
50                text_parts.push(text.clone());
51                content.push(json!({
52                    "type": "text",
53                    "bytes": text.len(),
54                }));
55            }
56            ToolContent::Image { data, mime_type } => {
57                let bytes = decode_bounded(data, &mut decoded_bytes)?;
58                let artifact =
59                    materialize_artifact(tool_name, None, mime_type, &bytes, context).await?;
60                text_parts.push(format!(
61                    "[Image: {mime_type}, {} bytes, artifact: {}]",
62                    bytes.len(),
63                    artifact.path.display()
64                ));
65                if model_image_mime_type(mime_type) {
66                    images.push(Attachment::new(bytes.clone(), mime_type.clone()));
67                }
68                content.push(json!({
69                    "type": "image",
70                    "mimeType": mime_type,
71                    "bytes": bytes.len(),
72                    "sha256": artifact.sha256.clone(),
73                    "path": artifact.path.clone(),
74                    "attachedToModel": model_image_mime_type(mime_type),
75                }));
76                artifacts.push(artifact.value("image", None));
77            }
78            ToolContent::Resource { resource } => {
79                let mut projection = ResourceProjection {
80                    text_parts: &mut text_parts,
81                    images: &mut images,
82                    artifacts: &mut artifacts,
83                    content: &mut content,
84                    decoded_bytes: &mut decoded_bytes,
85                };
86                project_resource(tool_name, resource, context, &mut projection).await?;
87            }
88        }
89    }
90
91    if text_parts.is_empty() {
92        if let Some(structured) = &result.structured_content {
93            text_parts.push(
94                serde_json::to_string_pretty(structured)
95                    .context("failed to format MCP structured content")?,
96            );
97        }
98    }
99
100    let mut metadata = serde_json::Map::new();
101    if let Some(structured) = &result.structured_content {
102        metadata.insert("structuredContent".to_string(), structured.clone());
103    }
104    if let Some(meta) = &result.meta {
105        metadata.insert("_meta".to_string(), meta.clone());
106    }
107    if !content.is_empty() {
108        metadata.insert("content".to_string(), Value::Array(content));
109    }
110    if !artifacts.is_empty() {
111        metadata.insert("artifacts".to_string(), Value::Array(artifacts));
112    }
113    metadata.insert("isError".to_string(), Value::Bool(result.is_error));
114
115    let text = text_parts.join("\n");
116    let output = if result.is_error {
117        ToolOutput::error(text)
118    } else {
119        ToolOutput::success(text)
120    }
121    .with_metadata(json!({ "mcp": Value::Object(metadata) }))
122    .with_images(images);
123    Ok(output)
124}
125
126struct ResourceProjection<'a> {
127    text_parts: &'a mut Vec<String>,
128    images: &'a mut Vec<Attachment>,
129    artifacts: &'a mut Vec<Value>,
130    content: &'a mut Vec<Value>,
131    decoded_bytes: &'a mut usize,
132}
133
134async fn project_resource(
135    tool_name: &str,
136    resource: &ResourceContent,
137    context: &ToolContext,
138    projection: &mut ResourceProjection<'_>,
139) -> Result<()> {
140    let mut descriptor = json!({
141        "type": "resource",
142        "uri": resource.uri,
143        "mimeType": resource.mime_type,
144        "hasText": resource.text.is_some(),
145        "hasBlob": resource.blob.is_some(),
146    });
147
148    if let Some(text) = &resource.text {
149        projection.text_parts.push(text.clone());
150        descriptor["textBytes"] = json!(text.len());
151    }
152    if let Some(blob) = &resource.blob {
153        let bytes = decode_bounded(blob, projection.decoded_bytes)?;
154        let mime_type = resource
155            .mime_type
156            .as_deref()
157            .unwrap_or("application/octet-stream");
158        let artifact =
159            materialize_artifact(tool_name, Some(&resource.uri), mime_type, &bytes, context)
160                .await?;
161        projection.text_parts.push(format!(
162            "[Resource: {}, {mime_type}, {} bytes, artifact: {}]",
163            resource.uri,
164            bytes.len(),
165            artifact.path.display()
166        ));
167        if model_image_mime_type(mime_type) {
168            projection
169                .images
170                .push(Attachment::new(bytes.clone(), mime_type.to_string()));
171        }
172        descriptor["blobBytes"] = json!(bytes.len());
173        descriptor["sha256"] = json!(artifact.sha256.clone());
174        descriptor["path"] = json!(artifact.path.clone());
175        descriptor["attachedToModel"] = json!(model_image_mime_type(mime_type));
176        projection
177            .artifacts
178            .push(artifact.value("resource", Some(&resource.uri)));
179    } else if resource.text.is_none() {
180        projection
181            .text_parts
182            .push(format!("[Resource: {}]", resource.uri));
183    }
184    projection.content.push(descriptor);
185    Ok(())
186}
187
188fn decode_bounded(data: &str, decoded_total: &mut usize) -> Result<Vec<u8>> {
189    let max_encoded = MAX_ARTIFACT_BYTES.div_ceil(3) * 4 + 4;
190    if data.len() > max_encoded {
191        return Err(anyhow!(
192            "MCP base64 artifact exceeds the {} MiB per-item limit",
193            MAX_ARTIFACT_BYTES / (1024 * 1024)
194        ));
195    }
196    let bytes = BASE64_STANDARD
197        .decode(data)
198        .context("MCP tool returned invalid base64 artifact data")?;
199    if bytes.len() > MAX_ARTIFACT_BYTES {
200        return Err(anyhow!(
201            "MCP decoded artifact exceeds the {} MiB per-item limit",
202            MAX_ARTIFACT_BYTES / (1024 * 1024)
203        ));
204    }
205    *decoded_total = decoded_total
206        .checked_add(bytes.len())
207        .ok_or_else(|| anyhow!("MCP artifact byte count overflowed"))?;
208    if *decoded_total > MAX_TOTAL_ARTIFACT_BYTES {
209        return Err(anyhow!(
210            "MCP decoded artifacts exceed the {} MiB per-call limit",
211            MAX_TOTAL_ARTIFACT_BYTES / (1024 * 1024)
212        ));
213    }
214    Ok(bytes)
215}
216
217fn validate_json_size(label: &str, value: Option<&Value>, limit: usize) -> Result<()> {
218    let Some(value) = value else {
219        return Ok(());
220    };
221    let bytes = serde_json::to_vec(value)
222        .with_context(|| format!("failed to encode MCP {label} for bounded projection"))?;
223    if bytes.len() > limit {
224        return Err(anyhow!(
225            "MCP {label} exceeds the {} MiB projection limit",
226            limit / (1024 * 1024)
227        ));
228    }
229    Ok(())
230}
231
232struct MaterializedArtifact {
233    path: PathBuf,
234    media_type: String,
235    size: usize,
236    sha256: String,
237}
238
239impl MaterializedArtifact {
240    fn value(&self, kind: &str, source_uri: Option<&str>) -> Value {
241        json!({
242            "kind": kind,
243            "path": self.path,
244            "mediaType": self.media_type,
245            "size": self.size,
246            "sha256": self.sha256,
247            "sourceUri": source_uri,
248        })
249    }
250}
251
252async fn materialize_artifact(
253    tool_name: &str,
254    _source_uri: Option<&str>,
255    media_type: &str,
256    bytes: &[u8],
257    context: &ToolContext,
258) -> Result<MaterializedArtifact> {
259    let digest = sha256::digest(bytes);
260    let session = context.session_id.as_deref().unwrap_or("anonymous");
261    let root = artifact_root(context)
262        .join(sanitize_segment(session))
263        .join(sanitize_segment(tool_name));
264    tokio::fs::create_dir_all(&root).await.with_context(|| {
265        format!(
266            "failed to create MCP artifact directory '{}'",
267            root.display()
268        )
269    })?;
270    #[cfg(unix)]
271    {
272        use std::os::unix::fs::PermissionsExt;
273        tokio::fs::set_permissions(&root, std::fs::Permissions::from_mode(0o700))
274            .await
275            .with_context(|| {
276                format!(
277                    "failed to secure MCP artifact directory '{}'",
278                    root.display()
279                )
280            })?;
281    }
282    let path = root.join(format!("{digest}.{}", extension_for_media_type(media_type)));
283    match tokio::fs::symlink_metadata(&path).await {
284        Ok(_) => verify_existing_artifact(&path, bytes, &digest).await?,
285        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
286            write_private_file(&path, bytes, &digest).await?;
287        }
288        Err(error) => {
289            return Err(error)
290                .with_context(|| format!("failed to inspect MCP artifact '{}'", path.display()));
291        }
292    }
293    prune_artifact_directory(&root, &path).await;
294    Ok(MaterializedArtifact {
295        path,
296        media_type: media_type.to_string(),
297        size: bytes.len(),
298        sha256: digest,
299    })
300}
301
302async fn prune_artifact_directory(directory: &Path, current: &Path) {
303    let Ok(mut entries) = tokio::fs::read_dir(directory).await else {
304        return;
305    };
306    let mut files = Vec::new();
307    while let Ok(Some(entry)) = entries.next_entry().await {
308        let Ok(metadata) = entry.metadata().await else {
309            continue;
310        };
311        if !metadata.is_file() {
312            continue;
313        }
314        let modified = metadata
315            .modified()
316            .unwrap_or(std::time::SystemTime::UNIX_EPOCH);
317        files.push((entry.path(), metadata.len(), modified));
318    }
319    files.sort_by_key(|(_, _, modified)| *modified);
320    let mut count = files.len();
321    let mut total = files
322        .iter()
323        .fold(0_u64, |sum, (_, size, _)| sum.saturating_add(*size));
324    for (path, size, _) in files {
325        if count <= MAX_CACHED_ARTIFACTS_PER_TOOL && total <= MAX_CACHED_ARTIFACT_BYTES_PER_TOOL {
326            break;
327        }
328        if path == current {
329            continue;
330        }
331        if tokio::fs::remove_file(&path).await.is_ok() {
332            count = count.saturating_sub(1);
333            total = total.saturating_sub(size);
334        }
335    }
336}
337
338fn artifact_root(context: &ToolContext) -> PathBuf {
339    #[cfg(test)]
340    {
341        context.workspace.join(".a3s-test-mcp-artifacts")
342    }
343    #[cfg(not(test))]
344    {
345        let _ = context;
346        dirs::cache_dir()
347            .unwrap_or_else(std::env::temp_dir)
348            .join("a3s-code")
349            .join("mcp-artifacts")
350    }
351}
352
353async fn verify_existing_artifact(path: &Path, bytes: &[u8], digest: &str) -> Result<()> {
354    let metadata = tokio::fs::symlink_metadata(path)
355        .await
356        .with_context(|| format!("failed to inspect MCP artifact '{}'", path.display()))?;
357    if !metadata.file_type().is_file() {
358        return Err(anyhow!(
359            "existing MCP artifact '{}' is not a regular file",
360            path.display()
361        ));
362    }
363    let existing = tokio::fs::read(path)
364        .await
365        .with_context(|| format!("failed to verify MCP artifact '{}'", path.display()))?;
366    if existing.len() != bytes.len() || sha256::digest(&existing) != digest {
367        return Err(anyhow!(
368            "existing MCP artifact '{}' does not match its content digest",
369            path.display()
370        ));
371    }
372    Ok(())
373}
374
375async fn write_private_file(path: &Path, bytes: &[u8], digest: &str) -> Result<()> {
376    let mut options = tokio::fs::OpenOptions::new();
377    options.write(true).create_new(true);
378    #[cfg(unix)]
379    options.mode(0o600);
380    match options.open(path).await {
381        Ok(mut file) => {
382            if let Err(error) = file.write_all(bytes).await {
383                drop(file);
384                let _ = tokio::fs::remove_file(path).await;
385                return Err(error)
386                    .with_context(|| format!("failed to write MCP artifact '{}'", path.display()));
387            }
388            if let Err(error) = file.flush().await {
389                drop(file);
390                let _ = tokio::fs::remove_file(path).await;
391                return Err(error)
392                    .with_context(|| format!("failed to flush MCP artifact '{}'", path.display()));
393            }
394            #[cfg(unix)]
395            {
396                use std::os::unix::fs::PermissionsExt;
397                tokio::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
398                    .await
399                    .with_context(|| {
400                        format!("failed to secure MCP artifact '{}'", path.display())
401                    })?;
402            }
403            Ok(())
404        }
405        Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
406            verify_existing_artifact(path, bytes, digest).await
407        }
408        Err(error) => Err(error)
409            .with_context(|| format!("failed to create MCP artifact '{}'", path.display())),
410    }
411}
412
413fn sanitize_segment(value: &str) -> String {
414    let sanitized = value
415        .chars()
416        .map(|character| {
417            if character.is_ascii_alphanumeric() || matches!(character, '-' | '_') {
418                character
419            } else {
420                '_'
421            }
422        })
423        .take(96)
424        .collect::<String>();
425    if sanitized.is_empty() {
426        "unknown".to_string()
427    } else {
428        sanitized
429    }
430}
431
432fn extension_for_media_type(media_type: &str) -> &'static str {
433    match media_type.split(';').next().unwrap_or(media_type).trim() {
434        "image/png" => "png",
435        "image/jpeg" => "jpg",
436        "image/gif" => "gif",
437        "image/webp" => "webp",
438        "image/svg+xml" => "svg",
439        "application/pdf" => "pdf",
440        "application/json" => "json",
441        "text/plain" => "txt",
442        "text/markdown" => "md",
443        _ => "bin",
444    }
445}
446
447fn model_image_mime_type(media_type: &str) -> bool {
448    matches!(
449        media_type.split(';').next().unwrap_or(media_type).trim(),
450        "image/png" | "image/jpeg" | "image/gif" | "image/webp"
451    )
452}
453
454/// Convert MCP tool result to a compact text representation.
455///
456/// Runtime tool execution uses [`project_tool_result`] so image bytes and
457/// structured data are retained. This helper remains for diagnostics and
458/// compatibility callers that explicitly need text only.
459pub fn tool_result_to_string(result: &CallToolResult) -> String {
460    let mut output = Vec::new();
461    for content in &result.content {
462        match content {
463            ToolContent::Text { text } => output.push(text.clone()),
464            ToolContent::Image { mime_type, .. } => {
465                output.push(format!("[Image: {mime_type}]"));
466            }
467            ToolContent::Resource { resource } => {
468                if let Some(text) = &resource.text {
469                    output.push(text.clone());
470                }
471                if resource.blob.is_some() || resource.text.is_none() {
472                    output.push(format!("[Resource: {}]", resource.uri));
473                }
474            }
475        }
476    }
477    if output.is_empty() {
478        if let Some(structured) = &result.structured_content {
479            return serde_json::to_string_pretty(structured)
480                .unwrap_or_else(|_| structured.to_string());
481        }
482    }
483    output.join("\n")
484}
485
486#[cfg(test)]
487mod tests {
488    use super::*;
489    use crate::mcp::protocol::CallToolResult;
490
491    #[tokio::test]
492    async fn projects_structured_content_and_materializes_an_image() {
493        let temp = tempfile::tempdir().unwrap();
494        let context =
495            ToolContext::new(temp.path().to_path_buf()).with_session_id("mcp-result-test");
496        let png = b"\x89PNG\r\n\x1a\nfixture";
497        let result = CallToolResult {
498            content: vec![ToolContent::Image {
499                data: BASE64_STANDARD.encode(png),
500                mime_type: "image/png".to_string(),
501            }],
502            structured_content: Some(json!({"text": "A3S"})),
503            is_error: false,
504            meta: Some(json!({"requestId": "one"})),
505        };
506
507        let output = project_tool_result("mcp__use_ocr__ocr_extract", &result, &context)
508            .await
509            .unwrap();
510        assert!(output.success);
511        assert_eq!(output.images.len(), 1);
512        assert_eq!(output.images[0].data, png);
513        assert_eq!(
514            output.metadata.as_ref().unwrap()["mcp"]["structuredContent"]["text"],
515            "A3S"
516        );
517        assert_eq!(
518            output.metadata.as_ref().unwrap()["mcp"]["_meta"]["requestId"],
519            "one"
520        );
521        let path = output.metadata.as_ref().unwrap()["mcp"]["artifacts"][0]["path"]
522            .as_str()
523            .unwrap();
524        assert_eq!(std::fs::read(path).unwrap(), png);
525    }
526
527    #[tokio::test]
528    async fn rejects_oversized_or_invalid_base64_without_losing_error_status() {
529        let temp = tempfile::tempdir().unwrap();
530        let context = ToolContext::new(temp.path().to_path_buf());
531        let invalid = CallToolResult {
532            content: vec![ToolContent::Image {
533                data: "***".to_string(),
534                mime_type: "image/png".to_string(),
535            }],
536            ..CallToolResult::default()
537        };
538        assert!(project_tool_result("mcp__test__image", &invalid, &context)
539            .await
540            .is_err());
541    }
542
543    #[tokio::test]
544    async fn rejects_a_preexisting_artifact_that_does_not_match_its_digest() {
545        let temp = tempfile::tempdir().unwrap();
546        let context = ToolContext::new(temp.path().to_path_buf());
547        let png = b"\x89PNG\r\n\x1a\nexpected";
548        let digest = sha256::digest(png);
549        let root = artifact_root(&context)
550            .join("anonymous")
551            .join("mcp__test__image");
552        tokio::fs::create_dir_all(&root).await.unwrap();
553        tokio::fs::write(root.join(format!("{digest}.png")), b"different")
554            .await
555            .unwrap();
556        let result = CallToolResult {
557            content: vec![ToolContent::Image {
558                data: BASE64_STANDARD.encode(png),
559                mime_type: "image/png".to_string(),
560            }],
561            ..CallToolResult::default()
562        };
563
564        let error = project_tool_result("mcp__test__image", &result, &context)
565            .await
566            .expect_err("content-addressed artifacts must reject a conflicting file");
567        assert!(
568            error
569                .to_string()
570                .contains("does not match its content digest"),
571            "{error:#}"
572        );
573    }
574}
a3s_code_core/mcp/result.rs

a3s_code_core/mcp/
result.rs
