From d3a1788445df4d6fb0e6f123fb4b749554efe587 Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 5 Oct 2026 17:08:39 +0900 Subject: [PATCH] fix: pin WIP cancellation to the active operation provider Resolve WIP_CANCEL_ROUTING from MR 01a10ac8-c63a-7e33-bc27-07a52f87abe9. Retain the dispatched handler instance and selected operation per active execution; forward cancellation only to that contributor, not unrelated providers or shared-handler aliases. Keep cancellation terminal confirmation and failure classification unchanged. Validation: contributed_cancellation tests (2 passed), all WIP tests (23 passed), root cargo check, cargo test -p worker -- --test-threads=1 (713 unit and 96 integration passed), cargo fmt --all -- --check, and git diff --check HEAD. Static/dynamic pending-call regressions cover unrelated cancel errors, shared handler duplication, exact dynamic handler identity, unknown IDs, terminal cleanup and cancelled audit. --- crates/worker/src/wip.rs | 298 ++++++++++++++++++++++++++++++--- docs/design/worker-wip-mode.md | 2 + 2 files changed, 277 insertions(+), 23 deletions(-) diff --git a/crates/worker/src/wip.rs b/crates/worker/src/wip.rs index dbc456b4f..703ac040e 100644 --- a/crates/worker/src/wip.rs +++ b/crates/worker/src/wip.rs @@ -131,6 +131,16 @@ pub trait WipOperationHandler: Send + Sync { async fn cancel(&self, _context: &ToolExecutionContext) -> Result<(), ToolError> { Ok(()) } + + /// Cancel the selected operation on this exact handler instance. Providers + /// that share a cancellation boundary may keep implementing `cancel`. + async fn cancel_operation( + &self, + _operation: &str, + context: &ToolExecutionContext, + ) -> Result<(), ToolError> { + self.cancel(context).await + } } /// One route contribution. Native projections replace compatibility projections @@ -249,11 +259,15 @@ impl WipOperationHandler for OperationDispatchHandler { handler.call(operation, arguments, context).await } - async fn cancel(&self, context: &ToolExecutionContext) -> Result<(), ToolError> { - for handler in self.handlers.values() { - handler.cancel(context).await?; - } - Ok(()) + async fn cancel_operation( + &self, + operation: &str, + context: &ToolExecutionContext, + ) -> Result<(), ToolError> { + let handler = self.handlers.get(operation).ok_or_else(|| { + ToolError::InvalidArgument("WIP cancellation operation is not published".into()) + })?; + handler.cancel_operation(operation, context).await } } @@ -1027,6 +1041,7 @@ impl WipHost { }) } + #[cfg(test)] async fn call( &self, request: CallOperationRequest, @@ -1035,6 +1050,15 @@ impl WipHost { let projection = self.projection(&request.target.path).ok_or_else(|| { WipOperationError::Protocol(unpublished_path_error(&request.target.path)) })?; + self.call_projection(projection, request, context).await + } + + async fn call_projection( + &self, + projection: WipProjection, + request: CallOperationRequest, + context: WipCallContext, + ) -> Result { let expected_object_validator = projection.object.validator.as_deref(); match ( expected_object_validator, @@ -1103,13 +1127,6 @@ impl WipHost { .call(&request.operation, &request.arguments, context) .await } - - async fn cancel(&self, path: &str, context: &ToolExecutionContext) -> Result<(), ToolError> { - let projection = self.projection(path).ok_or_else(|| { - ToolError::InvalidArgument("WIP cancellation target is no longer published".into()) - })?; - projection.handler.cancel(context).await - } } fn unpublished_path_error(path: &str) -> ProtocolError { @@ -1367,6 +1384,12 @@ struct ClientState { session: SessionId, } +#[derive(Clone)] +struct ActiveOperation { + operation: String, + handler: Arc, +} + pub struct WipRuntime { endpoint: Endpoint, wire_limits: Limits, @@ -1374,7 +1397,7 @@ pub struct WipRuntime { security_context: SecurityContext, state: Arc>, host: Arc, - active: Arc>>, + active: Arc>>, audit: Arc>>, metrics: Arc, } @@ -1576,10 +1599,6 @@ impl WipRuntime { .operation_round_trips .fetch_add(1, Ordering::Relaxed); let execution_id = execution.execution_id(); - self.active - .lock() - .unwrap_or_else(|error| error.into_inner()) - .insert(execution_id.clone(), path.clone()); let mut guard = DispatchedCallGuard::new( Arc::clone(&self.state), Arc::clone(&self.active), @@ -1621,14 +1640,19 @@ impl WipRuntime { } async fn cancel(&self, execution: &ToolExecutionContext) -> Result<(), ToolError> { - let path = self + let active = self .active .lock() .unwrap_or_else(|error| error.into_inner()) .get(&execution.execution_id()) .cloned(); - match path { - Some(path) => self.host.cancel(&path, execution).await, + match active { + Some(active) => { + active + .handler + .cancel_operation(&active.operation, execution) + .await + } None => Ok(()), } } @@ -1724,7 +1748,29 @@ impl WipRuntime { execution, security_context: self.security_context.as_str().to_string(), }; - match self.host.call(request_value.clone(), context).await { + let result = match self.host.projection(&request_value.target.path) { + Some(projection) => { + // Pin the exact resolved instance used for dispatch. A dynamic + // resolver may return a fresh stateful handler on each lookup. + self.active + .lock() + .unwrap_or_else(|error| error.into_inner()) + .insert( + context.execution.execution_id(), + ActiveOperation { + operation: request_value.operation.clone(), + handler: Arc::clone(&projection.handler), + }, + ); + self.host + .call_projection(projection, request_value.clone(), context) + .await + } + None => Err(WipOperationError::Protocol(unpublished_path_error( + &request_value.target.path, + ))), + }; + match result { Ok(output) => { let response = encode_call_operation_response( &request_value, @@ -1807,7 +1853,7 @@ impl WipRuntime { struct DispatchedCallGuard { state: Arc>, - active: Arc>>, + active: Arc>>, request_id: RequestId, execution_id: String, armed: bool, @@ -1816,7 +1862,7 @@ struct DispatchedCallGuard { impl DispatchedCallGuard { fn new( state: Arc>, - active: Arc>>, + active: Arc>>, request_id: RequestId, execution_id: String, ) -> Self { @@ -4526,6 +4572,212 @@ mod tests { runtime } + struct UnrelatedCancellation { + cancellations: Arc, + } + + #[async_trait] + impl WipOperationHandler for UnrelatedCancellation { + async fn call( + &self, + _: &str, + _: &BTreeMap, + _: WipCallContext, + ) -> Result { + panic!("unrelated provider must not execute"); + } + + async fn cancel(&self, _: &ToolExecutionContext) -> Result<(), ToolError> { + self.cancellations.fetch_add(1, Ordering::SeqCst); + Err(ToolError::InvalidArgument( + "unknown execution in unrelated provider".into(), + )) + } + } + + struct PendingCancellation { + started: Arc, + cancelled: tokio::sync::Notify, + was_called: AtomicBool, + cancellations: Arc, + } + + #[async_trait] + impl WipOperationHandler for PendingCancellation { + async fn call( + &self, + _: &str, + _: &BTreeMap, + _: WipCallContext, + ) -> Result { + self.was_called.store(true, Ordering::SeqCst); + self.started.notify_one(); + self.cancelled.notified().await; + Err(WipOperationError::Cancelled( + "provider confirmed cancellation".to_string().into(), + )) + } + + async fn cancel(&self, _: &ToolExecutionContext) -> Result<(), ToolError> { + // A second resolution of the same path would produce another instance. + assert!( + self.was_called.load(Ordering::SeqCst), + "cancel must use the dispatched instance" + ); + self.cancellations.fetch_add(1, Ordering::SeqCst); + self.cancelled.notify_one(); + Ok(()) + } + } + + struct PendingCancellationResolver { + started: Arc, + cancellations: Arc, + } + + impl WipDynamicOperationResolver for PendingCancellationResolver { + fn handler(&self, _: &str) -> Arc { + Arc::new(PendingCancellation { + started: Arc::clone(&self.started), + cancelled: tokio::sync::Notify::new(), + was_called: AtomicBool::new(false), + cancellations: Arc::clone(&self.cancellations), + }) + } + } + + struct CancellationItemResolver { + unrelated: Arc, + } + + impl WipDynamicItemResolver for CancellationItemResolver { + fn resolve(&self, item: &str) -> Option { + if item != "A-1" { + return None; + } + let projection = contribution_projection(Arc::new(UnrelatedCancellation { + cancellations: Arc::clone(&self.unrelated), + })); + Some(WipDynamicItem { + object: projection.object, + handler: projection.handler, + }) + } + } + + async fn assert_contributed_cancellation_is_selected_once(dynamic: bool) { + let started = Arc::new(tokio::sync::Notify::new()); + let cancellations = Arc::new(AtomicUsize::new(0)); + let unrelated = Arc::new(AtomicUsize::new(0)); + let mut registry = WipMountRegistry::new(); + registry.allocate_namespace("asset", "assets").unwrap(); + let projection = contribution_projection(Arc::new(UnrelatedCancellation { + cancellations: Arc::clone(&unrelated), + })); + let mut descriptor = projection.descriptor.clone(); + // Two names share one provider; read sorts first and fails cancellation. + descriptor.operations = vec![ + contributed_operation("write"), + contributed_operation("write_again"), + ]; + if dynamic { + let mut collection = projection.clone(); + collection.route = "/assets".into(); + collection.object.name = "assets".into(); + collection.interface = "test.asset/collection/v1".into(); + collection.object.interfaces = vec![collection.interface.clone()]; + registry.mount(collection).unwrap(); + registry + .mount_dynamic(WipDynamicMount { + collection_route: "/assets".into(), + capability: "asset:item".into(), + interface: projection.interface.clone(), + descriptor: projection.descriptor, + interface_validator: projection.interface_validator, + resolver: Arc::new(CancellationItemResolver { + unrelated: Arc::clone(&unrelated), + }), + }) + .unwrap(); + registry + .contribute_dynamic_operations(WipDynamicOperationContribution { + collection_route: "/assets".into(), + contributor: "management".into(), + interface: projection.interface, + descriptor, + resolver: Arc::new(PendingCancellationResolver { + started: Arc::clone(&started), + cancellations: Arc::clone(&cancellations), + }), + }) + .unwrap(); + } else { + registry.mount(projection).unwrap(); + registry + .contribute_operations(WipOperationContribution { + route: "/assets/A-1".into(), + contributor: "management".into(), + interface: "test.asset/item/v1".into(), + descriptor, + handler: PendingCancellationResolver { + started: Arc::clone(&started), + cancellations: Arc::clone(&cancellations), + } + .handler("A-1"), + }) + .unwrap(); + } + let runtime = + WipRuntime::new(WipHost::new(registry), SecurityContext::new("worker-a"), 0).unwrap(); + runtime + .discover("/assets/A-1".into(), 0, false) + .await + .unwrap(); + runtime + .inspect("test.asset/item/v1".into(), false) + .await + .unwrap(); + let execution = ToolExecutionContext::direct(); + let call = runtime.call( + "/assets/A-1".into(), + "test.asset/item/v1".into(), + "write".into(), + json!({}), + execution.clone(), + ); + let cancel = async { + started.notified().await; + // Unknown executions must not invoke any provider. + runtime + .cancel(&ToolExecutionContext::new("unknown", "other-batch", 0)) + .await + .unwrap(); + runtime.cancel(&execution).await.unwrap(); + }; + let (result, ()) = tokio::time::timeout(std::time::Duration::from_secs(2), async { + tokio::join!(call, cancel) + }) + .await + .expect("selected provider must receive cancellation"); + assert!(matches!(result, Err(ToolError::Cancelled(_)))); + assert_eq!(unrelated.load(Ordering::SeqCst), 0); + assert_eq!(cancellations.load(Ordering::SeqCst), 1); + assert!(runtime.active.lock().unwrap().is_empty()); + assert_eq!(runtime.audit()[0].outcome, WipAuditOutcome::Cancelled); + runtime.cancel(&execution).await.unwrap(); + assert_eq!(cancellations.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn static_contributed_cancellation_ignores_unrelated_and_shared_handlers() { + assert_contributed_cancellation_is_selected_once(false).await; + } + + #[tokio::test] + async fn dynamic_contributed_cancellation_pins_the_dispatched_instance() { + assert_contributed_cancellation_is_selected_once(true).await; + } + #[tokio::test] async fn cancelled_unknown_and_disconnected_calls_have_distinct_terminal_states() { let cancelled = diff --git a/docs/design/worker-wip-mode.md b/docs/design/worker-wip-mode.md index 9f044b2f5..ed3019ba4 100644 --- a/docs/design/worker-wip-mode.md +++ b/docs/design/worker-wip-mode.md @@ -124,6 +124,8 @@ The adapter keeps these outcomes distinct: - **disconnected/timeout after dispatch**: `wip-client` records an unknown outcome; - **operation outcome unknown**: an executing provider fails without a rollback/non-execution guarantee, a response cannot be completed after dispatch, or the call future is force-closed. +Cooperative cancellation is scoped to the active execution's selected operation and exact resolved handler instance. Static/dynamic dispatchers forward `cancel_operation` only to that operation's provider, once per cancellation request, never to unrelated contributors or every alias backed by a shared handler. The runtime pins the dispatcher used for execution rather than resolving the Object again on cancellation, since dynamic resolvers may return stateful instances. Providers may implement `cancel_operation` for operation-specific control or the shared `cancel` default. Terminal completion or a dropped execution future releases the active binding; unknown/finished execution IDs do not invoke any provider. A cancellation request is not itself terminal confirmation. + An unknown mutating outcome must never be retried automatically. The model must inspect external state or ask for operator guidance. Retrieval failures may be explicitly refreshed. ## Measurement