Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: order SubWorker cleanup before Workdir release

Browse Source
This commit is contained in:
Hare
2026-09-06 03:12:55 +09:00
parent e7079e223f
commit 0f8d61188a
5 changed files with 443 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workdir/src/scope.rs
+229 -19
View File
@@ -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<SessionValidity>,
cleanup_pending: Arc<AtomicBool>,
close_lock: Arc<tokio::sync::Mutex<()>>,
}
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<HashMap<u64, ActiveWriteLease>>,
next_lease_id: AtomicU64,
owned_commands: Arc<Mutex<HashSet<String>>>,
pending_command_events: Arc<Mutex<HashMap<String, Vec<CommandEvent>>>>,
starting_tool_calls: Arc<Mutex<HashSet<String>>>,
forwarded_terminals: Arc<Mutex<HashSet<String>>>,
command_events: broadcast::Sender<CommandEvent>,
closes_source: bool,
}
@@ -380,12 +410,42 @@ impl ScopedWorkdirSession {
}
}
fn publish_terminal_if_missing(
&self,
command_id: &str,
status: CommandStatus,
exit_code: Option<i32>,
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("scoped command set mutex poisoned")
.insert(handle.0.clone());
let _ = self.command_events.send(CommandEvent::Started {
.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");
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<broadcast::Receiver<CommandEvent>>,
owned_commands: Arc<Mutex<HashSet<String>>>,
pending_command_events: Arc<Mutex<HashMap<String, Vec<CommandEvent>>>>,
starting_tool_calls: Arc<Mutex<HashSet<String>>>,
forwarded_terminals: Arc<Mutex<HashSet<String>>>,
sender: broadcast::Sender<CommandEvent>,
) -> Option<tokio::task::JoinHandle<()>> {
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<CommandEvent>,
forwarded_terminals: &Mutex<HashSet<String>>,
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();
crates/worker/src/controller.rs
+35 -2
View File
@@ -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<C, St>(
// 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"));
}
}
crates/worker/src/feature/builtin/manage_workdir.rs
+119 -5
View File
@@ -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<dyn Fn() -> Pin<Box<dyn Future<Output = std::io::Result<()>> + 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<dyn WorkspaceClient>,
before_workdir_release: Option<BeforeWorkdirRelease>,
}
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<dyn WorkspaceClient>) -> Self {
Self { client }
Self {
client,
before_workdir_release: None,
}
}
pub(crate) fn with_before_workdir_release(
client: Arc<dyn WorkspaceClient>,
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<dyn WorkspaceClient>,
before_workdir_release: Option<BeforeWorkdirRelease>,
}
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<dyn WorkspaceClient>) -> Self {
Self { client }
Self {
client,
before_workdir_release: None,
}
}
fn with_before_workdir_release(
mut self,
before_workdir_release: Option<BeforeWorkdirRelease>,
) -> 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::<WorkdirAttachInput>(input_json)?),
WorkdirOperation::Detach => {
let _input = parse_input::<WorkdirDetachInput>(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::<WorkdirListInput>(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());
}
}
crates/worker/src/spawn/registry.rs
+47 -7
View File
@@ -679,13 +679,6 @@ impl SpawnedWorkerRegistry {
.unwrap_or_default()
}
pub(crate) fn reclaim_internal_scope(&self, worker_name: &str) -> io::Result<bool> {
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<bool> {
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<bool> {
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::<Vec<_>>();
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(&registry, 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();
crates/worker/src/spawn/tool.rs
+5 -2
View File
@@ -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) {
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 reclaim delegated scope after Internal SubWorker failure"
"failed to close parent-owned Workdir tools after Internal SubWorker failure"
);
}
});
}
}
let message = format!(