MCP: progress notifications for long scans, and cancellation

CHANGELOG.md

@@ -18,6 +18,14 @@ its heading and collects entries. The date and the link go on with the tag.
 
 ## [Unreleased]
 
+### Added
+
+- **The MCP server reports progress and can be cancelled.** Discover, port
+  scans and sweeps send progress notifications to clients that ask for them,
+  so a long sweep no longer looks stuck. A client that cancels a call now
+  stops the scan, instead of it running to the end and holding one of the
+  server's sixteen request slots.
+
 ### Changed
 
 - **Windows lists the desktop app's publisher as Felix Stubner.** Installed

crates/netscli-mcp/src/server.rs

@@ -1,9 +1,11 @@
 mod dispatch;
 mod errors;
+mod inflight;
 #[cfg(feature = "pcap")]
 mod jobs;
 mod limits;
 mod operations;
+mod progress;
 mod protocol;
 mod schemas;
 mod targets;
@@ -81,6 +83,7 @@ pub async fn run_server() -> anyhow::Result<()> {
     let mut reader = BufReader::new(stdin);
     let state = Arc::new(Mutex::new(ServerState::default()));
     let limiter = Arc::new(Semaphore::new(MAX_CONCURRENT_REQUESTS));
+    let inflight = Arc::new(inflight::InFlight::default());
 
     // Responses are funnelled through one channel to a single writer task.
     // Two concurrent handlers must never interleave bytes on stdout, and a
@@ -164,6 +167,11 @@ pub async fn run_server() -> anyhow::Result<()> {
                     guard.initialized = true;
                 }
             }
+            if request.method == "notifications/cancelled"
+                && inflight.cancel(request.params.as_ref())
+            {
+                tracing::info!("request cancelled by the client");
+            }
             continue;
         }
 
@@ -184,31 +192,42 @@ pub async fn run_server() -> anyhow::Result<()> {
         // `Option<Option<_>>` distinguishes absent from explicit null; a
         // notification never reaches here, so the outer layer is always Some.
         let id = request.id.clone().flatten();
-        handlers.spawn(async move {
+        let key = id.as_ref().map(inflight::key);
+        let progress = progress::Progress::for_request(request.params.as_ref(), &tx);
+        let inflight_for_task = Arc::clone(&inflight);
+        let key_for_task = key.clone();
+        let task = async move {
             let _permit = permit;
             // A hard ceiling on the whole call. Per-probe timeouts are
             // bounded but their product was not, and a request that never
             // returns holds its permit forever.
-            let response =
-                match tokio::time::timeout(MAX_REQUEST_DURATION, handle_request(state, request))
-                    .await
-                {
-                    Ok(response) => response,
-                    Err(_) => {
-                        tracing::warn!(
-                            timeout_s = MAX_REQUEST_DURATION.as_secs(),
-                            "request exceeded the maximum duration"
-                        );
-                        JsonRpcResponse::request_timeout(id, MAX_REQUEST_DURATION.as_secs())
-                    }
-                };
+            let call = progress::scope(progress, handle_request(state, request));
+            let response = match tokio::time::timeout(MAX_REQUEST_DURATION, call).await {
+                Ok(response) => response,
+                Err(_) => {
+                    tracing::warn!(
+                        timeout_s = MAX_REQUEST_DURATION.as_secs(),
+                        "request exceeded the maximum duration"
+                    );
+                    JsonRpcResponse::request_timeout(id, MAX_REQUEST_DURATION.as_secs())
+                }
+            };
             match serde_json::to_string(&response) {
                 Ok(s) => {
                     let _ = tx.send(s);
                 }
                 Err(e) => tracing::error!(error = %e, "failed to serialize response"),
             }
-        });
+            if let Some(key) = key_for_task {
+                inflight_for_task.finish(&key);
+            }
+        };
+        match key {
+            Some(key) => inflight.track(key, || handlers.spawn(task)),
+            None => {
+                handlers.spawn(task);
+            }
+        }
     }
 
     // EOF: the client is gone, so nothing is waiting for these answers.

crates/netscli-mcp/src/server/inflight.rs

