From 4b62b523f14e18440aaddfa0613081afaf1e0a9f Mon Sep 17 00:00:00 2001 From: daniel Date: Wed, 26 Aug 2026 20:40:25 +0100 Subject: [PATCH 1/2] fix(acp): await child background work --- src/protocols/acp.rs | 631 +++++++++++++++++++++++++++++++++++++--- src/protocols/acp/v2.rs | 277 +++++++++++++++--- src/runtime.rs | 176 +++++++++-- src/runtime/tests.rs | 51 ++++ 4 files changed, 1045 insertions(+), 90 deletions(-) diff --git a/src/protocols/acp.rs b/src/protocols/acp.rs index 9db5d58..559f713 100644 --- a/src/protocols/acp.rs +++ b/src/protocols/acp.rs @@ -336,6 +336,7 @@ struct SessionHandle { token: u64, commands: mpsc::Sender, background_jobs: BackgroundJobs, + structured_completion: bool, tasks: TaskManagerHandle, } @@ -344,6 +345,8 @@ struct RegisteredSession { token: u64, session_id: agentkit_acp::SessionId, integration: Arc, + background_jobs: BackgroundJobs, + tasks: TaskManagerHandle, commands: mpsc::WeakSender, actor: AbortHandle, completed: watch::Receiver, @@ -468,6 +471,7 @@ impl SessionRegistry { async fn shutdown_with_timeout(&self, limit: Duration) { let (sessions, v2_sessions) = self.close_gate_and_snapshot(); for session in &sessions { + session.background_jobs.cancel_all(); let _ = session.integration.interrupt_session(&session.session_id); } for session in &v2_sessions { @@ -478,6 +482,7 @@ impl SessionRegistry { for mut session in sessions.iter().cloned() { let registry = self.clone(); closing.spawn(async move { + cancel_background_jobs(&session.tasks, &session.background_jobs).await; if let Some(commands) = session.commands.upgrade() { let (reply, acknowledged) = oneshot::channel(); if commands.send(Command::Close { reply }).await.is_ok() { @@ -785,6 +790,7 @@ impl Server { let config_options = config_options(¤t, reasoning_effort, &catalog); let background_jobs = driver.background_jobs.clone(); let tasks = driver.tasks.clone(); + let structured_completion = driver.structured_completion; let canonical_transcript = driver.canonical_transcript; let (tx, rx) = mpsc::channel(8); let actor = SessionActor { @@ -793,6 +799,8 @@ impl Server { binding, driver: driver.driver, tasks: driver.tasks, + background_jobs: background_jobs.clone(), + structured_completion, adapter: driver.adapter, catalog, commands: rx, @@ -819,6 +827,8 @@ impl Server { token, session_id: session_id.clone(), integration: Arc::clone(&self.integration), + background_jobs: background_jobs.clone(), + tasks: tasks.clone(), commands: tx.downgrade(), actor: actor_task.abort_handle(), completed: completion, @@ -848,6 +858,7 @@ impl Server { token, commands: tx, background_jobs, + structured_completion, tasks, }, ); @@ -887,9 +898,25 @@ impl Server { // Interrupt out of band because the actor may currently be inside // `driver.next()`. The queued marker preserves command ordering once // that call settles. + let (sender, background_jobs, tasks, structured_completion) = self + .sessions + .lock() + .expect("ACP session map poisoned") + .get(¬ification.session_id) + .map(|session| { + ( + session.commands.clone(), + session.background_jobs.clone(), + session.tasks.clone(), + session.structured_completion, + ) + }) + .ok_or_else(|| AcpRuntimeError::SessionNotFound(notification.session_id.to_string()))?; + if structured_completion { + cancel_background_jobs(&tasks, &background_jobs).await; + } self.integration .interrupt_session(¬ification.session_id)?; - let sender = self.sender(¬ification.session_id).await?; sender .send(Command::Cancel) .await @@ -902,13 +929,14 @@ impl Server { ) -> Result { // Closing uses the same out-of-band interrupt, then waits for the // actor to reach and acknowledge the serialized close boundary. - self.integration.interrupt_session(&request.session_id)?; let session = self .sessions .lock() .expect("ACP session map poisoned") .remove(&request.session_id) .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; + cancel_background_jobs(&session.tasks, &session.background_jobs).await; + self.integration.interrupt_session(&request.session_id)?; let (tx, rx) = oneshot::channel(); session .commands @@ -996,6 +1024,8 @@ struct SessionActor { binding: SessionBindingGuard, driver: LoopDriver, tasks: TaskManagerHandle, + background_jobs: BackgroundJobs, + structured_completion: bool, adapter: SelectableAdapter, catalog: Vec, commands: mpsc::Receiver, @@ -1010,6 +1040,8 @@ async fn session_actor(actor: SessionActor) { binding, mut driver, tasks, + background_jobs, + structured_completion, adapter, catalog, mut commands, @@ -1025,7 +1057,15 @@ async fn session_actor(actor: SessionActor) { biased; command = commands.recv() => match command { Some(Command::Prompt { request, reply }) => { - let result = drive_prompt(&session_id, &integration, &mut driver, request).await; + let result = drive_prompt( + &session_id, + &integration, + &mut driver, + request, + &tasks, + &background_jobs, + structured_completion, + ).await; let _ = reply.send(result); } // The server already interrupted the shared controller; this @@ -1036,7 +1076,7 @@ async fn session_actor(actor: SessionActor) { let _ = reply.send(result); } Some(Command::Close { reply }) => { - clean_up_session(&session_id, &mut driver, &tasks).await; + clean_up_session(&session_id, &mut driver, &tasks, &background_jobs).await; // A close acknowledgement means the actor-owned binding is // already gone, so callers can immediately reuse the id. drop(binding.take()); @@ -1044,7 +1084,7 @@ async fn session_actor(actor: SessionActor) { break; } None => { - clean_up_session(&session_id, &mut driver, &tasks).await; + clean_up_session(&session_id, &mut driver, &tasks, &background_jobs).await; break; }, }, @@ -1071,20 +1111,24 @@ async fn session_actor(actor: SessionActor) { // Task events remain queued while a prompt is being driven, so a // completion cannot be lost in the prompt-to-idle transition. event = tasks.next_event() => match event { - Some(TaskEvent::Completed(snapshot, _)) - if snapshot.kind == agentkit_task_manager::TaskKind::Background => - { - let result = drive_unsolicited( - &session_id, - &integration, - &mut driver, - &turn_states, - &mut next_autonomous_turn_id, - ).await; - if let Err(error) = result { - eprintln!("autonomous ACP continuation failed for {session_id}: {error}"); + Some(TaskEvent::Completed(snapshot, _)) => { + background_jobs.acknowledge_terminal(&snapshot.call_id); + if snapshot.kind == agentkit_task_manager::TaskKind::Background { + let result = drive_unsolicited( + &session_id, + &integration, + &mut driver, + &turn_states, + &mut next_autonomous_turn_id, + ).await; + if let Err(error) = result { + eprintln!("autonomous ACP continuation failed for {session_id}: {error}"); + } } } + Some(TaskEvent::Cancelled(snapshot) | TaskEvent::Failed(snapshot, _)) => { + background_jobs.acknowledge_terminal(&snapshot.call_id); + } Some(_) => {} None => break, } @@ -1191,16 +1235,89 @@ pub(super) fn set_config( ))) } +pub(super) async fn settle_background_jobs( + tasks: &TaskManagerHandle, + background_jobs: &BackgroundJobs, +) -> Result { + let initial = background_jobs.activity(); + let mut observed_activity = false; + loop { + let activity = background_jobs.activity(); + let running = tasks + .list_running() + .await + .into_iter() + .any(|task| task.kind == agentkit_task_manager::TaskKind::Background); + let stable = background_jobs.activity(); + if stable.generation != activity.generation { + observed_activity = true; + continue; + } + observed_activity |= stable.background_started > initial.background_started + || stable.active + || running + || stable.unacknowledged_terminals; + if !stable.active && !running && !stable.unacknowledged_terminals { + return Ok(observed_activity); + } + tokio::select! { + biased; + _ = background_jobs.activity_after(stable.generation) => {} + event = tasks.next_event(), if running || stable.unacknowledged_terminals => { + match event { + Some( + TaskEvent::Completed(snapshot, _) + | TaskEvent::Cancelled(snapshot) + | TaskEvent::Failed(snapshot, _), + ) => { + background_jobs.acknowledge_terminal(&snapshot.call_id); + } + Some(_) => {} + None => { + return Err(AcpRuntimeError::Loop( + "background task event stream closed before quiescence".into(), + )); + } + } + } + } + } +} + +pub(super) async fn cancel_background_jobs( + tasks: &TaskManagerHandle, + background_jobs: &BackgroundJobs, +) { + background_jobs.cancel_all(); + let _ = timeout( + Duration::from_millis(100), + background_jobs.wait_for_quiescence(), + ) + .await; + for task in tasks.list_running().await { + if task.kind == agentkit_task_manager::TaskKind::Background { + let _ = tasks.cancel(task.id).await; + } + } +} + pub(super) async fn clean_up_session( session_id: &agentkit_acp::SessionId, driver: &mut LoopDriver, tasks: &TaskManagerHandle, + background_jobs: &BackgroundJobs, ) { + cancel_background_jobs(tasks, background_jobs).await; for task in tasks.list_running().await { if let Err(error) = tasks.cancel(task.id).await { eprintln!("failed to cancel ACP task for {session_id}: {error}"); } } + let _ = tokio::time::timeout( + Duration::from_secs(1), + background_jobs.wait_for_quiescence(), + ) + .await; if let Err(error) = driver.cancel_pending_approvals().await { eprintln!("failed to cancel ACP approvals for {session_id}: {error}"); } @@ -1251,14 +1368,46 @@ async fn drive_prompt( integration: &AcpIntegration, driver: &mut LoopDriver, request: PromptRequest, + tasks: &TaskManagerHandle, + background_jobs: &BackgroundJobs, + structured_completion: bool, ) -> Result { + if structured_completion { + let _ = settle_background_jobs(tasks, background_jobs).await?; + } + background_jobs.begin_turn(); let items = integration.input_port().prompt_to_items(&request)?; driver .submit_input(items) .map_err(|error| record_acp_loop_failure(session_id, &error))?; - drive_until_pause(session_id, integration, driver, true) - .await? - .ok_or_else(|| AcpRuntimeError::Loop("prompt ended without a response".into())) + let response = match drive_until_pause( + session_id, + integration, + driver, + true, + structured_completion.then_some((tasks, background_jobs)), + ) + .await + { + Ok(Some(response)) => response, + Ok(None) => { + return Err(AcpRuntimeError::Loop( + "prompt ended without a response".into(), + )); + } + Err(error) => { + if structured_completion { + cancel_background_jobs(tasks, background_jobs).await; + let _ = settle_background_jobs(tasks, background_jobs).await; + } + return Err(error); + } + }; + if structured_completion && response.stop_reason == StopReason::Cancelled { + cancel_background_jobs(tasks, background_jobs).await; + let _ = settle_background_jobs(tasks, background_jobs).await?; + } + Ok(response) } async fn drive_unsolicited( @@ -1292,7 +1441,7 @@ async fn drive_autonomous( integration: &AcpIntegration, driver: &mut LoopDriver, ) -> Result<(), AcpRuntimeError> { - let _ = drive_until_pause(session_id, integration, driver, false).await?; + let _ = drive_until_pause(session_id, integration, driver, false, None).await?; Ok(()) } @@ -1301,6 +1450,7 @@ async fn drive_until_pause( integration: &AcpIntegration, driver: &mut LoopDriver, answer_prompt: bool, + structured: Option<(&TaskManagerHandle, &BackgroundJobs)>, ) -> Result, AcpRuntimeError> { loop { let step = match driver.next().await { @@ -1321,10 +1471,21 @@ async fn drive_until_pause( return Ok(None); } let reason = agentkit_acp::finish_reason_to_stop_reason(&result.finish_reason)?; + if let Some((tasks, background_jobs)) = structured + && settle_background_jobs(tasks, background_jobs).await? + { + continue; + } return Ok(Some(PromptResponse::new(reason))); } LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => { integration.flush_session_updates(session_id).await?; + if answer_prompt + && let Some((tasks, background_jobs)) = structured + && settle_background_jobs(tasks, background_jobs).await? + { + continue; + } return Ok(answer_prompt.then(|| PromptResponse::new(StopReason::EndTurn))); } LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => continue, @@ -1599,7 +1760,7 @@ async fn drain_client_messages( } #[cfg(test)] -mod tests { +pub(super) mod tests { use std::{ collections::VecDeque, sync::atomic::{AtomicBool, AtomicUsize, Ordering}, @@ -1614,9 +1775,13 @@ mod tests { Agent, LoopError, ModelAdapter, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, TurnRequest, }; - use agentkit_task_manager::{AsyncTaskManager, RoutingDecision, TaskManager}; + use agentkit_task_manager::{ + AsyncTaskManager, RoutingDecision, TaskLaunchRequest, TaskManager, TaskStartContext, + }; use agentkit_tools_core::{ - Tool, ToolAnnotations, ToolContext, ToolName, ToolRegistry, ToolResult, ToolSpec, + AllowAllPermissions, BasicToolExecutor, OwnedToolContext, Tool, ToolAnnotations, + ToolContext, ToolExecutor, ToolName, ToolRegistry, ToolRequest, ToolResult, ToolSource, + ToolSpec, }; use async_trait::async_trait; use tokio::{ @@ -1645,9 +1810,10 @@ mod tests { } #[tokio::test] - async fn shutdown_sends_close_and_rejects_new_registration() { + async fn shutdown_force_cancels_background_before_close_and_rejects_registration() { let registry = SessionRegistry::new(); let integration = shutdown_test_integration(); + let (background_jobs, tasks) = start_non_cooperative_background("shutdown-call").await; let session_id = agentkit_acp::SessionId::new("close-me"); let token = registry.next_token(); let (commands, mut received) = mpsc::channel(1); @@ -1667,6 +1833,8 @@ mod tests { token, session_id, integration: Arc::clone(&integration), + background_jobs: background_jobs.clone(), + tasks: tasks.clone(), commands: weak_commands, actor: actor.abort_handle(), completed: completion, @@ -1676,6 +1844,8 @@ mod tests { registry.shutdown_with_timeout(Duration::from_secs(1)).await; assert!(closed.load(Ordering::SeqCst)); + assert!(tasks.list_running().await.is_empty()); + assert!(!background_jobs.activity().active); let late_token = registry.next_token(); let (late_commands, _late_received) = mpsc::channel(1); @@ -1690,6 +1860,8 @@ mod tests { token: late_token, session_id: agentkit_acp::SessionId::new("too-late"), integration, + background_jobs: BackgroundJobs::default(), + tasks: AsyncTaskManager::new().handle(), commands: late_commands.downgrade(), actor: late_actor.abort_handle(), completed: late_completion, @@ -1712,6 +1884,8 @@ mod tests { token, session_id: agentkit_acp::SessionId::new("stuck"), integration: shutdown_test_integration(), + background_jobs: BackgroundJobs::default(), + tasks: AsyncTaskManager::new().handle(), commands: commands.downgrade(), actor: actor.abort_handle(), completed: completion, @@ -1963,19 +2137,19 @@ mod tests { integration.unbind_session(&session_id).unwrap(); } - struct ScriptAdapter { - turns: Arc, - user_items_seen: Arc, - notification_items_seen: Arc, + pub(super) struct ScriptAdapter { + pub(super) turns: Arc, + pub(super) user_items_seen: Arc, + pub(super) notification_items_seen: Arc, } - struct ScriptSession { + pub(super) struct ScriptSession { turns: Arc, user_items_seen: Arc, notification_items_seen: Arc, } - struct ScriptTurn { + pub(super) struct ScriptTurn { events: VecDeque, } @@ -2129,10 +2303,10 @@ mod tests { } } - struct BlockingTool { - spec: ToolSpec, - entered: Arc, - release: Arc, + pub(super) struct BlockingTool { + pub(super) spec: ToolSpec, + pub(super) entered: Arc, + pub(super) release: Arc, } #[async_trait] @@ -2161,6 +2335,259 @@ mod tests { } } + struct FinishBackgroundOnDrop { + jobs: BackgroundJobs, + call_id: String, + } + + impl Drop for FinishBackgroundOnDrop { + fn drop(&mut self) { + self.jobs.finish_for_test(&self.call_id); + } + } + + struct NonCooperativeTool { + spec: ToolSpec, + entered: Arc, + jobs: BackgroundJobs, + } + + #[async_trait] + impl Tool for NonCooperativeTool { + fn spec(&self) -> &ToolSpec { + &self.spec + } + + async fn invoke( + &self, + request: ToolRequest, + _ctx: &mut ToolContext<'_>, + ) -> Result { + let _finish = FinishBackgroundOnDrop { + jobs: self.jobs.clone(), + call_id: request.call_id.to_string(), + }; + self.entered.store(true, Ordering::SeqCst); + std::future::pending::<()>().await; + unreachable!() + } + } + + struct FailingBackgroundTool { + spec: ToolSpec, + jobs: BackgroundJobs, + } + + #[async_trait] + impl Tool for FailingBackgroundTool { + fn spec(&self) -> &ToolSpec { + &self.spec + } + + async fn invoke( + &self, + request: ToolRequest, + _ctx: &mut ToolContext<'_>, + ) -> Result { + let _finish = FinishBackgroundOnDrop { + jobs: self.jobs.clone(), + call_id: request.call_id.to_string(), + }; + Err(agentkit_tools_core::ToolError::Unavailable( + "fast failure".into(), + )) + } + } + + fn background_task_context(executor: Arc) -> TaskStartContext { + TaskStartContext { + executor, + tool_context: OwnedToolContext { + session_id: agentkit_core::SessionId::new("background-session"), + turn_id: agentkit_core::TurnId::new("background-turn"), + metadata: MetadataMap::new(), + permissions: Arc::new(AllowAllPermissions), + resources: Arc::new(()), + cancellation: None, + execution_scope: None, + approved_request: None, + }, + } + } + + fn background_task_request(call_id: &str, tool_name: &str) -> TaskLaunchRequest { + TaskLaunchRequest::plain( + None, + ToolRequest { + call_id: ToolCallId::new(call_id), + tool_name: ToolName::new(tool_name), + input: json!({}), + session_id: agentkit_core::SessionId::new("background-session"), + turn_id: agentkit_core::TurnId::new("background-turn"), + metadata: MetadataMap::new(), + }, + ) + } + + async fn start_non_cooperative_background( + call_id: &str, + ) -> (BackgroundJobs, TaskManagerHandle) { + let jobs = BackgroundJobs::default(); + jobs.register_foreground_for_test(call_id); + assert_eq!(jobs.detach(call_id), Some(DetachRegistration::Registered)); + let entered = Arc::new(AtomicBool::new(false)); + let tools = ToolRegistry::new().with(NonCooperativeTool { + spec: ToolSpec { + name: ToolName::new("non-cooperative"), + description: "ignores cooperative cancellation".into(), + input_schema: json!({"type": "object"}), + output_schema: None, + annotations: ToolAnnotations::default(), + metadata: MetadataMap::new(), + }, + entered: Arc::clone(&entered), + jobs: jobs.clone(), + }); + let executor: Arc = Arc::new(BasicToolExecutor::new([ + Arc::new(tools) as Arc + ])); + let manager = + AsyncTaskManager::new().routing(|_request: &ToolRequest| RoutingDecision::Background); + let tasks = manager.handle(); + manager + .start_task( + background_task_request(call_id, "non-cooperative"), + background_task_context(executor), + ) + .await + .unwrap(); + timeout(Duration::from_secs(1), async { + while !entered.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .expect("non-cooperative task did not start"); + (jobs, tasks) + } + + #[tokio::test] + async fn depth_zero_cancel_leaves_detached_background_running() { + let (jobs, tasks) = start_non_cooperative_background("root-call").await; + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let integration = AcpIntegration::builder() + .name("root-cancel-test") + .approval_resolver(AutoDenyResolver) + .build() + .unwrap(); + let session_id = agentkit_acp::SessionId::new("root-cancel-session"); + let (client, _messages) = AcpClientHandle::channel(); + integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + AgentkitSessionId::new("root-cancel-session"), + client, + )) + .unwrap(); + let server = Server::new(runtime, integration, SessionRegistry::new()); + let (commands, mut received) = mpsc::channel(1); + server.sessions.lock().unwrap().insert( + session_id.clone(), + SessionHandle { + token: 1, + commands, + background_jobs: jobs.clone(), + structured_completion: false, + tasks: tasks.clone(), + }, + ); + + server + .cancel(CancelNotification::new(session_id)) + .await + .unwrap(); + assert!(matches!(received.recv().await, Some(Command::Cancel))); + assert!(!jobs.is_cancelled_for_test("root-call")); + assert!(!tasks.list_running().await.is_empty()); + + cancel_background_jobs(&tasks, &jobs).await; + timeout( + Duration::from_secs(1), + settle_background_jobs(&tasks, &jobs), + ) + .await + .unwrap() + .unwrap(); + } + + #[tokio::test] + async fn failed_background_terminal_settles_after_fast_completion() { + let jobs = BackgroundJobs::default(); + jobs.register_foreground_for_test("failed-call"); + assert_eq!( + jobs.detach("failed-call"), + Some(DetachRegistration::Registered) + ); + let tools = ToolRegistry::new().with(FailingBackgroundTool { + spec: ToolSpec { + name: ToolName::new("fast-failure"), + description: "fails immediately".into(), + input_schema: json!({"type": "object"}), + output_schema: None, + annotations: ToolAnnotations::default(), + metadata: MetadataMap::new(), + }, + jobs: jobs.clone(), + }); + let executor: Arc = Arc::new(BasicToolExecutor::new([ + Arc::new(tools) as Arc + ])); + let manager = + AsyncTaskManager::new().routing(|_request: &ToolRequest| RoutingDecision::Background); + let tasks = manager.handle(); + manager + .start_task( + background_task_request("failed-call", "fast-failure"), + background_task_context(executor), + ) + .await + .unwrap(); + + assert!( + timeout( + Duration::from_secs(1), + settle_background_jobs(&tasks, &jobs) + ) + .await + .expect("failed background terminal did not settle") + .unwrap() + ); + assert!(tasks.list_running().await.is_empty()); + } + + #[tokio::test] + async fn non_cooperative_background_is_force_cancelled_before_settlement() { + let (jobs, tasks) = start_non_cooperative_background("stuck-call").await; + + timeout( + Duration::from_secs(1), + cancel_background_jobs(&tasks, &jobs), + ) + .await + .expect("bounded background cancellation hung"); + assert!(tasks.list_running().await.is_empty()); + assert!( + timeout( + Duration::from_secs(1), + settle_background_jobs(&tasks, &jobs) + ) + .await + .expect("forced cancellation did not settle") + .unwrap() + ); + } + #[tokio::test] async fn cancelled_driver_returns_normal_cancelled_prompt_response() { let integration = Arc::new( @@ -2196,6 +2623,9 @@ mod tests { .await .unwrap(); + let task_manager = AsyncTaskManager::new(); + let tasks = task_manager.handle(); + let background_jobs = BackgroundJobs::default(); let response = drive_prompt( &acp_session_id, &integration, @@ -2206,6 +2636,9 @@ mod tests { agentkit_acp::TextContent::new("cancel me"), )], ), + &tasks, + &background_jobs, + false, ) .await .expect("cancellation must be an ACP response, not an RPC error"); @@ -2214,6 +2647,131 @@ mod tests { drain.abort(); } + #[tokio::test] + async fn structured_prompt_waits_for_background_completion_and_synthesis() { + let turns = Arc::new(AtomicUsize::new(0)); + let user_items_seen = Arc::new(AtomicUsize::new(0)); + let notification_items_seen = Arc::new(AtomicUsize::new(0)); + let entered = Arc::new(AtomicBool::new(false)); + let release = Arc::new(Notify::new()); + let integration = Arc::new( + AcpIntegration::builder() + .name("structured-test") + .approval_resolver(AutoDenyResolver) + .build() + .unwrap(), + ); + let acp_session_id = agentkit_acp::SessionId::new("s-structured-session"); + let agentkit_session_id = AgentkitSessionId::new("s-structured-session"); + let (client, mut messages) = AcpClientHandle::channel(); + integration + .bind_session(AcpSessionBinding::new( + acp_session_id.clone(), + agentkit_session_id.clone(), + client, + )) + .unwrap(); + let drain = tokio::spawn(async move { + while let Some(message) = messages.recv().await { + if let AcpClientMessage::Flush { response } = message { + let _ = response.send(()); + } + } + }); + + let task_manager = AsyncTaskManager::new() + .routing(|_request: &agentkit_tools_core::ToolRequest| RoutingDecision::Foreground); + let tasks = task_manager.handle(); + let tools = ToolRegistry::new().with(BlockingTool { + spec: ToolSpec { + name: ToolName::new(agentkit_tool_compose::COMPOSE_TOOL_NAME), + description: "controlled compose tool".into(), + input_schema: json!({"type": "object", "additionalProperties": false}), + output_schema: None, + annotations: ToolAnnotations::default(), + metadata: MetadataMap::new(), + }, + entered: Arc::clone(&entered), + release: Arc::clone(&release), + }); + let mut driver = Agent::builder() + .model(ScriptAdapter { + turns: Arc::clone(&turns), + user_items_seen, + notification_items_seen, + }) + .add_tool_source(tools) + .task_manager(task_manager) + .observer(integration.as_ref().clone()) + .build() + .unwrap() + .start(SessionConfig::new(agentkit_session_id).without_cache()) + .await + .unwrap(); + let background_jobs = BackgroundJobs::default(); + let request = PromptRequest::new( + acp_session_id.clone(), + vec![agentkit_acp::ContentBlock::Text( + agentkit_acp::TextContent::new("start one background call"), + )], + ); + let prompt = drive_prompt( + &acp_session_id, + &integration, + &mut driver, + request, + &tasks, + &background_jobs, + true, + ); + tokio::pin!(prompt); + + timeout(Duration::from_secs(1), async { + tokio::select! { + response = &mut prompt => panic!("structured prompt resolved before its task started: {response:?}"), + _ = async { + while !entered.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + assert!(detach_compose_call( + &tasks, + &background_jobs, + "background-call", + ).await); + background_jobs.register_foreground_for_test("background-call"); + while turns.load(Ordering::SeqCst) < 2 { + tokio::task::yield_now().await; + } + } => {} + } + }) + .await + .expect("structured prompt did not reach its provisional response"); + assert!( + timeout(Duration::from_millis(20), &mut prompt) + .await + .is_err() + ); + + // The compose guard can disappear before the task manager publishes its + // result. Structured completion must wait through that handoff. + background_jobs.finish_for_test("background-call"); + assert!( + timeout(Duration::from_millis(20), &mut prompt) + .await + .is_err(), + "structured prompt crossed the terminal publication handoff early" + ); + release.notify_one(); + let response = timeout(Duration::from_secs(1), &mut prompt) + .await + .expect("structured prompt did not synthesize the background result") + .unwrap(); + assert_eq!(response.stop_reason, StopReason::EndTurn); + assert_eq!(turns.load(Ordering::SeqCst), 3); + drain.abort(); + } + #[tokio::test] async fn foreground_compose_detaches_out_of_band_and_completes_autonomously() { let turns = Arc::new(AtomicUsize::new(0)); @@ -2297,6 +2855,8 @@ mod tests { binding: SessionBindingGuard::new(Arc::clone(&integration), acp_session_id.clone()), driver, tasks: tasks.clone(), + background_jobs: background_jobs.clone(), + structured_completion: false, adapter: SelectableAdapter::new(crate::ProviderKind::OpenAiSubscription, "gpt-5.4") .unwrap(), catalog: Vec::new(), @@ -2347,6 +2907,7 @@ mod tests { completed[0].kind, agentkit_task_manager::TaskKind::Background ); + background_jobs.finish_for_test("background-call"); let notification = timeout(Duration::from_secs(1), async { loop { let notification = updates_rx.recv().await.expect("update stream closed"); diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index 54129c0..ee6ecba 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -364,6 +364,7 @@ struct SessionHandle { integration: AcpSessionHandle, busy: Arc, background_jobs: BackgroundJobs, + structured_completion: bool, tasks: TaskManagerHandle, } @@ -609,6 +610,7 @@ impl Server { let canonical_transcript = driver.canonical_transcript; let background_jobs = driver.background_jobs.clone(); let tasks = driver.tasks.clone(); + let structured_completion = driver.structured_completion; let mcp_events = self.runtime.subscribe_mcp(session_id.to_string()); let (tx, rx) = mpsc::channel(8); let busy = Arc::new(AtomicBool::new(false)); @@ -621,6 +623,8 @@ impl Server { sink, driver: driver.driver, tasks: driver.tasks, + background_jobs: background_jobs.clone(), + structured_completion, adapter: driver.adapter, catalog, commands: rx, @@ -643,11 +647,20 @@ impl Server { } }); let interrupt_handle = handle.clone(); - let interrupt = Arc::new(move || interrupt_handle.interrupt()); + let interrupt_background_jobs = background_jobs.clone(); + let interrupt = Arc::new(move || { + interrupt_background_jobs.cancel_all(); + interrupt_handle.interrupt(); + }); let weak = tx.downgrade(); + let close_background_jobs = background_jobs.clone(); + let close_tasks = tasks.clone(); let close = Arc::new(move || { + let close_background_jobs = close_background_jobs.clone(); + let close_tasks = close_tasks.clone(); let weak = weak.clone(); Box::pin(async move { + super::cancel_background_jobs(&close_tasks, &close_background_jobs).await; if let Some(commands) = weak.upgrade() { let (reply, acknowledged) = oneshot::channel(); if commands.send(Command::Close { reply }).await.is_ok() { @@ -684,6 +697,7 @@ impl Server { integration: handle, busy, background_jobs, + structured_completion, tasks, }, ); @@ -770,13 +784,23 @@ impl Server { &self, notification: wire::CancelSessionNotification, ) -> Result<(), AcpRuntimeError> { - let handle = self + let session = self .sessions .lock() .expect("ACP v2 session map poisoned") .get(¬ification.session_id) - .map(|session| session.integration.clone()); - if let Some(handle) = handle { + .map(|session| { + ( + session.integration.clone(), + session.background_jobs.clone(), + session.tasks.clone(), + session.structured_completion, + ) + }); + if let Some((handle, background_jobs, tasks, structured_completion)) = session { + if structured_completion { + super::cancel_background_jobs(&tasks, &background_jobs).await; + } handle.interrupt(); } Ok(()) @@ -792,6 +816,7 @@ impl Server { .expect("ACP v2 session map poisoned") .remove(&request.session_id) .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; + super::cancel_background_jobs(&session.tasks, &session.background_jobs).await; session.integration.close(); let (reply, acknowledged) = oneshot::channel(); session @@ -862,6 +887,8 @@ struct SessionActor { sink: ResponseReplacementSink, driver: LoopDriver, tasks: TaskManagerHandle, + background_jobs: BackgroundJobs, + structured_completion: bool, adapter: SelectableAdapter, catalog: Vec, commands: mpsc::Receiver, @@ -878,6 +905,8 @@ async fn session_actor(actor: SessionActor) sink, mut driver, tasks, + background_jobs, + structured_completion, adapter, catalog, mut commands, @@ -896,6 +925,9 @@ async fn session_actor(actor: SessionActor) &mut driver, command, &sink, + &tasks, + &background_jobs, + structured_completion, ) .await; busy.store(false, Ordering::Release); @@ -909,14 +941,14 @@ async fn session_actor(actor: SessionActor) } Some(Command::Close { reply }) => { let v1_id = agentkit_acp::SessionId::new(session_id.to_string()); - super::clean_up_session(&v1_id, &mut driver, &tasks).await; + super::clean_up_session(&v1_id, &mut driver, &tasks, &background_jobs).await; drop(binding.take()); let _ = reply.send(()); break; } None => { let v1_id = agentkit_acp::SessionId::new(session_id.to_string()); - super::clean_up_session(&v1_id, &mut driver, &tasks).await; + super::clean_up_session(&v1_id, &mut driver, &tasks, &background_jobs).await; break; } }, @@ -939,20 +971,24 @@ async fn session_actor(actor: SessionActor) } } event = tasks.next_event() => match event { - Some(TaskEvent::Completed(snapshot, _)) - if snapshot.kind == agentkit_task_manager::TaskKind::Background => - { - if let Err(error) = drive_autonomous( - &session_id, - &integration, - &handle, - &busy, - &mut driver, - &sink, - ).await { + Some(TaskEvent::Completed(snapshot, _)) => { + background_jobs.acknowledge_terminal(&snapshot.call_id); + if snapshot.kind == agentkit_task_manager::TaskKind::Background + && let Err(error) = drive_autonomous( + &session_id, + &integration, + &handle, + &busy, + &mut driver, + &sink, + ).await + { eprintln!("ACP v2 autonomous turn failed for {session_id}: {error}"); } } + Some(TaskEvent::Cancelled(snapshot) | TaskEvent::Failed(snapshot, _)) => { + background_jobs.acknowledge_terminal(&snapshot.call_id); + } Some(_) => {} None => break, } @@ -967,12 +1003,23 @@ async fn prepare_prompt( driver: &mut LoopDriver, command: PromptCommand, sink: &impl AcpSessionUpdateSink, + tasks: &TaskManagerHandle, + background_jobs: &BackgroundJobs, + structured_completion: bool, ) -> Result<(), AcpRuntimeError> { let PromptCommand { request, cancellation_generation, reply, } = command; + if structured_completion + && let Err(error) = super::settle_background_jobs(tasks, background_jobs).await + { + handle.stop_injection_turn(); + let _ = reply.send(Err(error)); + return Ok(()); + } + background_jobs.begin_turn(); let prepared = integration.prompt_to_items(&request).and_then(|items| { driver .submit_input(items) @@ -1006,8 +1053,14 @@ async fn prepare_prompt( session_id, wire::StateUpdate::Running(wire::RunningStateUpdate::new()), )?; - let stop_reason = match drive_prompt(session_id, driver, handle, cancellation_generation) - .await + let stop_reason = match drive_prompt( + session_id, + driver, + handle, + cancellation_generation, + structured_completion.then_some((tasks, background_jobs)), + ) + .await { Ok(stop_reason) => stop_reason, Err(_) @@ -1018,10 +1071,18 @@ async fn prepare_prompt( wire::StopReason::Cancelled } Err(error) => { + if structured_completion { + super::cancel_background_jobs(tasks, background_jobs).await; + let _ = super::settle_background_jobs(tasks, background_jobs).await; + } terminalize_running_error(session_id, integration, handle, sink, &error).await?; return Err(error); } }; + if structured_completion && stop_reason == wire::StopReason::Cancelled { + super::cancel_background_jobs(tasks, background_jobs).await; + let _ = super::settle_background_jobs(tasks, background_jobs).await?; + } let _ = integration.flush_session_updates(session_id).await; integration.finish_prompt(session_id); send_state( @@ -1031,6 +1092,10 @@ async fn prepare_prompt( ) } .await; + if result.is_err() && structured_completion { + super::cancel_background_jobs(tasks, background_jobs).await; + let _ = super::settle_background_jobs(tasks, background_jobs).await; + } integration.finish_prompt(session_id); handle.stop_injection_turn(); result @@ -1071,6 +1136,7 @@ async fn drive_prompt( driver: &mut LoopDriver, control: &C, cancellation_generation: u64, + structured: Option<(&TaskManagerHandle, &BackgroundJobs)>, ) -> Result where S: ModelSession + Send + 'static, @@ -1102,6 +1168,11 @@ where } return Err(AcpRuntimeError::Loop("model turn failed".into())); } + if let Some((tasks, background_jobs)) = structured + && super::settle_background_jobs(tasks, background_jobs).await? + { + continue; + } match control.handle_injection_boundary(driver, true).await { Ok(AcpInjectionBoundary::Delivered | AcpInjectionBoundary::Continue) => { continue; @@ -1122,6 +1193,11 @@ where } } LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => { + if let Some((tasks, background_jobs)) = structured + && super::settle_background_jobs(tasks, background_jobs).await? + { + continue; + } match control.handle_injection_boundary(driver, true).await { Ok(AcpInjectionBoundary::Delivered | AcpInjectionBoundary::Continue) => { continue; @@ -1129,7 +1205,9 @@ where Ok(AcpInjectionBoundary::Stopped) => { return Ok(wire::StopReason::Cancelled); } - Ok(AcpInjectionBoundary::Finished) => return Ok(wire::StopReason::EndTurn), + Ok(AcpInjectionBoundary::Finished) => { + return Ok(wire::StopReason::EndTurn); + } Err(error) => { control.stop_injection_turn(); if control.is_cancelled_since(cancellation_generation) { @@ -1188,22 +1266,22 @@ async fn drive_autonomous( session_id, wire::StateUpdate::Running(wire::RunningStateUpdate::new()), )?; - let stop_reason = match drive_prompt(session_id, driver, handle, cancellation_generation) - .await - { - Ok(stop_reason) => stop_reason, - Err(_) - if handle - .cancellation_handle() - .is_cancelled_since(cancellation_generation) => - { - wire::StopReason::Cancelled - } - Err(error) => { - terminalize_running_error(session_id, integration, handle, sink, &error).await?; - return Err(error); - } - }; + let stop_reason = + match drive_prompt(session_id, driver, handle, cancellation_generation, None).await { + Ok(stop_reason) => stop_reason, + Err(_) + if handle + .cancellation_handle() + .is_cancelled_since(cancellation_generation) => + { + wire::StopReason::Cancelled + } + Err(error) => { + terminalize_running_error(session_id, integration, handle, sink, &error) + .await?; + return Err(error); + } + }; let _ = integration.flush_session_updates(session_id).await; send_state( sink, @@ -1750,6 +1828,8 @@ pub(crate) fn component( #[cfg(test)] mod tests { + use std::sync::atomic::AtomicUsize; + use serde_json::json; use agent_client_protocol::schema::MaybeUndefined; @@ -1758,8 +1838,15 @@ mod tests { Agent, ModelAdapter, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, TurnRequest, TurnResult, }; + use agentkit_task_manager::{AsyncTaskManager, RoutingDecision, TaskManager}; + use agentkit_tools_core::{ToolAnnotations, ToolName, ToolRegistry, ToolSpec}; + use tokio::{ + sync::Notify, + time::{Duration, timeout}, + }; use super::*; + use crate::protocols::acp::tests::{BlockingTool, ScriptAdapter}; #[derive(Clone, Default)] struct RecordingSink { @@ -2412,6 +2499,9 @@ mod tests { response.await.unwrap().unwrap().send(()).unwrap(); }; + let task_manager = AsyncTaskManager::new(); + let tasks = task_manager.handle(); + let background_jobs = BackgroundJobs::default(); let (result, ()) = tokio::join!( prepare_prompt( &session_id, @@ -2420,6 +2510,9 @@ mod tests { &mut driver, command, &sink, + &tasks, + &background_jobs, + false, ), acknowledge, ); @@ -2459,6 +2552,114 @@ mod tests { )); } + #[tokio::test] + async fn structured_prompt_waits_for_background_synthesis_and_consumes_completion() { + let turns = Arc::new(AtomicUsize::new(0)); + let entered = Arc::new(AtomicBool::new(false)); + let release = Arc::new(Notify::new()); + let task_manager = AsyncTaskManager::new() + .routing(|_request: &agentkit_tools_core::ToolRequest| RoutingDecision::Foreground); + let tasks = task_manager.handle(); + let tools = ToolRegistry::new().with(BlockingTool { + spec: ToolSpec { + name: ToolName::new(agentkit_tool_compose::COMPOSE_TOOL_NAME), + description: "controlled compose tool".into(), + input_schema: json!({"type": "object", "additionalProperties": false}), + output_schema: None, + annotations: ToolAnnotations::default(), + metadata: MetadataMap::new(), + }, + entered: Arc::clone(&entered), + release: Arc::clone(&release), + }); + let mut driver = Agent::builder() + .model(ScriptAdapter { + turns: Arc::clone(&turns), + user_items_seen: Arc::new(AtomicUsize::new(0)), + notification_items_seen: Arc::new(AtomicUsize::new(0)), + }) + .add_tool_source(tools) + .task_manager(task_manager) + .build() + .unwrap() + .start(SessionConfig::new(SessionId::new("v2-structured-loop")).without_cache()) + .await + .unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "start background")]) + .unwrap(); + + let integration = AcpIntegration::default(); + let session_id = wire::SessionId::new("v2-structured"); + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + SessionId::new("v2-structured-loop"), + RecordingSink::default(), + )) + .unwrap(); + handle.prepare_injection_turn(); + handle.start_injection_turn(); + let generation = handle.cancellation_handle().generation(); + let background_jobs = BackgroundJobs::default(); + let prompt = drive_prompt( + &session_id, + &mut driver, + &handle, + generation, + Some((&tasks, &background_jobs)), + ); + tokio::pin!(prompt); + + timeout(Duration::from_secs(1), async { + tokio::select! { + result = &mut prompt => panic!("structured v2 prompt resolved early: {result:?}"), + _ = async { + while !entered.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + assert!(super::super::detach_compose_call( + &tasks, + &background_jobs, + "background-call", + ).await); + background_jobs.register_foreground_for_test("background-call"); + while turns.load(Ordering::SeqCst) < 2 { + tokio::task::yield_now().await; + } + } => {} + } + }) + .await + .expect("structured v2 prompt did not reach its provisional response"); + assert!( + timeout(Duration::from_millis(20), &mut prompt) + .await + .is_err() + ); + + background_jobs.finish_for_test("background-call"); + assert!( + timeout(Duration::from_millis(20), &mut prompt) + .await + .is_err(), + "structured v2 prompt crossed the terminal publication handoff early" + ); + release.notify_one(); + let reason = timeout(Duration::from_secs(1), &mut prompt) + .await + .expect("structured v2 prompt did not synthesize") + .unwrap(); + assert_eq!(reason, wire::StopReason::EndTurn); + assert_eq!(turns.load(Ordering::SeqCst), 3); + assert!( + timeout(Duration::from_millis(20), tasks.next_event()) + .await + .is_err() + ); + handle.stop_injection_turn(); + } + #[tokio::test] async fn finish_error_stops_before_delivering_pending_steer() { let (mut driver, turns) = test_driver(TestOutcome::FinishError, "finish-error").await; @@ -2472,6 +2673,7 @@ mod tests { &mut driver, &control, 0, + None, ) .await; @@ -2507,7 +2709,7 @@ mod tests { .submit_input(vec![Item::text(ItemKind::User, "cancel")]) .unwrap(); - let result = drive_prompt(&session_id, &mut driver, &handle, generation).await; + let result = drive_prompt(&session_id, &mut driver, &handle, generation, None).await; assert_eq!(result.unwrap(), wire::StopReason::Cancelled); } @@ -2525,6 +2727,7 @@ mod tests { &mut driver, &control, 0, + None, ) .await; diff --git a/src/runtime.rs b/src/runtime.rs index 1f03560..a28fe69 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -28,6 +28,7 @@ use agentkit_tools_core::{ }; use async_trait::async_trait; use serde_json::{Value, json}; +use tokio::sync::watch; use tokio_util::sync::CancellationToken; use crate::{ @@ -218,6 +219,7 @@ pub(crate) struct AcpDriver { pub driver: LoopDriver, pub tasks: TaskManagerHandle, pub background_jobs: BackgroundJobs, + pub structured_completion: bool, pub adapter: SelectableAdapter, pub canonical_transcript: Vec, } @@ -954,6 +956,7 @@ impl Runtime { driver, tasks, background_jobs, + structured_completion: self.base_depth > 0, adapter, canonical_transcript, }; @@ -1074,6 +1077,7 @@ struct BackgroundJob { cancellation_relay: Option, detached: bool, manual_detach: bool, + terminal_published: bool, } #[derive(Default)] @@ -1081,6 +1085,10 @@ struct BackgroundJobState { running: HashMap, pending_cancellations: std::collections::HashSet, pending_detaches: std::collections::HashSet, + unacknowledged_terminals: std::collections::HashSet, + generation: u64, + background_started: u64, + cancel_all: bool, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -1089,13 +1097,107 @@ pub(crate) enum DetachRegistration { AlreadyDetached, } -#[derive(Clone, Default)] -pub(crate) struct BackgroundJobs(Arc>); +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct BackgroundActivity { + pub generation: u64, + pub active: bool, + pub background_started: u64, + pub unacknowledged_terminals: bool, +} + +#[derive(Clone)] +pub(crate) struct BackgroundJobs { + state: Arc>, + activity: watch::Sender, +} + +impl Default for BackgroundJobs { + fn default() -> Self { + let (activity, _) = watch::channel(0); + Self { + state: Arc::new(Mutex::new(BackgroundJobState::default())), + activity, + } + } +} impl BackgroundJobs { + fn changed(&self, jobs: &mut BackgroundJobState) { + jobs.generation = jobs.generation.wrapping_add(1); + self.activity.send_replace(jobs.generation); + } + + pub(crate) fn activity(&self) -> BackgroundActivity { + self.state.lock().map_or( + BackgroundActivity { + generation: *self.activity.borrow(), + active: false, + background_started: 0, + unacknowledged_terminals: false, + }, + |jobs| BackgroundActivity { + generation: jobs.generation, + active: !jobs.running.is_empty(), + background_started: jobs.background_started, + unacknowledged_terminals: !jobs.unacknowledged_terminals.is_empty(), + }, + ) + } + + pub(crate) async fn activity_after(&self, generation: u64) -> BackgroundActivity { + let mut receiver = self.activity.subscribe(); + loop { + let current = self.activity(); + if current.generation != generation { + return current; + } + if receiver.changed().await.is_err() { + return self.activity(); + } + } + } + + pub(crate) async fn wait_for_quiescence(&self) { + loop { + let activity = self.activity(); + if !activity.active { + return; + } + let _ = self.activity_after(activity.generation).await; + } + } + + pub(crate) fn acknowledge_terminal(&self, call_id: &agentkit_core::ToolCallId) { + if let Ok(mut jobs) = self.state.lock() { + let changed = if let Some(job) = jobs.running.get_mut(call_id) { + !std::mem::replace(&mut job.terminal_published, true) + } else { + jobs.unacknowledged_terminals.remove(call_id) + }; + if changed { + self.changed(&mut jobs); + } + } + } + + pub(crate) fn begin_turn(&self) { + if let Ok(mut jobs) = self.state.lock() { + jobs.cancel_all = false; + } + } + + pub(crate) fn cancel_all(&self) { + if let Ok(mut jobs) = self.state.lock() { + jobs.cancel_all = true; + for job in jobs.running.values() { + job.controller.interrupt(); + } + } + } + fn cancel_running(&self, call_id: &str) -> bool { let call_id = agentkit_core::ToolCallId::new(call_id); - let Ok(jobs) = self.0.lock() else { + let Ok(jobs) = self.state.lock() else { return false; }; let Some(job) = jobs.running.get(&call_id) else { @@ -1107,7 +1209,7 @@ impl BackgroundJobs { pub fn cancel(&self, call_id: &str) -> bool { let call_id = agentkit_core::ToolCallId::new(call_id); - let Ok(mut jobs) = self.0.lock() else { + let Ok(mut jobs) = self.state.lock() else { return false; }; if let Some(job) = jobs.running.get(&call_id) { @@ -1123,13 +1225,14 @@ impl BackgroundJobs { pub(crate) fn detach(&self, call_id: &str) -> Option { let call_id = agentkit_core::ToolCallId::new(call_id); - let Ok(mut jobs) = self.0.lock() else { + let Ok(mut jobs) = self.state.lock() else { return None; }; if let Some(job) = jobs.running.get_mut(&call_id) { if job.manual_detach { return Some(DetachRegistration::AlreadyDetached); } + let newly_detached = !job.detached; job.detached = true; job.manual_detach = true; if job @@ -1142,6 +1245,10 @@ impl BackgroundJobs { job.controller.interrupt(); return None; } + if newly_detached { + jobs.background_started = jobs.background_started.wrapping_add(1); + self.changed(&mut jobs); + } return Some(DetachRegistration::Registered); } Some(if jobs.pending_detaches.insert(call_id) { @@ -1153,7 +1260,7 @@ impl BackgroundJobs { pub(crate) fn restore_foreground(&self, call_id: &str) { let call_id = agentkit_core::ToolCallId::new(call_id); - let Ok(mut jobs) = self.0.lock() else { + let Ok(mut jobs) = self.state.lock() else { return; }; let Some(job) = jobs.running.get_mut(&call_id) else { @@ -1174,7 +1281,7 @@ impl BackgroundJobs { } fn propagate_foreground_cancellation(&self, call_id: &agentkit_core::ToolCallId) { - let Ok(jobs) = self.0.lock() else { + let Ok(jobs) = self.state.lock() else { return; }; if let Some(job) = jobs.running.get(call_id) @@ -1185,39 +1292,67 @@ impl BackgroundJobs { } fn finish(&self, call_id: &agentkit_core::ToolCallId) { - if let Ok(mut jobs) = self.0.lock() { - if let Some(job) = jobs.running.remove(call_id) - && let Some(relay) = job.cancellation_relay - { - relay.abort(); + if let Ok(mut jobs) = self.state.lock() { + if let Some(job) = jobs.running.remove(call_id) { + if job.detached && !job.terminal_published { + jobs.unacknowledged_terminals.insert(call_id.clone()); + } + if let Some(relay) = job.cancellation_relay { + relay.abort(); + } } jobs.pending_cancellations.remove(call_id); jobs.pending_detaches.remove(call_id); + self.changed(&mut jobs); } } #[cfg(test)] pub(crate) fn register_foreground_for_test(&self, call_id: &str) { - if let Ok(mut jobs) = self.0.lock() { + if let Ok(mut jobs) = self.state.lock() { let call_id = agentkit_core::ToolCallId::new(call_id); let manual_detach = jobs.pending_detaches.remove(&call_id); + let controller = CancellationController::new(); + if jobs.cancel_all { + controller.interrupt(); + } jobs.running.insert( call_id, BackgroundJob { - controller: CancellationController::new(), + controller, foreground_cancellation: None, cancellation_relay: None, detached: manual_detach, manual_detach, + terminal_published: false, }, ); + if manual_detach { + jobs.background_started = jobs.background_started.wrapping_add(1); + } + self.changed(&mut jobs); } } + #[cfg(test)] + pub(crate) fn finish_for_test(&self, call_id: &str) { + self.finish(&agentkit_core::ToolCallId::new(call_id)); + } + + #[cfg(test)] + pub(crate) fn is_cancelled_for_test(&self, call_id: &str) -> bool { + let call_id = agentkit_core::ToolCallId::new(call_id); + self.state.lock().is_ok_and(|jobs| { + jobs.running + .get(&call_id) + .is_some_and(|job| job.controller.handle().is_cancelled_since(0)) + }) + } + #[cfg(test)] pub(crate) fn is_detached_for_test(&self, call_id: &str) -> bool { let call_id = agentkit_core::ToolCallId::new(call_id); - self.0.lock().is_ok_and(|jobs| { + self.state.lock().is_ok_and(|jobs| { jobs.running.get(&call_id).is_some_and(|job| job.detached) || jobs.pending_detaches.contains(&call_id) }) @@ -1361,8 +1496,8 @@ impl BackgroundableCompose { if let Some(scope) = &mut ctx.execution_scope { scope.cancellation = Some(cancellation); } - if let Ok(mut jobs) = self.background_jobs.0.lock() { - if jobs.pending_cancellations.remove(call_id) { + if let Ok(mut jobs) = self.background_jobs.state.lock() { + if jobs.pending_cancellations.remove(call_id) || jobs.cancel_all { controller.interrupt(); } let manual_detach = jobs.pending_detaches.remove(call_id); @@ -1382,8 +1517,13 @@ impl BackgroundableCompose { cancellation_relay: None, detached, manual_detach, + terminal_published: false, }, ); + if detached { + jobs.background_started = jobs.background_started.wrapping_add(1); + } + self.background_jobs.changed(&mut jobs); } if let Some(cancellation) = foreground_cancellation { let jobs = self.background_jobs.clone(); @@ -1393,7 +1533,7 @@ impl BackgroundableCompose { jobs.propagate_foreground_cancellation(&relay_call_id); }) .abort_handle(); - if let Ok(mut jobs) = self.background_jobs.0.lock() + if let Ok(mut jobs) = self.background_jobs.state.lock() && let Some(job) = jobs.running.get_mut(call_id) { job.cancellation_relay = Some(relay); diff --git a/src/runtime/tests.rs b/src/runtime/tests.rs index ebb8230..13d265b 100644 --- a/src/runtime/tests.rs +++ b/src/runtime/tests.rs @@ -870,6 +870,57 @@ fn duplicate_detach_does_not_take_rollback_ownership() { assert!(jobs.is_detached_for_test("duplicate-call")); } +#[test] +fn background_terminal_publication_is_acknowledged_by_call_id() { + let jobs = BackgroundJobs::default(); + + jobs.register_foreground_for_test("finished-first"); + assert_eq!( + jobs.detach("finished-first"), + Some(DetachRegistration::Registered) + ); + jobs.finish_for_test("finished-first"); + assert!(jobs.activity().unacknowledged_terminals); + + jobs.acknowledge_terminal(&ToolCallId::new("other-call")); + assert!(jobs.activity().unacknowledged_terminals); + jobs.acknowledge_terminal(&ToolCallId::new("finished-first")); + assert!(!jobs.activity().unacknowledged_terminals); + + jobs.register_foreground_for_test("published-first"); + assert_eq!( + jobs.detach("published-first"), + Some(DetachRegistration::Registered) + ); + jobs.acknowledge_terminal(&ToolCallId::new("published-first")); + jobs.finish_for_test("published-first"); + assert!(!jobs.activity().unacknowledged_terminals); +} + +#[test] +fn cancel_all_covers_running_and_late_background_registration() { + let jobs = BackgroundJobs::default(); + let initial = jobs.activity(); + jobs.register_foreground_for_test("running"); + assert!(jobs.activity().active); + + jobs.cancel_all(); + assert!(jobs.is_cancelled_for_test("running")); + jobs.register_foreground_for_test("late"); + assert!(jobs.is_cancelled_for_test("late")); + + jobs.finish_for_test("running"); + jobs.finish_for_test("late"); + let quiescent = jobs.activity(); + assert!(!quiescent.active); + assert!(quiescent.generation > initial.generation); + + jobs.begin_turn(); + jobs.register_foreground_for_test("next-turn"); + assert!(!jobs.is_cancelled_for_test("next-turn")); + jobs.finish_for_test("next-turn"); +} + #[test] fn system_prompt_guides_compose_and_subagent_hygiene() { let root = tempfile::tempdir().unwrap(); From e6243861c14d82a4d0fb5a9ca8c7825def70795e Mon Sep 17 00:00:00 2001 From: daniel Date: Wed, 26 Aug 2026 20:50:13 +0100 Subject: [PATCH 2/2] fix: satisfy ACP completion checks --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/protocols/acp/v2.rs | 1 + 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7223e79..b729171 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2517,7 +2517,7 @@ dependencies = [ [[package]] name = "kit" -version = "0.1.94" +version = "0.1.95" dependencies = [ "a2a-protocol-client", "a2a-protocol-server", diff --git a/Cargo.toml b/Cargo.toml index 18749dc..e50a30b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "kit" -version = "0.1.94" +version = "0.1.95" edition = "2024" rust-version = "1.94.0" publish = false diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index ee6ecba..87b3a4c 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -996,6 +996,7 @@ async fn session_actor(actor: SessionActor) } } +#[allow(clippy::too_many_arguments)] async fn prepare_prompt( session_id: &wire::SessionId, integration: &AcpIntegration,