Skip to main content

dnslib/mcp/
helpers.rs

1//! MCP helper functions — error and result formatting.
2
3use std::future::Future;
4
5use miette::Diagnostic;
6use rmcp::{ErrorData as McpError, model::*};
7use serde_json::json;
8use tracing::Instrument;
9
10use crate::core::error::{Error, Result};
11
12/// Build the structured JSON `data` payload that accompanies an MCP error.
13///
14/// Includes the miette diagnostic code, help text, and a `retryable` hint
15/// so AI agent clients can decide whether to retry the call without parsing
16/// the human-readable message.
17fn error_data(e: &Error) -> serde_json::Value {
18    let code = e.code().map(|c| c.to_string());
19    let help = e.help().map(|h| h.to_string());
20    json!({
21        "code": code,
22        "help": help,
23        "retryable": is_retryable(e),
24    })
25}
26
27/// True when the error represents a transient failure the client may retry.
28fn is_retryable(e: &Error) -> bool {
29    matches!(
30        e,
31        Error::Network(_) | Error::Mcp { .. } | Error::InvalidJson(_)
32    )
33}
34
35/// Build the human-readable MCP error message: `[code] <display>\n\nhelp: <help>`.
36fn error_message(e: &Error) -> String {
37    let mut msg = e.to_string();
38    if let Some(code) = e.code() {
39        msg = format!("[{code}] {msg}");
40    }
41    if let Some(help) = e.help() {
42        msg = format!("{msg}\n\nhelp: {help}");
43    }
44    msg
45}
46
47/// Convert a crate `Error` into an MCP protocol error with structured data.
48///
49/// The `data` payload exposes the diagnostic code, help text, and a
50/// `retryable` flag for client-side handling.
51pub fn mcp_err(e: Error) -> McpError {
52    let msg = error_message(&e);
53    let data = error_data(&e);
54    McpError::internal_error(msg, Some(data))
55}
56
57/// Wrap a JSON value into a successful MCP call result.
58pub fn json_result(value: serde_json::Value) -> CallToolResult {
59    // `unwrap_or_else` is a safe fallback, not a panic: any `serde_json::Value`
60    // round-trips through `to_string()`, so pretty-printing failure degrades
61    // gracefully instead of bringing the MCP transport down.
62    CallToolResult::success(vec![Content::text(
63        serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string()),
64    )])
65}
66
67/// Wrap a plain text string into a successful MCP call result.
68pub fn text_result(s: String) -> CallToolResult {
69    CallToolResult::success(vec![Content::text(s)])
70}
71
72/// Build a tool-level error result (`isError: true`) carrying both the
73/// human-readable message and the structured data payload.
74fn error_tool_result(e: &Error) -> CallToolResult {
75    let msg = error_message(e);
76    let data = error_data(e);
77    let body = json!({
78        "error": msg,
79        "data": data,
80    });
81    let text = serde_json::to_string_pretty(&body).unwrap_or_else(|_| body.to_string());
82    CallToolResult::error(vec![Content::text(text)])
83}
84
85/// The outcome label recorded on a tool's span.
86fn outcome_for(e: &Error) -> &'static str {
87    match e {
88        Error::PolicyViolation { .. } | Error::Forbidden { .. } => "permission_denied",
89        _ => "error",
90    }
91}
92
93/// Run a permission check (or `.and(...)`-chained checks), then a future,
94/// and wrap the JSON result. Errors are surfaced as a `CallToolResult` with
95/// `is_error: true` (per MCP spec) rather than as a transport-level `McpError`,
96/// so an AI agent client can recover without aborting its loop.
97pub async fn run_json<F>(tool: &str, check: Result<()>, fut: F) -> CallToolResult
98where
99    F: Future<Output = Result<serde_json::Value>>,
100{
101    let span = tracing::info_span!("mcp_tool", tool = tool, outcome = tracing::field::Empty);
102    async move {
103        match check {
104            Err(e) => {
105                tracing::warn!(tool = tool, error = %e, "MCP tool error");
106                tracing::Span::current().record("outcome", outcome_for(&e));
107                error_tool_result(&e)
108            }
109            Ok(()) => match fut.await {
110                Ok(v) => {
111                    tracing::Span::current().record("outcome", "success");
112                    json_result(v)
113                }
114                Err(e) => {
115                    tracing::warn!(tool = tool, error = %e, "MCP tool error");
116                    tracing::Span::current().record("outcome", outcome_for(&e));
117                    error_tool_result(&e)
118                }
119            },
120        }
121    }
122    .instrument(span)
123    .await
124}
125
126/// Like [`run_json`] but wraps a plain text result.
127pub async fn run_text<F>(tool: &str, check: Result<()>, fut: F) -> CallToolResult
128where
129    F: Future<Output = Result<String>>,
130{
131    let span = tracing::info_span!("mcp_tool", tool = tool, outcome = tracing::field::Empty);
132    async move {
133        match check {
134            Err(e) => {
135                tracing::warn!(tool = tool, error = %e, "MCP tool error");
136                tracing::Span::current().record("outcome", outcome_for(&e));
137                error_tool_result(&e)
138            }
139            Ok(()) => match fut.await {
140                Ok(s) => {
141                    tracing::Span::current().record("outcome", "success");
142                    text_result(s)
143                }
144                Err(e) => {
145                    tracing::warn!(tool = tool, error = %e, "MCP tool error");
146                    tracing::Span::current().record("outcome", outcome_for(&e));
147                    error_tool_result(&e)
148                }
149            },
150        }
151    }
152    .instrument(span)
153    .await
154}