From 0f8d61188aad1842f14e5ebf94a3be9f48564225 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 03:12:55 +0900 Subject: [PATCH] fix: order SubWorker cleanup before Workdir release --- crates/workdir/src/scope.rs | 254 ++++++++++++++++-- crates/worker/src/controller.rs | 37 ++- .../src/feature/builtin/manage_workdir.rs | 124 ++++++++- crates/worker/src/spawn/registry.rs | 54 +++- crates/worker/src/spawn/tool.rs | 17 +- 5 files changed, 443 insertions(+), 43 deletions(-) diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index cf6db8cc..c48c1541 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -70,6 +70,9 @@ impl WorkdirToolBroker { child_write_leases: Mutex::new(HashMap::new()), next_lease_id: AtomicU64::new(1), owned_commands: Arc::new(Mutex::new(HashSet::new())), + pending_command_events: Arc::new(Mutex::new(HashMap::new())), + starting_tool_calls: Arc::new(Mutex::new(HashSet::new())), + forwarded_terminals: Arc::new(Mutex::new(HashSet::new())), command_events, closes_source: true, }); @@ -107,6 +110,7 @@ pub struct WorkdirScopeLease { pub capabilities: WorkdirSessionCapabilities, validity: Arc, cleanup_pending: Arc, + close_lock: Arc>, } impl std::fmt::Debug for WorkdirScopeLease { @@ -136,6 +140,10 @@ impl WorkdirScopeLease { } pub async fn close(&self) -> Result<(), WorkdirError> { + let _close_guard = self.close_lock.lock().await; + if !self.cleanup_pending.load(Ordering::Acquire) { + return Ok(()); + } self.validity.active.store(false, Ordering::Release); let command_ids = self .broker @@ -167,9 +175,28 @@ impl WorkdirScopeLease { }) .await; match (cancel, terminal) { - (_, Ok(_)) - | (Ok(()), Err(WorkdirError::UnknownCommand(_))) + (_, Ok(output)) => { + self.broker.authority.publish_terminal_if_missing( + &command_id, + output.status, + output.exit_code, + output.next_cursor.unwrap_or(output.content.len()) as u64, + ); + self.broker + .authority + .owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .remove(&command_id); + } + (Ok(()), Err(WorkdirError::UnknownCommand(_))) | (Err(WorkdirError::UnknownCommand(_)), Err(WorkdirError::UnknownCommand(_))) => { + self.broker.authority.publish_terminal_if_missing( + &command_id, + CommandStatus::Cancelled, + None, + 0, + ); self.broker .authority .owned_commands @@ -271,6 +298,9 @@ struct ScopedWorkdirSession { child_write_leases: Mutex>, next_lease_id: AtomicU64, owned_commands: Arc>>, + pending_command_events: Arc>>>, + starting_tool_calls: Arc>>, + forwarded_terminals: Arc>>, command_events: broadcast::Sender, closes_source: bool, } @@ -380,12 +410,42 @@ impl ScopedWorkdirSession { } } + fn publish_terminal_if_missing( + &self, + command_id: &str, + status: CommandStatus, + exit_code: Option, + offset: u64, + ) { + publish_owned_command_event( + &self.command_events, + &self.forwarded_terminals, + CommandEvent::Terminal { + command_id: command_id.to_string(), + status, + exit_code, + stdout_end_offset: offset, + stderr_end_offset: 0, + observed_at_ms: unix_timestamp_ms(), + }, + ); + } + fn ensure_parent_write_available(&self, path: &FsPath) -> Result<(), WorkdirError> { let mut leases = self .child_write_leases .lock() .expect("Workdir tool scope lease mutex poisoned"); - leases.retain(|_, lease| lease.validity.upgrade().is_some_and(|v| v.is_active())); + leases.retain(|_, lease| { + lease + .validity + .upgrade() + .is_some_and(|validity| validity.is_active()) + || lease + .cleanup_pending + .upgrade() + .is_some_and(|pending| pending.load(Ordering::Acquire)) + }); if leases.values().any(|lease| { lease.rules.iter().any(|rule| { rule.permission == WorkdirToolScopePermission::Write @@ -583,10 +643,16 @@ impl ScopedWorkdirSession { ); } let owned_commands = Arc::new(Mutex::new(HashSet::new())); + let pending_command_events = Arc::new(Mutex::new(HashMap::new())); + let starting_tool_calls = Arc::new(Mutex::new(HashSet::new())); + let forwarded_terminals = Arc::new(Mutex::new(HashSet::new())); let (command_events, _) = broadcast::channel(64); let event_forwarder = forward_owned_command_events( self.source.subscribe_command_events(), owned_commands.clone(), + pending_command_events.clone(), + starting_tool_calls.clone(), + forwarded_terminals.clone(), command_events.clone(), ) .map(|handle| Arc::new(Mutex::new(Some(handle)))); @@ -599,6 +665,9 @@ impl ScopedWorkdirSession { child_write_leases: Mutex::new(HashMap::new()), next_lease_id: AtomicU64::new(1), owned_commands, + pending_command_events, + starting_tool_calls, + forwarded_terminals, command_events, closes_source: false, }); @@ -612,6 +681,7 @@ impl ScopedWorkdirSession { capabilities, validity, cleanup_pending, + close_lock: Arc::new(tokio::sync::Mutex::new(())), }) } } @@ -687,16 +757,59 @@ impl WorkdirSession for ScopedWorkdirSession { None => self.cwd.clone(), }); } - let handle = self.source.start_command(request).await?; - self.owned_commands + if let Some(tool_call_id) = &tool_call_id { + self.starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .insert(tool_call_id.clone()); + } + let handle = match self.source.start_command(request).await { + Ok(handle) => handle, + Err(error) => { + if let Some(tool_call_id) = &tool_call_id { + self.starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .remove(tool_call_id); + } + return Err(error); + } + }; + let mut owned = self + .owned_commands .lock() - .expect("scoped command set mutex poisoned") - .insert(handle.0.clone()); - let _ = self.command_events.send(CommandEvent::Started { - command_id: handle.0.clone(), - tool_call_id, - observed_at_ms: unix_timestamp_ms(), - }); + .expect("scoped command set mutex poisoned"); + owned.insert(handle.0.clone()); + if let Some(tool_call_id) = &tool_call_id { + self.starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .remove(tool_call_id); + } + let pending = self + .pending_command_events + .lock() + .expect("pending scoped command event mutex poisoned") + .remove(&handle.0) + .unwrap_or_default(); + drop(owned); + if !pending + .iter() + .any(|event| matches!(event, CommandEvent::Started { .. })) + { + publish_owned_command_event( + &self.command_events, + &self.forwarded_terminals, + CommandEvent::Started { + command_id: handle.0.clone(), + tool_call_id, + observed_at_ms: unix_timestamp_ms(), + }, + ); + } + for event in pending { + publish_owned_command_event(&self.command_events, &self.forwarded_terminals, event); + } Ok(handle) } @@ -713,6 +826,12 @@ impl WorkdirSession for ScopedWorkdirSession { let command_id = request.handle.0.clone(); let output = self.source.command_output(request).await?; if !matches!(output.status, CommandStatus::Running) { + self.publish_terminal_if_missing( + &command_id, + output.status, + output.exit_code, + output.next_cursor.unwrap_or(output.content.len()) as u64, + ); self.owned_commands .lock() .expect("scoped command set mutex poisoned") @@ -848,6 +967,9 @@ impl WorkdirSession for ReadOnlyWorkdirSession { fn forward_owned_command_events( receiver: Option>, owned_commands: Arc>>, + pending_command_events: Arc>>>, + starting_tool_calls: Arc>>, + forwarded_terminals: Arc>>, sender: broadcast::Sender, ) -> Option> { let mut receiver = receiver?; @@ -858,22 +980,73 @@ fn forward_owned_command_events( Err(broadcast::error::RecvError::Lagged(_)) => continue, Err(broadcast::error::RecvError::Closed) => break, }; - let command_id = match &event { - CommandEvent::Started { .. } => continue, - CommandEvent::Output { command_id, .. } - | CommandEvent::Terminal { command_id, .. } => command_id.clone(), - }; - let owned = owned_commands + let command_id = command_event_id(&event).to_string(); + let mut owned = owned_commands .lock() - .expect("scoped command set mutex poisoned") - .contains(&command_id); - if owned { - let _ = sender.send(event); + .expect("scoped command set mutex poisoned"); + if !owned.contains(&command_id) { + let claimed = matches!( + &event, + CommandEvent::Started { + tool_call_id: Some(tool_call_id), + .. + } if starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .contains(tool_call_id) + ); + if !claimed { + continue; + } + owned.insert(command_id.clone()); + pending_command_events + .lock() + .expect("pending scoped command event mutex poisoned") + .entry(command_id) + .or_default() + .push(event); + continue; } + let mut pending = pending_command_events + .lock() + .expect("pending scoped command event mutex poisoned"); + if let Some(events) = pending.get_mut(&command_id) { + if events.len() < 64 { + events.push(event); + } + continue; + } + drop(pending); + drop(owned); + publish_owned_command_event(&sender, &forwarded_terminals, event); } })) } +fn command_event_id(event: &CommandEvent) -> &str { + match event { + CommandEvent::Started { command_id, .. } + | CommandEvent::Output { command_id, .. } + | CommandEvent::Terminal { command_id, .. } => command_id, + } +} + +fn publish_owned_command_event( + sender: &broadcast::Sender, + forwarded_terminals: &Mutex>, + event: CommandEvent, +) { + if let CommandEvent::Terminal { command_id, .. } = &event + && !forwarded_terminals + .lock() + .expect("forwarded terminal command mutex poisoned") + .insert(command_id.clone()) + { + return; + } + let _ = sender.send(event); +} + fn unix_timestamp_ms() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -1382,6 +1555,43 @@ mod tests { first.close().await.unwrap(); } + #[tokio::test] + async fn fast_command_keeps_started_output_terminal_event_order() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("work")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(request("work", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + let mut events = child.subscribe_command_events().unwrap(); + + let output = run_command(&child.tool_session(), "printf fast-output", "fast-command").await; + assert_eq!(output.content, "fast-output"); + + let mut kinds = Vec::new(); + let mut streamed = String::new(); + while kinds.last().is_none_or(|kind| *kind != "terminal") { + let event = tokio::time::timeout(std::time::Duration::from_secs(1), events.recv()) + .await + .expect("fast command event timeout") + .expect("fast command event channel"); + match event { + CommandEvent::Started { .. } => kinds.push("started"), + CommandEvent::Output { content, .. } => { + kinds.push("output"); + streamed.push_str(&content); + } + CommandEvent::Terminal { .. } => kinds.push("terminal"), + } + } + assert_eq!(kinds.first(), Some(&"started")); + assert_eq!(kinds.last(), Some(&"terminal")); + assert!(kinds.contains(&"output")); + assert!(streamed.contains("fast-output")); + child.close().await.unwrap(); + } + #[tokio::test] async fn closing_scope_cancels_and_terminalizes_owned_commands() { let root = TempDir::new().unwrap(); diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 214b1d55..a73a175a 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1101,8 +1101,15 @@ where "manage Workdir tools require Backend Workspace API authority", )); } + let child_registry = spawned_registry.clone(); feature_registry.add_module( - crate::feature::builtin::manage_workdir::manage_workdir_feature(workspace_client), + crate::feature::builtin::manage_workdir::ManageWorkdirFeature::with_before_workdir_release( + workspace_client, + Arc::new(move || { + let child_registry = child_registry.clone(); + Box::pin(async move { child_registry.shutdown_internal().await }) + }), + ), ); } if feature_config.workspace_worker_discovery.enabled { @@ -1726,7 +1733,16 @@ async fn controller_loop( // Memory/Workdir teardown so they cannot observe a partially closed Worker. worker.stop_feature_runtime("controller shutdown").await; - if let Some(session) = worker.workdir_session() + let child_cleanup_succeeded = match spawned_registry.shutdown_internal().await { + Ok(()) => true, + Err(error) => { + tracing::warn!(%error, "Internal SubWorker cleanup failed before Workdir shutdown"); + false + } + }; + + if child_cleanup_succeeded + && let Some(session) = worker.workdir_session() && let Err(error) = session.close().await { tracing::warn!(%error, "Workdir session close failed"); @@ -2604,4 +2620,21 @@ mod tests { other => panic!("expected compact rejection error, got {other:?}"), } } + + #[test] + fn controller_shutdown_orders_child_cleanup_before_workdir_close() { + let source = include_str!("controller.rs"); + let shutdown_start = source + .rfind("worker.stop_feature_runtime(\"controller shutdown\")") + .expect("controller shutdown block"); + let shutdown = &source[shutdown_start..]; + let children = shutdown + .find("spawned_registry.shutdown_internal().await") + .expect("Internal SubWorker cleanup"); + let workdir = shutdown + .find("session.close().await") + .expect("parent Workdir close"); + assert!(children < workdir); + assert!(shutdown.contains("if child_cleanup_succeeded")); + } } diff --git a/crates/worker/src/feature/builtin/manage_workdir.rs b/crates/worker/src/feature/builtin/manage_workdir.rs index ab252835..24cbc0a9 100644 --- a/crates/worker/src/feature/builtin/manage_workdir.rs +++ b/crates/worker/src/feature/builtin/manage_workdir.rs @@ -5,6 +5,8 @@ //! endpoints, credentials, materializer handles, and operation sessions stay //! behind [`WorkspaceClient`]. +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use agen::tool::{Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput}; @@ -52,16 +54,43 @@ const LIST_DESCRIPTION: &str = "List persistent Workdirs in the current Workspac const CREATE_DESCRIPTION: &str = "Materialize a persistent Workdir on a selected Runtime from a Workspace repository and optional selector. This does not change this Worker's attachment; use WorkdirAttach explicitly after creation."; const ATTACH_DESCRIPTION: &str = "Attach this Worker to one existing Workdir. The Backend enforces one active Workdir per Worker and one active Worker per Workdir, then opens an ephemeral operation session."; const DETACH_DESCRIPTION: &str = "Detach this Worker from its active Workdir and release Workdir occupancy. Any ephemeral operation session is closed."; +pub(crate) type BeforeWorkdirRelease = + Arc Pin> + Send>> + Send + Sync>; + const DELETE_DESCRIPTION: &str = "Request removal of one persistent Workdir by id through durable Backend Workspace authority. The input includes only the Workdir id and a bounded reason. The result reports removed, retained, or attention_required without exposing operation-table or provider internals."; -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct ManageWorkdirFeature { client: Arc, + before_workdir_release: Option, +} + +impl std::fmt::Debug for ManageWorkdirFeature { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ManageWorkdirFeature") + .field("client_kind", &self.client.kind()) + .field("release_guard", &self.before_workdir_release.is_some()) + .finish() + } } impl ManageWorkdirFeature { pub fn new(client: Arc) -> Self { - Self { client } + Self { + client, + before_workdir_release: None, + } + } + + pub(crate) fn with_before_workdir_release( + client: Arc, + before_workdir_release: BeforeWorkdirRelease, + ) -> Self { + Self { + client, + before_workdir_release: Some(before_workdir_release), + } } } @@ -81,7 +110,8 @@ impl FeatureModule for ManageWorkdirFeature { } fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { - let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()); + let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()) + .with_before_workdir_release(self.before_workdir_release.clone()); for (name, definition) in [ ( LIST_TOOL, @@ -142,9 +172,20 @@ impl FeatureModule for ManageWorkdirFeature { } } -#[derive(Clone, Debug)] +#[derive(Clone)] struct WorkspaceHttpWorkdirBackend { client: Arc, + before_workdir_release: Option, +} + +impl std::fmt::Debug for WorkspaceHttpWorkdirBackend { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkspaceHttpWorkdirBackend") + .field("client_kind", &self.client.kind()) + .field("release_guard", &self.before_workdir_release.is_some()) + .finish() + } } /// Worker-local Workdir handle whose operation authority remains in the Workspace Backend. @@ -327,7 +368,18 @@ impl WorkdirSession for WorkspaceAttachedWorkdirSession { impl WorkspaceHttpWorkdirBackend { fn new(client: Arc) -> Self { - Self { client } + Self { + client, + before_workdir_release: None, + } + } + + fn with_before_workdir_release( + mut self, + before_workdir_release: Option, + ) -> Self { + self.before_workdir_release = before_workdir_release; + self } fn workspace_id(&self) -> Result<&str, ToolError> { @@ -510,6 +562,13 @@ impl Tool for WorkspaceHttpWorkdirTool { .attach(parse_input::(input_json)?), WorkdirOperation::Detach => { let _input = parse_input::(input_json)?; + if let Some(before_release) = &self.backend.before_workdir_release { + before_release().await.map_err(|error| { + ToolError::ExecutionFailed(format!( + "stop Internal SubWorkers before Workdir detach: {error}" + )) + })?; + } self.backend.detach() } WorkdirOperation::Delete => self @@ -705,6 +764,7 @@ struct WorkdirDeleteInput { #[cfg(test)] mod tests { use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; use super::*; use crate::feature::{FeatureModule, FeatureRegistryBuilder}; @@ -1259,4 +1319,58 @@ mod tests { assert!(client.requests().is_empty()); assert!(parse_input::(r#"{"path":"/tmp"}"#).is_err()); } + + #[tokio::test] + async fn detach_stops_internal_subworkers_before_backend_release() { + let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({ + "workspace_id": "workspace/test", + "workdir_id": "wd-attached", + "attached": false + }))])); + let cleanup_calls = Arc::new(AtomicUsize::new(0)); + let cleanup_calls_for_guard = cleanup_calls.clone(); + let before_release: BeforeWorkdirRelease = Arc::new(move || { + let cleanup_calls = cleanup_calls_for_guard.clone(); + Box::pin(async move { + cleanup_calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + }); + let tool = WorkspaceHttpWorkdirTool { + backend: WorkspaceHttpWorkdirBackend::new(client.clone()) + .with_before_workdir_release(Some(before_release)), + operation: WorkdirOperation::Detach, + }; + + tool.execute("{}", ToolExecutionContext::default()) + .await + .unwrap(); + + assert_eq!(cleanup_calls.load(Ordering::SeqCst), 1); + assert_eq!(client.requests().len(), 1); + assert_eq!( + client.requests()[0].path, + "/api/w/workspace%2Ftest/workers/self/workdir-attachment" + ); + } + + #[tokio::test] + async fn detach_does_not_release_backend_when_child_cleanup_fails() { + let client = Arc::new(RecordingWorkspaceClient::new(Vec::new())); + let before_release: BeforeWorkdirRelease = + Arc::new(|| Box::pin(async { Err(std::io::Error::other("child cleanup failed")) })); + let tool = WorkspaceHttpWorkdirTool { + backend: WorkspaceHttpWorkdirBackend::new(client.clone()) + .with_before_workdir_release(Some(before_release)), + operation: WorkdirOperation::Detach, + }; + + let error = tool + .execute("{}", ToolExecutionContext::default()) + .await + .unwrap_err(); + + assert!(error.to_string().contains("stop Internal SubWorkers")); + assert!(client.requests().is_empty()); + } } diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index c270243f..3db26267 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -679,13 +679,6 @@ impl SpawnedWorkerRegistry { .unwrap_or_default() } - pub(crate) fn reclaim_internal_scope(&self, worker_name: &str) -> io::Result { - let record = self.get_internal(worker_name).ok_or_else(|| { - io::Error::new(io::ErrorKind::NotFound, "internal SubWorker not found") - })?; - self.reclaim_record_scope(&record) - } - fn reclaim_record_scope(&self, record: &InternalSpawnedWorkerRecord) -> io::Result { if !record.claim_scope_reclaim() { return Ok(false); @@ -705,6 +698,35 @@ impl SpawnedWorkerRegistry { result } + pub(crate) async fn close_internal_scope(&self, name: &str) -> io::Result { + let Some(record) = self.get_internal(name) else { + return Ok(false); + }; + record + .workdir_tool_scope + .close() + .await + .map_err(|error| io::Error::other(error.to_string()))?; + self.reclaim_record_scope(&record) + } + + pub(crate) async fn shutdown_internal(&self) -> io::Result<()> { + let names = self + .internal_records + .lock() + .expect("internal Worker registry lock poisoned") + .iter() + .map(|record| record.worker_name.clone()) + .collect::>(); + let mut first_error = None; + for name in names { + if let Err(error) = self.remove_internal(&name).await { + first_error.get_or_insert(error); + } + } + first_error.map_or(Ok(()), Err) + } + /// Stop one direct Internal SubWorker and discard its registry/scope state. /// /// The child actor must acknowledge its stop before the registry is removed. @@ -1236,6 +1258,24 @@ mod tests { } } + #[tokio::test] + async fn parent_shutdown_stops_all_internal_workers_before_returning() { + let registry = registry(); + for name in ["first", "second"] { + let (record, _events) = record(name, InternalWorkerVisibility::ParentClient).await; + record + .session + .force_status(InternalWorkerSessionStatus::Running); + install_record(®istry, record); + } + + registry.shutdown_internal().await.unwrap(); + + assert!(registry.list_internal().is_empty()); + assert!(registry.get_internal("first").is_none()); + assert!(registry.get_internal("second").is_none()); + } + #[tokio::test] async fn running_worker_is_stopped_before_removal() { let registry = registry(); diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 3091642e..294983b0 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -532,13 +532,16 @@ impl Tool for SubWorkerSpawnTool { InternalWorkerSessionStatus::Failed | InternalWorkerSessionStatus::Stopped ) { if let Some(registry) = registry.upgrade() { - if let Err(error) = registry.reclaim_internal_scope(&child_name) { - tracing::warn!( - child_name, - %error, - "failed to reclaim delegated scope after Internal SubWorker failure" - ); - } + let child_name = child_name.clone(); + tokio::spawn(async move { + if let Err(error) = registry.close_internal_scope(&child_name).await { + tracing::warn!( + child_name, + %error, + "failed to close parent-owned Workdir tools after Internal SubWorker failure" + ); + } + }); } } let message = format!(