zai_rs/mcp/
responses.rs

1//! Typed responses and MCP content decoding.
2
3use std::{collections::BTreeMap, fmt, ops::Deref};
4
5use rmcp::model::CallToolResult;
6use serde::{Deserialize, Serialize, de::DeserializeOwned};
7use serde_json::Value;
8
9use crate::{
10    ZaiResult,
11    client::error::{ZaiError, codes},
12};
13
14/// Text returned by repository and vision tools.
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16pub struct McpTextResponse {
17    text: String,
18}
19
20impl McpTextResponse {
21    pub(crate) fn new(text: String) -> Self {
22        Self { text }
23    }
24
25    /// Borrow the returned text.
26    pub fn text(&self) -> &str {
27        &self.text
28    }
29
30    /// Consume the response and return its text.
31    pub fn into_text(self) -> String {
32        self.text
33    }
34}
35
36impl Deref for McpTextResponse {
37    type Target = str;
38
39    fn deref(&self) -> &Self::Target {
40        self.text()
41    }
42}
43
44impl fmt::Display for McpTextResponse {
45    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
46        formatter.write_str(&self.text)
47    }
48}
49
50/// One result returned by web search.
51#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
52pub struct WebSearchResult {
53    /// Result title.
54    pub title: String,
55    /// Result URL.
56    pub link: String,
57    /// Search-generated page summary.
58    pub content: String,
59    /// Upstream reference identifier, when supplied.
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub refer: Option<String>,
62    /// Source site or publisher name, when supplied.
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub media: Option<String>,
65    /// Source-site icon URL, when supplied.
66    #[serde(skip_serializing_if = "Option::is_none")]
67    pub icon: Option<String>,
68    /// Publication date in the upstream service's original representation.
69    #[serde(skip_serializing_if = "Option::is_none")]
70    pub publish_date: Option<String>,
71    /// Additional fields returned by newer versions of the upstream service.
72    #[serde(flatten)]
73    pub extra: BTreeMap<String, Value>,
74}
75
76/// Typed response from `web_search_prime`.
77#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
78pub struct WebSearchResponse {
79    /// Search results in upstream ranking order.
80    pub results: Vec<WebSearchResult>,
81}
82
83impl WebSearchResponse {
84    /// Consume the response and return the result list.
85    pub fn into_results(self) -> Vec<WebSearchResult> {
86        self.results
87    }
88}
89
90impl Deref for WebSearchResponse {
91    type Target = [WebSearchResult];
92
93    fn deref(&self) -> &Self::Target {
94        &self.results
95    }
96}
97
98/// Typed response from `webReader`.
99#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
100pub struct WebReaderResponse {
101    /// Page title, or an empty string when the upstream response omits it.
102    #[serde(default)]
103    pub title: String,
104    /// Page description, or an empty string when unavailable.
105    #[serde(default)]
106    pub description: String,
107    /// Canonical or requested page URL reported by the reader.
108    #[serde(default)]
109    pub url: String,
110    /// Extracted page content in the requested format.
111    #[serde(default)]
112    pub content: String,
113    /// Image labels mapped to their source URLs.
114    #[serde(default)]
115    pub images: BTreeMap<String, String>,
116    /// Reader metadata whose shape is controlled by the upstream service.
117    #[serde(default)]
118    pub metadata: BTreeMap<String, Value>,
119    /// Additional externally sourced reader data.
120    #[serde(default)]
121    pub external: BTreeMap<String, Value>,
122    /// Unknown top-level fields retained for forward compatibility.
123    #[serde(flatten)]
124    pub extra: BTreeMap<String, Value>,
125}
126
127pub(crate) fn text_response(result: CallToolResult) -> ZaiResult<McpTextResponse> {
128    extract_text(result).map(McpTextResponse::new)
129}
130
131pub(crate) fn web_search_response(result: CallToolResult) -> ZaiResult<WebSearchResponse> {
132    let results = decode_json_text::<Vec<WebSearchResult>>(result)?;
133    Ok(WebSearchResponse { results })
134}
135
136pub(crate) fn web_reader_response(result: CallToolResult) -> ZaiResult<WebReaderResponse> {
137    decode_json_text(result)
138}
139
140fn decode_json_text<T: DeserializeOwned>(result: CallToolResult) -> ZaiResult<T> {
141    check_tool_error(&result)?;
142    if let Some(structured) = result.structured_content {
143        return Ok(serde_json::from_value(structured)?);
144    }
145    let text = collect_text(&result)?;
146    let text = decode_string_layer(text);
147    Ok(serde_json::from_str(&text)?)
148}
149
150fn extract_text(result: CallToolResult) -> ZaiResult<String> {
151    check_tool_error(&result)?;
152    if let Some(structured) = result.structured_content {
153        return match structured {
154            Value::String(text) => Ok(text),
155            other => Ok(serde_json::to_string(&other)?),
156        };
157    }
158    collect_text(&result).map(decode_string_layer)
159}
160
161fn collect_text(result: &CallToolResult) -> ZaiResult<String> {
162    let mut texts = result
163        .content
164        .iter()
165        .filter_map(|content| content.as_text().map(|text| text.text.as_str()));
166    let first = texts.next().ok_or_else(|| ZaiError::Unknown {
167        code: codes::SDK_EXTERNAL_TOOL,
168        message: "MCP tool response did not contain text content".to_owned(),
169    })?;
170    let mut text = String::from(first);
171    for additional in texts {
172        text.push('\n');
173        text.push_str(additional);
174    }
175    Ok(text)
176}
177
178fn decode_string_layer(text: String) -> String {
179    match serde_json::from_str::<Value>(&text) {
180        Ok(Value::String(decoded)) => decoded,
181        _ => text,
182    }
183}
184
185fn check_tool_error(result: &CallToolResult) -> ZaiResult<()> {
186    if result.is_error != Some(true) {
187        return Ok(());
188    }
189    let message = collect_text(result)
190        .or_else(|_| {
191            result
192                .structured_content
193                .as_ref()
194                .map(serde_json::to_string)
195                .transpose()?
196                .ok_or_else(|| ZaiError::Unknown {
197                    code: codes::SDK_EXTERNAL_TOOL,
198                    message: "MCP tool returned an error without content".to_owned(),
199                })
200        })
201        .unwrap_or_else(|error| error.to_string());
202    Err(ZaiError::Unknown {
203        code: codes::SDK_EXTERNAL_TOOL,
204        message: decode_string_layer(message),
205    })
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use serde_json::json;
212
213    fn result(value: Value) -> CallToolResult {
214        serde_json::from_value(value).unwrap()
215    }
216
217    #[test]
218    fn unwraps_json_encoded_text_content() {
219        let value = json!({
220            "content": [{"type": "text", "text": "\"hello\""}],
221            "isError": false
222        });
223        assert_eq!(text_response(result(value)).unwrap().text(), "hello");
224    }
225
226    #[test]
227    fn surfaces_mcp_error_text_as_sdk_error() {
228        let value = json!({
229            "content": [{"type": "text", "text": "invalid input"}],
230            "isError": true
231        });
232        let error = text_response(result(value)).unwrap_err();
233        assert!(error.to_string().contains("invalid input"));
234    }
235
236    #[test]
237    fn parses_captured_search_encoding() {
238        let value = json!({
239            "content": [{
240                "type": "text",
241                "text": "\"[{\\\"title\\\":\\\"Docs\\\",\\\"link\\\":\\\"https://docs.rs\\\",\\\"content\\\":\\\"Rust docs\\\",\\\"refer\\\":\\\"ref_1\\\"}]\""
242            }],
243            "isError": false
244        });
245        let response = web_search_response(result(value)).unwrap();
246        assert_eq!(response[0].title, "Docs");
247        assert_eq!(response[0].refer.as_deref(), Some("ref_1"));
248    }
249
250    #[test]
251    fn parses_captured_reader_encoding() {
252        let value = json!({
253            "content": [{
254                "type": "text",
255                "text": "\"{\\\"title\\\":\\\"Page\\\",\\\"url\\\":\\\"https://example.com\\\",\\\"content\\\":\\\"Body\\\",\\\"images\\\":{},\\\"metadata\\\":{},\\\"external\\\":{}}\""
256            }],
257            "isError": false
258        });
259        let response = web_reader_response(result(value)).unwrap();
260        assert_eq!(response.title, "Page");
261        assert_eq!(response.content, "Body");
262    }
263
264    #[test]
265    fn structured_content_has_priority() {
266        let value = json!({
267            "content": [{"type": "text", "text": "ignored"}],
268            "structuredContent": [{
269                "title": "Structured",
270                "link": "https://example.com",
271                "content": "Body"
272            }],
273            "isError": false
274        });
275        let response = web_search_response(result(value)).unwrap();
276        assert_eq!(response[0].title, "Structured");
277    }
278
279    #[test]
280    fn structured_error_never_becomes_success() {
281        let value = json!({
282            "structuredContent": {"message": "quota exhausted"},
283            "isError": true
284        });
285        let error = text_response(result(value)).unwrap_err();
286        assert!(error.to_string().contains("quota exhausted"));
287    }
288
289    #[test]
290    fn accepts_direct_json_and_unknown_response_fields() {
291        let search = json!({
292            "content": [{
293                "type": "text",
294                "text": "[{\"title\":\"Docs\",\"link\":\"https://docs.rs\",\"content\":\"Body\",\"icon\":\"icon.png\",\"futureField\":42}]"
295            }]
296        });
297        let response = web_search_response(result(search)).unwrap();
298        assert_eq!(response[0].icon.as_deref(), Some("icon.png"));
299        assert_eq!(response[0].extra.get("futureField"), Some(&json!(42)));
300
301        let reader = json!({
302            "content": [{
303                "type": "text",
304                "text": "{\"title\":\"Page\",\"content\":\"Body\",\"links\":[\"https://example.com\"]}"
305            }]
306        });
307        let response = web_reader_response(result(reader)).unwrap();
308        assert_eq!(
309            response.extra.get("links").unwrap()[0],
310            "https://example.com"
311        );
312    }
313
314    #[test]
315    fn joins_multiple_text_blocks_in_order() {
316        let value = json!({
317            "content": [
318                {"type": "text", "text": "first"},
319                {"type": "image", "data": "AA==", "mimeType": "image/png"},
320                {"type": "text", "text": "second"}
321            ]
322        });
323        assert_eq!(
324            text_response(result(value)).unwrap().text(),
325            "first\nsecond"
326        );
327    }
328
329    #[test]
330    fn missing_text_is_an_error() {
331        let value = json!({"content": [], "isError": false});
332        assert!(text_response(result(value)).is_err());
333    }
334}
zai_rs/mcp/responses.rs

zai_rs/mcp/
responses.rs
