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

a3s_code_core/mcp/
result.rs
