From 3e833b529569be5ada51d060012663c99ceedebf Mon Sep 17 00:00:00 2001 From: Hare Date: Thu, 20 Aug 2026 11:19:02 +0900 Subject: [PATCH] fix: finalize stopped subworker sessions --- crates/protocol/src/lib.rs | 29 ++ crates/tools/src/edit.rs | 13 +- crates/tools/src/lib.rs | 2 +- crates/tools/src/tracker.rs | 92 +++++- crates/tools/src/write.rs | 3 + crates/tui/src/app.rs | 113 ++++++- .../src/feature/builtin/manage_worker.rs | 96 +++++- crates/worker/src/internal_worker.rs | 35 +- crates/worker/src/spawn/comm_tools.rs | 29 +- crates/worker/src/spawn/registry.rs | 304 +++++++++++++++++- crates/worker/src/spawn/tool.rs | 2 + web/workspace/src/lib/generated/protocol.ts | 2 +- .../src/lib/workspace/console/model.test.ts | 108 ++++++- .../src/lib/workspace/console/model.ts | 25 ++ 14 files changed, 816 insertions(+), 37 deletions(-) diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index d44427cf..19e73b50 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -530,6 +530,15 @@ pub enum Event { revision: u64, event: Box, }, + /// Terminal removal fence for one parent-owned Internal Worker session. + /// + /// Clients discard the matching child and descendants, then ignore later + /// nested events for this identity until an authoritative snapshot replaces + /// the projection. + InternalWorkerRemoved { + worker: InternalWorkerRef, + revision: u64, + }, /// Server-side segment log rotated to a fresh `SegmentStart`. /// /// Fires on compaction and on auto-fork when the store head drifts @@ -1802,6 +1811,26 @@ mod tests { } } + #[test] + fn internal_worker_removal_roundtrip_preserves_terminal_fence() { + let event = Event::InternalWorkerRemoved { + worker: InternalWorkerRef { + session_id: "session-1".into(), + name: "research".into(), + parent_session_id: Some("parent-session".into()), + kind: InternalWorkerKind::SubWorker, + }, + revision: 8, + }; + let json = serde_json::to_string(&event).unwrap(); + let decoded: Event = serde_json::from_str(&json).unwrap(); + assert!(matches!( + decoded, + Event::InternalWorkerRemoved { worker, revision } + if worker.session_id == "session-1" && revision == 8 + )); + } + #[test] fn legacy_snapshot_defaults_internal_workers_to_empty() { let snapshot: Event = serde_json::from_value(serde_json::json!({ diff --git a/crates/tools/src/edit.rs b/crates/tools/src/edit.rs index a6a4869a..17e87dae 100644 --- a/crates/tools/src/edit.rs +++ b/crates/tools/src/edit.rs @@ -72,13 +72,20 @@ impl Tool for EditTool { }) .await .map_err(ToolsError::from)?; - self.tracker.record_workdir_hash(&path, result.content_hash); + let replacements = result.replacements; + self.tracker.record_workdir_edit( + &path, + result.content_hash, + replacements, + params.new_string.lines().count(), + params.old_string.lines().count(), + ); let summary = format!( "Edited {} ({} replacement{})", path, - result.replacements, - if result.replacements == 1 { "" } else { "s" } + replacements, + if replacements == 1 { "" } else { "s" } ); let preview = make_preview(¶ms.new_string, ¶ms.new_string); diff --git a/crates/tools/src/lib.rs b/crates/tools/src/lib.rs index 319ee71f..df140220 100644 --- a/crates/tools/src/lib.rs +++ b/crates/tools/src/lib.rs @@ -27,7 +27,7 @@ pub use error::ToolsError; pub use glob::glob_tool; pub use grep::grep_tool; pub use read::read_tool; -pub use tracker::Tracker; +pub use tracker::{ChangeStat, Tracker}; pub use view_image::view_image_tool; pub use web::{web_fetch_tool, web_search_tool}; pub use write::write_tool; diff --git a/crates/tools/src/tracker.rs b/crates/tools/src/tracker.rs index 40a771be..5c284bcb 100644 --- a/crates/tools/src/tracker.rs +++ b/crates/tools/src/tracker.rs @@ -119,12 +119,22 @@ fn normalize_path_lexically(path: &Path) -> PathBuf { normalized } +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ChangeStat { + pub added: u64, + pub deleted: u64, +} + #[derive(Debug, Default)] struct Inner { /// Hash of each file's last observed contents, keyed by canonical path. hashes: HashMap, + /// Line count paired with observations that included the file content. + line_counts: HashMap, /// LRU list of touched files. Front = most recently touched. recency: VecDeque, + /// Successful Write/Edit mutations attributed to this session's tools. + change_stat: ChangeStat, } /// Canonical-path keyed tracker of file observations and their recency. @@ -187,8 +197,27 @@ impl Tracker { } } - pub fn record_workdir_content(&self, path: &workdir::WorkdirPath, bytes: &[u8]) { - self.record_workdir_hash(path, hash_bytes(bytes)); + pub fn record_workdir_content(&self, path: &workdir::WorkdirPath, content: &[u8]) { + let key = PathBuf::from(path.as_str()); + let hash = hash_bytes(content); + let line_count = String::from_utf8_lossy(content).lines().count(); + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + inner.line_counts.insert(key.clone(), line_count); + inner.hashes.insert(key.clone(), hash); + inner.recency.retain(|candidate| candidate != &key); + inner.recency.push_front(key); + if inner.recency.len() > RECENCY_CAPACITY { + inner.recency.pop_back(); + } + } + + pub fn observed_workdir_line_count(&self, path: &workdir::WorkdirPath) -> Option { + self.inner + .lock() + .unwrap_or_else(|e| e.into_inner()) + .line_counts + .get(Path::new(path.as_str())) + .copied() } pub fn record_workdir_hash(&self, path: &workdir::WorkdirPath, hash: workdir::ContentHash) { @@ -202,6 +231,50 @@ impl Tracker { } } + /// Record a successful, session-attributable source mutation. + /// + /// Callers supply line counts derived from the exact replacement accepted + /// by a Write/Edit tool. Bash and external process mutations are excluded + /// because this tracker cannot attribute them to one tool operation + /// authoritatively. + pub fn record_change(&self, added: usize, deleted: usize) { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + inner.change_stat.added = inner.change_stat.added.saturating_add(added as u64); + inner.change_stat.deleted = inner.change_stat.deleted.saturating_add(deleted as u64); + } + + pub fn record_workdir_edit( + &self, + path: &workdir::WorkdirPath, + hash: workdir::ContentHash, + replacements: usize, + added_lines_per_replacement: usize, + deleted_lines_per_replacement: usize, + ) { + let added = added_lines_per_replacement.saturating_mul(replacements); + let deleted = deleted_lines_per_replacement.saturating_mul(replacements); + let key = PathBuf::from(path.as_str()); + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + inner.change_stat.added = inner.change_stat.added.saturating_add(added as u64); + inner.change_stat.deleted = inner.change_stat.deleted.saturating_add(deleted as u64); + if let Some(line_count) = inner.line_counts.get_mut(&key) { + *line_count = line_count.saturating_sub(deleted).saturating_add(added); + } + inner.hashes.insert(key.clone(), hash); + inner.recency.retain(|candidate| candidate != &key); + inner.recency.push_front(key); + if inner.recency.len() > RECENCY_CAPACITY { + inner.recency.pop_back(); + } + } + + pub fn change_stat(&self) -> ChangeStat { + self.inner + .lock() + .unwrap_or_else(|e| e.into_inner()) + .change_stat + } + pub fn expected_workdir_hash( &self, path: &workdir::WorkdirPath, @@ -458,6 +531,21 @@ mod tests { } } + #[test] + fn change_stat_saturates_and_accumulates_tracked_mutations() { + let tracker = Tracker::new(); + tracker.record_change(7, 3); + tracker.record_change(5, 2); + + assert_eq!( + tracker.change_stat(), + ChangeStat { + added: 12, + deleted: 5, + } + ); + } + #[tokio::test] async fn mutation_guard_blocks_equivalent_paths_until_drop() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/tools/src/write.rs b/crates/tools/src/write.rs index cec2ad05..ff9ff09d 100644 --- a/crates/tools/src/write.rs +++ b/crates/tools/src/write.rs @@ -50,6 +50,7 @@ impl Tool for WriteTool { Err(error) => return Err(ToolsError::from(error).into()), }; + let old_line_count = self.tracker.observed_workdir_line_count(&path).unwrap_or(0); let outcome = self .session .write(WriteRequest { @@ -60,6 +61,8 @@ impl Tool for WriteTool { .await .map_err(ToolsError::from)?; + self.tracker + .record_change(params.content.lines().count(), old_line_count); self.tracker .record_workdir_content(&path, params.content.as_bytes()); diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index 473b7a91..a742b1bd 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -1,4 +1,4 @@ -use std::collections::VecDeque; +use std::collections::{HashMap, VecDeque}; use std::path::Path; use std::time::{Duration, Instant}; @@ -283,6 +283,8 @@ pub struct App { /// Presentation-only Internal Worker projections keyed by session identity. /// They are rendered in separate sub-panes and never mixed into `blocks`. pub internal_workers: Vec, + /// Terminal child-session fences, reset only by an authoritative snapshot. + removed_internal_workers: HashMap, pub scroll: Scroll, pub mode: Mode, pub cache: FileCache, @@ -361,6 +363,7 @@ impl App { blocks: Vec::new(), run_error_messages: Vec::new(), internal_workers: Vec::new(), + removed_internal_workers: HashMap::new(), scroll: Scroll::default(), mode: Mode::Normal, cache: FileCache::new(), @@ -1318,6 +1321,9 @@ impl App { revision, event, } => self.apply_internal_worker_event(worker, revision, *event), + Event::InternalWorkerRemoved { worker, revision } => { + self.remove_internal_worker(worker, revision) + } Event::Status { status } => { self.rewind_refresh_fence = false; self.set_worker_status(status); @@ -2002,6 +2008,7 @@ impl App { .into_iter() .map(Self::internal_worker_view_from_snapshot) .collect(); + self.removed_internal_workers.clear(); } fn internal_worker_view_from_snapshot(snapshot: InternalWorkerSnapshot) -> InternalWorkerView { @@ -2029,6 +2036,12 @@ impl App { revision: u64, event: Event, ) { + if self + .removed_internal_workers + .contains_key(&worker.session_id) + { + return; + } let index = self .internal_workers .iter() @@ -2051,6 +2064,26 @@ impl App { let _ = target.app.handle_worker_event(event); } + fn remove_internal_worker(&mut self, worker: InternalWorkerRef, revision: u64) { + let Some(index) = self + .internal_workers + .iter() + .position(|candidate| candidate.worker.session_id == worker.session_id) + else { + self.removed_internal_workers + .entry(worker.session_id) + .and_modify(|current| *current = (*current).max(revision)) + .or_insert(revision); + return; + }; + if revision <= self.internal_workers[index].revision { + return; + } + self.internal_workers.remove(index); + self.removed_internal_workers + .insert(worker.session_id, revision); + } + fn restore_snapshot( &mut self, entries: &[serde_json::Value], @@ -3542,6 +3575,84 @@ mod completion_flow_tests { ); } + #[test] + fn terminal_internal_worker_removal_drops_descendants_and_fences_late_events() { + let mut app = App::new("parent".into()); + let worker = InternalWorkerRef { + session_id: "child-session".into(), + name: "child".into(), + parent_session_id: Some("parent-session".into()), + kind: protocol::InternalWorkerKind::SubWorker, + }; + let nested = InternalWorkerRef { + session_id: "grandchild-session".into(), + name: "grandchild".into(), + parent_session_id: Some("child-session".into()), + kind: protocol::InternalWorkerKind::SubWorker, + }; + app.handle_worker_event(Event::InternalWorker { + worker: worker.clone(), + revision: 2, + event: Box::new(Event::InternalWorker { + worker: nested, + revision: 1, + event: Box::new(Event::TextDone { + text: "nested".into(), + }), + }), + }); + assert_eq!(app.internal_workers.len(), 1); + assert_eq!(app.internal_workers[0].app.internal_workers.len(), 1); + + app.handle_worker_event(Event::InternalWorkerRemoved { + worker: worker.clone(), + revision: 3, + }); + app.handle_worker_event(Event::InternalWorker { + worker, + revision: 4, + event: Box::new(Event::TextDone { + text: "late".into(), + }), + }); + + assert!(app.internal_workers.is_empty()); + app.handle_worker_event(Event::Snapshot { + greeting: test_greeting(), + entries: Vec::new(), + status: WorkerStatus::Idle, + in_flight: Default::default(), + internal_workers: Vec::new(), + }); + assert!(app.internal_workers.is_empty()); + assert!(app.removed_internal_workers.is_empty()); + } + + #[test] + fn stale_internal_worker_removal_keeps_newer_projection() { + let mut app = App::new("parent".into()); + let worker = InternalWorkerRef { + session_id: "child-session".into(), + name: "child".into(), + parent_session_id: Some("parent-session".into()), + kind: protocol::InternalWorkerKind::SubWorker, + }; + app.handle_worker_event(Event::InternalWorker { + worker: worker.clone(), + revision: 4, + event: Box::new(Event::TextDone { + text: "current".into(), + }), + }); + app.handle_worker_event(Event::InternalWorkerRemoved { + worker, + revision: 3, + }); + + assert_eq!(app.internal_workers.len(), 1); + assert_eq!(app.internal_workers[0].revision, 4); + } + #[test] fn snapshot_authoritatively_replaces_internal_worker_views() { let mut app = App::new("parent".into()); diff --git a/crates/worker/src/feature/builtin/manage_worker.rs b/crates/worker/src/feature/builtin/manage_worker.rs index 54b97524..c48383a7 100644 --- a/crates/worker/src/feature/builtin/manage_worker.rs +++ b/crates/worker/src/feature/builtin/manage_worker.rs @@ -16,7 +16,7 @@ use crate::feature::{ FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureModule, ServiceDeclaration, ServiceId, ToolContribution, ToolDeclaration, }; -use crate::spawn::registry::SpawnedWorkerRegistry; +use crate::spawn::registry::{SpawnedWorkerRegistry, SubWorkerStopSummary}; use crate::worker::{ WorkspaceClient, WorkspaceClientError, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, @@ -138,14 +138,19 @@ impl WorkerControlService for WorkspaceWorkerControlService { let registry = self.registry.as_ref().ok_or_else(|| { WorkspaceClientError::Request("unknown Worker or permission not granted".to_string()) })?; - registry + let summary = registry .remove_internal(name) .await - .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; + .map_err(|error| WorkspaceClientError::Request(error.to_string()))? + .ok_or_else(|| { + WorkspaceClientError::Request( + "unknown Worker or permission not granted".to_string(), + ) + })?; Ok(WorkspaceResponse { status: 200, - body: serde_json::json!({ "subject": { "kind": "sub_worker", "name": name } }) - .to_string(), + body: serde_json::to_string(&summary) + .map_err(|error| WorkspaceClientError::Request(error.to_string()))?, }) } @@ -839,6 +844,15 @@ fn tool_output( response.status, response.body ))); } + if operation == WorkerOperation::Stop + && let Ok(summary) = serde_json::from_str::(&response.body) + { + return Ok(ToolOutput { + summary: render_subworker_stop_summary(&summary), + content: Some(response.body), + attachments: Vec::new(), + }); + } Ok(ToolOutput { summary: format!("{} completed", operation.tool_name()), content: Some(response.body), @@ -846,6 +860,37 @@ fn tool_output( }) } +fn render_subworker_stop_summary(summary: &SubWorkerStopSummary) -> String { + let tools = if summary.tool_counts.is_empty() { + "No tool calls".to_string() + } else { + summary + .tool_counts + .iter() + .map(|tool| format!("{} {}", tool.count, tool.name)) + .collect::>() + .join(", ") + }; + let elapsed = format_elapsed(summary.elapsed_ms); + let changes = summary + .change_stat + .as_ref() + .map(|stat| format!("+{}/-{} Changes · ", stat.added, stat.deleted)) + .unwrap_or_default(); + format!("SubWorkerStop - done\n {tools}\n {changes}{elapsed}",) +} + +fn format_elapsed(elapsed_ms: u64) -> String { + let seconds = elapsed_ms / 1_000; + let minutes = seconds / 60; + let seconds = seconds % 60; + if minutes > 0 { + format!("{minutes}m {seconds}s") + } else { + format!("{seconds}s") + } +} + fn definition( operation: WorkerOperation, control: Arc, @@ -1252,6 +1297,47 @@ mod tests { assert!(client.removals.lock().unwrap().is_empty()); } + #[test] + fn subworker_stop_output_is_compact_and_keeps_typed_evidence() { + let summary = SubWorkerStopSummary { + session_id: "session-1".to_string(), + display_name: "research".to_string(), + outcome: crate::spawn::registry::SubWorkerFinalOutcome::Done, + elapsed_ms: 78_000, + tool_counts: vec![ + crate::spawn::registry::SubWorkerToolCount { + name: "Read".to_string(), + count: 26, + }, + crate::spawn::registry::SubWorkerToolCount { + name: "Grep".to_string(), + count: 5, + }, + ], + change_stat: Some(crate::spawn::registry::SubWorkerChangeStat { + added: 215, + deleted: 148, + source: "tracked_write_edit_tools".to_string(), + }), + }; + let response = WorkspaceResponse { + status: 200, + body: serde_json::to_string(&summary).unwrap(), + }; + + let output = tool_output(WorkerOperation::Stop, response).unwrap(); + + assert_eq!( + output.summary, + "SubWorkerStop - done\n 26 Read, 5 Grep\n +215/-148 Changes · 1m 18s" + ); + assert_eq!( + serde_json::from_str::(output.content.as_deref().unwrap()) + .unwrap(), + summary + ); + } + #[test] fn worker_inputs_reject_paths_and_parent_traversal() { assert!(authority_id("https://runtime.example", "runtime_id").is_err()); diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index 0a21d0d2..9b00dc19 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -294,6 +294,8 @@ pub(crate) struct InternalWorkerSessionHandle { last_error: Arc>>, child_registry: Option>, sink: SegmentLogSink, + #[cfg(test)] + fail_stop: Arc, } impl InternalWorkerSessionHandle { @@ -319,6 +321,9 @@ impl InternalWorkerSessionHandle { #[cfg(test)] pub(crate) fn publish_test_entry(&self, entry: LogEntry) { + self.store + .append(self.session_id, self.segment_id, &entry) + .expect("append test Internal Worker entry"); self.sink.publish(entry); } @@ -419,7 +424,23 @@ impl InternalWorkerSessionHandle { } } + #[cfg(test)] + pub(crate) fn force_status(&self, status: InternalWorkerSessionStatus) { + self.status + .store(status.encode(), std::sync::atomic::Ordering::Release); + } + + #[cfg(test)] + pub(crate) fn force_stop_failure(&self) { + self.fail_stop + .store(true, std::sync::atomic::Ordering::Release); + } + pub(crate) async fn stop(&self) -> Result<(), InternalWorkerSessionError> { + #[cfg(test)] + if self.fail_stop.load(std::sync::atomic::Ordering::Acquire) { + return Err(InternalWorkerSessionError::Unavailable); + } let prior = self.status.swap( InternalWorkerSessionStatus::Stopping.encode(), std::sync::atomic::Ordering::AcqRel, @@ -582,6 +603,8 @@ pub(crate) async fn prepare_internal_worker_session( last_error: last_error.clone(), child_registry, sink, + #[cfg(test)] + fail_stop: Arc::new(std::sync::atomic::AtomicBool::new(false)), }; tokio::spawn(async move { @@ -890,8 +913,17 @@ pub(crate) fn test_internal_worker_session( let session_id = session_store::new_session_id(); let segment_id = session_store::new_segment_id(); let (command_tx, mut command_rx) = tokio::sync::mpsc::channel(1); - tokio::spawn(async move { while command_rx.recv().await.is_some() {} }); let (event_tx, _) = broadcast::channel(256); + let command_event_tx = event_tx.clone(); + tokio::spawn(async move { + while let Some(command) = command_rx.recv().await { + if let InternalWorkerSessionCommand::Stop(done_tx) = command { + let _ = command_event_tx.send(Event::Shutdown); + let _ = done_tx.send(()); + break; + } + } + }); let sink = SegmentLogSink::new(); spawn_internal_log_event_bridge(sink.clone(), event_tx.clone()); let handle = InternalWorkerSessionHandle { @@ -909,6 +941,7 @@ pub(crate) fn test_internal_worker_session( last_error: Arc::new(Mutex::new(None)), child_registry: None, sink, + fail_stop: Arc::new(std::sync::atomic::AtomicBool::new(false)), }; (handle, event_tx) } diff --git a/crates/worker/src/spawn/comm_tools.rs b/crates/worker/src/spawn/comm_tools.rs index c0fd4d55..ca1d16a1 100644 --- a/crates/worker/src/spawn/comm_tools.rs +++ b/crates/worker/src/spawn/comm_tools.rs @@ -170,20 +170,27 @@ impl Tool for SubWorkerStopTool { ) -> Result { let input: NameInput = serde_json::from_str(input_json) .map_err(|e| ToolError::InvalidArgument(format!("invalid SubWorkerStop input: {e}")))?; - if let Some(record) = self.registry.get_internal(&input.name) { - record.session.stop().await.map_err(|error| { - ToolError::ExecutionFailed(format!("stop `{}`: {error}", input.name)) - })?; - self.registry - .remove_internal(&input.name) - .await - .map_err(|error| ToolError::ExecutionFailed(error.to_string()))?; + if let Some(summary) = self + .registry + .remove_internal(&input.name) + .await + .map_err(|error| ToolError::ExecutionFailed(error.to_string()))? + { return Ok(ToolOutput { summary: format!( - "stopped worker `{}` and reclaimed delegated scope", - input.name + "SubWorkerStop - done\n {} tool kind{}\n {}ms", + summary.tool_counts.len(), + if summary.tool_counts.len() == 1 { + "" + } else { + "s" + }, + summary.elapsed_ms, + ), + content: Some( + serde_json::to_string(&summary) + .map_err(|error| ToolError::ExecutionFailed(error.to_string()))?, ), - content: None, attachments: Vec::new(), }); } diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 0f9042cf..c800b01d 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -7,17 +7,20 @@ //! Parent registry drop closes all session handles and synchronously returns delegated Write deny //! rules to the parent scope. -use std::collections::HashSet; +use std::collections::{BTreeMap, HashSet}; use std::io; use std::sync::{ Arc, Mutex, atomic::{AtomicBool, AtomicU64, Ordering}, }; +use std::time::Instant; + +use serde::{Deserialize, Serialize}; use manifest::{Permission, ScopeRule, SharedScope}; use protocol::{Event, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot}; use session_store::{ - WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerStoreError, + LoggedItem, WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerStoreError, }; use tokio::sync::broadcast; use tracing::warn; @@ -27,6 +30,39 @@ use crate::internal_worker::{InternalWorkerSessionHandle, InternalWorkerVisibili use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord}; use crate::runtime::worker_allocation; +const STOP_SUMMARY_TOOL_LIMIT: usize = 16; +const STOP_SUMMARY_TOOL_NAME_LIMIT: usize = 64; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum SubWorkerFinalOutcome { + Done, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct SubWorkerToolCount { + pub name: String, + pub count: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct SubWorkerChangeStat { + pub added: u64, + pub deleted: u64, + pub source: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct SubWorkerStopSummary { + pub session_id: String, + pub display_name: String, + pub outcome: SubWorkerFinalOutcome, + pub elapsed_ms: u64, + pub tool_counts: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub change_stat: Option, +} + #[derive(Clone)] pub(crate) struct InternalSpawnedWorkerRecord { pub worker_name: String, @@ -35,8 +71,13 @@ pub(crate) struct InternalSpawnedWorkerRecord { #[cfg(test)] pub installed_tools: Arc<[String]>, pub session: InternalWorkerSessionHandle, + change_tracker: Option, + started_at: Instant, + stop_lock: Arc>, scope_reclaimed: Arc, protocol_revision: Arc, + protocol_emit_lock: Arc>, + protocol_terminal: Arc, forwarding_started: Arc, } @@ -47,6 +88,7 @@ impl InternalSpawnedWorkerRecord { workdir_delegation: WorkdirDelegation, #[cfg(test)] installed_tools: Vec, session: InternalWorkerSessionHandle, + change_tracker: Option, ) -> Self { Self { worker_name, @@ -55,12 +97,64 @@ impl InternalSpawnedWorkerRecord { #[cfg(test)] installed_tools: installed_tools.into(), session, + change_tracker, + started_at: Instant::now(), + stop_lock: Arc::new(tokio::sync::Mutex::new(())), scope_reclaimed: Arc::new(AtomicBool::new(false)), protocol_revision: Arc::new(AtomicU64::new(0)), + protocol_emit_lock: Arc::new(Mutex::new(())), + protocol_terminal: Arc::new(AtomicBool::new(false)), forwarding_started: Arc::new(AtomicBool::new(false)), } } + fn stop_summary(&self) -> SubWorkerStopSummary { + let mut counts = BTreeMap::::new(); + for entry in self.session.entries() { + if let session_store::LogEntry::AssistantItem { + item: LoggedItem::ToolCall { name, .. }, + .. + } = entry + { + let count = counts.entry(bounded_tool_name(&name)).or_default(); + *count = count.saturating_add(1); + } + } + let mut tool_counts = counts + .into_iter() + .map(|(name, count)| SubWorkerToolCount { name, count }) + .collect::>(); + tool_counts.sort_by(|left, right| { + right + .count + .cmp(&left.count) + .then_with(|| left.name.cmp(&right.name)) + }); + tool_counts.truncate(STOP_SUMMARY_TOOL_LIMIT); + + let change_stat = self.change_tracker.as_ref().and_then(|tracker| { + let stat = tracker.change_stat(); + (stat.added > 0 || stat.deleted > 0).then(|| SubWorkerChangeStat { + added: stat.added, + deleted: stat.deleted, + source: "tracked_write_edit_tools".to_string(), + }) + }); + + SubWorkerStopSummary { + session_id: self.session.session_id_string(), + display_name: self.worker_name.clone(), + outcome: SubWorkerFinalOutcome::Done, + elapsed_ms: self + .started_at + .elapsed() + .as_millis() + .min(u128::from(u64::MAX)) as u64, + tool_counts, + change_stat, + } + } + fn claim_scope_reclaim(&self) -> bool { !self.scope_reclaimed.swap(true, Ordering::AcqRel) } @@ -277,12 +371,20 @@ impl SpawnedWorkerRegistry { }; let worker = record.protocol_ref(Some(parent_session_id)); let protocol_revision = record.protocol_revision.clone(); + let protocol_emit_lock = record.protocol_emit_lock.clone(); + let protocol_terminal = record.protocol_terminal.clone(); let mut child_rx = record.session.subscribe_events(); tokio::spawn(async move { loop { match child_rx.recv().await { Ok(event) => { let shutdown = matches!(event, Event::Shutdown); + let _emit_guard = protocol_emit_lock + .lock() + .unwrap_or_else(|error| error.into_inner()); + if protocol_terminal.load(Ordering::Acquire) { + break; + } let revision = protocol_revision.fetch_add(1, Ordering::AcqRel) + 1; let _ = parent_tx.send(Event::InternalWorker { worker: worker.clone(), @@ -294,6 +396,12 @@ impl SpawnedWorkerRegistry { } } Err(broadcast::error::RecvError::Lagged(skipped)) => { + let _emit_guard = protocol_emit_lock + .lock() + .unwrap_or_else(|error| error.into_inner()); + if protocol_terminal.load(Ordering::Acquire) { + break; + } let revision = protocol_revision.fetch_add(1, Ordering::AcqRel) + 1; let _ = parent_tx.send(Event::InternalWorker { worker: worker.clone(), @@ -385,13 +493,34 @@ impl SpawnedWorkerRegistry { result } + /// Stop one direct Internal SubWorker and discard its registry/scope state. + /// + /// The child actor must acknowledge its stop before the registry is removed. + /// After scope reclamation and removal, `InternalWorkerRemoved` is published + /// exactly once as the parent-stream terminal fence. Callers only receive + /// `Done` after all authoritative cleanup succeeds. pub(crate) async fn remove_internal( &self, worker_name: &str, - ) -> io::Result> { - if let Some(record) = self.get_internal(worker_name) { - self.reclaim_record_scope(&record)?; + ) -> io::Result> { + let Some(record) = self.get_internal(worker_name) else { + return Ok(None); + }; + let _stop_guard = record.stop_lock.lock().await; + let still_registered = self.get_internal(worker_name).is_some_and(|current| { + current.session.session_id_string() == record.session.session_id_string() + }); + if !still_registered { + return Ok(None); } + + record + .session + .stop() + .await + .map_err(|error| io::Error::other(error.to_string()))?; + let summary = record.stop_summary(); + self.reclaim_record_scope(&record)?; let removed = { let mut records = self.internal_records.lock().map_err(|_| { @@ -402,14 +531,41 @@ impl SpawnedWorkerRegistry { })?; let removed = records .iter() - .position(|record| record.worker_name == worker_name) + .position(|candidate| { + candidate.worker_name == worker_name + && candidate.session.session_id_string() + == record.session.session_id_string() + }) .map(|index| records.remove(index)); if removed.is_some() { names.remove(worker_name); } removed }; - Ok(removed) + if removed.is_some() { + self.publish_internal_removal(&record); + } + Ok(removed.map(|_| summary)) + } + + fn publish_internal_removal(&self, record: &InternalSpawnedWorkerRecord) { + if record.session.visibility() != InternalWorkerVisibility::ParentClient { + return; + } + let Some((parent_tx, parent_session_id)) = self.parent_protocol.lock().unwrap().clone() + else { + return; + }; + let _emit_guard = record + .protocol_emit_lock + .lock() + .unwrap_or_else(|error| error.into_inner()); + record.protocol_terminal.store(true, Ordering::Release); + let revision = record.protocol_revision.fetch_add(1, Ordering::AcqRel) + 1; + let _ = parent_tx.send(Event::InternalWorkerRemoved { + worker: record.protocol_ref(Some(parent_session_id)), + revision, + }); } } @@ -508,6 +664,17 @@ fn record_from_worker_state(child: &WorkerSpawnedChild) -> io::Result String { + let mut bounded = name + .chars() + .take(STOP_SUMMARY_TOOL_NAME_LIMIT) + .collect::(); + if name.chars().count() > STOP_SUMMARY_TOOL_NAME_LIMIT { + bounded.push('…'); + } + bounded +} + fn store_error_to_io(error: WorkerStoreError) -> io::Error { io::Error::other(error) } @@ -520,7 +687,7 @@ mod tests { use session_store::LogEntry; use super::*; - use crate::internal_worker::test_internal_worker_session; + use crate::internal_worker::{InternalWorkerSessionStatus, test_internal_worker_session}; fn registry() -> Arc { let scope = Scope::from_config(&ScopeConfig { @@ -577,6 +744,7 @@ mod tests { delegation, Vec::new(), session, + None, ), sender, ) @@ -669,4 +837,124 @@ mod tests { ); assert!(registry.internal_worker_snapshots().is_empty()); } + + fn install_record(registry: &SpawnedWorkerRegistry, record: InternalSpawnedWorkerRecord) { + registry + .internal_names + .lock() + .unwrap() + .insert(record.worker_name.clone()); + registry.internal_records.lock().unwrap().push(record); + } + + #[tokio::test] + async fn stop_removes_internal_worker_and_returns_bounded_summary() { + let registry = registry(); + let (parent_tx, mut parent_rx) = broadcast::channel(32); + registry.attach_parent_protocol(parent_tx, "parent-session".into()); + let tracker = tools::Tracker::new(); + tracker.record_change(12, 4); + let (mut record, _events) = record("child", InternalWorkerVisibility::ParentClient).await; + record.change_tracker = Some(tracker); + for (index, name) in ["Read", "Read", "Grep"].into_iter().enumerate() { + record.session.publish_test_entry(LogEntry::AssistantItem { + ts: index as u64, + item: LoggedItem::ToolCall { + call_id: format!("call-{index}"), + name: name.to_string(), + arguments: "{}".to_string(), + }, + }); + } + registry.start_protocol_forwarding(record.clone()); + install_record(®istry, record); + + let summary = registry.remove_internal("child").await.unwrap().unwrap(); + + assert_eq!(summary.display_name, "child"); + assert_eq!(summary.outcome, SubWorkerFinalOutcome::Done); + assert_eq!( + summary.tool_counts, + vec![ + SubWorkerToolCount { + name: "Read".to_string(), + count: 2, + }, + SubWorkerToolCount { + name: "Grep".to_string(), + count: 1, + }, + ] + ); + assert_eq!( + summary.change_stat, + Some(SubWorkerChangeStat { + added: 12, + deleted: 4, + source: "tracked_write_edit_tools".to_string(), + }) + ); + assert!(registry.get_internal("child").is_none()); + let terminal_revision = loop { + if let Event::InternalWorkerRemoved { worker, revision } = + parent_rx.recv().await.unwrap() + { + assert_eq!(worker.session_id, summary.session_id); + assert!(revision > 0); + break revision; + } + }; + assert!(registry.remove_internal("child").await.unwrap().is_none()); + while let Ok(Ok(event)) = + tokio::time::timeout(Duration::from_millis(20), parent_rx.recv()).await + { + assert!(!matches!(event, Event::InternalWorkerRemoved { .. })); + if let Event::InternalWorker { revision, .. } = event { + assert!(revision > terminal_revision); + } + } + } + + #[tokio::test] + async fn running_worker_is_stopped_before_removal() { + let registry = registry(); + let (record, _events) = record("running", InternalWorkerVisibility::ParentClient).await; + record + .session + .force_status(InternalWorkerSessionStatus::Running); + install_record(®istry, record); + + let summary = registry.remove_internal("running").await.unwrap().unwrap(); + + assert_eq!(summary.outcome, SubWorkerFinalOutcome::Done); + assert!(registry.get_internal("running").is_none()); + } + + #[tokio::test] + async fn stop_failure_keeps_registry_and_emits_no_removal() { + let registry = registry(); + let (parent_tx, mut parent_rx) = broadcast::channel(8); + registry.attach_parent_protocol(parent_tx, "parent-session".into()); + let (record, _events) = record("child", InternalWorkerVisibility::ParentClient).await; + record.session.force_stop_failure(); + install_record(®istry, record); + + let error = registry.remove_internal("child").await.unwrap_err(); + + assert!(error.to_string().contains("unavailable")); + assert!(registry.get_internal("child").is_some()); + assert!(matches!( + parent_rx.try_recv(), + Err(broadcast::error::TryRecvError::Empty) + )); + } + + #[tokio::test] + async fn read_only_summary_omits_unavailable_change_stat() { + let tracker = tools::Tracker::new(); + let (mut record, _events) = record("reader", InternalWorkerVisibility::ParentClient).await; + record.change_tracker = Some(tracker); + + assert_eq!(record.stop_summary().change_stat, None); + } } diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index b4c117cd..41e0866c 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -481,6 +481,7 @@ impl Tool for SubWorkerSpawnTool { .map_err(|error| { ToolError::ExecutionFailed(format!("install Internal Worker features: {error}")) })?; + let child_change_tracker = child.tracker().cloned(); #[cfg(test)] let installed_tools = child .engine() @@ -587,6 +588,7 @@ impl Tool for SubWorkerSpawnTool { #[cfg(test)] installed_tools, session.clone(), + child_change_tracker, ); if let Err(error) = name_reservation.commit(record) { let _ = session.stop().await; diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index 471f4dab..cf8b298c 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -178,4 +178,4 @@ in_flight?: InFlightSnapshot, * Parent-owned Internal Worker sessions visible to this client. * Service-private Internal Workers are deliberately excluded. */ -internal_workers?: Array, } } | { "event": "internal_worker", "data": { worker: InternalWorkerRef, revision: number, event: Event, } } | { "event": "segment_rotated", "data": { entry: unknown, } } | { "event": "status", "data": { status: WorkerStatus, } } | { "event": "completions", "data": { kind: CompletionKind, entries: Array, } } | { "event": "rewind_targets", "data": { head_entries: number, targets: Array, } } | { "event": "rewind_applied", "data": { entries: Array, input: Array, summary: RewindSummary, } } | { "event": "workers_listed", "data": { workers: unknown, } } | { "event": "worker_restored", "data": { result: unknown, } } | { "event": "peer_registered", "data": { result: unknown, } } | { "event": "alert", "data": Alert } | { "event": "memory_worker", "data": MemoryWorkerEvent } | { "event": "compact_start" } | { "event": "compact_done", "data": { new_segment_id: string, } } | { "event": "compact_failed", "data": { error: string, } } | { "event": "shutdown" }; +internal_workers?: Array, } } | { "event": "internal_worker", "data": { worker: InternalWorkerRef, revision: number, event: Event, } } | { "event": "internal_worker_removed", "data": { worker: InternalWorkerRef, revision: number, } } | { "event": "segment_rotated", "data": { entry: unknown, } } | { "event": "status", "data": { status: WorkerStatus, } } | { "event": "completions", "data": { kind: CompletionKind, entries: Array, } } | { "event": "rewind_targets", "data": { head_entries: number, targets: Array, } } | { "event": "rewind_applied", "data": { entries: Array, input: Array, summary: RewindSummary, } } | { "event": "workers_listed", "data": { workers: unknown, } } | { "event": "worker_restored", "data": { result: unknown, } } | { "event": "peer_registered", "data": { result: unknown, } } | { "event": "alert", "data": Alert } | { "event": "memory_worker", "data": MemoryWorkerEvent } | { "event": "compact_start" } | { "event": "compact_done", "data": { new_segment_id: string, } } | { "event": "compact_failed", "data": { error: string, } } | { "event": "shutdown" }; diff --git a/web/workspace/src/lib/workspace/console/model.test.ts b/web/workspace/src/lib/workspace/console/model.test.ts index 041fdbfb..dbd701b9 100644 --- a/web/workspace/src/lib/workspace/console/model.test.ts +++ b/web/workspace/src/lib/workspace/console/model.test.ts @@ -1287,7 +1287,10 @@ Deno.test("Internal Worker output stays separate and revision-fenced", () => { }]); assertEquals(projection.lines, []); assertEquals(projection.internalWorkers.length, 1); - assertEquals(projection.internalWorkers[0].console.lines[0].body, "child output"); + assertEquals( + projection.internalWorkers[0].console.lines[0].body, + "child output", + ); projection = projector.append([{ eventId: "2", @@ -1364,15 +1367,112 @@ Deno.test("parent snapshot authoritatively replaces Internal Worker projections" }, }]); const projection = projector.append([{ eventId: "snapshot", event }]); - assertEquals(projection.internalWorkers.map((worker) => worker.worker.session_id), [ - "replacement", - ]); + assertEquals( + projection.internalWorkers.map((worker) => worker.worker.session_id), + [ + "replacement", + ], + ); const childLines = projection.internalWorkers[0].console.lines; assertEquals(childLines.length, 1); assertEquals(new Set(childLines.map((line) => line.id)).size, 1); assertEquals(childLines[0].kind, "tool"); }); +Deno.test("terminal Internal Worker removal drops descendants and fences late events", () => { + const worker = { + session_id: "child-session", + name: "child", + parent_session_id: "parent-session", + kind: "sub_worker" as const, + }; + const nestedWorker = { + session_id: "grandchild-session", + name: "grandchild", + parent_session_id: "child-session", + kind: "sub_worker" as const, + }; + const projector = createConsoleProjector(); + let projection = projector.append([{ + eventId: "child", + event: { + event: "internal_worker", + data: { + worker, + revision: 2, + event: { + event: "internal_worker", + data: { + worker: nestedWorker, + revision: 1, + event: { event: "text_done", data: { text: "nested" } }, + }, + }, + }, + }, + }]); + assertEquals(projection.internalWorkers.length, 1); + assertEquals( + projection.internalWorkers[0].console.internalWorkers.length, + 1, + ); + + projection = projector.append([{ + eventId: "removed", + event: { + event: "internal_worker_removed", + data: { worker, revision: 3 }, + }, + }, { + eventId: "late", + event: { + event: "internal_worker", + data: { + worker, + revision: 4, + event: { event: "text_done", data: { text: "must stay removed" } }, + }, + }, + }]); + assertEquals(projection.internalWorkers, []); + + const snapshot = snapshotEvent("/repo"); + projection = projector.append([{ eventId: "snapshot", event: snapshot }]); + assertEquals(projection.internalWorkers, []); + assertEquals(projection.removedInternalWorkers, {}); +}); + +Deno.test("stale Internal Worker removal cannot discard a newer projection", () => { + const worker = { + session_id: "child-session", + name: "child", + parent_session_id: "parent-session", + kind: "sub_worker" as const, + }; + const projector = createConsoleProjector(); + projector.append([{ + eventId: "current", + event: { + event: "internal_worker", + data: { + worker, + revision: 4, + event: { event: "text_done", data: { text: "current" } }, + }, + }, + }]); + + const projection = projector.append([{ + eventId: "stale-removal", + event: { + event: "internal_worker_removed", + data: { worker, revision: 3 }, + }, + }]); + assertEquals(projection.internalWorkers.length, 1); + assertEquals(projection.internalWorkers[0].revision, 4); +}); + Deno.test("snapshot restores TaskStore state from system history", () => { const taskSnapshot = `[Session TaskStore snapshot]\n\n\`\`\`json\n{\n "tasks": [{"taskid": 3, "status": "pending", "subject": "Restored", "description": "From compaction"}]\n}\n\`\`\``; diff --git a/web/workspace/src/lib/workspace/console/model.ts b/web/workspace/src/lib/workspace/console/model.ts index 2e632fff..cd7d6e44 100644 --- a/web/workspace/src/lib/workspace/console/model.ts +++ b/web/workspace/src/lib/workspace/console/model.ts @@ -94,6 +94,8 @@ export type ConsoleProjection = { cwd: string | null; lastEventId: string | null; internalWorkers: InternalWorkerProjection[]; + /** Terminal child-session fences, reset only by an authoritative snapshot. */ + removedInternalWorkers: Record; }; export type ConsoleTimelineLineSelection = { @@ -179,6 +181,7 @@ export function emptyConsoleProjection(): ConsoleProjection { cwd: null, lastEventId: null, internalWorkers: [], + removedInternalWorkers: {}, }; } @@ -286,6 +289,7 @@ export function applyProtocolEvent( cwd: projection.cwd, lastEventId: envelope.eventId, internalWorkers: [...projection.internalWorkers], + removedInternalWorkers: { ...projection.removedInternalWorkers }, }; const event = envelope.event; @@ -406,9 +410,16 @@ export function applyProtocolEvent( next.internalWorkers = (event.data.internal_workers ?? []).map((worker) => projectInternalWorkerSnapshot(worker, envelope.eventId, next.cwd) ); + next.removedInternalWorkers = {}; break; } case "internal_worker": { + if ( + Object.hasOwn( + next.removedInternalWorkers, + event.data.worker.session_id, + ) + ) break; const existingIndex = next.internalWorkers.findIndex((worker) => worker.worker.session_id === event.data.worker.session_id ); @@ -433,6 +444,19 @@ export function applyProtocolEvent( else next.internalWorkers.push(updated); break; } + case "internal_worker_removed": { + const existingIndex = next.internalWorkers.findIndex((worker) => + worker.worker.session_id === event.data.worker.session_id + ); + const existingRevision = existingIndex >= 0 + ? next.internalWorkers[existingIndex].revision + : 0; + if (event.data.revision <= existingRevision) break; + next.removedInternalWorkers[event.data.worker.session_id] = + event.data.revision; + if (existingIndex >= 0) next.internalWorkers.splice(existingIndex, 1); + break; + } case "status": next.status = event.data.status; break; @@ -1294,6 +1318,7 @@ function snapshotProjectionFromEntries( cwd, lastEventId: eventId, internalWorkers: [], + removedInternalWorkers: {}, }; entries.forEach((entry, index) => applyLogEntry(projection, `${eventId}-snapshot-${index}`, entry)