From 8df6cdd2dcc279383666797a52b0a178a7c53d81 Mon Sep 17 00:00:00 2001 From: wsp1911 Date: Tue, 28 Jul 2026 19:57:41 +0800 Subject: [PATCH] fix(agentic): stabilize Task subagent execution - isolate coordinator futures behind Tokio task boundaries - align send_input and cancel permissions with the agent_id contract - propagate Tool cancellation through background subagent startup --- .../src/agentic/coordination/coordinator.rs | 128 +++++++++++++++++- .../tools/implementations/task/execution.rs | 104 ++++++++------ .../agentic/tools/implementations/task/mod.rs | 16 +-- .../tools/implementations/task/tests.rs | 21 +++ 4 files changed, 209 insertions(+), 60 deletions(-) diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index 64bb4ea7a7..1c3ec4c382 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -7869,13 +7869,34 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet &self, request: SubagentExecutionRequest, timeout_seconds: Option, + // Tool cancellation is narrower than parent-turn cancellation: round + // injection cancels the Tool while keeping the dialog turn alive. + tool_cancellation_token: Option, ) -> BitFunResult { + if tool_cancellation_token + .as_ref() + .is_some_and(CancellationToken::is_cancelled) + { + return Err(BitFunError::Cancelled( + "Background subagent start was cancelled".to_string(), + )); + } let request = self .resolve_hidden_subagent_execution_request(request) .await?; let mut request = self .prepare_hidden_subagent_execution_request(request) .await?; + if tool_cancellation_token + .as_ref() + .is_some_and(CancellationToken::is_cancelled) + { + self.cleanup_prepared_hidden_subagent_session_if_unsubmitted(&request) + .await; + return Err(BitFunError::Cancelled( + "Background subagent start was cancelled".to_string(), + )); + } let subagent_dialog_turn_id = request.ensure_dialog_turn_id(); let subagent_session_id = request .target_session_id() @@ -7942,6 +7963,22 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet let task_pk = registered_task.task_pk; let bg_task_id = registered_task.bg_task_id; let agent_id = registered_task.agent_id; + if tool_cancellation_token + .as_ref() + .is_some_and(CancellationToken::is_cancelled) + { + if let Err(error) = self.background_subagent_outcomes.discard(task_pk).await { + warn!( + "Failed to discard cancelled background task start: task_pk={}, error={}", + task_pk, error + ); + } + self.cleanup_prepared_hidden_subagent_session_if_unsubmitted(&request) + .await; + return Err(BitFunError::Cancelled( + "Background subagent start was cancelled".to_string(), + )); + } let parent_cancel_token = self .execution_engine .cancel_token_for_dialog_turn(&subagent_parent_info.dialog_turn_id) @@ -7980,8 +8017,33 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet let background_subagent_outcomes = self.background_subagent_outcomes.clone(); tokio::spawn(async move { - let result = match parent_cancel_token.as_ref() { - Some(token) => { + let result = match (parent_cancel_token, tool_cancellation_token) { + (Some(parent_token), Some(tool_token)) => { + let received = Self::await_hidden_subagent_receiver(receiver); + tokio::pin!(received); + tokio::select! { + _ = parent_token.cancelled() => { + scheduler_for_cancel + .request_hidden_subagent_cancellation(&cancel_handle) + .await; + Self::await_hidden_subagent_cancellation( + &mut received, + SUBAGENT_TIMEOUT_GRACE_PERIOD, + ).await + }, + _ = tool_token.cancelled() => { + scheduler_for_cancel + .request_hidden_subagent_cancellation(&cancel_handle) + .await; + Self::await_hidden_subagent_cancellation( + &mut received, + SUBAGENT_TIMEOUT_GRACE_PERIOD, + ).await + }, + result = &mut received => result, + } + } + (Some(token), None) | (None, Some(token)) => { let received = Self::await_hidden_subagent_receiver(receiver); tokio::pin!(received); tokio::select! { @@ -7997,7 +8059,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet result = &mut received => result, } } - None => Self::await_hidden_subagent_receiver(receiver).await, + (None, None) => Self::await_hidden_subagent_receiver(receiver).await, }; if suppress_delivery.load(Ordering::SeqCst) { background_subagent_tasks.remove(&task_pk); @@ -8024,10 +8086,23 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet let execution_cancel_token = CancellationToken::new(); let background_cancel_token_for_bridge = background_cancel_token.clone(); let execution_cancel_token_for_bridge = execution_cancel_token.clone(); - let cancel_bridge_handle = match parent_cancel_token { - Some(parent_cancel_token) => tokio::spawn(async move { + let cancel_bridge_handle = match (parent_cancel_token, tool_cancellation_token) { + (Some(parent_token), Some(tool_token)) => tokio::spawn(async move { + tokio::select! { + _ = parent_token.cancelled() => { + execution_cancel_token_for_bridge.cancel(); + } + _ = tool_token.cancelled() => { + execution_cancel_token_for_bridge.cancel(); + } + _ = background_cancel_token_for_bridge.cancelled() => { + execution_cancel_token_for_bridge.cancel(); + } + } + }), + (Some(token), None) | (None, Some(token)) => tokio::spawn(async move { tokio::select! { - _ = parent_cancel_token.cancelled() => { + _ = token.cancelled() => { execution_cancel_token_for_bridge.cancel(); } _ = background_cancel_token_for_bridge.cancelled() => { @@ -8035,7 +8110,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet } } }), - None => tokio::spawn(async move { + (None, None) => tokio::spawn(async move { background_cancel_token_for_bridge.cancelled().await; execution_cancel_token_for_bridge.cancel(); }), @@ -9335,6 +9410,7 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; + use tokio_util::sync::CancellationToken; #[test] fn worktree_execution_root_is_a_legacy_alias_for_project_storage() { @@ -9395,6 +9471,44 @@ mod tests { ); } + #[tokio::test] + async fn background_subagent_start_honors_an_already_cancelled_tool() { + let (coordinator, _session_manager) = test_coordinator(); + let cancellation_token = CancellationToken::new(); + cancellation_token.cancel(); + let request = SubagentExecutionRequest { + task_description: "should not start".to_string(), + context_mode: SubagentContextMode::Fresh, + target_session_id: None, + subagent_type: Some("Explore".to_string()), + logical_subagent_type: Some("Explore".to_string()), + continuation_policy: SessionContinuationPolicy::Reusable, + model_binding_policy: SessionModelBindingPolicy::Mutable, + workspace_path: None, + model_id: None, + inherit_parent_model: false, + subagent_parent_info: SubagentParentInfo { + session_id: "parent-session".to_string(), + dialog_turn_id: "parent-turn".to_string(), + tool_call_id: "task-tool".to_string(), + }, + context: HashMap::new(), + permission_runtime_ceiling: PermissionRuntimeCeiling::default(), + delegation_policy: DelegationPolicy::top_level().spawn_child(), + external_generation_lease: None, + }; + + let error = coordinator + .start_background_subagent(request, None, Some(cancellation_token)) + .await + .expect_err("a cancelled Tool must not start a background subagent"); + + assert!(matches!( + error, + crate::util::errors::BitFunError::Cancelled(_) + )); + } + #[test] fn session_reference_artifact_stems_extend_only_for_collisions() { let references = vec![ diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/task/execution.rs b/src/crates/assembly/core/src/agentic/tools/implementations/task/execution.rs index 6ced66639a..c52dd5bbfb 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/task/execution.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/task/execution.rs @@ -755,28 +755,36 @@ impl TaskTool { session_id, dialog_turn_id, }; - let background_result = coordinator - .start_background_subagent( - SubagentExecutionRequest { - task_description: prepared_prompt, - context_mode, - target_session_id, - subagent_type, - logical_subagent_type, - continuation_policy, - model_binding_policy, - workspace_path: effective_workspace_path, - model_id, - inherit_parent_model, - subagent_parent_info: parent_info, - context: subagent_context.unwrap_or_default(), - permission_runtime_ceiling, - delegation_policy: context.delegation_policy().spawn_child(), - external_generation_lease, - }, - timeout_seconds, - ) - .await?; + let request = SubagentExecutionRequest { + task_description: prepared_prompt, + context_mode, + target_session_id, + subagent_type, + logical_subagent_type, + continuation_policy, + model_binding_policy, + workspace_path: effective_workspace_path, + model_id, + inherit_parent_model, + subagent_parent_info: parent_info, + context: subagent_context.unwrap_or_default(), + permission_runtime_ceiling, + delegation_policy: context.delegation_policy().spawn_child(), + external_generation_lease, + }; + let coordinator = coordinator.clone(); + // The Tool future may be dropped on round injection. Keep its token in + // the spawned task so a detached background start still self-cancels. + let cancellation_token = context.cancellation_token().cloned(); + let background_result = tokio::spawn(async move { + coordinator + .start_background_subagent(request, timeout_seconds, cancellation_token) + .await + }) + .await + .map_err(|error| { + BitFunError::tool(format!("Background subagent task failed to join: {error}")) + })??; Ok(vec![ToolResult::Result { data: json!({ @@ -850,29 +858,35 @@ impl TaskTool { model_id, inherit_parent_model ); - let execution_result = coordinator - .execute_subagent( - SubagentExecutionRequest { - task_description: prepared_prompt.clone(), - context_mode, - target_session_id: target_session_id.clone(), - subagent_type: subagent_type.clone(), - logical_subagent_type: logical_subagent_type.clone(), - continuation_policy, - model_binding_policy, - workspace_path: effective_workspace_path.clone(), - model_id: model_id.clone(), - inherit_parent_model, - subagent_parent_info: parent_info, - context: subagent_context.clone().unwrap_or_default(), - permission_runtime_ceiling: permission_runtime_ceiling.clone(), - delegation_policy: context.delegation_policy().spawn_child(), - external_generation_lease: external_generation_lease.clone(), - }, - context.cancellation_token(), - timeout_seconds, - ) - .await; + let request = SubagentExecutionRequest { + task_description: prepared_prompt.clone(), + context_mode, + target_session_id: target_session_id.clone(), + subagent_type: subagent_type.clone(), + logical_subagent_type: logical_subagent_type.clone(), + continuation_policy, + model_binding_policy, + workspace_path: effective_workspace_path.clone(), + model_id: model_id.clone(), + inherit_parent_model, + subagent_parent_info: parent_info, + context: subagent_context.clone().unwrap_or_default(), + permission_runtime_ceiling: permission_runtime_ceiling.clone(), + delegation_policy: context.delegation_policy().spawn_child(), + external_generation_lease: external_generation_lease.clone(), + }; + let coordinator = coordinator.clone(); + let cancellation_token = context.cancellation_token().cloned(); + let execution_timeout = timeout_seconds; + let execution_result = tokio::spawn(async move { + coordinator + .execute_subagent(request, cancellation_token.as_ref(), execution_timeout) + .await + }) + .await + .map_err(|error| { + BitFunError::tool(format!("Foreground subagent task failed to join: {error}")) + })?; match execution_result { Ok(result) => { diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/task/mod.rs b/src/crates/assembly/core/src/agentic/tools/implementations/task/mod.rs index 8910e68c28..a01cf241f2 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/task/mod.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/task/mod.rs @@ -210,19 +210,19 @@ impl Tool for TaskTool { .unwrap_or("fork_context") .to_string(), TaskAction::SendInput => input - .get("session_id") + .get("agent_id") .and_then(Value::as_str) .map(str::trim) - .filter(|session_id| !session_id.is_empty()) - .map(|session_id| format!("send_input:{session_id}")) - .ok_or_else(|| BitFunError::validation("session_id is required".to_string()))?, + .filter(|agent_id| !agent_id.is_empty()) + .map(|agent_id| format!("send_input:{agent_id}")) + .ok_or_else(|| BitFunError::validation("agent_id is required".to_string()))?, TaskAction::Cancel => input - .get("session_id") + .get("agent_id") .and_then(Value::as_str) .map(str::trim) - .filter(|session_id| !session_id.is_empty()) - .map(|session_id| format!("cancel:{session_id}")) - .ok_or_else(|| BitFunError::validation("session_id is required".to_string()))?, + .filter(|agent_id| !agent_id.is_empty()) + .map(|agent_id| format!("cancel:{agent_id}")) + .ok_or_else(|| BitFunError::validation("agent_id is required".to_string()))?, }; Ok(vec![PermissionIntent::new("task", vec![resource])]) } diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/task/tests.rs b/src/crates/assembly/core/src/agentic/tools/implementations/task/tests.rs index ba45632c5a..a362f2273e 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/task/tests.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/task/tests.rs @@ -628,6 +628,27 @@ async fn validate_input_infers_send_input_without_action_when_agent_id_present() assert!(validation.result); } +#[test] +fn permission_intents_follow_the_agent_id_contract() { + let tool = TaskTool::new(); + let context = test_tool_context("agentic"); + + for (action, expected_resource) in [("send_input", "send_input:a1"), ("cancel", "cancel:a1")] { + let intents = tool + .permission_intents(&json!({ "action": action, "agent_id": "a1" }), &context) + .expect("agent_id should produce a permission intent"); + + assert_eq!(intents.len(), 1); + assert_eq!(intents[0].action, "task"); + assert_eq!(intents[0].resources, vec![expected_resource]); + } + + let error = tool + .permission_intents(&json!({ "action": "send_input" }), &context) + .expect_err("send_input without agent_id should be rejected"); + assert!(error.to_string().contains("agent_id is required")); +} + #[tokio::test] async fn validate_input_rejects_send_input_with_subagent_type() { let validation = TaskTool::new()