Skip to main content

aether_core/mcp/
tool_bridge.rs

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