Skip to main content

stakpak_shared/models/integrations/
mcp.rs

1use rmcp::model::{Annotated, CallToolResult, Content, RawContent};
2
3use crate::models::integrations::openai::{ChatMessage, MessageContent, ToolCallResultStatus};
4
5pub trait CallToolResultExt {
6    /// Create a success result with a simple text message
7    fn cancel(content: Option<&Vec<Annotated<RawContent>>>) -> Self;
8    fn get_status(&self) -> ToolCallResultStatus;
9    fn get_status_from_chat_message(message: &ChatMessage) -> ToolCallResultStatus;
10}
11
12impl CallToolResultExt for CallToolResult {
13    fn cancel(content: Option<&Vec<Annotated<RawContent>>>) -> Self {
14        let mut error_content = vec![Content::text("TOOL_CALL_CANCELLED")];
15
16        // Convert RawContent to Content by extracting text and creating new Content
17        if let Some(content) = content {
18            for raw in content {
19                if let Some(text_content) = raw.as_text() {
20                    error_content.push(Content::text(&text_content.text));
21                }
22            }
23        }
24
25        CallToolResult::error(error_content)
26    }
27
28    fn get_status(&self) -> ToolCallResultStatus {
29        if self.is_error == Some(true) {
30            // Check if any content contains the cancellation message
31            let is_cancelled = self.content.iter().any(|content| {
32                if let Some(text_content) = content.raw.as_text() {
33                    text_content.text.contains("TOOL_CALL_CANCELLED")
34                } else {
35                    false
36                }
37            });
38
39            if is_cancelled {
40                ToolCallResultStatus::Cancelled
41            } else {
42                ToolCallResultStatus::Error
43            }
44        } else {
45            ToolCallResultStatus::Success
46        }
47    }
48
49    fn get_status_from_chat_message(message: &ChatMessage) -> ToolCallResultStatus {
50        match message
51            .content
52            .as_ref()
53            .unwrap_or(&MessageContent::String(String::new()))
54            .to_string()
55            .contains("TOOL_CALL_CANCELLED")
56            || message
57                .content
58                .as_ref()
59                .unwrap_or(&MessageContent::String(String::new()))
60                .to_string()
61                .contains("cancelled")
62        {
63            true => ToolCallResultStatus::Cancelled,
64            false => ToolCallResultStatus::Success,
65        }
66    }
67}