Refactor MCP code

Author: sudeeptarlekar

application/apps/indexer/Cargo.toml

@@ -9,13 +9,13 @@ members = [
   "addons/bufread",
   "indexer_base",
   "merging",
+  "mcp",
   "parsers",
   "plugins_host",
   "processor",
   "session",
   "sources",
   "stypes",
-  "mcp",
 ]
 
 [workspace.dependencies]

application/apps/indexer/mcp/Cargo.toml

@@ -11,6 +11,7 @@ schemars = "1.0"
 dirs.workspace = true
 anyhow.workspace = true
 log.workspace = true
+tokio.workspace = true
 
 [dev-dependencies]
 tempfile.workspace = true

application/apps/indexer/mcp/src/ai_config.rs

@@ -20,6 +20,15 @@ impl AiConfig {
             .unwrap_or_default()
     }
 
+    pub fn new(model: String, api_key: String, endpoint: String, timeout_secs: u64) -> Self {
+        AiConfig {
+            model,
+            api_key,
+            endpoint,
+            timeout_secs,
+        }
+    }
+
     pub fn is_valid(&self) -> bool {
         !self.model.is_empty() && !self.api_key.is_empty() && !self.endpoint.is_empty()
     }

application/apps/indexer/mcp/src/ai_manager.rs

@@ -1,16 +1,94 @@
 mod ai_config;
-mod ai_manager;
 mod parameters;
+mod tasks;
 mod tools;
 mod utils;
 
 use anyhow::Result;
+use rmcp::{
+    handler::server::{ServerHandler, router::tool::ToolRouter},
+    model::{ServerCapabilities, ServerInfo},
+    tool_handler,
+};
+use tokio::sync::mpsc::{Receiver, Sender};
 
 use ai_config::AiConfig;
-use ai_manager::ChipmunkAI;
+use tasks::Task;
 
-async fn start() -> Result<()> {
-    let config = AiConfig::init();
-    let service = ChipmunkAI::new(config);
-    Ok(())
+/// ChipmunkAI is responsible for receiving the prompt from the UI and the processing
+/// the prompt with the LLM agent. If agent invokes some tool then
+pub struct ChipmunkAI {
+    pub config: AiConfig,
+
+    // Channel for sending the Task to Chipmunk Core
+    pub task_tx: Sender<Task>,
+
+    //Channel for receiving the prompt from the User
+    pub prompt_rx: Receiver<String>,
+
+    // Channel for sending the UI updates to Chipmunk Core.
+    // This channel will receive the updated from AI agent,
+    // e.g. UI messages to display, thinking..., Searching for tools...
+    pub ui_tx: Sender<String>,
+
+    pub tool_router: ToolRouter<Self>,
+}
+
+#[tool_handler]
+impl ServerHandler for ChipmunkAI {
+    fn get_info(&self) -> ServerInfo {
+        ServerInfo {
+            instructions: Some("Chipmunk AI".to_string()),
+            capabilities: ServerCapabilities::builder()
+                .enable_tools()
+                .enable_prompts()
+                .enable_resources()
+                .build(),
+            ..Default::default()
+        }
+    }
+}
+
+/// ServerCommunication is responsible to receiving the tasks from
+/// AI. Once task is received then Chipmunk takes the necessary action on the tasks,
+/// apply those tasks to the UI.
+pub struct ServerCommunication {
+    pub task_rx: Receiver<Task>,
+}
+
+impl ServerCommunication {
+    pub fn new(task_rx: Receiver<Task>) -> Self {
+        Self { task_rx }
+    }
+}
+
+/// ClientCommuinication is responsible for communicating the tasks from
+/// UI to the AI agent. With this user can send the prompt to AI agent and
+/// AI agent can send the response over mpsc receiver to render on the UI
+pub struct ClientCommunication {
+    pub ui_rx: Receiver<String>,
+    pub prompt_tx: Sender<String>,
+}
+
+impl ClientCommunication {
+    pub fn new(ui_receiver: Receiver<String>, prompt_sender: Sender<String>) -> Self {
+        Self {
+            ui_rx: ui_receiver,
+            prompt_tx: prompt_sender,
+        }
+    }
+}
+
+pub fn start() -> Result<(ChipmunkAI, ServerCommunication, ClientCommunication)> {
+    let (task_tx, task_rx) = tokio::sync::mpsc::channel::<Task>(10);
+    let (prompt_tx, prompt_rx) = tokio::sync::mpsc::channel::<String>(1);
+    let (ui_tx, ui_rx) = tokio::sync::mpsc::channel::<String>(1);
+
+    let server_comm = ServerCommunication::new(task_rx);
+    let client_comm = ClientCommunication::new(ui_rx, prompt_tx);
+
+    let config = AiConfig::default();
+    let manager = ChipmunkAI::new(config, prompt_rx, task_tx, ui_tx);
+
+    Ok((manager, server_comm, client_comm))
 }

application/apps/indexer/mcp/src/lib.rs

@@ -0,0 +1,18 @@
+use anyhow::Result;
+
+use mcp_server::start;
+use rmcp::{ServiceExt, transport::io::stdio};
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let (ai_manager, _server_comm, _client_comm) = start()?;
+
+    let service = ai_manager
+        .serve(stdio())
+        .await
+        .inspect_err(|err| eprintln!("Error while starting the AI Server; {err:?}"))?;
+
+    service.waiting().await?;
+    Ok(())
+}
+