@@ -0,0 +1,94 @@
+//! Requests still running, so a client can cancel one.
+//!
+//! MCP clients send `notifications/cancelled` with the `requestId` of a call
+//! they no longer want, for example when the user stops a model mid-answer.
+//! The server used to ignore it: a sweep carried on to the end, holding one of
+//! the sixteen request slots and sending probes nobody would read. Now the
+//! request's task is aborted, and, as the protocol asks, no response is sent
+//! for it.
+
+use std::collections::HashMap;
+use std::sync::{Mutex, MutexGuard};
+
+use serde_json::Value;
+use tokio::task::AbortHandle;
+
+#[derive(Default)]
+pub(super) struct InFlight(Mutex<HashMap<String, AbortHandle>>);
+
+/// One key per JSON-RPC id. The id's JSON text keeps `1` and `"1"` apart,
+/// which JSON-RPC treats as different ids.
+pub(super) fn key(id: &Value) -> String {
+    id.to_string()
+}
+
+impl InFlight {
+    fn map(&self) -> MutexGuard<'_, HashMap<String, AbortHandle>> {
+        self.0
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
+    /// Track a request while it is spawned. The lock is held across `spawn`
+    /// so a task that finishes at once cannot call `finish` before its handle
+    /// is recorded, which would leave the handle behind for good.
+    pub(super) fn track(&self, key: String, spawn: impl FnOnce() -> AbortHandle) {
+        let mut map = self.map();
+        let handle = spawn();
+        map.insert(key, handle);
+    }
+
+    /// The request answered; stop tracking it.
+    pub(super) fn finish(&self, key: &str) {
+        self.map().remove(key);
+    }
+
+    /// Handle `notifications/cancelled`. Returns whether a running request
+    /// was stopped; an unknown or finished id is ignored, as the protocol
+    /// expects, since the response may already be on its way.
+    pub(super) fn cancel(&self, params: Option<&Value>) -> bool {
+        let Some(id) = params.and_then(|p| p.get("requestId")) else {
+            return false;
+        };
+        match self.map().remove(&key(id)) {
+            Some(handle) => {
+                handle.abort();
+                true
+            }
+            None => false,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+    use std::time::Duration;
+
+    #[tokio::test]
+    async fn cancelling_a_running_request_aborts_it() {
+        let inflight = InFlight::default();
+        let mut set = tokio::task::JoinSet::new();
+        inflight.track(key(&json!(5)), || {
+            set.spawn(async { tokio::time::sleep(Duration::from_secs(60)).await })
+        });
+        assert!(inflight.cancel(Some(&json!({"requestId": 5, "reason": "user stopped"}))));
+        let joined = tokio::time::timeout(Duration::from_secs(1), set.join_next()).await;
+        let outcome = joined.expect("cancel must end the task promptly").unwrap();
+        assert!(outcome.unwrap_err().is_cancelled());
+    }
+
+    #[test]
+    fn unknown_or_finished_ids_are_ignored() {
+        let inflight = InFlight::default();
+        assert!(!inflight.cancel(Some(&json!({"requestId": 9}))));
+        assert!(!inflight.cancel(Some(&json!({}))));
+        assert!(!inflight.cancel(None));
+    }
+
+    #[test]
+    fn numeric_and_string_ids_are_different_requests() {
+        assert_ne!(key(&json!(1)), key(&json!("1")));
+    }
+}

crates/netscli-mcp/src/server/operations.rs

@@ -1,4 +1,5 @@
 use super::errors::RpcError;
+use super::progress;
 use super::schemas::{
     clamp_concurrency, clamp_timeout_ms, normalize_ports, validate_subnet, DiscoverParams,
     DnsParams, PingHostParams, ScanParams, SweepParams,
@@ -75,7 +76,21 @@ pub(super) async fn op_discover(p: DiscoverParams) -> Result<Vec<netscli_core::H
     let cfg = ops_config(concurrency, p.timeout);
     let ops = netscli_core::Ops::new(cfg);
     let (_subnet, hosts) = ops
-        .discover_ipv4(Some(subnet), p.resolve_hostnames.unwrap_or(false))
+        .discover_ipv4_with_progress(
+            Some(subnet),
+            p.resolve_hostnames.unwrap_or(false),
+            progress::callback(|e: &netscli_core::DiscoverProgress| {
+                use netscli_core::DiscoverPhase;
+                let (start, span, what) = match e.phase {
+                    DiscoverPhase::Ping => (0.0, 0.8, "pinging"),
+                    DiscoverPhase::Resolve => (0.8, 0.2, "naming hosts"),
+                };
+                (
+                    progress::phase_fraction(start, span, e.completed, e.total),
+                    format!("{what}: {} of {}, {} found", e.completed, e.total, e.found),
+                )
+            }),
+        )
         .await
         .map_err(|e| RpcError::ToolError(e.to_string()))?;
     Ok(hosts)
@@ -91,10 +106,21 @@ pub(super) async fn op_scan_ports(
     // see `ensure_host_allowed` for why the two can differ.
     let ip = ensure_host_allowed(&p.host, netscli_core::DEFAULT_DNS_TIMEOUT_MS).await?;
     let ops = netscli_core::Ops::new(cfg);
+    let progress = progress::callback(|e: &netscli_core::PortScanProgress| {
+        (
+            progress::phase_fraction(0.0, 1.0, e.completed, e.total),
+            format!(
+                "{} of {} ports, {} open",
+                e.completed, e.total, e.open_found
+            ),
+        )
+    });
     let scan = if p.udp.unwrap_or(false) {
-        ops.scan_udp_ports(&ip.to_string(), ports).await
+        ops.scan_udp_ports_with_progress(&ip.to_string(), ports, progress)
+            .await
     } else {
-        ops.scan_ports(&ip.to_string(), ports).await
+        ops.scan_ports_with_progress(&ip.to_string(), ports, progress)
+            .await
     };
     let (_ip, mut res) = scan.map_err(|e| RpcError::ToolError(e.to_string()))?;
     if !p.include_closed.unwrap_or(false) {
@@ -144,8 +170,27 @@ pub(super) async fn op_sweep(p: SweepParams) -> Result<Vec<netscli_core::SweepEn
     let concurrency = clamp_concurrency(p.max_concurrent, netscli_core::DEFAULT_CONCURRENCY);
     let cfg = ops_config(concurrency, p.timeout);
     let ops = netscli_core::Ops::new(cfg);
+    // Pinging is most of a sweep's wait on a quiet network, scanning most
+    // of it on a busy one; the split only has to keep the count rising.
+    let progress = progress::callback(|e: &netscli_core::SweepProgress| {
+        use netscli_core::SweepPhase;
+        let (start, span, what) = match e.phase {
+            SweepPhase::DiscoverPing => (0.0, 0.4, "pinging"),
+            SweepPhase::DiscoverResolve => (0.4, 0.1, "naming hosts"),
+            SweepPhase::Scan => (0.5, 0.5, "scanning hosts"),
+        };
+        (
+            progress::phase_fraction(start, span, e.completed, e.total),
+            format!("{what}: {} of {}, {} found", e.completed, e.total, e.found),
+        )
+    });
     let (_subnet, res) = ops
-        .sweep_ipv4(Some(subnet), ports, p.resolve_hostnames.unwrap_or(false))
+        .sweep_ipv4_with_progress(
+            Some(subnet),
+            ports,
+            p.resolve_hostnames.unwrap_or(false),
+            progress,
+        )
         .await
         .map_err(|e| RpcError::ToolError(e.to_string()))?;
     Ok(res)
@@ -225,69 +270,4 @@ pub(super) async fn op_discover_mdns(
 }
 
 #[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn closed_and_filtered_ports_are_dropped_and_the_rest_kept() {
-        use netscli_core::PortStatus;
-        assert!(worth_reporting(PortStatus::Open));
-        assert!(worth_reporting(PortStatus::OpenFiltered));
-        assert!(worth_reporting(PortStatus::Error));
-        assert!(!worth_reporting(PortStatus::Closed));
-        assert!(!worth_reporting(PortStatus::Filtered));
-    }
-
-    #[test]
-    fn without_a_timeout_each_phase_keeps_the_core_default() {
-        let cfg = ops_config(64, None);
-        assert_eq!(cfg.concurrency, 64);
-        assert_eq!(cfg.ping_timeout_ms, netscli_core::DEFAULT_PING_TIMEOUT_MS);
-        assert_eq!(cfg.scan_timeout_ms, netscli_core::DEFAULT_SCAN_TIMEOUT_MS);
-        assert_eq!(cfg.dns_timeout_ms, netscli_core::DEFAULT_DNS_TIMEOUT_MS);
-    }
-
-    #[test]
-    fn a_timeout_from_the_client_applies_to_every_phase_within_bounds() {
-        let cfg = ops_config(64, Some(2_000));
-        assert_eq!(
-            (cfg.ping_timeout_ms, cfg.scan_timeout_ms, cfg.dns_timeout_ms),
-            (2_000, 2_000, 2_000)
-        );
-        assert_eq!(ops_config(64, Some(1)).scan_timeout_ms, 10);
-    }
-
-    #[test]
-    fn a_publicly_addressed_interface_is_refused_rather_than_scanned() {
-        // The case this whole change exists for. Before, the default was
-        // substituted inside `Ops` after every check had run, so a host whose
-        // LAN is publicly addressed -- some universities and ISPs still do
-        // this -- had its own subnet scanned with no policy consulted.
-        let err = resolve_and_check(None, || "203.0.113.0/24".to_string())
-            .expect_err("a public default must be refused, not silently scanned");
-        let message = err.to_string();
-        assert!(message.contains("203.0.113.0/24"), "got: {message}");
-        assert!(
-            message.contains("NETSCLI_MCP_ALLOW_PUBLIC_TARGETS"),
-            "got: {message}"
-        );
-    }
-
-    #[test]
-    fn an_ordinary_private_default_is_returned_unchanged() {
-        let subnet = resolve_and_check(None, || "192.168.1.0/24".to_string())
-            .expect("a private default must pass");
-        assert_eq!(subnet, "192.168.1.0/24");
-    }
-
-    #[test]
-    fn an_explicit_subnet_wins_over_the_default() {
-        // The default must not be consulted at all when one was supplied,
-        // or a bad default could override a good request.
-        let subnet = resolve_and_check(Some("10.1.2.0/24".to_string()), || {
-            panic!("the default must not be evaluated when a subnet was given")
-        })
-        .expect("an explicit private subnet must pass");
-        assert_eq!(subnet, "10.1.2.0/24");
-    }
-}
+mod tests;