aether_core/mcp/
tool_bridge.rs

1use std::path::{Path, PathBuf};
2
3use crate::events::{TaskOutcome, TaskOutcomeState};
4use mcp_utils::{
5    client::{CallToolError, SERVERNAME_DELIMITER},
6    display_meta::ToolResultMeta,
7};
8use rmcp::model::{CallToolRequestParams, CallToolResult, Task};
9use serde_json;
10
11use llm::{ToolCallError, ToolCallRequest, ToolCallResult};
12
13/// Maximum bytes for a tool result before spilling to disk.
14/// ~50K tokens at ~4 bytes/token.
15const TOOL_RESULT_MAX_BYTES: usize = 200_000;
16
17/// Size of the head preview included inline when a result spills to disk.
18const SPILLOVER_PREVIEW_BYTES: usize = 10_000;
19
20/// Convert a `ToolCallRequest` to `rmcp::CallToolRequestParams`
21pub fn tool_call_request_to_mcp(request: &ToolCallRequest) -> Result<CallToolRequestParams, String> {
22    let tool_name = request
23        .name
24        .split_once(SERVERNAME_DELIMITER)
25        .map_or_else(|| request.name.clone(), |(_, tool_name)| tool_name.to_string());
26
27    // Parse arguments from JSON string
28    let arguments = serde_json::from_str::<serde_json::Value>(&request.arguments)
29        .map_err(|e| format!("Invalid tool arguments: {e}"))?
30        .as_object()
31        .cloned();
32
33    let mut params = CallToolRequestParams::new(tool_name);
34    if let Some(args) = arguments {
35        params = params.with_arguments(args);
36    }
37    Ok(params)
38}
39
40/// Convert an rmcp `CallToolResult` and request to `ToolCallResult` or `ToolCallError`,
41/// extracting any `_meta` metadata from structured content.
42pub fn mcp_result_to_tool_call_result(
43    request: &ToolCallRequest,
44    mcp_result: rmcp::model::CallToolResult,
45) -> Result<(ToolCallResult, Option<ToolResultMeta>), ToolCallError> {
46    if mcp_result.is_error.unwrap_or(false) {
47        let error_msg = mcp_result.content.first().map_or_else(
48            || "Unknown error".to_string(),
49            |content| {
50                content.as_text().map_or_else(
51                    || serde_json::to_string(content).unwrap_or_else(|_| "Unknown error".to_string()),
52                    |text| text.text.clone(),
53                )
54            },
55        );
56        Err(ToolCallError {
57            id: request.id.clone(),
58            name: request.name.clone(),
59            arguments: Some(request.arguments.clone()),
60            error: format!("Tool execution error: {error_msg}"),
61        })
62    } else {
63        let (result_value, result_meta) = extract_result_and_meta(mcp_result.structured_content, &mcp_result.content);
64        // YAML is ~18% more token-efficient than JSON for LLM consumption
65        let yaml = serde_yml::to_string(&result_value).unwrap_or_else(|_| result_value.to_string());
66        let result_str = maybe_spillover(&request.id, yaml, TOOL_RESULT_MAX_BYTES, &spillover_dir());
67        Ok((
68            ToolCallResult {
69                id: request.id.clone(),
70                name: request.name.clone(),
71                arguments: request.arguments.clone(),
72                result: result_str,
73            },
74            result_meta,
75        ))
76    }
77}
78
79pub fn map_task_result_to_outcome(
80    request: ToolCallRequest,
81    task: Task,
82    outcome: Result<CallToolResult, CallToolError>,
83) -> TaskOutcome {
84    let state = match convert_tool_result(&request, outcome) {
85        Ok((result, result_meta)) => TaskOutcomeState::Completed { result, result_meta },
86        Err(error) => TaskOutcomeState::Failed { error },
87    };
88    TaskOutcome { request, task_id: task.task_id, state }
89}
90
91pub fn convert_tool_result(
92    request: &ToolCallRequest,
93    outcome: Result<CallToolResult, CallToolError>,
94) -> Result<(ToolCallResult, Option<ToolResultMeta>), ToolCallError> {
95    outcome
96        .map_err(|error| ToolCallError::from_request(request, error.to_string()))
97        .and_then(|mcp_result| mcp_result_to_tool_call_result(request, mcp_result))
98}
99
100fn spillover_dir() -> PathBuf {
101    std::env::temp_dir().join("aether-tool-output")
102}
103
104/// If `result` exceeds `max_bytes`, write the full output to disk and return
105/// a head preview with a pointer to the file. Otherwise return unchanged.
106fn maybe_spillover(tool_id: &str, result: String, max_bytes: usize, dir: &Path) -> String {
107    if result.len() <= max_bytes {
108        return result;
109    }
110
111    if let Err(e) = std::fs::create_dir_all(dir) {
112        tracing::warn!("Failed to create tool-output dir: {e}");
113        return result;
114    }
115
116    let file_path = dir.join(format!("{tool_id}.txt"));
117
118    if let Err(e) = std::fs::write(&file_path, &result) {
119        tracing::warn!("Failed to write spillover file: {e}");
120        return result;
121    }
122
123    let preview_end = result.floor_char_boundary(SPILLOVER_PREVIEW_BYTES);
124    let preview = &result[..preview_end];
125    let total_bytes = result.len();
126
127    format!(
128        "<preview>\n{preview}\n</preview>\n\n[Tool result too large ({total_bytes} bytes). Full output saved to {path}. Use grep, read, or tail to explore the full result.]",
129        path = file_path.display()
130    )
131}
132
133fn extract_result_and_meta(
134    structured_content: Option<serde_json::Value>,
135    content: &[rmcp::model::ContentBlock],
136) -> (serde_json::Value, Option<ToolResultMeta>) {
137    if let Some(mut val) = structured_content {
138        let result_meta = extract_result_meta(&mut val);
139        (val, result_meta)
140    } else {
141        let fallback = content.first().map_or_else(
142            || serde_json::Value::String("No result".to_string()),
143            |c| serde_json::to_value(c).unwrap_or(serde_json::Value::String("Serialization error".to_string())),
144        );
145        (fallback, None)
146    }
147}
148
149fn extract_result_meta(value: &mut serde_json::Value) -> Option<ToolResultMeta> {
150    let obj = value.as_object_mut()?;
151    let parsed: ToolResultMeta = {
152        let meta = obj.get("_meta")?.as_object()?;
153        serde_json::from_value(serde_json::Value::Object(meta.clone())).ok()?
154    };
155
156    let meta_empty = {
157        let meta = obj.get_mut("_meta")?.as_object_mut()?;
158        for key in ["display", "file_diff", "plan"] {
159            meta.remove(key);
160        }
161        meta.is_empty()
162    };
163
164    if meta_empty {
165        obj.remove("_meta");
166    }
167
168    Some(parsed)
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174    use mcp_utils::display_meta::PlanMetaStatus;
175    use rmcp::model::{CallToolResult as McpCallToolResult, ContentBlock};
176    use serde::Serialize;
177    use serde_json::json;
178
179    fn req() -> ToolCallRequest {
180        ToolCallRequest { id: "call_123".into(), name: "test_tool".into(), arguments: "{}".into() }
181    }
182
183    fn call_structured(structured: serde_json::Value) -> (ToolCallResult, Option<ToolResultMeta>) {
184        let mut mcp = McpCallToolResult::structured(structured);
185        mcp.content = vec![];
186        mcp_result_to_tool_call_result(&req(), mcp).unwrap()
187    }
188
189    fn extract_preview(result: &str) -> &str {
190        let start = result.find("<preview>\n").unwrap() + "<preview>\n".len();
191        let end = result.find("\n</preview>").unwrap();
192        &result[start..end]
193    }
194
195    #[test]
196    fn test_extracts_and_strips_meta() {
197        let structured = json!({
198            "status": "success", "file_path": "/test/file.rs",
199            "_meta": { "display": { "title": "Read file", "value": "file.rs, 50 lines" } }
200        });
201        let mut mcp = McpCallToolResult::structured(structured);
202        mcp.content = vec![ContentBlock::text("plain text fallback")];
203        let (result, meta) = mcp_result_to_tool_call_result(&req(), mcp).unwrap();
204
205        assert!(!result.result.contains("_meta"));
206        assert!(result.result.contains("success"));
207        let rm = meta.expect("meta should be present");
208        assert_eq!(rm.display.title, "Read file");
209        assert_eq!(rm.display.value, "file.rs, 50 lines");
210        assert!(rm.file_diff.is_none());
211    }
212
213    #[test]
214    fn test_extracts_meta_with_file_diff() {
215        let (result, meta) = call_structured(json!({
216            "status": "success",
217            "_meta": {
218                "display": { "title": "Edit file", "value": "main.rs" },
219                "file_diff": { "path": "/tmp/main.rs", "old_text": "old content", "new_text": "new content" }
220            }
221        }));
222        assert!(!result.result.contains("_meta"));
223        let rm = meta.expect("meta should be present");
224        assert_eq!(rm.display.title, "Edit file");
225        let fd = rm.file_diff.expect("file_diff should be present");
226        assert_eq!(fd.path, "/tmp/main.rs");
227        assert_eq!(fd.old_text.as_deref(), Some("old content"));
228        assert_eq!(fd.new_text, "new content");
229    }
230
231    #[test]
232    fn test_extracts_known_meta_and_preserves_unknown_meta_keys() {
233        let (result, meta) = call_structured(json!({
234            "status": "success",
235            "_meta": {
236                "display": { "title": "Edit file", "value": "main.rs" },
237                "file_diff": { "path": "/tmp/main.rs", "old_text": "old", "new_text": "new" },
238                "trace_id": "trace-123", "duration_ms": 18
239            }
240        }));
241        let rm = meta.expect("meta should be present");
242        assert_eq!(rm.display.title, "Edit file");
243        assert!(rm.file_diff.is_some());
244        for absent in ["display:", "file_diff:"] {
245            assert!(!result.result.contains(absent));
246        }
247        for present in ["trace_id:", "trace-123", "duration_ms:", "18"] {
248            assert!(result.result.contains(present));
249        }
250    }
251
252    #[test]
253    fn test_malformed_meta_returns_none() {
254        let (result, meta) = call_structured(json!({
255            "status": "success",
256            "_meta": { "display": "not a valid ToolDisplayMeta" }
257        }));
258        assert!(meta.is_none());
259        assert!(result.result.contains("not a valid ToolDisplayMeta"));
260    }
261
262    #[test]
263    fn test_no_meta_passes_through_unchanged() {
264        let (result, meta) = call_structured(json!({"status": "success", "data": "hello"}));
265        assert!(result.result.contains("success"));
266        assert!(result.result.contains("hello"));
267        assert!(meta.is_none());
268    }
269
270    #[test]
271    fn test_tool_call_result_falls_back_to_content() {
272        let mcp = McpCallToolResult::success(vec![ContentBlock::text("plain text result")]);
273        let (result, meta) = mcp_result_to_tool_call_result(&req(), mcp).unwrap();
274        assert!(result.result.contains("plain text result"));
275        assert!(meta.is_none());
276    }
277
278    #[test]
279    fn test_extracts_meta_with_plan() {
280        let (result, meta) = call_structured(json!({
281            "status": "success",
282            "_meta": {
283                "display": { "title": "Todo", "value": "Research AI agents" },
284                "plan": { "entries": [
285                    { "content": "Research AI agents", "status": "in_progress" },
286                    { "content": "Write tests", "status": "pending" }
287                ]}
288            }
289        }));
290        assert!(!result.result.contains("_meta"));
291        let rm = meta.expect("meta should be present");
292        assert_eq!(rm.display.title, "Todo");
293        let plan = rm.plan.expect("plan should be present");
294        assert_eq!(plan.entries.len(), 2);
295        assert_eq!(plan.entries[0].content, "Research AI agents");
296        assert_eq!(plan.entries[0].status, PlanMetaStatus::InProgress);
297        assert_eq!(plan.entries[1].status, PlanMetaStatus::Pending);
298    }
299
300    /// Regression: verifies `#[serde(rename = "_meta")]` preserves the key under camelCase,
301    /// and that omitting the rename breaks extraction.
302    #[test]
303    fn test_meta_camel_case_serde_round_trip() {
304        #[derive(Serialize)]
305        #[serde(rename_all = "camelCase")]
306        struct GoodResult {
307            file_path: String,
308            total_lines: usize,
309            #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
310            _meta: Option<serde_json::Value>,
311        }
312
313        #[derive(Serialize)]
314        #[serde(rename_all = "camelCase")]
315        struct BrokenResult {
316            file_path: String,
317            #[serde(skip_serializing_if = "Option::is_none")]
318            _meta: Option<serde_json::Value>,
319        }
320
321        let display_meta = json!({
322            "display": { "title": "Read file", "value": "file.rs, 50 lines" }
323        });
324
325        let good = serde_json::to_value(&GoodResult {
326            file_path: "/test/file.rs".into(),
327            total_lines: 50,
328            _meta: Some(display_meta.clone()),
329        })
330        .unwrap();
331        assert!(good.get("_meta").is_some(), "expected `_meta` key, got: {good}");
332        let (stripped, meta) = extract_result_and_meta(Some(good), &[]);
333        let rm = meta.expect("meta should be extracted");
334        assert_eq!(rm.display.title, "Read file");
335        assert_eq!(rm.display.value, "file.rs, 50 lines");
336        assert!(stripped.get("_meta").is_none());
337
338        let broken =
339            serde_json::to_value(&BrokenResult { file_path: "/test/file.rs".into(), _meta: Some(display_meta) })
340                .unwrap();
341        assert!(broken.get("_meta").is_none(), "should be mangled by camelCase");
342        assert!(broken.get("meta").is_some());
343        let (_, meta) = extract_result_and_meta(Some(broken), &[]);
344        assert!(meta.is_none(), "extraction should fail when _meta is mangled");
345    }
346
347    #[test]
348    fn test_tool_call_result_handles_text_error_without_sdk_debug_output() {
349        let mcp = McpCallToolResult::error(vec![ContentBlock::text("Error: file not found")]);
350        let err = mcp_result_to_tool_call_result(&req(), mcp).unwrap_err();
351        assert_eq!(err.error, "Tool execution error: Error: file not found");
352    }
353
354    #[test]
355    fn test_tool_call_result_serializes_non_text_error_content() {
356        let image = serde_json::from_value(serde_json::json!({
357            "type": "image",
358            "data": "aW1hZ2U=",
359            "mimeType": "image/png"
360        }))
361        .unwrap();
362        let mcp = McpCallToolResult::error(vec![image]);
363        let err = mcp_result_to_tool_call_result(&req(), mcp).unwrap_err();
364        assert_eq!(err.error, r#"Tool execution error: {"type":"image","data":"aW1hZ2U=","mimeType":"image/png"}"#);
365    }
366
367    #[test]
368    fn test_result_is_yaml_format() {
369        let (result, _) = call_structured(json!({
370            "status": "success",
371            "files": [{"name": "Cargo.toml", "path": "./Cargo.toml"}, {"name": "src", "path": "./src"}],
372            "totalCount": 2
373        }));
374        let r = &result.result;
375        for expected in ["status: success", "totalCount: 2", "- name:"] {
376            assert!(r.contains(expected), "expected '{expected}' in YAML, got: {r}");
377        }
378        assert!(!r.starts_with('{'), "expected YAML, not JSON: {r}");
379    }
380
381    #[test]
382    fn test_serde_yml_produces_yaml_not_json() {
383        let yaml = serde_yml::to_string(&json!({"key": "value"})).unwrap();
384        assert!(yaml.contains("key:") && yaml.contains("value") && !yaml.starts_with('{'));
385    }
386
387    #[test]
388    fn test_spillover_small_input_unchanged() {
389        let dir = tempfile::tempdir().unwrap();
390        let input = "hello world".to_string();
391        assert_eq!(maybe_spillover("id", input.clone(), 1000, dir.path()), input);
392    }
393
394    #[test]
395    fn test_spillover_large_input_writes_file() {
396        let dir = tempfile::tempdir().unwrap();
397        let large = "x".repeat(5000);
398        let result = maybe_spillover("test_large", large.clone(), 1000, dir.path());
399        for expected in ["<preview>", "</preview>", "Tool result too large", "5000 bytes", "test_large.txt"] {
400            assert!(result.contains(expected), "missing '{expected}' in: {result}");
401        }
402        let on_disk = std::fs::read_to_string(dir.path().join("test_large.txt")).unwrap();
403        assert_eq!(on_disk, large);
404    }
405
406    #[test]
407    fn test_spillover_preview_content() {
408        let dir = tempfile::tempdir().unwrap();
409        let large = format!("HEAD_{}{}", "z".repeat(SPILLOVER_PREVIEW_BYTES + 5000), "TAIL");
410        let result = maybe_spillover("id", large, 1000, dir.path());
411        assert!(result.contains("HEAD_"));
412        assert!(!result.contains("TAIL"));
413    }
414
415    #[test]
416    fn test_spillover_preserves_utf8_boundaries() {
417        let dir = tempfile::tempdir().unwrap();
418        let large = format!("{}{}", "\u{1F600}".repeat(300), "a".repeat(5000));
419        let result = maybe_spillover("id", large, 100, dir.path());
420        assert!(extract_preview(&result).chars().count() > 0);
421    }
422
423    #[test]
424    fn test_mcp_result_spills_large_output() {
425        let request =
426            ToolCallRequest { id: "spill_integration".into(), name: "big_tool".into(), arguments: "{}".into() };
427        let mut mcp = McpCallToolResult::structured(json!({"data": "x".repeat(TOOL_RESULT_MAX_BYTES + 1000)}));
428        mcp.content = vec![];
429        let (result, _) = mcp_result_to_tool_call_result(&request, mcp).unwrap();
430        for expected in ["<preview>", "Tool result too large", "spill_integration.txt"] {
431            assert!(result.result.contains(expected));
432        }
433    }
434}
aether_core/mcp/tool_bridge.rs

aether_core/mcp/
tool_bridge.rs