application/apps/indexer/mcp/src/main.rs

@@ -1,6 +1,14 @@
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
+#[derive(Debug, Default, JsonSchema, Serialize, Deserialize)]
+pub struct ChartFilter {
+    pub value: String,
+    pub is_regex: bool,
+    pub ignore_case: bool,
+    pub is_word: bool,
+}
+
 #[derive(Debug, Default, JsonSchema, Serialize, Deserialize)]
 pub struct SearchFilter {
     pub value: String,
@@ -10,6 +18,11 @@ pub struct SearchFilter {
 }
 
 #[derive(Debug, Default, JsonSchema, Serialize, Deserialize)]
-pub struct FilterParameter {
+pub struct SearchFilterParameter {
     pub filters: Vec<SearchFilter>,
 }
+
+#[derive(Debug, Default, JsonSchema, Serialize, Deserialize)]
+pub struct ChartFilterParameter {
+    pub filters: Vec<ChartFilter>,
+}

application/apps/indexer/mcp/src/parameters.rs

@@ -0,0 +1,5 @@
+use crate::parameters::{ChartFilter, SearchFilter};
+pub enum Task {
+    ApplyFilter { filters: Vec<SearchFilter> },
+    ApplyChart { filter: Vec<ChartFilter> },
+}

application/apps/indexer/mcp/src/tasks.rs

@@ -1,23 +1,34 @@
+use anyhow::Result;
 use log::info;
 use rmcp::{
     ErrorData as McpError,
     handler::server::wrapper::Parameters,
     model::{CallToolResult, Content},
     tool, tool_router,
 };
+use tokio::sync::mpsc::{Receiver, Sender};
 
 use super::{
+    ChipmunkAI, ClientCommunication, ServerCommunication,
     ai_config::AiConfig,
-    ai_manager::ChipmunkAI,
-    parameters::{FilterParameter, SearchFilter},
+    parameters::{SearchFilter, SearchFilterParameter},
+    tasks::Task,
 };
 
 #[tool_router]
 impl ChipmunkAI {
     #[allow(dead_code)]
-    pub fn new(config: AiConfig) -> Self {
+    pub fn new(
+        config: AiConfig,
+        prompt_rx: Receiver<String>,
+        task_tx: Sender<Task>,
+        ui_tx: Sender<String>,
+    ) -> Self {
         Self {
             config,
+            prompt_rx,
+            ui_tx,
+            task_tx,
             tool_router: Self::tool_router(),
         }
     }
@@ -61,9 +72,9 @@ When the user provides natural language instructions, interpret them as follows:
 - "regex pattern \\d+" → set is_regex: true
 - "find ERROR, WARNING, and CRITICAL" → three separate filters
 "#)]
-    async fn apply_filters(
+    async fn apply_search_filters(
         &self,
-        Parameters(param): Parameters<FilterParameter>,
+        Parameters(param): Parameters<SearchFilterParameter>,
     ) -> Result<CallToolResult, McpError> {
         let filters = param
             .filters