From fb85df8f111522e3c1e17581c3d74fbfd841ec3c Mon Sep 17 00:00:00 2001 From: ScriptedAlchemy Date: Sun, 18 Oct 2026 20:47:39 +0000 Subject: [PATCH] fix(mcp): preserve live graph request cancellation --- .../src/invocation/dispatch.rs | 16 ++- .../src/invocation/graph_tool.rs | 13 +- .../src/request_cancellation.rs | 124 +++++++++++++++++- .../src/daemon/invocation_executor.rs | 33 ++++- 4 files changed, 164 insertions(+), 22 deletions(-) diff --git a/crates/tracedecay-daemon-service/src/invocation/dispatch.rs b/crates/tracedecay-daemon-service/src/invocation/dispatch.rs index 02d486f67a..b6bc595321 100644 --- a/crates/tracedecay-daemon-service/src/invocation/dispatch.rs +++ b/crates/tracedecay-daemon-service/src/invocation/dispatch.rs @@ -1,6 +1,7 @@ //! The daemon invocation dispatcher: `DaemonInvocationService::invoke`. use super::*; +use tracedecay_contracts::CancellationSignal; use tracedecay_runtime_core::cancellation::CancellationToken; use tracedecay_runtime_core::path_safety::canonical_existing_identity; use tracedecay_tool_catalog::ApplicationSurfaceOperation; @@ -785,13 +786,26 @@ impl DaemonInvocationService { let Some(owner) = graph_tool_owner else { return missing_registered_owner_problem(publication, request_id); }; + let signal = match self.request_cancellations.signal(&request_id, &request_cancellation) { + Some(signal) => signal, + None => match CancellationSignal::active(cancellation.token_id.as_str()) { + Ok(signal) => signal, + Err(_) => return DaemonInvocationResponse::problem( + request_id, + DaemonInvocationProblem::InvalidRequest, + ), + }, + }; + if let CancellationState::Cancelled { requested_at } = cancellation.state { + signal.cancel(requested_at); + } Box::pin(execute_graph_tool( request_id, owner, surface_operation, arguments, deadline, - cancellation, + signal, request_cancellation, )) .await diff --git a/crates/tracedecay-daemon-service/src/invocation/graph_tool.rs b/crates/tracedecay-daemon-service/src/invocation/graph_tool.rs index c215591889..ab7880061a 100644 --- a/crates/tracedecay-daemon-service/src/invocation/graph_tool.rs +++ b/crates/tracedecay-daemon-service/src/invocation/graph_tool.rs @@ -9,8 +9,7 @@ use std::sync::Arc; use serde_json::{Map, Value}; use tracedecay_contracts::graph_tool::GraphToolCompletionV1; use tracedecay_contracts::{ - ApplicationProblem, CancellationContext, CancellationSignal, CancellationState, Deadline, - RequestId, ResolvedScope, now_micros, + ApplicationProblem, CancellationSignal, Deadline, RequestId, ResolvedScope, now_micros, }; use tracedecay_daemon_protocol::{ DaemonInvocationOutcome, DaemonInvocationProblem, DaemonInvocationResponse, @@ -92,21 +91,15 @@ pub(super) async fn execute_graph_tool( operation: ApplicationSurfaceOperation, arguments: Map, deadline: Deadline, - cancellation: CancellationContext, + signal: CancellationSignal, request_cancellation: tracedecay_runtime_core::cancellation::CancellationToken, ) -> DaemonInvocationResponse { - let (Ok(typed_request_id), Ok(signal)) = ( - RequestId::new(request_id.clone()), - CancellationSignal::active(cancellation.token_id.as_str()), - ) else { + let Ok(typed_request_id) = RequestId::new(request_id.clone()) else { return DaemonInvocationResponse::problem( request_id, DaemonInvocationProblem::InvalidRequest, ); }; - if let CancellationState::Cancelled { requested_at } = &cancellation.state { - signal.cancel(*requested_at); - } let execution = owner.port.execute(GraphToolInvocationV1 { operation, arguments, diff --git a/crates/tracedecay-daemon-service/src/request_cancellation.rs b/crates/tracedecay-daemon-service/src/request_cancellation.rs index eee4ff1dc2..90e27ad813 100644 --- a/crates/tracedecay-daemon-service/src/request_cancellation.rs +++ b/crates/tracedecay-daemon-service/src/request_cancellation.rs @@ -12,6 +12,7 @@ use std::collections::BTreeMap; use std::sync::{Arc, Mutex, MutexGuard}; use std::time::{Duration, Instant}; +use tracedecay_contracts::{CancellationSignal, clock::now_micros}; use tracedecay_runtime_core::cancellation::CancellationToken; const PENDING_CAPACITY: usize = 1_024; @@ -21,11 +22,17 @@ const COMPLETED_TTL: Duration = Duration::from_mins(1); #[derive(Default)] struct State { - active: BTreeMap, + active: BTreeMap, pending: BTreeMap, completed: BTreeMap, } +#[derive(Clone)] +struct ActiveRequest { + token: CancellationToken, + signal: Option, +} + /// The cancellation table of one invocation service generation. /// /// Clones share the same table. A [`Lease`] retains the table it was @@ -55,6 +62,37 @@ impl RequestCancellationRegistryV1 { } pub fn register(&self, request_id: &str) -> Option { + self.register_control(request_id, None) + } + + /// Retains an in-process caller's live signal alongside its request lease. + /// Reconstructing it from its wire context would require an asynchronous + /// relay before a blocking worker could observe caller cancellation. + pub fn register_with_signal( + &self, + request_id: &str, + signal: CancellationSignal, + ) -> Option { + self.register_control(request_id, Some(signal)) + } + + pub(crate) fn signal( + &self, + request_id: &str, + token: &CancellationToken, + ) -> Option { + self.lock() + .active + .get(request_id) + .filter(|request| request.token.is_same_token(token)) + .and_then(|request| request.signal.clone()) + } + + fn register_control( + &self, + request_id: &str, + signal: Option, + ) -> Option { let token = CancellationToken::for_application_request(request_id); let mut state = self.lock(); expire_ephemeral(&mut state, Instant::now()); @@ -67,8 +105,17 @@ impl RequestCancellationRegistryV1 { state.completed.remove(request_id); if state.pending.remove(request_id).is_some() { token.cancel(); + if let Some(signal) = &signal { + signal.cancel(now_micros()); + } } - state.active.insert(request_id.to_owned(), token.clone()); + state.active.insert( + request_id.to_owned(), + ActiveRequest { + token: token.clone(), + signal, + }, + ); Some(Lease { state: Arc::clone(&self.state), request_id: request_id.to_owned(), @@ -81,9 +128,12 @@ impl RequestCancellationRegistryV1 { let mut state = self.lock(); let now = Instant::now(); expire_ephemeral(&mut state, now); - if let Some(token) = state.active.get(request_id).cloned() { + if let Some(request) = state.active.get(request_id).cloned() { drop(state); - token.cancel(); + request.token.cancel(); + if let Some(signal) = request.signal { + signal.cancel(now_micros()); + } true } else if state.completed.contains_key(request_id) { true @@ -128,7 +178,7 @@ impl Drop for Lease { if state .active .get(&self.request_id) - .is_some_and(|token| token.is_same_token(&self.token)) + .is_some_and(|request| request.token.is_same_token(&self.token)) { state.active.remove(&self.request_id); record_completed(&mut state, self.request_id.clone(), Instant::now()); @@ -141,6 +191,70 @@ mod tests { use std::time::{Duration, Instant}; use super::{COMPLETED_CAPACITY, RequestCancellationRegistryV1, State, record_completed}; + use tracedecay_contracts::{CancellationSignal, clock::now_micros}; + + #[test] + fn controlled_request_shares_live_cancellation_without_polling_a_relay() { + let registry = RequestCancellationRegistryV1::default(); + let signal = CancellationSignal::active("cancel.controlled").unwrap(); + let lease = registry + .register_with_signal("request.controlled", signal.clone()) + .unwrap(); + let worker = registry + .signal("request.controlled", &lease.token()) + .unwrap(); + + assert!(signal.cancel(now_micros())); + assert!(worker.is_cancelled()); + + let foreign = RequestCancellationRegistryV1::default() + .register("request.controlled") + .unwrap(); + assert!( + registry + .signal("request.controlled", &foreign.token()) + .is_none() + ); + drop(lease); + let retry = registry.register("request.controlled").unwrap(); + assert!( + registry + .signal("request.controlled", &retry.token()) + .is_none() + ); + assert!(!retry.token().is_cancelled()); + } + + #[test] + fn daemon_cancel_reaches_the_controlled_signal_and_preserves_request_isolation() { + let registry = RequestCancellationRegistryV1::default(); + let signal = CancellationSignal::active("cancel.controlled").unwrap(); + let lease = registry + .register_with_signal("request.controlled", signal.clone()) + .unwrap(); + let other_signal = CancellationSignal::active("cancel.other").unwrap(); + let other = registry + .register_with_signal("request.other", other_signal.clone()) + .unwrap(); + + assert!(registry.cancel("request.controlled")); + assert!(signal.is_cancelled()); + assert!(lease.token().is_cancelled()); + assert!(!other_signal.is_cancelled()); + assert!(!other.token().is_cancelled()); + } + + #[test] + fn pending_cancel_reaches_a_later_controlled_registration() { + let registry = RequestCancellationRegistryV1::default(); + assert!(!registry.cancel("request.controlled")); + let signal = CancellationSignal::active("cancel.controlled").unwrap(); + let lease = registry + .register_with_signal("request.controlled", signal.clone()) + .unwrap(); + assert!(signal.is_cancelled()); + assert!(lease.token().is_cancelled()); + } #[test] fn pre_registration_cancellation_is_retained_and_cleanup_is_exact() { diff --git a/crates/tracedecay/src/daemon/invocation_executor.rs b/crates/tracedecay/src/daemon/invocation_executor.rs index 87b343f20a..d886c35740 100644 --- a/crates/tracedecay/src/daemon/invocation_executor.rs +++ b/crates/tracedecay/src/daemon/invocation_executor.rs @@ -13,8 +13,8 @@ use tracedecay_tool_catalog::ApplicationSurfaceOperation; use tracedecay_application::operation_stream::OperationRequestControls; use tracedecay_daemon_protocol::DAEMON_TOOL_RESPONSE_GRACE; use tracedecay_daemon_service::{ - DaemonInvocationOperation, DaemonInvocationProblem, ProjectRuntimeRequestLeaseV1, - WorkApplicationOutcomeV1, + DaemonInvocationOperation, DaemonInvocationPayload, DaemonInvocationProblem, + ProjectRuntimeRequestLeaseV1, WorkApplicationOutcomeV1, }; use tracedecay_domain::errors::{Result, TraceDecayError}; use tracedecay_runtime_core::native_runtime; @@ -285,7 +285,7 @@ impl InProcessDaemonInvocationExecutor { &self.store_administration, Some(&self.project_path), request, - None, + self.admitted_cancellation.clone(), ) .await } @@ -659,10 +659,29 @@ impl tracedecay_daemon_protocol::DaemonInvocationExecutor for InProcessDaemonInv stage: CancellationStage::BeforeAdmission, }, )?; - let executor = self.clone(); + let mut executor = self.clone(); let admitted_cancellation = self.admitted_cancellation.clone(); let request_cancellations = self.invocation.service.request_cancellations().clone(); let unavailable = |_| tracedecay_daemon_protocol::DaemonInvocationError::Unavailable; + // The graph owner must observe the caller's live signal directly: + // relaying through two native tasks lets a blocking file walk + // continue after the transport has acknowledged cancellation. + let cancellation_lease = if executor.admitted_cancellation.is_none() + && matches!(&request.payload, DaemonInvocationPayload::GraphTool { .. }) + { + let Some(lease) = request_cancellations + .register_with_signal(&request.request_id, cancellation.clone()) + else { + return Ok(DaemonInvocationResponse::problem( + request.request_id, + DaemonInvocationProblem::InvalidRequest, + )); + }; + executor.admitted_cancellation = Some(lease.token()); + Some(lease) + } else { + None + }; match native_runtime::NativeTask::spawn(async move { let request_id = request.request_id.clone(); let invocation = @@ -671,7 +690,7 @@ impl tracedecay_daemon_protocol::DaemonInvocationExecutor for InProcessDaemonInv tracing::trace_span!("daemon.invocation.invoke_once"), )) .map_err(unavailable)?; - settle_in_process_invocation( + let result = settle_in_process_invocation( &request_cancellations, &request_id, invocation, @@ -681,7 +700,9 @@ impl tracedecay_daemon_protocol::DaemonInvocationExecutor for InProcessDaemonInv policy, DAEMON_TOOL_RESPONSE_GRACE, ) - .await + .await; + drop(cancellation_lease); + result }) .map_err(unavailable)? .await