intent_engine/mcp/
ws_client.rs

1// WebSocket client for MCP → Dashboard communication
2// Handles registration and keep-alive for MCP server instances
3
4use anyhow::{Context, Result};
5use futures_util::{SinkExt, StreamExt};
6use serde::{Deserialize, Serialize};
7use std::path::PathBuf;
8use std::time::Duration;
9use tokio_tungstenite::{connect_async, tungstenite::Message};
10
11/// Project information sent to Dashboard
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct ProjectInfo {
14    pub path: String,
15    pub name: String,
16    pub db_path: String,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub agent: Option<String>,
19}
20
21/// Message types sent by MCP client
22#[derive(Debug, Serialize)]
23#[serde(tag = "type")]
24enum McpMessage {
25    #[serde(rename = "register")]
26    Register { project: ProjectInfo },
27    #[serde(rename = "ping")]
28    Ping,
29}
30
31/// Response types from Dashboard
32#[derive(Debug, Deserialize)]
33#[serde(tag = "type")]
34enum DashboardResponse {
35    #[serde(rename = "registered")]
36    Registered { success: bool },
37    #[serde(rename = "pong")]
38    Pong,
39}
40
41/// Start WebSocket client connection to Dashboard
42/// This replaces the Registry-based registration mechanism
43pub async fn connect_to_dashboard(
44    project_path: PathBuf,
45    db_path: PathBuf,
46    agent: Option<String>,
47) -> Result<()> {
48    // Extract project name from path
49    let project_name = project_path
50        .file_name()
51        .and_then(|n| n.to_str())
52        .unwrap_or("unknown")
53        .to_string();
54
55    // Normalize paths to handle symlinks
56    let normalized_project_path = project_path
57        .canonicalize()
58        .unwrap_or_else(|_| project_path.clone());
59    let normalized_db_path = db_path.canonicalize().unwrap_or_else(|_| db_path.clone());
60
61    // Create project info
62    let project_info = ProjectInfo {
63        path: normalized_project_path.to_string_lossy().to_string(),
64        name: project_name,
65        db_path: normalized_db_path.to_string_lossy().to_string(),
66        agent,
67    };
68
69    // Connect to Dashboard WebSocket
70    let url = "ws://127.0.0.1:11391/ws/mcp";
71    let (ws_stream, _) = connect_async(url)
72        .await
73        .context("Failed to connect to Dashboard WebSocket")?;
74
75    tracing::debug!("Connected to Dashboard at {}", url);
76
77    let (mut write, mut read) = ws_stream.split();
78
79    // Send registration message
80    let register_msg = McpMessage::Register {
81        project: project_info.clone(),
82    };
83    let register_json = serde_json::to_string(&register_msg)?;
84    write
85        .send(Message::Text(register_json))
86        .await
87        .context("Failed to send register message")?;
88
89    // Wait for registration confirmation
90    if let Some(Ok(Message::Text(text))) = read.next().await {
91        match serde_json::from_str::<DashboardResponse>(&text) {
92            Ok(DashboardResponse::Registered { success: true }) => {
93                tracing::debug!("Successfully registered with Dashboard");
94            },
95            Ok(DashboardResponse::Registered { success: false }) => {
96                anyhow::bail!("Dashboard rejected registration");
97            },
98            _ => {
99                tracing::debug!("Unexpected response during registration: {}", text);
100            },
101        }
102    }
103
104    // Spawn ping task
105    let mut write_clone = write;
106    tokio::spawn(async move {
107        let mut interval = tokio::time::interval(Duration::from_secs(30));
108
109        loop {
110            interval.tick().await;
111
112            let ping_msg = McpMessage::Ping;
113            if let Ok(ping_json) = serde_json::to_string(&ping_msg) {
114                if write_clone.send(Message::Text(ping_json)).await.is_err() {
115                    tracing::warn!("Failed to send ping - Dashboard connection lost");
116                    break;
117                }
118            }
119        }
120    });
121
122    // Spawn read task to handle pongs and other messages
123    tokio::spawn(async move {
124        while let Some(Ok(msg)) = read.next().await {
125            match msg {
126                Message::Text(text) => match serde_json::from_str::<DashboardResponse>(&text) {
127                    Ok(DashboardResponse::Pong) => {
128                        tracing::debug!("Received pong from Dashboard");
129                    },
130                    _ => {
131                        tracing::debug!("Received message from Dashboard: {}", text);
132                    },
133                },
134                Message::Close(_) => {
135                    tracing::info!("Dashboard closed connection");
136                    break;
137                },
138                _ => {},
139            }
140        }
141    });
142
143    Ok(())
144}
intent_engine/mcp/ws_client.rs

intent_engine/mcp/
ws_client.rs
