outrig_cli/
rig_tool.rs

1//! MCP -> Rig dynamic-tool adapter.
2//!
3//! [`McpToolAdapter`] wraps an MCP-discovered tool as a [`rig::tool::ToolDyn`]
4//! so the agent loop can hand it directly to a Rig `Agent`. The original
5//! tool name (used on the MCP wire) lives next to the sanitized
6//! `<server>__<tool>` name (the form the LLM sees, produced by
7//! [`outrig::sanitize_tool_name`]), so a single adapter knows both ends of
8//! the dispatch.
9
10use std::sync::Arc;
11
12use rig::completion::ToolDefinition;
13use rig::tool::{ToolDyn, ToolError};
14use rig::wasm_compat::WasmBoxedFuture;
15use serde_json::Value;
16
17use crate::error::Result;
18use outrig::{McpClient, McpToolResult};
19
20/// A Rig dynamic-tool view of one MCP-discovered tool.
21///
22/// `openai_name` is what the LLM sees and what
23/// [`ToolDyn::name`] returns. `mcp_tool_name` is the original name from
24/// `tools/list`, used on the wire when dispatching back into the MCP server.
25/// `client` is shared (via `Arc`) so multiple adapters fronting the same MCP
26/// server reuse one connection.
27#[derive(Debug, Clone)]
28pub struct McpToolAdapter {
29    pub openai_name: String,
30    pub mcp_tool_name: String,
31    pub description: String,
32    pub input_schema: Value,
33    pub result_cap_bytes: usize,
34    pub client: Arc<McpClient>,
35}
36
37impl McpToolAdapter {
38    /// Build one adapter per tool the server advertises. The server's local
39    /// name (from [`McpClient::name`]) becomes the prefix.
40    pub async fn from_client_tools(
41        client: Arc<McpClient>,
42        result_cap_bytes: usize,
43    ) -> Result<Vec<McpToolAdapter>> {
44        let tools = client.list_tools().await?;
45        let server_name = client.name().to_string();
46        Ok(tools
47            .into_iter()
48            .map(|t| McpToolAdapter {
49                openai_name: outrig::sanitize_tool_name(&server_name, &t.name),
50                mcp_tool_name: t.name,
51                description: t.description.unwrap_or_default(),
52                input_schema: t.input_schema,
53                result_cap_bytes,
54                client: client.clone(),
55            })
56            .collect())
57    }
58}
59
60/// Error wrapper that carries an MCP tool-call failure into rig's
61/// `ToolError::ToolCallError(Box<dyn Error>)` channel without leaking our
62/// concrete `OutrigError` type into rig's API.
63#[derive(Debug, thiserror::Error)]
64#[error("{0}")]
65struct McpAdapterError(String);
66
67pub fn truncate_for_llm(result: &str, max: usize) -> String {
68    if result.is_empty() || result.len() <= max {
69        return result.to_string();
70    }
71    if max == 0 {
72        return String::new();
73    }
74
75    let original_len = result.len();
76    let mut cut = max.saturating_sub(truncation_marker(original_len, max, 0).len());
77    loop {
78        cut = floor_char_boundary(result, cut.min(result.len()));
79        let marker = truncation_marker(original_len, max, cut);
80        if marker.len() >= max {
81            return truncate_marker(&marker, max);
82        }
83
84        let content_budget = max - marker.len();
85        if cut <= content_budget {
86            let mut truncated = String::with_capacity(cut + marker.len());
87            truncated.push_str(&result[..cut]);
88            truncated.push_str(&marker);
89            debug_assert!(truncated.len() <= max);
90            return truncated;
91        }
92
93        cut = content_budget;
94    }
95}
96
97fn adapt_tool_result(result: McpToolResult, max: usize) -> std::result::Result<String, ToolError> {
98    let content_text = truncate_for_llm(&result.content_text, max);
99    if result.is_error {
100        Err(ToolError::ToolCallError(Box::new(McpAdapterError(
101            content_text,
102        ))))
103    } else {
104        Ok(content_text)
105    }
106}
107
108fn truncation_marker(original_len: usize, max: usize, kept: usize) -> String {
109    let dropped = original_len.saturating_sub(kept);
110    format!(
111        concat!(
112            "\n\n[outrig: tool result truncated]\n",
113            "  original size: {original_len} bytes\n",
114            "  max:           {max} bytes\n",
115            "  kept:          first {kept} bytes; trailing {dropped} bytes dropped.\n\n",
116            "  This tool result was larger than the configured max. Your next call\n",
117            "  should narrow the query: use head/tail/grep/--max-count, scope a\n",
118            "  directory or line range, or call a more specific tool. Re-running\n",
119            "  the same call will produce the same truncation.",
120        ),
121        original_len = original_len,
122        max = max,
123        kept = kept,
124        dropped = dropped,
125    )
126}
127
128fn truncate_marker(marker: &str, max: usize) -> String {
129    let cut = floor_char_boundary(marker, max.min(marker.len()));
130    marker[..cut].to_string()
131}
132
133fn floor_char_boundary(s: &str, mut index: usize) -> usize {
134    while !s.is_char_boundary(index) {
135        index -= 1;
136    }
137    index
138}
139
140impl ToolDyn for McpToolAdapter {
141    fn name(&self) -> String {
142        self.openai_name.clone()
143    }
144
145    fn definition(&self, _prompt: String) -> WasmBoxedFuture<'_, ToolDefinition> {
146        Box::pin(async move {
147            ToolDefinition {
148                name: self.openai_name.clone(),
149                description: self.description.clone(),
150                parameters: self.input_schema.clone(),
151            }
152        })
153    }
154
155    fn call(&self, args: String) -> WasmBoxedFuture<'_, std::result::Result<String, ToolError>> {
156        Box::pin(async move {
157            let parsed: Value = if args.is_empty() {
158                Value::Null
159            } else {
160                serde_json::from_str(&args)?
161            };
162
163            let result = self
164                .client
165                .call_tool(&self.mcp_tool_name, parsed)
166                .await
167                .map_err(|e| ToolError::ToolCallError(Box::new(McpAdapterError(e.to_string()))))?;
168
169            adapt_tool_result(result, self.result_cap_bytes)
170        })
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    const MAX: usize = 1024;
179
180    #[test]
181    fn truncate_for_llm_leaves_empty_and_under_cap_results_unchanged() {
182        assert_eq!(truncate_for_llm("", MAX), "");
183        assert_eq!(truncate_for_llm("short", MAX), "short");
184    }
185
186    #[test]
187    fn truncate_for_llm_leaves_exact_cap_result_unchanged() {
188        let input = "a".repeat(MAX);
189
190        let output = truncate_for_llm(&input, MAX);
191
192        assert_eq!(output, input);
193    }
194
195    #[test]
196    fn truncate_for_llm_caps_one_byte_over_with_marker() {
197        let input = "a".repeat(MAX + 1);
198
199        let output = truncate_for_llm(&input, MAX);
200
201        assert!(output.len() <= MAX, "output len: {}", output.len());
202        assert!(output.contains("[outrig: tool result truncated]"));
203        assert!(output.contains("original size: 1025 bytes"));
204        assert!(output.contains("max:           1024 bytes"));
205        assert!(output.ends_with("produce the same truncation."));
206    }
207
208    #[test]
209    fn truncate_for_llm_caps_large_result_with_original_size_and_hint() {
210        let input = "x".repeat(5 * 1024 * 1024);
211
212        let output = truncate_for_llm(&input, 4096);
213
214        assert!(output.len() <= 4096, "output len: {}", output.len());
215        assert!(output.contains("original size: 5242880 bytes"));
216        assert!(output.contains("max:           4096 bytes"));
217        assert!(output.contains("should narrow the query"));
218    }
219
220    #[test]
221    fn truncate_for_llm_keeps_valid_utf8_at_boundary() {
222        let input = format!("{}{}", "a".repeat(900), "🙂".repeat(200));
223
224        let output = truncate_for_llm(&input, MAX);
225
226        assert!(output.len() <= MAX, "output len: {}", output.len());
227        assert!(output.is_char_boundary(output.len()));
228        assert!(output.contains("[outrig: tool result truncated]"));
229    }
230
231    #[test]
232    fn adapt_tool_result_truncates_success_content() {
233        let result = McpToolResult {
234            content_text: "a".repeat(MAX + 1),
235            is_error: false,
236        };
237
238        let output = adapt_tool_result(result, MAX).expect("success result");
239
240        assert!(output.len() <= MAX, "output len: {}", output.len());
241        assert!(output.contains("[outrig: tool result truncated]"));
242    }
243
244    #[test]
245    fn adapt_tool_result_truncates_error_content() {
246        let result = McpToolResult {
247            content_text: "e".repeat(MAX + 1),
248            is_error: true,
249        };
250
251        let err = adapt_tool_result(result, MAX).expect_err("error result");
252        let ToolError::ToolCallError(source) = err else {
253            panic!("expected tool-call error");
254        };
255        let msg = source.to_string();
256
257        assert!(msg.len() <= MAX, "error len: {}", msg.len());
258        assert!(msg.contains("[outrig: tool result truncated]"));
259    }
260}
outrig_cli/rig_tool.rs

outrig_cli/
rig_tool.rs
