feat(agent): let AgentWait yield to steering

src/crates/assembly/core/src/agentic/coordination/background_outcomes.rs

@@ -42,6 +42,7 @@ impl BackgroundSubagentOutcome {
 pub(crate) enum BackgroundSubagentWaitStatus {
     Completed,
     TimedOut,
+    Steered,
     NoMatchingTasks,
 }
 
@@ -50,6 +51,7 @@ impl BackgroundSubagentWaitStatus {
         match self {
             Self::Completed => "completed",
             Self::TimedOut => "timed_out",
+            Self::Steered => "steered",
             Self::NoMatchingTasks => "no_matching_tasks",
         }
     }
@@ -207,6 +209,7 @@ impl BackgroundSubagentOutcomeStore {
         timeout: Duration,
         delivered_parent_dialog_turn_id: &str,
         cancellation_token: Option<&CancellationToken>,
+        round_injection_preemption_token: Option<&CancellationToken>,
     ) -> BitFunResult<BackgroundSubagentWaitResult> {
         self.reconcile_stale_running_tasks(parent_session_id)
             .await?;
@@ -229,10 +232,38 @@ impl BackgroundSubagentOutcomeStore {
         let mut debounce_deadline = None;
 
         loop {
+            if cancellation_token.is_some_and(CancellationToken::is_cancelled) {
+                return Err(BitFunError::cancelled(
+                    "AgentWait was cancelled".to_string(),
+                ));
+            }
+
             let notified = self.changes.notified();
             tokio::pin!(notified);
 
             let available = self.collect_available(&selected_task_pks).await?;
+            if round_injection_preemption_token.is_some_and(CancellationToken::is_cancelled) {
+                if available.outcomes.is_empty() {
+                    return Ok(wait_result(
+                        BackgroundSubagentWaitStatus::Steered,
+                        Vec::new(),
+                        available.pending_bg_task_ids,
+                    ));
+                }
+                if let Some(result) = self
+                    .claim_result(
+                        parent_session_id,
+                        delivered_parent_dialog_turn_id,
+                        BackgroundSubagentWaitStatus::Steered,
+                        available,
+                    )
+                    .await?
+                {
+                    return Ok(result);
+                }
+                debounce_deadline = None;
+                continue;
+            }
             if !available.outcomes.is_empty() && available.pending_bg_task_ids.is_empty() {
                 if let Some(result) = self
                     .claim_result(
@@ -283,22 +314,24 @@ impl BackgroundSubagentOutcomeStore {
                 continue;
             }
 
-            match cancellation_token {
-                Some(token) => {
-                    tokio::select! {
-                        _ = token.cancelled() => {
-                            return Err(BitFunError::cancelled("AgentWait was cancelled".to_string()));
-                        }
-                        _ = &mut notified => {}
-                        _ = sleep_until(wake_at) => {}
+            tokio::select! {
+                biased;
+                _ = async {
+                    match cancellation_token {
+                        Some(token) => token.cancelled().await,
+                        None => std::future::pending::<()>().await,
                     }
+                } => {
+                    return Err(BitFunError::cancelled("AgentWait was cancelled".to_string()));
                 }
-                None => {
-                    tokio::select! {
-                        _ = &mut notified => {}
-                        _ = sleep_until(wake_at) => {}
+                _ = async {
+                    match round_injection_preemption_token {
+                        Some(token) => token.cancelled().await,
+                        None => std::future::pending::<()>().await,
                     }
-                }
+                } => {}
+                _ = &mut notified => {}
+                _ = sleep_until(wake_at) => {}
             }
         }
     }
@@ -699,6 +732,7 @@ mod tests {
             Duration::from_millis(50),
             "wait-turn",
             None,
+            None,
         )
         .await
         .expect("recover persisted background result");

src/crates/assembly/core/src/agentic/coordination/coordinator.rs

@@ -11480,6 +11480,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet
         timeout: Duration,
         delivered_parent_dialog_turn_id: &str,
         cancellation_token: Option<&CancellationToken>,
+        round_injection_preemption_token: Option<&CancellationToken>,
     ) -> BitFunResult<BackgroundSubagentWaitResult> {
         self.background_subagent_outcomes
             .wait_for(
@@ -11489,6 +11490,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet
                 timeout,
                 delivered_parent_dialog_turn_id,
                 cancellation_token,
+                round_injection_preemption_token,
             )
             .await
     }
@@ -17432,6 +17434,7 @@ mod tests {
                 Duration::from_millis(10),
                 "wait-turn-1",
                 None,
+                None,
             )
             .await
             .expect("AgentWait should collect the completed outcome");
@@ -17451,6 +17454,7 @@ mod tests {
                 Duration::from_millis(10),
                 "wait-turn-2",
                 None,
+                None,
             )
             .await
             .expect("a consumed outcome should not be delivered twice");
@@ -17555,6 +17559,7 @@ mod tests {
                 Duration::from_millis(10),
                 "wait-turn",
                 None,
+                None,
             )
             .await
             .expect("AgentWait should collect a prior-turn outcome in the same session");
@@ -17596,6 +17601,7 @@ mod tests {
                 Duration::from_millis(1),
                 "wait-turn-1",
                 None,
+                None,
             )
             .await
             .expect("all selector timeout should return partial results");
@@ -17613,6 +17619,7 @@ mod tests {
                 Duration::from_millis(10),
                 "wait-turn-2",
                 None,
+                None,
             )
             .await
             .expect("returned results should be consumed");
@@ -17650,6 +17657,7 @@ mod tests {
                 Duration::from_secs(6),
                 "wait-turn",
                 None,
+                None,
             )
             .await
             .expect("any selector should return after the result debounce");
@@ -17697,6 +17705,7 @@ mod tests {
                 Duration::from_secs(10),
                 "cancelled-wait-turn",
                 Some(&cancellation),
+                None,
             )
             .await
             .expect_err("cancelled AgentWait should not return a partial result");
@@ -17710,6 +17719,7 @@ mod tests {
                 Duration::from_millis(10),
                 "retry-wait-turn",
                 None,
+                None,
             )
             .await
             .expect("a cancelled wait must not consume the completed outcome");
@@ -17737,6 +17747,7 @@ mod tests {
                 Duration::from_millis(1),
                 "wait-turn",
                 None,
+                None,
             )
             .await
             .expect("AgentWait timeout should be returned normally");
@@ -17746,6 +17757,95 @@ mod tests {
         assert_eq!(result.pending_bg_task_ids, vec![registered.bg_task_id]);
     }
 
+    #[tokio::test]
+    async fn steered_agent_wait_returns_pending_tasks_without_cancelling_them() {
+        let (coordinator, _) = test_coordinator();
+        let registered = register_test_background_task(
+            &coordinator,
+            "parent-session",
+            "parent-turn",
+            "subagent-session",
+        )
+        .await;
+        let preemption = tokio_util::sync::CancellationToken::new();
+        preemption.cancel();
+
+        let result = coordinator
+            .wait_for_background_subagent_outcomes(
+                "parent-session",
+                std::slice::from_ref(&registered.bg_task_id),
+                BackgroundSubagentWaitMode::All,
+                Duration::from_secs(10),
+                "steered-wait-turn",
+                None,
+                Some(&preemption),
+            )
+            .await
+            .expect("steering should end AgentWait normally");
+
+        assert_eq!(result.status.as_str(), "steered");
+        assert!(result.outcomes.is_empty());
+        assert_eq!(result.pending_bg_task_ids, vec![registered.bg_task_id]);
+    }
+
+    #[tokio::test]
+    async fn steered_agent_wait_returns_and_consumes_available_partial_results() {
+        let (coordinator, _) = test_coordinator();
+        let completed_task = register_test_background_task(
+            &coordinator,
+            "parent-session",
+            "parent-turn",
+            "subagent-session-completed",
+        )
+        .await;
+        let pending_task = register_test_background_task(
+            &coordinator,
+            "parent-session",
+            "parent-turn",
+            "subagent-session-pending",
+        )
+        .await;
+        let completed = super::SubagentResult::completed("done".to_string());
+        coordinator
+            .background_subagent_outcomes
+            .complete(completed_task.task_pk, Ok(&completed))
+            .await;
+        let preemption = tokio_util::sync::CancellationToken::new();
+        preemption.cancel();
+
+        let result = coordinator
+            .wait_for_background_subagent_outcomes(
+                "parent-session",
+                &[],
+                BackgroundSubagentWaitMode::All,
+                Duration::from_secs(10),
+                "steered-wait-turn",
+                None,
+                Some(&preemption),
+            )
+            .await
+            .expect("steering should return collected outcomes");
+
+        assert_eq!(result.status.as_str(), "steered");
+        assert_eq!(result.outcomes.len(), 1);
+        assert_eq!(result.outcomes[0].bg_task_id, completed_task.bg_task_id);
+        assert_eq!(result.pending_bg_task_ids, vec![pending_task.bg_task_id]);
+
+        let retry = coordinator
+            .wait_for_background_subagent_outcomes(
+                "parent-session",
+                std::slice::from_ref(&completed_task.bg_task_id),
+                BackgroundSubagentWaitMode::All,
+                Duration::from_millis(10),
+                "retry-wait-turn",
+                None,
+                None,
+            )
+            .await
+            .expect("a steered wait should consume returned outcomes");
+        assert_eq!(retry.status.as_str(), "no_matching_tasks");
+    }
+
     #[test]
     fn external_subagent_surfaces_use_logical_id_instead_of_runtime_generation_key() {
         let runtime_type = "external_subagent_runtime:generation-hash";

src/crates/assembly/core/src/agentic/coordination/mod.rs

@@ -17,7 +17,7 @@ pub use turn_outcome::*;
 
 pub(crate) use background_outcomes::{
     BackgroundSubagentOutcome, BackgroundSubagentOutcomeStore, BackgroundSubagentWaitMode,
-    BackgroundSubagentWaitResult,
+    BackgroundSubagentWaitResult, BackgroundSubagentWaitStatus,
 };
 pub(crate) use coordination_store::DirectChildAgentRecord;
 

src/crates/assembly/core/src/agentic/execution/round_executor.rs

@@ -1129,6 +1129,7 @@ impl RoundExecutor {
                     Some(format!("round-budget-{}", round_id)),
                     self.computer_use_host(),
                     CancellationToken::new(),
+                    None,
                 );
 
             // Execute tools — convert pipeline-level Err into per-tool error results

src/crates/assembly/core/src/agentic/tools/framework.rs

@@ -134,6 +134,12 @@ pub trait Tool: Send + Sync {
         false
     }
 
+    /// Whether a pending round injection may end this tool normally so the
+    /// current turn can advance without cancelling the underlying work.
+    fn round_injection_yieldable(&self) -> bool {
+        false
+    }
+
     /// Whether to support streaming output
     fn supports_streaming(&self) -> bool {
         false

src/crates/assembly/core/src/agentic/tools/implementations/agent_wait_tool.rs

@@ -1,6 +1,6 @@
 use crate::agentic::coordination::{
     get_global_coordinator, BackgroundSubagentOutcome, BackgroundSubagentWaitMode,
-    BackgroundSubagentWaitResult,
+    BackgroundSubagentWaitResult, BackgroundSubagentWaitStatus,
 };
 use crate::agentic::tools::framework::{
     PermissionIntent, Tool, ToolRenderOptions, ToolResult, ToolUseContext, ValidationResult,
@@ -93,15 +93,21 @@ impl AgentWaitTool {
     }
 
     fn assistant_result(result: &BackgroundSubagentWaitResult) -> String {
+        let finished = if result.status == BackgroundSubagentWaitStatus::Steered {
+            "AgentWait ended early because user steering arrived. Background agents continue running."
+                .to_string()
+        } else {
+            format!("AgentWait finished with status {}.", result.status.as_str())
+        };
         if result.outcomes.is_empty() {
             return format!(
-                "AgentWait finished with status {}. Pending background task IDs: {}.",
-                result.status.as_str(),
+                "{} Pending background task IDs: {}.",
+                finished,
                 result.pending_bg_task_ids.join(", ")
             );
         }
 
-        let mut message = format!("AgentWait finished with status {}.", result.status.as_str());
+        let mut message = finished;
         for outcome in &result.outcomes {
             message.push_str(&format!(
                 "\n<result bg_task_id=\"{}\" agent_id=\"{}\" status=\"{}\">",
@@ -138,6 +144,10 @@ impl Tool for AgentWaitTool {
         true
     }
 
+    fn round_injection_yieldable(&self) -> bool {
+        true
+    }
+
     async fn description(&self) -> BitFunResult<String> {
         Ok("Wait for background agent results.
 Wait for every selected task to complete. The tool also returns when `timeout_seconds` has elapsed.".to_string())
@@ -230,6 +240,7 @@ Wait for every selected task to complete. The tool also returns when `timeout_se
                 Duration::from_secs(request.timeout_seconds),
                 dialog_turn_id,
                 context.cancellation_token(),
+                context.round_injection_preemption_token(),
             )
             .await?;
         let data = json!({
@@ -250,6 +261,13 @@ mod tests {
     use super::{AgentWaitTool, DEFAULT_TIMEOUT_SECONDS, MAX_TIMEOUT_SECONDS, MIN_TIMEOUT_SECONDS};
     use crate::agentic::tools::framework::Tool;
 
+    #[test]
+    fn agent_wait_owns_timeout_and_yields_to_round_injection() {
+        let tool = AgentWaitTool::new();
+        assert!(tool.manages_own_execution_timeout());
+        assert!(tool.round_injection_yieldable());
+    }
+
     #[test]
     fn missing_or_empty_task_ids_are_tolerated_by_the_parser() {
         let request = AgentWaitTool::parse_request(&serde_json::json!({})).expect("valid request");