From db3727500a641e98477155224efc06277c5477c0 Mon Sep 17 00:00:00 2001
From: Matt Toohey <contact@matttoohey.com>
Date: Wed, 18 Mar 2026 12:46:16 +1100
Subject: [PATCH] feat(project_mcp): handle per-request MCP cancellation in
 start_repo_session

When Goose's 5-minute tool timeout fires, it sends a CancelledNotification
to the MCP server. rmcp 0.17 automatically cancels the RequestContext's
CancellationToken for that request. Previously, the start_repo_session
handler only watched self.cancel_token (the parent project session token)
and never noticed per-request cancellation, leaving child sessions running
as orphans after timeout.

Add the per-request CancellationToken as a parameter to start_repo_session
using rmcp's FromContextPart extractor (which extracts RequestContext.ct).
Add a new branch in the polling tokio::select! loop that watches this token
and cancels the child session via SessionRegistry::cancel when it fires.

The existing self.cancel_token branch is preserved for parent project
session cancellation. add_project_repo is not modified as it lacks a
similar long-running polling loop.
---
 apps/staged/src-tauri/src/project_mcp.rs | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/apps/staged/src-tauri/src/project_mcp.rs b/apps/staged/src-tauri/src/project_mcp.rs
index fb4297f3..41da93bd 100644
--- a/apps/staged/src-tauri/src/project_mcp.rs
+++ b/apps/staged/src-tauri/src/project_mcp.rs
@@ -134,6 +134,7 @@ impl ProjectToolsHandler {
     async fn start_repo_session(
         &self,
         Parameters(p): Parameters<StartRepoSessionParams>,
+        request_ct: CancellationToken,
     ) -> String {
         log::debug!(
             "[project_mcp] start_repo_session called: repo={:?} subpath={:?} expected_outcome={:?} return_info={:?} provider={:?} instructions={:?}",
@@ -448,8 +449,9 @@ impl ProjectToolsHandler {
         }
 
         // Poll until the session reaches a terminal state.
-        // Also watch the parent project session's cancellation token so we
-        // don't loop forever if the project session is cancelled while waiting.
+        // Watch both the parent project session's cancellation token and the
+        // per-request cancellation token (fired by rmcp when the MCP client
+        // sends a CancelledNotification, e.g. on Goose's 5-minute timeout).
         loop {
             tokio::select! {
                 _ = tokio::time::sleep(Duration::from_secs(2)) => {}
@@ -463,6 +465,20 @@ impl ProjectToolsHandler {
                     })
                     .to_string();
                 }
+                _ = request_ct.cancelled() => {
+                    // The MCP client cancelled this specific tool call (e.g.
+                    // Goose's per-tool timeout fired). Cancel the child session
+                    // so it doesn't continue running as an orphan.
+                    log::info!(
+                        "[project_mcp] per-request cancellation received for session {session_id}"
+                    );
+                    self.registry.cancel(&session_id);
+                    return serde_json::json!({
+                        "outcome": "cancelled",
+                        "output": "",
+                    })
+                    .to_string();
+                }
             }
             match self.store.get_session(&session_id) {
                 Ok(Some(s)) if s.status != SessionStatus::Running => {