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

a3s_code_core/mcp/
result.rs
