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

fix: finalize stopped subworker sessions

Browse Source
This commit is contained in:
Hare
2026-08-20 11:19:02 +09:00
parent 9194f0a1ba
commit 3e833b5295
14 changed files with 816 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/protocol/src/lib.rs
+29
View File
@@ -530,6 +530,15 @@ pub enum Event {
revision: u64,
event: Box<Event>,
},
/// 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!({
crates/tools/src/edit.rs
+10 -3
View File
@@ -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(&params.new_string, &params.new_string);
crates/tools/src/lib.rs
+1 -1
View File
@@ -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;
crates/tools/src/tracker.rs
+90 -2
View File
@@ -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<PathBuf, ContentHash>,
/// Line count paired with observations that included the file content.
line_counts: HashMap<PathBuf, usize>,
/// LRU list of touched files. Front = most recently touched.
recency: VecDeque<PathBuf>,
/// 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<usize> {
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();
crates/tools/src/write.rs
+3
View File
@@ -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());
crates/tui/src/app.rs
+112 -1
View File
@@ -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<InternalWorkerView>,
/// Terminal child-session fences, reset only by an authoritative snapshot.
removed_internal_workers: HashMap<String, u64>,
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());
crates/worker/src/feature/builtin/manage_worker.rs
+91 -5
View File
@@ -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::<SubWorkerStopSummary>(&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::<Vec<_>>()
.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<I: JsonSchema + 'static>(
operation: WorkerOperation,
control: Arc<dyn WorkerControlService>,
@@ -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::<SubWorkerStopSummary>(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());
crates/worker/src/internal_worker.rs
+34 -1
View File
@@ -294,6 +294,8 @@ pub(crate) struct InternalWorkerSessionHandle {
last_error: Arc<Mutex<Option<String>>>,
child_registry: Option<Arc<SpawnedWorkerRegistry>>,
sink: SegmentLogSink,
#[cfg(test)]
fail_stop: Arc<std::sync::atomic::AtomicBool>,
}
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)
}
crates/worker/src/spawn/comm_tools.rs
+16 -9
View File
@@ -170,20 +170,27 @@ impl Tool for SubWorkerStopTool {
) -> Result<ToolOutput, ToolError> {
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
if let Some(summary) = self
.registry
.remove_internal(&input.name)
.await
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
.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(),
});
}
crates/worker/src/spawn/registry.rs
+296 -8
View File
@@ -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<SubWorkerToolCount>,
#[serde(skip_serializing_if = "Option::is_none")]
pub change_stat: Option<SubWorkerChangeStat>,
}
#[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<tools::Tracker>,
started_at: Instant,
stop_lock: Arc<tokio::sync::Mutex<()>>,
scope_reclaimed: Arc<AtomicBool>,
protocol_revision: Arc<AtomicU64>,
protocol_emit_lock: Arc<Mutex<()>>,
protocol_terminal: Arc<AtomicBool>,
forwarding_started: Arc<AtomicBool>,
}
@@ -47,6 +88,7 @@ impl InternalSpawnedWorkerRecord {
workdir_delegation: WorkdirDelegation,
#[cfg(test)] installed_tools: Vec<String>,
session: InternalWorkerSessionHandle,
change_tracker: Option<tools::Tracker>,
) -> 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::<String, u64>::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::<Vec<_>>();
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<Option<InternalSpawnedWorkerRecord>> {
if let Some(record) = self.get_internal(worker_name) {
self.reclaim_record_scope(&record)?;
) -> io::Result<Option<SubWorkerStopSummary>> {
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<SpawnedWor
})
}
fn bounded_tool_name(name: &str) -> String {
let mut bounded = name
.chars()
.take(STOP_SUMMARY_TOOL_NAME_LIMIT)
.collect::<String>();
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<SpawnedWorkerRegistry> {
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(&registry, 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(&registry, 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(&registry, 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);
}
}
crates/worker/src/spawn/tool.rs
+2
View File
@@ -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;
web/workspace/src/lib/generated/protocol.ts
+1 -1
View File
@@ -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<InternalWorkerSnapshot>, } } | { "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<CompletionEntry>, } } | { "event": "rewind_targets", "data": { head_entries: number, targets: Array<RewindTarget>, } } | { "event": "rewind_applied", "data": { entries: Array<unknown>, input: Array<Segment>, 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<InternalWorkerSnapshot>, } } | { "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<CompletionEntry>, } } | { "event": "rewind_targets", "data": { head_entries: number, targets: Array<RewindTarget>, } } | { "event": "rewind_applied", "data": { entries: Array<unknown>, input: Array<Segment>, 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" };
web/workspace/src/lib/workspace/console/model.test.ts
+103 -3
View File
@@ -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), [
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\`\`\``;
web/workspace/src/lib/workspace/console/model.ts
+25
View File
@@ -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<string, number>;
};
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)