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

Merge remote-tracking branch 'origin/develop' into develop

Browse Source
This commit is contained in:
Hare
2026-08-21 05:37:57 +09:00
parent 614157424f a21ef31ee7
commit 402543d617
61 changed files with 7101 additions and 523 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+130 -2
View File
@@ -28,8 +28,14 @@ use crate::worker::{
WorkerRunResult,
};
use protocol::{
AlertLevel, AlertSource, ErrorCode, Event, Method, RewindTargetId, RunResult, Segment,
TurnResult, WorkerStatus,
AlertLevel, AlertSource, CommandEvent as ProtocolCommandEvent,
CommandSnapshot as ProtocolCommandSnapshot, CommandStatus as ProtocolCommandStatus,
CommandStream as ProtocolCommandStream, CommandStreamSlice as ProtocolCommandStreamSlice,
ErrorCode, Event, Method, RewindTargetId, RunResult, Segment, TurnResult, WorkerStatus,
};
use workdir::{
CommandEvent as WorkdirCommandEvent, CommandSnapshot as WorkdirCommandSnapshot,
CommandStatus as WorkdirCommandStatus, CommandStream as WorkdirCommandStream, WorkdirSession,
};
// ---------------------------------------------------------------------------
@@ -424,6 +430,9 @@ impl WorkerController {
Some(method_tx.downgrade()),
)
.await?;
if let Some(session) = fs_for_view.as_ref() {
wire_workdir_command_events(session, &in_flight);
}
// Intake role Workers self-terminate only after a successful
// TicketIntakeReady turn has fully settled back to Idle. The request
@@ -498,6 +507,125 @@ impl WorkerController {
}
}
pub(crate) fn wire_workdir_command_events(
session: &Arc<dyn WorkdirSession>,
in_flight: &InFlightEvents,
) {
in_flight.replace_command_snapshot(protocol_command_snapshots(session.as_ref()));
let Some(mut events) = session.subscribe_command_events() else {
return;
};
// Keep only a weak reference in the observer task. Holding the session
// strongly here would keep its broadcast sender alive forever and prevent
// the receiver from observing closure during Worker teardown.
let session = Arc::downgrade(session);
let in_flight = in_flight.clone();
tokio::spawn(async move {
loop {
match events.recv().await {
Ok(event) => in_flight.publish_command_event(protocol_command_event(event)),
Err(broadcast::error::RecvError::Lagged(_)) => {
let Some(session) = session.upgrade() else {
break;
};
in_flight
.replace_command_snapshot(protocol_command_snapshots(session.as_ref()));
}
Err(broadcast::error::RecvError::Closed) => break,
}
}
});
}
fn protocol_command_snapshots(session: &dyn WorkdirSession) -> Vec<ProtocolCommandSnapshot> {
session
.command_snapshot()
.into_iter()
.map(protocol_command_snapshot)
.collect()
}
fn protocol_command_snapshot(snapshot: WorkdirCommandSnapshot) -> ProtocolCommandSnapshot {
ProtocolCommandSnapshot {
command_id: snapshot.command_id,
tool_call_id: snapshot.tool_call_id,
status: protocol_command_status(snapshot.status),
started_at_ms: snapshot.started_at_ms,
observed_at_ms: snapshot.observed_at_ms,
last_output_at_ms: snapshot.last_output_at_ms,
stdout: ProtocolCommandStreamSlice {
start_offset: snapshot.stdout.start_offset,
end_offset: snapshot.stdout.end_offset,
content: snapshot.stdout.content,
truncated: snapshot.stdout.truncated,
},
stderr: ProtocolCommandStreamSlice {
start_offset: snapshot.stderr.start_offset,
end_offset: snapshot.stderr.end_offset,
content: snapshot.stderr.content,
truncated: snapshot.stderr.truncated,
},
exit_code: snapshot.exit_code,
}
}
fn protocol_command_event(event: WorkdirCommandEvent) -> ProtocolCommandEvent {
match event {
WorkdirCommandEvent::Started {
command_id,
tool_call_id,
observed_at_ms,
} => ProtocolCommandEvent::Started {
command_id,
tool_call_id,
observed_at_ms,
},
WorkdirCommandEvent::Output {
command_id,
stream,
start_offset,
end_offset,
content,
observed_at_ms,
} => ProtocolCommandEvent::Output {
command_id,
stream: match stream {
WorkdirCommandStream::Stdout => ProtocolCommandStream::Stdout,
WorkdirCommandStream::Stderr => ProtocolCommandStream::Stderr,
},
start_offset,
end_offset,
content,
observed_at_ms,
},
WorkdirCommandEvent::Terminal {
command_id,
status,
exit_code,
stdout_end_offset,
stderr_end_offset,
observed_at_ms,
} => ProtocolCommandEvent::Terminal {
command_id,
status: protocol_command_status(status),
exit_code,
stdout_end_offset,
stderr_end_offset,
observed_at_ms,
},
}
}
fn protocol_command_status(status: WorkdirCommandStatus) -> ProtocolCommandStatus {
match status {
WorkdirCommandStatus::Running => ProtocolCommandStatus::Running,
WorkdirCommandStatus::Completed => ProtocolCommandStatus::Completed,
WorkdirCommandStatus::Failed => ProtocolCommandStatus::Failed,
WorkdirCommandStatus::TimedOut => ProtocolCommandStatus::TimedOut,
WorkdirCommandStatus::Cancelled => ProtocolCommandStatus::Cancelled,
}
}
/// Wire the per-event broadcast bridges on the Worker's Engine. Each callback
/// re-publishes a worker-level signal as a `protocol::Event` on `event_tx`
/// so subscribers (TUI, socket clients) get a single typed stream.
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/feature/builtin/ticket.rs
+1
View File
@@ -1781,6 +1781,7 @@ provider = "github"
assert!(request.contains("\"title\":\"HTTP ticket\""));
let response_body = serde_json::to_string(&TicketRef {
id: "01TEST".to_string(),
human_key: None,
slug: "http-ticket".to_string(),
status: ticket::TicketStatus::Open,
})
crates/worker/src/in_flight.rs
+157 -2
View File
@@ -1,9 +1,14 @@
use std::sync::{Arc, Mutex, MutexGuard};
use protocol::{Event, InFlightBlock, InFlightSnapshot, InFlightToolCallState};
use protocol::{
CommandEvent, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice, Event,
InFlightBlock, InFlightSnapshot, InFlightToolCallState,
};
use session_store::{LoggedContentPart, LoggedItem};
use tokio::sync::broadcast;
const COMMAND_SNAPSHOT_STREAM_BYTES: usize = 32 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct InFlightBlockId(u64);
@@ -17,6 +22,7 @@ pub struct InFlightEvents {
pub(crate) struct InFlightInner {
next_block_id: u64,
blocks: Vec<TrackedBlock>,
commands: Vec<CommandSnapshot>,
}
#[derive(Debug, Clone)]
@@ -46,6 +52,7 @@ impl InFlightEvents {
inner: Arc::new(Mutex::new(InFlightInner {
next_block_id: 1,
blocks: Vec::new(),
commands: Vec::new(),
})),
event_tx,
}
@@ -201,6 +208,15 @@ impl InFlightEvents {
f()
}
pub(crate) fn publish_command_event(&self, event: CommandEvent) {
self.lock().apply_command_event(&event);
let _ = self.event_tx.send(Event::Command { event });
}
pub(crate) fn replace_command_snapshot(&self, commands: Vec<CommandSnapshot>) {
self.lock().commands = commands;
}
pub(crate) fn clear(&self) {
let mut inner = self.lock();
inner.clear();
@@ -224,6 +240,92 @@ impl InFlightInner {
.find(|block| block.block_id() == block_id)
}
fn apply_command_event(&mut self, event: &CommandEvent) {
match event {
CommandEvent::Started {
command_id,
tool_call_id,
observed_at_ms,
} => {
self.commands
.retain(|command| command.command_id != *command_id);
self.commands.push(CommandSnapshot {
command_id: command_id.clone(),
tool_call_id: tool_call_id.clone(),
status: CommandStatus::Running,
started_at_ms: *observed_at_ms,
observed_at_ms: *observed_at_ms,
last_output_at_ms: None,
stdout: CommandStreamSlice::default(),
stderr: CommandStreamSlice::default(),
exit_code: None,
});
}
CommandEvent::Output {
command_id,
stream,
start_offset,
end_offset,
content,
observed_at_ms,
} => {
let command = match self
.commands
.iter_mut()
.find(|command| command.command_id == *command_id)
{
Some(command) => command,
None => {
self.commands.push(CommandSnapshot {
command_id: command_id.clone(),
tool_call_id: None,
status: CommandStatus::Running,
started_at_ms: *observed_at_ms,
observed_at_ms: *observed_at_ms,
last_output_at_ms: Some(*observed_at_ms),
stdout: CommandStreamSlice::default(),
stderr: CommandStreamSlice::default(),
exit_code: None,
});
self.commands.last_mut().expect("command was inserted")
}
};
command.observed_at_ms = *observed_at_ms;
command.last_output_at_ms = Some(*observed_at_ms);
let target = match stream {
CommandStream::Stdout => &mut command.stdout,
CommandStream::Stderr => &mut command.stderr,
};
if target.end_offset != *start_offset {
target.content.clear();
target.start_offset = *start_offset;
target.truncated = *start_offset > 0;
}
target.content.push_str(content);
target.end_offset = *end_offset;
if target.content.len() > COMMAND_SNAPSHOT_STREAM_BYTES {
let mut cut = target.content.len() - COMMAND_SNAPSHOT_STREAM_BYTES;
while cut < target.content.len() && !target.content.is_char_boundary(cut) {
cut += 1;
}
target.content.drain(..cut);
target.start_offset = target
.end_offset
.saturating_sub(target.content.len() as u64);
target.truncated = true;
}
}
CommandEvent::Terminal { command_id, .. } => {
// Terminal state is delivered as a live protocol event. It is
// no longer in-flight snapshot state, and removing it here
// also prevents queued output from an aborted turn from
// surviving the subsequent terminal event after `clear()`.
self.commands
.retain(|command| command.command_id != *command_id);
}
}
}
fn clear_for_committed_item(&mut self, item: &LoggedItem) {
match item {
LoggedItem::Message { role, content }
@@ -273,14 +375,16 @@ impl InFlightInner {
.iter()
.filter_map(TrackedBlock::to_snapshot_block)
.collect(),
commands: self.commands.clone(),
}
}
fn clear(&mut self) -> bool {
if self.blocks.is_empty() {
if self.blocks.is_empty() && self.commands.is_empty() {
false
} else {
self.blocks.clear();
self.commands.clear();
true
}
}
@@ -583,6 +687,57 @@ mod tests {
);
}
#[test]
fn command_events_are_bounded_and_recoverable_from_snapshot() {
let (event_tx, _) = broadcast::channel(16);
let mut rx = event_tx.subscribe();
let in_flight = InFlightEvents::new(event_tx);
in_flight.publish_command_event(CommandEvent::Started {
command_id: "command-1".into(),
tool_call_id: Some("tool-1".into()),
observed_at_ms: 100,
});
in_flight.publish_command_event(CommandEvent::Output {
command_id: "command-1".into(),
stream: CommandStream::Stdout,
start_offset: 0,
end_offset: 5,
content: "ready".into(),
observed_at_ms: 110,
});
let guard = in_flight.snapshot_guard();
let snapshot = snapshot_from_guard(&guard);
assert_eq!(snapshot.commands.len(), 1);
assert_eq!(snapshot.commands[0].tool_call_id.as_deref(), Some("tool-1"));
assert_eq!(snapshot.commands[0].stdout.content, "ready");
assert_eq!(snapshot.commands[0].status, CommandStatus::Running);
drop(guard);
assert!(matches!(
rx.try_recv().unwrap(),
Event::Command {
event: CommandEvent::Started { .. }
}
));
assert!(matches!(
rx.try_recv().unwrap(),
Event::Command {
event: CommandEvent::Output { .. }
}
));
in_flight.publish_command_event(CommandEvent::Terminal {
command_id: "command-1".into(),
status: CommandStatus::TimedOut,
exit_code: None,
stdout_end_offset: 5,
stderr_end_offset: 0,
observed_at_ms: 200,
});
let guard = in_flight.snapshot_guard();
assert!(snapshot_from_guard(&guard).commands.is_empty());
}
#[test]
fn clear_discards_uncommitted_blocks_without_protocol_event() {
let (event_tx, _) = broadcast::channel(16);
crates/worker/src/internal_worker.rs
+38 -2
View File
@@ -17,7 +17,7 @@ use session_store::{LogEntry, SegmentId, SessionId, Store, StoreError, TraceEntr
use tokio::sync::broadcast;
use uuid::Uuid;
use crate::controller::wire_event_bridges_on_engine;
use crate::controller::{wire_event_bridges_on_engine, wire_workdir_command_events};
use crate::feature::FeatureRegistryBuilder;
use crate::in_flight::{InFlightEvents, snapshot_from_guard};
use crate::ipc::alerter::Alerter;
@@ -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,
@@ -555,6 +576,9 @@ pub(crate) async fn prepare_internal_worker_session(
spawn_internal_log_event_bridge(sink.clone(), event_tx.clone());
let alerter = Alerter::new(event_tx.clone());
let in_flight = InFlightEvents::new(event_tx.clone());
if let Some(session) = worker.workdir_session() {
wire_workdir_command_events(session, &in_flight);
}
let actor_in_flight = in_flight.clone();
worker.attach_alerter(alerter.clone());
worker.attach_event_tx(event_tx.clone());
@@ -582,6 +606,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 +916,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 +944,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
+18 -11
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
.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(),
});
}
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;
crates/worker/tests/controller_test.rs
+178 -2
View File
@@ -12,8 +12,8 @@ use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use session_store::{CombinedStore, FsWorkerStore};
use session_store::{FsStore, LogEntry};
use workdir::{
CommandRequest, LocalWorkdirSession, Workdir, WorkdirError, WorkdirSessionCapabilities,
WorkdirSessionHandle,
CommandOutputRequest, CommandRequest, LocalWorkdirSession, Workdir, WorkdirError,
WorkdirSessionCapabilities, WorkdirSessionHandle,
};
use worker::{
@@ -232,6 +232,7 @@ async fn shutdown_closes_bound_workdir_session() {
command: "sleep 30".to_owned(),
timeout_secs: 60,
output_limit: 1024,
tool_call_id: None,
})
.await
.unwrap();
@@ -253,6 +254,180 @@ async fn shutdown_closes_bound_workdir_session() {
));
}
#[tokio::test]
async fn controller_projects_workdir_command_events_and_snapshot_state() {
let (mut worker, pwd) = make_worker_with_pwd(MockClient::new(simple_text_events())).await;
let session: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::materialized_bound(
Workdir::new("controller-command-observation-workdir"),
pwd.clone(),
pwd,
worker.scope().clone(),
WorkdirSessionCapabilities::ALL,
));
worker.bind_workdir_session(Some(Arc::clone(&session)));
let handle = spawn_controller(worker).await;
let mut events = handle.subscribe();
let command = session
.start_command(CommandRequest {
command: "printf ready; sleep 0.3; printf done".to_owned(),
timeout_secs: 5,
output_limit: 1024,
tool_call_id: Some("tool-command-1".into()),
})
.await
.unwrap();
let mut saw_started = false;
let mut saw_output = false;
while !saw_output {
let event = tokio::time::timeout(std::time::Duration::from_secs(2), events.recv())
.await
.expect("command event should arrive")
.unwrap();
match event {
Event::Command {
event:
protocol::CommandEvent::Started {
command_id,
tool_call_id,
..
},
} => {
assert_eq!(command_id, command.0);
assert_eq!(tool_call_id.as_deref(), Some("tool-command-1"));
saw_started = true;
}
Event::Command {
event:
protocol::CommandEvent::Output {
command_id,
stream: protocol::CommandStream::Stdout,
content,
..
},
} if command_id == command.0 && content.contains("ready") => saw_output = true,
_ => {}
}
}
assert!(saw_started);
let Event::Snapshot { in_flight, .. } = handle.snapshot_event() else {
panic!("worker snapshot expected");
};
assert_eq!(in_flight.commands.len(), 1);
assert_eq!(in_flight.commands[0].command_id, command.0);
assert_eq!(in_flight.commands[0].stdout.content, "ready");
assert_eq!(
in_flight.commands[0].status,
protocol::CommandStatus::Running
);
let saw_terminal = drain_until(&mut events, std::time::Duration::from_secs(2), |event| {
matches!(
event,
Event::Command {
event: protocol::CommandEvent::Terminal {
command_id,
status: protocol::CommandStatus::Completed,
exit_code: Some(0),
..
}
} if command_id == &command.0
)
})
.await;
assert!(saw_terminal, "completed command event should arrive");
let output = session
.command_output(CommandOutputRequest {
handle: command,
cursor: 0,
limit: 1024,
wait: true,
})
.await
.unwrap();
assert_eq!(output.status, workdir::CommandStatus::Completed);
let (entries, _) = handle.sink.subscribe_with_snapshot();
let durable_history = serde_json::to_string(&entries).unwrap();
assert!(
!durable_history.contains("ready") && !durable_history.contains("done"),
"operational command chunks must not be appended to Worker history: {durable_history}"
);
handle.send(Method::Shutdown).await.unwrap();
}
#[tokio::test]
async fn controller_refreshes_command_snapshot_after_high_output_provider_lag() {
let (mut worker, pwd) = make_worker_with_pwd(MockClient::new(simple_text_events())).await;
let session: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::materialized_bound(
Workdir::new("controller-command-lag-recovery-workdir"),
pwd.clone(),
pwd,
worker.scope().clone(),
WorkdirSessionCapabilities::ALL,
));
worker.bind_workdir_session(Some(Arc::clone(&session)));
let handle = spawn_controller(worker).await;
// Local command telemetry uses 8 KiB chunks and a 256-event channel. One
// synchronous file-poll burst with 300 chunks deterministically makes the
// worker-side receiver observe `Lagged` before this command terminates.
let command = session
.start_command(CommandRequest {
command: "dd if=/dev/zero bs=8192 count=300 2>/dev/null | tr '\\0' x; sleep 5"
.to_owned(),
timeout_secs: 10,
output_limit: 1024,
tool_call_id: Some("tool-high-output".into()),
})
.await
.unwrap();
let expected_end_offset = 300_u64 * 8192;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(3);
let recovered = loop {
let Event::Snapshot { in_flight, .. } = handle.snapshot_event() else {
panic!("worker snapshot expected");
};
if let Some(snapshot) = in_flight
.commands
.iter()
.find(|snapshot| snapshot.command_id == command.0)
&& snapshot.stdout.end_offset >= expected_end_offset
{
break snapshot.clone();
}
assert!(
tokio::time::Instant::now() < deadline,
"timed out waiting for lag recovery snapshot"
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
};
assert_eq!(recovered.tool_call_id.as_deref(), Some("tool-high-output"));
assert_eq!(recovered.status, protocol::CommandStatus::Running);
assert!(recovered.stdout.truncated);
assert!(recovered.stdout.start_offset > 0);
assert_eq!(recovered.stdout.end_offset, expected_end_offset);
assert!(recovered.stdout.content.len() <= 32 * 1024);
assert!(recovered.stdout.content.bytes().all(|byte| byte == b'x'));
session.cancel_command(command.clone()).await.unwrap();
let output = session
.command_output(CommandOutputRequest {
handle: command,
cursor: 0,
limit: 1024,
wait: true,
})
.await
.unwrap();
assert_eq!(output.status, workdir::CommandStatus::Cancelled);
handle.send(Method::Shutdown).await.unwrap();
}
#[tokio::test]
async fn controller_startup_failure_closes_bound_workdir_session() {
let (mut worker, pwd) = make_worker_with_pwd(MockClient::new(simple_text_events())).await;
@@ -279,6 +454,7 @@ async fn controller_startup_failure_closes_bound_workdir_session() {
command: "printf unreachable".to_owned(),
timeout_secs: 5,
output_limit: 1024,
tool_call_id: None,
})
.await,
Err(WorkdirError::Unavailable(_))