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    // Validate project path - reject temporary directories (Defense Layer 2)
62    // This prevents test environments from polluting the Dashboard registry
63    let temp_dir = std::env::temp_dir();
64    if normalized_project_path.starts_with(&temp_dir) {
65        tracing::warn!(
66            "Skipping Dashboard registration for temporary path: {}",
67            normalized_project_path.display()
68        );
69        return Ok(()); // Silently skip, don't error - non-fatal for MCP server
70    }
71
72    // Create project info
73    let project_info = ProjectInfo {
74        path: normalized_project_path.to_string_lossy().to_string(),
75        name: project_name,
76        db_path: normalized_db_path.to_string_lossy().to_string(),
77        agent,
78    };
79
80    // Connect to Dashboard WebSocket
81    let url = "ws://127.0.0.1:11391/ws/mcp";
82    let (ws_stream, _) = connect_async(url)
83        .await
84        .context("Failed to connect to Dashboard WebSocket")?;
85
86    tracing::debug!("Connected to Dashboard at {}", url);
87
88    let (mut write, mut read) = ws_stream.split();
89
90    // Send registration message
91    let register_msg = McpMessage::Register {
92        project: project_info.clone(),
93    };
94    let register_json = serde_json::to_string(&register_msg)?;
95    write
96        .send(Message::Text(register_json))
97        .await
98        .context("Failed to send register message")?;
99
100    // Wait for registration confirmation
101    if let Some(Ok(Message::Text(text))) = read.next().await {
102        match serde_json::from_str::<DashboardResponse>(&text) {
103            Ok(DashboardResponse::Registered { success: true }) => {
104                tracing::debug!("Successfully registered with Dashboard");
105            },
106            Ok(DashboardResponse::Registered { success: false }) => {
107                anyhow::bail!("Dashboard rejected registration");
108            },
109            _ => {
110                tracing::debug!("Unexpected response during registration: {}", text);
111            },
112        }
113    }
114
115    // Spawn ping task
116    let mut write_clone = write;
117    tokio::spawn(async move {
118        let mut interval = tokio::time::interval(Duration::from_secs(30));
119
120        loop {
121            interval.tick().await;
122
123            let ping_msg = McpMessage::Ping;
124            if let Ok(ping_json) = serde_json::to_string(&ping_msg) {
125                if write_clone.send(Message::Text(ping_json)).await.is_err() {
126                    tracing::warn!("Failed to send ping - Dashboard connection lost");
127                    break;
128                }
129            }
130        }
131    });
132
133    // Spawn read task to handle pongs and other messages
134    tokio::spawn(async move {
135        while let Some(Ok(msg)) = read.next().await {
136            match msg {
137                Message::Text(text) => match serde_json::from_str::<DashboardResponse>(&text) {
138                    Ok(DashboardResponse::Pong) => {
139                        tracing::debug!("Received pong from Dashboard");
140                    },
141                    _ => {
142                        tracing::debug!("Received message from Dashboard: {}", text);
143                    },
144                },
145                Message::Close(_) => {
146                    tracing::info!("Dashboard closed connection");
147                    break;
148                },
149                _ => {},
150            }
151        }
152    });
153
154    Ok(())
155}