feat: Support custom messages when calling mcp

Author: L-jasmine

src/ai/mod.rs

@@ -682,6 +682,11 @@ impl ChatSession {
         Ok(response)
     }
 
+    pub fn get_tool_call_message(&self, tool_call: &llm::ToolCall) -> Option<String> {
+        let tool = self.tools.get_tool(tool_call.function.name.as_str())?;
+        Some(tool.call_mcp_message().to_string())
+    }
+
     pub async fn execute_tool(&mut self, tool_call: &llm::ToolCall) -> anyhow::Result<()> {
         use crate::ai::openai::tool::Tool;
 
@@ -775,6 +780,7 @@ pub async fn load_sse_tools(
         rmcp::service::RunningService<rmcp::RoleClient, rmcp::model::InitializeRequestParam>,
     >,
     mcp_servers_url: &str,
+    call_mcp_message: &str,
 ) -> anyhow::Result<()> {
     // load MCP
     let transport = SseClientTransport::start(mcp_servers_url).await?;
@@ -794,7 +800,11 @@ pub async fn load_sse_tools(
     for tool in tools {
         let server = client.peer().clone();
         log::info!("add tool: {}", tool.name);
-        tool_set.add_tool(McpToolAdapter::new(tool, server));
+        tool_set.add_tool(McpToolAdapter::new(
+            tool,
+            call_mcp_message.to_string(),
+            server,
+        ));
     }
     clients.push(client);
     Ok(())
@@ -806,6 +816,7 @@ pub async fn load_http_streamable_tools(
         rmcp::service::RunningService<rmcp::RoleClient, rmcp::model::InitializeRequestParam>,
     >,
     mcp_servers_url: &str,
+    call_mcp_message: &str,
 ) -> anyhow::Result<()> {
     // load MCP
     let transport = StreamableHttpClientTransport::from_uri(mcp_servers_url);
@@ -825,7 +836,11 @@ pub async fn load_http_streamable_tools(
     for tool in tools {
         let server = client.peer().clone();
         log::info!("add tool: {}", tool.name);
-        tool_set.add_tool(McpToolAdapter::new(tool, server));
+        tool_set.add_tool(McpToolAdapter::new(
+            tool,
+            call_mcp_message.to_string(),
+            server,
+        ));
     }
 
     clients.push(client);
@@ -857,7 +872,7 @@ async fn test_chat_session() {
     let mut clients = vec![];
 
     let mut tools = ToolSet::default();
-    load_http_streamable_tools(&mut tools, &mut clients, "http://localhost:8000/mcp")
+    load_http_streamable_tools(&mut tools, &mut clients, "http://localhost:8000/mcp", "")
         .await
         .unwrap();
 

src/ai/openai/tool.rs

@@ -17,12 +17,23 @@ pub trait Tool: Send + Sync {
 
 pub struct McpToolAdapter {
     tool: McpTool,
+    call_mcp_message: String,
     server: ServerSink,
 }
 
 impl McpToolAdapter {
-    pub fn new(tool: McpTool, server: ServerSink) -> Self {
-        Self { tool, server }
+    pub fn new(tool: McpTool, call_mcp_message: String, server: ServerSink) -> Self {
+        Self {
+            tool,
+            call_mcp_message,
+            server,
+        }
+    }
+}
+
+impl McpToolAdapter {
+    pub fn call_mcp_message(&self) -> &str {
+        &self.call_mcp_message
     }
 }
 

src/config.rs

@@ -22,6 +22,8 @@ pub struct MCPServerConfig {
     pub api_key: String,
     #[serde(rename = "type", default)]
     pub type_: MCPType,
+    #[serde(default)]
+    pub call_mcp_message: String,
 }
 
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]

src/main.rs

@@ -66,16 +66,25 @@ async fn routes(
             for server in &llm.mcp_server {
                 match server.type_ {
                     config::MCPType::SSE => {
-                        if let Err(e) =
-                            ai::load_sse_tools(&mut tool_set, clients, &server.server).await
+                        if let Err(e) = ai::load_sse_tools(
+                            &mut tool_set,
+                            clients,
+                            &server.server,
+                            &server.call_mcp_message,
+                        )
+                        .await
                         {
                             log::error!("Failed to load tools from {}: {}", &server.server, e);
                         }
                     }
                     config::MCPType::HttpStreamable => {
-                        if let Err(e) =
-                            ai::load_http_streamable_tools(&mut tool_set, clients, &server.server)
-                                .await
+                        if let Err(e) = ai::load_http_streamable_tools(
+                            &mut tool_set,
+                            clients,
+                            &server.server,
+                            &server.call_mcp_message,
+                        )
+                        .await
                         {
                             log::error!("Failed to load tools from {}: {}", &server.server, e);
                         }

src/services/ws/stable/llm.rs

@@ -204,6 +204,19 @@ pub async fn chat(
                 log::info!("llm functions: {:#?}", functions);
                 chat_session.add_assistant_tool_call(functions.clone());
                 for function in functions {
+                    if let Some(message) = chat_session.get_tool_call_message(&function) {
+                        log::info!("tool {} call message: {}", &function.function.name, message);
+                        if !message.is_empty() {
+                            let (tts_resp_tx, tts_resp_rx) = tokio::sync::mpsc::unbounded_channel();
+                            drop(tts_resp_tx);
+
+                            chunks_tx.send((message, tts_resp_rx)).map_err(|e| {
+                                anyhow::anyhow!(
+                                    "error sending tts chunks receiver for llm chunk: {e}"
+                                )
+                            })?;
+                        }
+                    }
                     chat_session.execute_tool(&function).await?
                 }
                 resp = chat_session.complete().await?;

src/services/ws/stable/mod.rs

@@ -210,6 +210,10 @@ async fn handle_tts_requests(mut chunks_rx: ChunksRx, session: &mut Session) ->
                 tts_chunk.len()
             );
 
+            if tts_chunk.is_empty() {
+                continue;
+            }
+
             session
                 .cmd_tx
                 .send(super::WsCommand::Audio(tts_chunk))