Add MCP server (v1 read-only): toggle, per-pane gate, panel UI

src-tauri/src/mcp.rs

@@ -0,0 +1,464 @@
+//! Embedded MCP server. Lets a Claude session running anywhere on the
+//! same machine — including inside one of tiletopia's own panes — inspect
+//! the workspace via Model Context Protocol.
+//!
+//! V1 surface (read-only):
+//!   resources: tiletopia://layout, tiletopia://panes, tiletopia://hosts
+//!   tools:     read_pane(leaf_id, last_lines?, after_seq?)
+//!              wait_for_idle(leaf_id, idle_ms?, timeout_ms?)
+//!
+//! Per-pane `mcpAllow` gate (default-deny) lives in the frontend tree;
+//! the frontend mirrors the gated subset into {@link McpState} via the
+//! `mcp_update_state` Tauri command. The MCP server only sees what the
+//! mirror exposes — no peeking around it.
+
+use std::collections::HashMap;
+use std::net::SocketAddr;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use anyhow::Result;
+use axum::{
+    body::Body,
+    http::{HeaderMap, HeaderValue, Request, StatusCode},
+    middleware::{self, Next},
+    response::Response,
+    Router,
+};
+use parking_lot::Mutex as PlMutex;
+use rmcp::{
+    handler::server::{router::tool::ToolRouter, wrapper::Parameters},
+    model::*,
+    schemars, tool, tool_handler, tool_router,
+    service::RequestContext,
+    transport::streamable_http_server::{
+        session::local::LocalSessionManager, tower::StreamableHttpService,
+    },
+    ErrorData as McpError, RoleServer, ServerHandler,
+};
+use serde::{Deserialize, Serialize};
+use serde_json::json;
+use tokio::{net::TcpListener, sync::RwLock, task::JoinHandle};
+use tokio_util::sync::CancellationToken;
+
+use crate::pty::{PaneId, PtyManager};
+
+// ----------------------------------------------------------------------------
+// Shared state mirrored from the frontend.
+// ----------------------------------------------------------------------------
+
+pub type LeafId = String;
+
+/// Cached snapshot the frontend pushes via `mcp_update_state` whenever the
+/// tree or hosts change. Source of truth for everything except scrollback,
+/// which the backend collects directly via {@link PtyManager}.
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct McpMirror {
+    /// Serialised layout tree (full structure, post-filtering happens
+    /// per-resource — see read_resource).
+    #[serde(default)]
+    pub layout_json: String,
+    /// Map of leaf id → pane metadata. Includes only leaves with
+    /// `mcpAllow === true` (frontend gates before mirroring).
+    #[serde(default)]
+    pub leaves: HashMap<LeafId, MirroredLeaf>,
+    /// Saved SSH hosts, password fields stripped.
+    #[serde(default)]
+    pub hosts: Vec<MirroredHost>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct MirroredLeaf {
+    pub pane_id: Option<PaneId>,
+    pub label: Option<String>,
+    pub shell_kind: String,
+    pub distro: Option<String>,
+    pub ssh_host_id: Option<String>,
+    #[serde(default)]
+    pub broadcast: bool,
+    #[serde(default)]
+    pub active: bool,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct MirroredHost {
+    pub id: String,
+    pub label: String,
+    pub hostname: String,
+    pub user: Option<String>,
+    pub port: Option<u16>,
+    #[serde(default)]
+    pub has_password: bool,
+}
+
+#[derive(Default)]
+pub struct McpState {
+    pub bearer_token: String,
+    pub mirror: McpMirror,
+}
+
+// ----------------------------------------------------------------------------
+// MCP service: tools + resources.
+// ----------------------------------------------------------------------------
+
+#[derive(Clone)]
+pub struct TileService {
+    ptys: Arc<PtyManager>,
+    state: Arc<RwLock<McpState>>,
+    tool_router: ToolRouter<Self>,
+}
+
+#[derive(Debug, Deserialize, schemars::JsonSchema)]
+pub struct ReadPaneArgs {
+    /// Stable leaf id from the tree (uuid-shaped). Must belong to a pane
+    /// the user has allow-listed for MCP access.
+    pub leaf_id: LeafId,
+    /// Return only the last N lines (default 200, hard cap 3000).
+    #[serde(default)]
+    pub last_lines: Option<usize>,
+    /// Only return bytes whose seq > this. Pair with the `__seq__` value
+    /// returned in a prior call for incremental polling.
+    #[serde(default)]
+    pub after_seq: Option<u64>,
+}
+
+#[derive(Debug, Deserialize, schemars::JsonSchema)]
+pub struct WaitForIdleArgs {
+    pub leaf_id: LeafId,
+    /// Required quiet window before declaring idle (default 500 ms).
+    #[serde(default)]
+    pub idle_ms: Option<u64>,
+    /// Hard timeout in ms; returns timeout=true after this (default 30s,
+    /// hard cap 5 min).
+    #[serde(default)]
+    pub timeout_ms: Option<u64>,
+}
+
+const READ_PANE_HARD_CAP_LINES: usize = 3000;
+const WAIT_TIMEOUT_HARD_CAP_MS: u64 = 5 * 60 * 1000;
+
+#[tool_router]
+impl TileService {
+    pub fn new(ptys: Arc<PtyManager>, state: Arc<RwLock<McpState>>) -> Self {
+        Self {
+            ptys,
+            state,
+            tool_router: Self::tool_router(),
+        }
+    }
+
+    /// Look up a leaf_id → pane_id under the MCP-allow gate.
+    async fn resolve_pane(&self, leaf_id: &str) -> Result<PaneId, McpError> {
+        let st = self.state.read().await;
+        let leaf = st.mirror.leaves.get(leaf_id).ok_or_else(|| {
+            McpError::invalid_params(
+                "unknown leaf_id (not visible to MCP; user may need to allow it)",
+                Some(json!({ "leaf_id": leaf_id })),
+            )
+        })?;
+        leaf.pane_id.ok_or_else(|| {
+            McpError::invalid_params(
+                "leaf has no live pane",
+                Some(json!({ "leaf_id": leaf_id })),
+            )
+        })
+    }
+
+    #[tool(description = "Read the recent scrollback of a terminal pane. \
+        Returns text plus a __seq__=N marker that can be passed back as \
+        after_seq for incremental polling.")]
+    async fn read_pane(
+        &self,
+        Parameters(args): Parameters<ReadPaneArgs>,
+    ) -> Result<CallToolResult, McpError> {
+        let pane_id = self.resolve_pane(&args.leaf_id).await?;
+
+        let ring = self.ptys.ring(pane_id).ok_or_else(|| {
+            McpError::internal_error(
+                "pane ring missing (pane may have just exited)",
+                Some(json!({ "leaf_id": args.leaf_id })),
+            )
+        })?;
+        let (bytes, seq) = {
+            let g = ring.lock();
+            g.snapshot()
+        };
+
+        // Trim by after_seq if provided: bytes in the ring beyond
+        // `after_seq` is `seq - after_seq`, clamped against ring size.
+        let start = match args.after_seq {
+            Some(prev) if seq > prev => {
+                let new_bytes = (seq - prev) as usize;
+                bytes.len().saturating_sub(new_bytes)
+            }
+            Some(_) => bytes.len(),
+            None => 0,
+        };
+        let tail = &bytes[start..];
+
+        let text = String::from_utf8_lossy(tail);
+        let cap = args
+            .last_lines
+            .map(|n| n.min(READ_PANE_HARD_CAP_LINES))
+            .unwrap_or(200);
+        let limited: String = if cap == 0 {
+            String::new()
+        } else {
+            let lines: Vec<&str> = text.lines().collect();
+            let start_line = lines.len().saturating_sub(cap);
+            lines[start_line..].join("\n")
+        };
+
+        Ok(CallToolResult::success(vec![
+            Content::text(limited),
+            Content::text(format!("__seq__={seq}")),
+        ]))
+    }
+
+    #[tool(description = "Block until a pane has been quiet (no output) \
+        for idle_ms, or timeout_ms elapses. Useful for command-completion \
+        synchronisation. Returns {idle:bool, seq:u64, elapsed_ms:u64}.")]
+    async fn wait_for_idle(
+        &self,
+        Parameters(args): Parameters<WaitForIdleArgs>,
+    ) -> Result<CallToolResult, McpError> {
+        let pane_id = self.resolve_pane(&args.leaf_id).await?;
+        let ring = self.ptys.ring(pane_id).ok_or_else(|| {
+            McpError::internal_error("pane ring missing", None)
+        })?;
+
+        let idle_target = Duration::from_millis(args.idle_ms.unwrap_or(500));
+        let timeout = Duration::from_millis(
+            args.timeout_ms
+                .unwrap_or(30_000)
+                .min(WAIT_TIMEOUT_HARD_CAP_MS),
+        );
+        let start = Instant::now();
+        let mut last_seq = ring.lock().snapshot().1;
+        let mut last_change = Instant::now();
+
+        loop {
+            // Sleep in small slices so we notice both incoming data and
+            // the overall timeout promptly.
+            tokio::time::sleep(Duration::from_millis(50)).await;
+            let now_seq = ring.lock().snapshot().1;
+            if now_seq != last_seq {
+                last_seq = now_seq;
+                last_change = Instant::now();
+            }
+            if last_change.elapsed() >= idle_target {
+                return Ok(CallToolResult::success(vec![Content::text(
+                    json!({
+                        "idle": true,
+                        "seq": last_seq,
+                        "elapsed_ms": start.elapsed().as_millis() as u64,
+                    })
+                    .to_string(),
+                )]));
+            }
+            if start.elapsed() >= timeout {
+                return Ok(CallToolResult::success(vec![Content::text(
+                    json!({
+                        "idle": false,
+                        "seq": last_seq,
+                        "elapsed_ms": start.elapsed().as_millis() as u64,
+                    })
+                    .to_string(),
+                )]));
+            }
+        }
+    }
+}
+
+#[tool_handler]
+impl ServerHandler for TileService {
+    fn get_info(&self) -> ServerInfo {
+        ServerInfo::new(
+            ServerCapabilities::builder()
+                .enable_tools()
+                .enable_resources()
+                .build(),
+        )
+        .with_server_info(Implementation::from_build_env())
+        .with_protocol_version(ProtocolVersion::V_2024_11_05)
+        .with_instructions(
+            "Tiletopia MCP (read-only v1). Resources: tiletopia://layout, \
+             tiletopia://panes, tiletopia://hosts. Tools: read_pane, \
+             wait_for_idle. Only panes the user has allow-listed are \
+             visible.",
+        )
+    }
+
+    async fn list_resources(
+        &self,
+        _r: Option<PaginatedRequestParams>,
+        _: RequestContext<RoleServer>,
+    ) -> Result<ListResourcesResult, McpError> {
+        Ok(ListResourcesResult {
+            resources: vec![
+                RawResource::new("tiletopia://layout", "layout").no_annotation(),
+                RawResource::new("tiletopia://panes", "panes").no_annotation(),
+                RawResource::new("tiletopia://hosts", "hosts").no_annotation(),
+            ],
+            next_cursor: None,
+            meta: None,
+        })
+    }
+
+    async fn read_resource(
+        &self,
+        req: ReadResourceRequestParams,
+        _: RequestContext<RoleServer>,
+    ) -> Result<ReadResourceResult, McpError> {
+        let state = self.state.read().await;
+        let body = match req.uri.as_str() {
+            "tiletopia://layout" => state.mirror.layout_json.clone(),
+            "tiletopia://panes" => {
+                serde_json::to_string(&state.mirror.leaves).unwrap_or_default()
+            }
+            "tiletopia://hosts" => {
+                serde_json::to_string(&state.mirror.hosts).unwrap_or_default()
+            }
+            other => {
+                return Err(McpError::resource_not_found(
+                    "resource_not_found",
+                    Some(json!({ "uri": other })),
+                ));
+            }
+        };
+        Ok(ReadResourceResult {
+            contents: vec![ResourceContents::text(body, req.uri)],
+        })
+    }
+
+    async fn list_resource_templates(
+        &self,
+        _r: Option<PaginatedRequestParams>,
+        _: RequestContext<RoleServer>,
+    ) -> Result<ListResourceTemplatesResult, McpError> {
+        Ok(ListResourceTemplatesResult {
+            resource_templates: vec![],
+            next_cursor: None,
+            meta: None,
+        })
+    }
+}
+
+// ----------------------------------------------------------------------------
+// HTTP wiring + bearer auth.
+// ----------------------------------------------------------------------------
+
+async fn bearer_auth(
+    axum::extract::State(expected): axum::extract::State<Arc<String>>,
+    headers: HeaderMap,
+    req: Request<Body>,
+    next: Next,
+) -> Result<Response, Response> {
+    let supplied = headers
+        .get(axum::http::header::AUTHORIZATION)
+        .and_then(|v| v.to_str().ok())
+        .and_then(|s| s.strip_prefix("Bearer "));
+    let ok = supplied
+        .map(|t| constant_time_eq(t.as_bytes(), expected.as_bytes()))
+        .unwrap_or(false);
+    if ok {
+        return Ok(next.run(req).await);
+    }
+
+    let mut resp = Response::builder()
+        .status(StatusCode::UNAUTHORIZED)
+        .body(Body::empty())
+        .unwrap();
+    resp.headers_mut().insert(
+        axum::http::header::WWW_AUTHENTICATE,
+        HeaderValue::from_static(r#"Bearer realm="tiletopia""#),
+    );
+    Err(resp)
+}
+
+fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
+    if a.len() != b.len() {
+        return false;
+    }
+    let mut d = 0u8;
+    for (x, y) in a.iter().zip(b) {
+        d |= x ^ y;
+    }
+    d == 0
+}
+
+// ----------------------------------------------------------------------------
+// Lifecycle.
+// ----------------------------------------------------------------------------
+
+pub struct RunningServer {
+    pub addr: SocketAddr,
+    pub token: String,
+    pub cancel: CancellationToken,
+    pub task: JoinHandle<()>,
+}
+
+#[derive(Default)]
+pub struct McpServerHandle(pub PlMutex<Option<RunningServer>>);
+
+pub async fn start_server(
+    ptys: Arc<PtyManager>,
+    state: Arc<RwLock<McpState>>,
+) -> Result<RunningServer> {
+    // 256-bit bearer token, hex-encoded.
+    use rand::RngCore;
+    let mut buf = [0u8; 32];
+    rand::rng().fill_bytes(&mut buf);
+    let token = hex::encode(buf);
+    state.write().await.bearer_token = token.clone();
+
+    let cancel = CancellationToken::new();
+
+    // Fresh service per session; cheap because we share state via Arcs.
+    let ptys_f = ptys.clone();
+    let state_f = state.clone();
+    let mcp_service = StreamableHttpService::new(
+        move || Ok(TileService::new(ptys_f.clone(), state_f.clone())),
+        LocalSessionManager::default().into(),
+        Default::default(),
+    );
+
+    let app = Router::new()
+        .nest_service("/mcp", mcp_service)
+        .layer(middleware::from_fn_with_state(
+            Arc::new(token.clone()),
+            bearer_auth,
+        ));
+
+    // Port 0 → OS picks. Recover via local_addr() before serving.
+    let listener = TcpListener::bind("127.0.0.1:0").await?;
+    let addr = listener.local_addr()?;
+
+    let cancel_inner = cancel.clone();
+    let task = tokio::spawn(async move {
+        let _ = axum::serve(listener, app)
+            .with_graceful_shutdown(async move {
+                cancel_inner.cancelled().await;
+            })
+            .await;
+    });
+
+    tracing::info!("MCP server listening on http://{addr}/mcp");
+    Ok(RunningServer {
+        addr,
+        token,
+        cancel,
+        task,
+    })
+}
+
+pub fn stop_server(handle: &McpServerHandle) {
+    if let Some(srv) = handle.0.lock().take() {
+        srv.cancel.cancel();
+        srv.task.abort();
+        tracing::info!("MCP server stopped");
+    }
+}