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

feat: stream internal subworker output through parent

Browse Source
This commit is contained in:
Hare
2026-08-19 09:24:14 +09:00
parent fe74d7c4b8
commit ce62e09919
18 changed files with 990 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/protocol/src/lib.rs
+110
View File
@@ -278,6 +278,47 @@ impl Method {
// Event (Worker → Client via Unix Socket broadcast)
// ---------------------------------------------------------------------------
/// Presentation category for an Internal Worker exposed through its parent's
/// protocol stream. Internal Workers never become independently addressable
/// protocol subjects.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(rename_all = "snake_case")]
pub enum InternalWorkerKind {
SubWorker,
}
/// Stable presentation identity for one parent-owned Internal Worker session.
/// `name` is display-only; `session_id` is the identity used by clients.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct InternalWorkerRef {
pub session_id: String,
pub name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub parent_session_id: Option<String>,
pub kind: InternalWorkerKind,
}
/// Reconnect state for one visible Internal Worker. The revision fences live
/// `Event::InternalWorker` updates that raced with parent snapshot assembly.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct InternalWorkerSnapshot {
pub worker: InternalWorkerRef,
pub revision: u64,
#[cfg_attr(feature = "typescript", ts(type = "Array<unknown>"))]
pub entries: Vec<serde_json::Value>,
#[serde(default)]
pub status: WorkerStatus,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
#[serde(default, skip_serializing_if = "InFlightSnapshot::is_empty")]
pub in_flight: InFlightSnapshot,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub internal_workers: Vec<InternalWorkerSnapshot>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(tag = "event", content = "data", rename_all = "snake_case")]
@@ -476,6 +517,18 @@ pub enum Event {
/// run but is not yet represented by committed snapshot entries.
#[serde(default, skip_serializing_if = "InFlightSnapshot::is_empty")]
in_flight: InFlightSnapshot,
/// Parent-owned Internal Worker sessions visible to this client.
/// Service-private Internal Workers are deliberately excluded.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
internal_workers: Vec<InternalWorkerSnapshot>,
},
/// A live event from a parent-owned Internal Worker. The payload reuses the
/// normal Worker event vocabulary while the wrapper carries stable origin
/// identity and a per-child revision fence.
InternalWorker {
worker: InternalWorkerRef,
revision: u64,
event: Box<Event>,
},
/// Server-side segment log rotated to a fresh `SegmentStart`.
///
@@ -1269,6 +1322,7 @@ mod tests {
},
status: WorkerStatus::Paused,
in_flight: InFlightSnapshot::default(),
internal_workers: Vec::new(),
};
let json = serde_json::to_string(&event).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
@@ -1322,6 +1376,7 @@ mod tests {
},
],
},
internal_workers: Vec::new(),
};
let json = serde_json::to_string(&event).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
@@ -1716,6 +1771,61 @@ mod tests {
}
}
#[test]
fn internal_worker_event_roundtrip_preserves_origin_and_payload() {
let event = Event::InternalWorker {
worker: InternalWorkerRef {
session_id: "session-1".into(),
name: "research".into(),
parent_session_id: Some("parent-session".into()),
kind: InternalWorkerKind::SubWorker,
},
revision: 7,
event: Box::new(Event::TextDone {
text: "result".into(),
}),
};
let json = serde_json::to_string(&event).unwrap();
let decoded: Event = serde_json::from_str(&json).unwrap();
match decoded {
Event::InternalWorker {
worker,
revision,
event,
} => {
assert_eq!(worker.session_id, "session-1");
assert_eq!(worker.parent_session_id.as_deref(), Some("parent-session"));
assert_eq!(revision, 7);
assert!(matches!(*event, Event::TextDone { ref text } if text == "result"));
}
other => panic!("expected internal Worker event, got {other:?}"),
}
}
#[test]
fn legacy_snapshot_defaults_internal_workers_to_empty() {
let snapshot: Event = serde_json::from_value(serde_json::json!({
"event": "snapshot",
"data": {
"entries": [],
"greeting": {
"worker_name": "parent",
"cwd": ".",
"provider": "provider",
"model": "model",
"scope_summary": "scope",
"tools": []
},
"status": "idle"
}
}))
.unwrap();
assert!(matches!(
snapshot,
Event::Snapshot { internal_workers, .. } if internal_workers.is_empty()
));
}
#[test]
fn worker_discovery_events_roundtrip() {
let events = [
crates/protocol/src/typescript.rs
+7 -3
View File
@@ -4,9 +4,10 @@ use ts_rs::{Config, TS};
use crate::{
Alert, AlertLevel, AlertSource, CompletionEntry, CompletionKind, ErrorCode, Event, Greeting,
InFlightBlock, InFlightSnapshot, InFlightToolCallState, InvokeKind, MemoryWorkerEvent, Method,
Permission, RewindSummary, RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment,
TurnResult, WorkerEvent, WorkerStatus,
InFlightBlock, InFlightSnapshot, InFlightToolCallState, InternalWorkerKind, InternalWorkerRef,
InternalWorkerSnapshot, InvokeKind, MemoryWorkerEvent, Method, Permission, RewindSummary,
RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, TurnResult, WorkerEvent,
WorkerStatus,
subscription::{
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
@@ -53,6 +54,9 @@ pub fn generated_protocol_types() -> String {
push_decl::<RewindSummary>(&cfg, &mut output);
push_decl::<InFlightBlock>(&cfg, &mut output);
push_decl::<InFlightSnapshot>(&cfg, &mut output);
push_decl::<InternalWorkerKind>(&cfg, &mut output);
push_decl::<InternalWorkerRef>(&cfg, &mut output);
push_decl::<InternalWorkerSnapshot>(&cfg, &mut output);
push_decl::<Greeting>(&cfg, &mut output);
push_decl::<Alert>(&cfg, &mut output);
push_decl::<MemoryWorkerEvent>(&cfg, &mut output);
crates/tui/src/app.rs
+155 -2
View File
@@ -4,8 +4,8 @@ use std::time::{Duration, Instant};
use protocol::{
AlertLevel, AlertSource, CompletionEntry, CompletionKind, ErrorCode, Event, InFlightBlock,
InFlightSnapshot, InFlightToolCallState, Method, RewindTarget, RunResult, Segment,
WorkerStatus,
InFlightSnapshot, InFlightToolCallState, InternalWorkerRef, InternalWorkerSnapshot, Method,
RewindTarget, RunResult, Segment, WorkerStatus,
};
use crate::block::{
@@ -227,6 +227,12 @@ impl ActionbarNotice {
}
}
pub struct InternalWorkerView {
pub worker: InternalWorkerRef,
pub revision: u64,
pub app: Box<App>,
}
pub struct App {
pub worker_name: String,
pub connected: bool,
@@ -274,6 +280,9 @@ pub struct App {
/// Turn/protocol errors retained when a real `SegmentStart` replaces the
/// replayable conversation rows during segment rotation.
run_error_messages: Vec<String>,
/// 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>,
pub scroll: Scroll,
pub mode: Mode,
pub cache: FileCache,
@@ -351,6 +360,7 @@ impl App {
quit_confirm: None,
blocks: Vec::new(),
run_error_messages: Vec::new(),
internal_workers: Vec::new(),
scroll: Scroll::default(),
mode: Mode::Normal,
cache: FileCache::new(),
@@ -1296,11 +1306,18 @@ impl App {
greeting,
status,
in_flight,
internal_workers,
} => {
self.rewind_refresh_fence = false;
self.restore_snapshot(&entries, greeting, in_flight);
self.replace_internal_worker_snapshots(internal_workers);
self.set_worker_status(status);
}
Event::InternalWorker {
worker,
revision,
event,
} => self.apply_internal_worker_event(worker, revision, *event),
Event::Status { status } => {
self.rewind_refresh_fence = false;
self.set_worker_status(status);
@@ -1980,6 +1997,60 @@ impl App {
/// LogEntry variant into the same blocks live events would have
/// produced. Followed by `Event::Entry` updates for anything
/// committed after the snapshot.
fn replace_internal_worker_snapshots(&mut self, snapshots: Vec<InternalWorkerSnapshot>) {
self.internal_workers = snapshots
.into_iter()
.map(Self::internal_worker_view_from_snapshot)
.collect();
}
fn internal_worker_view_from_snapshot(snapshot: InternalWorkerSnapshot) -> InternalWorkerView {
let mut app = App::new(snapshot.worker.name.clone());
app.restore_entries(&snapshot.entries, None);
app.apply_in_flight_snapshot(snapshot.in_flight);
app.set_worker_status(snapshot.status);
if let Some(error) = snapshot.error {
let _ = app.handle_worker_event(Event::Error {
code: protocol::ErrorCode::Internal,
message: error,
});
}
app.replace_internal_worker_snapshots(snapshot.internal_workers);
InternalWorkerView {
worker: snapshot.worker,
revision: snapshot.revision,
app: Box::new(app),
}
}
fn apply_internal_worker_event(
&mut self,
worker: InternalWorkerRef,
revision: u64,
event: Event,
) {
let index = self
.internal_workers
.iter()
.position(|candidate| candidate.worker.session_id == worker.session_id);
let target = if let Some(index) = index {
&mut self.internal_workers[index]
} else {
self.internal_workers.push(InternalWorkerView {
worker: worker.clone(),
revision: 0,
app: Box::new(App::new(worker.name.clone())),
});
self.internal_workers.last_mut().unwrap()
};
if revision <= target.revision {
return;
}
target.worker = worker;
target.revision = revision;
let _ = target.app.handle_worker_event(event);
}
fn restore_snapshot(
&mut self,
entries: &[serde_json::Value],
@@ -3276,6 +3347,7 @@ mod completion_flow_tests {
entries: vec![session_start_value],
status: WorkerStatus::Running,
in_flight: Default::default(),
internal_workers: Vec::new(),
});
assert!(matches!(app.worker_status, WorkerStatus::Running));
@@ -3321,6 +3393,7 @@ mod completion_flow_tests {
entries: vec![serde_json::to_value(run_errored).unwrap()],
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
});
let errors = app
@@ -3398,6 +3471,7 @@ mod completion_flow_tests {
},
],
},
internal_workers: Vec::new(),
});
app.handle_worker_event(Event::TextDelta { text: "lo".into() });
@@ -3434,6 +3508,83 @@ mod completion_flow_tests {
assert!(app.blocks.is_empty());
}
#[test]
fn internal_worker_events_project_by_session_without_mixing_parent_blocks() {
let mut app = App::new("parent".into());
let worker = InternalWorkerRef {
session_id: "child-session".into(),
name: "research".into(),
parent_session_id: Some("parent-session".into()),
kind: protocol::InternalWorkerKind::SubWorker,
};
app.handle_worker_event(Event::InternalWorker {
worker: worker.clone(),
revision: 2,
event: Box::new(Event::TextDelta {
text: "child output".into(),
}),
});
app.handle_worker_event(Event::InternalWorker {
worker,
revision: 1,
event: Box::new(Event::TextDelta {
text: "stale".into(),
}),
});
assert!(app.blocks.is_empty());
assert_eq!(app.internal_workers.len(), 1);
assert_eq!(app.internal_workers[0].revision, 2);
assert!(
app.internal_workers[0].app.blocks.iter().any(
|block| matches!(block, Block::AssistantText { text } if text == "child output")
)
);
}
#[test]
fn snapshot_authoritatively_replaces_internal_worker_views() {
let mut app = App::new("parent".into());
app.internal_workers.push(InternalWorkerView {
worker: InternalWorkerRef {
session_id: "old".into(),
name: "old".into(),
parent_session_id: None,
kind: protocol::InternalWorkerKind::SubWorker,
},
revision: 1,
app: Box::new(App::new("old".into())),
});
app.handle_worker_event(Event::Snapshot {
greeting: test_greeting(),
entries: Vec::new(),
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: vec![InternalWorkerSnapshot {
worker: InternalWorkerRef {
session_id: "replacement".into(),
name: "replacement".into(),
parent_session_id: Some("parent-session".into()),
kind: protocol::InternalWorkerKind::SubWorker,
},
revision: 4,
entries: Vec::new(),
status: WorkerStatus::Running,
error: None,
in_flight: Default::default(),
internal_workers: Vec::new(),
}],
});
assert_eq!(app.internal_workers.len(), 1);
assert_eq!(app.internal_workers[0].worker.session_id, "replacement");
assert_eq!(app.internal_workers[0].revision, 4);
assert_eq!(
app.internal_workers[0].app.worker_status,
WorkerStatus::Running
);
}
#[test]
fn live_system_item_notification_appends_notify_block() {
let mut app = App::new("test".into());
@@ -3552,6 +3703,7 @@ mod completion_flow_tests {
greeting,
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
});
assert_eq!(app.context_window, 123_000);
@@ -3749,6 +3901,7 @@ mod completion_flow_tests {
entries: assistant_item_entries,
status: WorkerStatus::Running,
in_flight: Default::default(),
internal_workers: Vec::new(),
});
let tasks = app.task_store.tasks();
crates/tui/src/console/mod.rs
+2
View File
@@ -2019,6 +2019,7 @@ mod tests {
entries: vec![],
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
});
app.handle_worker_event(Event::RewindApplied {
entries: vec![],
@@ -2045,6 +2046,7 @@ mod tests {
entries: vec![],
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
});
type_keys(&mut app, "draft");
crates/tui/src/dashboard/tests.rs
+2
View File
@@ -859,6 +859,7 @@ async fn ticket_queue_notification_sends_notify_when_socket_available() {
},
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
})
.await
.unwrap();
@@ -900,6 +901,7 @@ async fn send_notify_only_can_deliver_weak_notification_without_auto_run() {
},
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
})
.await
.unwrap();
crates/tui/src/ui.rs
+22
View File
@@ -381,6 +381,28 @@ pub fn compute_history(app: &App, width: u16) -> HistoryLayout {
i += 1;
}
for internal in &app.internal_workers {
logical.push((Line::from(""), false));
logical.push((
Line::from(vec![
Span::styled("SubWorker ", Style::default().bold()),
Span::raw(internal.worker.name.clone()),
Span::styled(
format!(" {:?}", internal.app.worker_status),
Style::default().fg(Color::DarkGray),
),
]),
false,
));
let child_width = width.saturating_sub(2).max(1);
let child_history = compute_history(&internal.app, child_width);
logical.extend(child_history.rows.into_iter().map(|row| {
let mut spans = vec![Span::raw(" ")];
spans.extend(row.line.spans);
(Line::from(spans), row.selectable)
}));
}
// Step 2: pre-wrap every logical line to char-based terminal rows so
// scroll math is exact. Track the logical → wrapped mapping so
// turn-start indices get translated into wrapped-row coordinates.
crates/tui/src/worker_list.rs
+1
View File
@@ -914,6 +914,7 @@ mod tests {
greeting: test_greeting(),
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
},
];
crates/worker-runtime/src/runtime.rs
+2
View File
@@ -1435,6 +1435,7 @@ impl Runtime {
},
status: protocol::WorkerStatus::Idle,
in_flight: protocol::InFlightSnapshot { blocks: Vec::new() },
internal_workers: Vec::new(),
})
}
@@ -3770,6 +3771,7 @@ mod tests {
},
status: protocol::WorkerStatus::Running,
in_flight: protocol::InFlightSnapshot { blocks: Vec::new() },
internal_workers: Vec::new(),
},
);
crates/worker/src/controller.rs
+6 -1
View File
@@ -48,6 +48,7 @@ pub struct WorkerHandle {
/// it on every new connection (Event::Snapshot) and forwards
/// subsequent commits (Event::Entry) on the receiver.
pub sink: SegmentLogSink,
spawned_registry: Arc<SpawnedWorkerRegistry>,
}
impl WorkerHandle {
@@ -84,6 +85,7 @@ impl WorkerHandle {
greeting: self.shared_state.greeting.clone(),
status: self.shared_state.get_status(),
in_flight,
internal_workers: self.spawned_registry.internal_worker_snapshots(),
};
(event, entry_rx)
}
@@ -413,6 +415,7 @@ impl WorkerController {
wire_event_bridges_on_engine(&mut worker, &event_tx, &alerter, &in_flight);
// === 3. Tool registration (builtin / memory / spawn-orchestration) ===
spawned_registry.attach_parent_protocol(event_tx.clone(), worker.session_id().to_string());
let fs_for_view = register_worker_tools(
&mut worker,
bash_output_dir,
@@ -460,6 +463,7 @@ impl WorkerController {
alerter: alerter.clone(),
in_flight: in_flight.clone(),
sink: worker.sink(),
spawned_registry: spawned_registry.clone(),
};
let socket_server = match transport {
@@ -502,7 +506,7 @@ impl WorkerController {
/// per-item history commit callback so every assistant / tool item
/// landing in `worker.history` becomes a singular `LogEntry::AssistantItem`
/// / `ToolResult` commit through the sync writer.
fn wire_event_bridges_on_engine<C, St>(
pub(crate) fn wire_event_bridges_on_engine<C, St>(
worker: &mut Worker<C, St>,
event_tx: &broadcast::Sender<Event>,
alerter: &Alerter,
@@ -1888,6 +1892,7 @@ mod tests {
},
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
})
.await
.ok()?;
crates/worker/src/discovery.rs
+6
View File
@@ -1494,6 +1494,7 @@ mod tests {
},
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
})
.await
.unwrap();
@@ -1526,6 +1527,7 @@ mod tests {
},
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
})
.await
.unwrap();
@@ -1614,6 +1616,7 @@ mod tests {
},
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
})
.await
.unwrap();
@@ -1637,6 +1640,7 @@ mod tests {
},
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
})
.await
.unwrap();
@@ -1738,6 +1742,7 @@ mod tests {
},
status: WorkerStatus::Paused,
in_flight: Default::default(),
internal_workers: Vec::new(),
})
.await
.unwrap();
@@ -1787,6 +1792,7 @@ mod tests {
},
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
})
.await;
});
crates/worker/src/internal_worker.rs
+203 -6
View File
@@ -12,10 +12,18 @@ use std::sync::{Arc, Mutex};
use llm_engine::timeline::event::UsageEvent;
use llm_engine::{Engine, llm_client::LlmClient};
use manifest::{Scope, WorkerManifest};
use protocol::{Event, InFlightSnapshot, WorkerStatus};
use session_store::{LogEntry, SegmentId, SessionId, Store, StoreError, TraceEntry};
use tokio::sync::broadcast;
use uuid::Uuid;
use crate::controller::wire_event_bridges_on_engine;
use crate::feature::FeatureRegistryBuilder;
use crate::in_flight::{InFlightEvents, snapshot_from_guard};
use crate::ipc::alerter::Alerter;
use crate::ipc::protocol_session::live_log_entry_event;
use crate::segment_log_sink::SegmentLogSink;
use crate::spawn::registry::SpawnedWorkerRegistry;
use crate::worker::{
Worker, WorkerError, WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext,
};
@@ -195,6 +203,20 @@ where
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum InternalWorkerVisibility {
/// Output may be projected only through the owning parent's protocol stream.
ParentClient,
/// Backend-owned helper output remains private to the service authority.
ServicePrivate,
}
impl Default for InternalWorkerVisibility {
fn default() -> Self {
Self::ServicePrivate
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum InternalWorkerSessionStatus {
Idle,
@@ -246,8 +268,18 @@ enum InternalWorkerSessionCommand {
/// Parent-owned handle for a long-lived Internal Worker session.
///
/// The handle exposes only typed turn, history, status, and stop operations. The underlying Worker,
/// Engine, ephemeral Store, and cancellation sender remain inside the actor task.
/// The handle exposes typed turn, history, status, presentation snapshot, event subscription, and
/// stop operations. The underlying Worker, Engine, and cancellation sender remain inside the actor
/// task; protocol access is consumed only by the owning parent registry.
#[derive(Debug, Clone)]
pub(crate) struct InternalWorkerSessionSnapshot {
pub entries: Vec<LogEntry>,
pub status: WorkerStatus,
pub error: Option<String>,
pub in_flight: InFlightSnapshot,
pub internal_workers: Vec<protocol::InternalWorkerSnapshot>,
}
#[derive(Clone)]
pub(crate) struct InternalWorkerSessionHandle {
command_tx: tokio::sync::mpsc::Sender<InternalWorkerSessionCommand>,
@@ -256,6 +288,12 @@ pub(crate) struct InternalWorkerSessionHandle {
session_id: SessionId,
segment_id: SegmentId,
state_changed: Arc<tokio::sync::Notify>,
in_flight: InFlightEvents,
event_tx: broadcast::Sender<Event>,
visibility: InternalWorkerVisibility,
last_error: Arc<Mutex<Option<String>>>,
child_registry: Option<Arc<SpawnedWorkerRegistry>>,
sink: SegmentLogSink,
}
impl InternalWorkerSessionHandle {
@@ -267,6 +305,54 @@ impl InternalWorkerSessionHandle {
InternalWorkerSessionStatus::decode(self.status.load(std::sync::atomic::Ordering::Acquire))
}
pub(crate) fn visibility(&self) -> InternalWorkerVisibility {
self.visibility
}
pub(crate) fn subscribe_events(&self) -> broadcast::Receiver<Event> {
self.event_tx.subscribe()
}
pub(crate) fn protocol_sender(&self) -> broadcast::Sender<Event> {
self.event_tx.clone()
}
#[cfg(test)]
pub(crate) fn publish_test_entry(&self, entry: LogEntry) {
self.sink.publish(entry);
}
#[cfg(test)]
pub(crate) fn emit_test_text_delta(&self, text: &str) {
let block_id = self.in_flight.start_text_block();
self.in_flight.text_delta(block_id, text.to_owned());
}
pub(crate) fn protocol_snapshot(&self) -> InternalWorkerSessionSnapshot {
let (entries, in_flight) = {
let guard = self.in_flight.snapshot_guard();
let (entries, _) = self.sink.subscribe_with_snapshot();
(entries, snapshot_from_guard(&guard))
};
InternalWorkerSessionSnapshot {
entries,
status: match self.status() {
InternalWorkerSessionStatus::Running => WorkerStatus::Running,
InternalWorkerSessionStatus::Idle => WorkerStatus::Idle,
InternalWorkerSessionStatus::Stopping
| InternalWorkerSessionStatus::Stopped
| InternalWorkerSessionStatus::Failed => WorkerStatus::Paused,
},
error: self.last_error.lock().unwrap().clone(),
in_flight,
internal_workers: self
.child_registry
.as_ref()
.map(|registry| registry.internal_worker_snapshots())
.unwrap_or_default(),
}
}
pub(crate) fn entries(&self) -> Vec<LogEntry> {
self.store
.read_all(self.session_id, self.segment_id)
@@ -305,8 +391,17 @@ impl InternalWorkerSessionHandle {
std::sync::atomic::Ordering::Release,
);
self.state_changed.notify_waiters();
let message = "internal Worker session actor is unavailable".to_owned();
*self.last_error.lock().unwrap() = Some(message.clone());
let _ = self.event_tx.send(Event::Error {
code: protocol::ErrorCode::Internal,
message,
});
return Err(InternalWorkerSessionError::Unavailable);
}
let _ = self.event_tx.send(Event::Status {
status: WorkerStatus::Running,
});
Ok(())
}
@@ -423,11 +518,48 @@ pub(crate) async fn spawn_internal_worker_session(
spawn_prepared_internal_worker_session(worker, store, input, None).await
}
fn spawn_internal_log_event_bridge(sink: SegmentLogSink, event_tx: broadcast::Sender<Event>) {
let (_, mut log_rx) = sink.subscribe_with_snapshot();
tokio::spawn(async move {
loop {
match log_rx.recv().await {
Ok(entry) => {
if let Some(event) = live_log_entry_event(entry) {
let _ = event_tx.send(event);
}
}
Err(broadcast::error::RecvError::Lagged(skipped)) => {
let _ = event_tx.send(Event::Error {
code: protocol::ErrorCode::Internal,
message: format!(
"internal Worker session-log output lagged by {skipped} entries; reconnect to resynchronize"
),
});
break;
}
Err(broadcast::error::RecvError::Closed) => break,
}
}
});
}
pub(crate) async fn prepare_internal_worker_session(
mut worker: Worker<Box<dyn LlmClient>, EphemeralSessionStore>,
store: EphemeralSessionStore,
visibility: InternalWorkerVisibility,
child_registry: Option<Arc<SpawnedWorkerRegistry>>,
on_turn_end: Option<Arc<dyn Fn(InternalWorkerSessionStatus) + Send + Sync>>,
) -> Result<InternalWorkerSessionHandle, InternalWorkerSessionError> {
let (event_tx, _event_rx) = broadcast::channel(256);
let sink = worker.sink();
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());
worker.attach_alerter(alerter.clone());
worker.attach_event_tx(event_tx.clone());
worker.attach_in_flight_events(in_flight.clone());
wire_event_bridges_on_engine(&mut worker, &event_tx, &alerter, &in_flight);
let session_id = worker.session_id();
let segment_id = worker.segment_id();
let (command_tx, mut command_rx) = tokio::sync::mpsc::channel(8);
@@ -435,6 +567,7 @@ pub(crate) async fn prepare_internal_worker_session(
InternalWorkerSessionStatus::Idle.encode(),
));
let state_changed = Arc::new(tokio::sync::Notify::new());
let last_error = Arc::new(Mutex::new(None));
let handle = InternalWorkerSessionHandle {
command_tx,
status: status.clone(),
@@ -442,6 +575,12 @@ pub(crate) async fn prepare_internal_worker_session(
session_id,
segment_id,
state_changed: state_changed.clone(),
in_flight,
event_tx: event_tx.clone(),
visibility,
last_error: last_error.clone(),
child_registry,
sink,
};
tokio::spawn(async move {
@@ -453,11 +592,25 @@ pub(crate) async fn prepare_internal_worker_session(
loop {
tokio::select! {
result = &mut run => {
let turn_status = match result {
Ok(_) => InternalWorkerSessionStatus::Idle,
Err(_) => InternalWorkerSessionStatus::Failed,
let (turn_status, error) = match result {
Ok(_) => (InternalWorkerSessionStatus::Idle, None),
Err(error) => (
InternalWorkerSessionStatus::Failed,
Some(error.to_string()),
),
};
status.store(turn_status.encode(), std::sync::atomic::Ordering::Release);
if let Some(message) = error {
*last_error.lock().unwrap() = Some(message.clone());
let _ = event_tx.send(Event::Error {
code: protocol::ErrorCode::Internal,
message,
});
} else {
let _ = event_tx.send(Event::Status {
status: WorkerStatus::Idle,
});
}
if let Some(callback) = &on_turn_end {
callback(turn_status);
}
@@ -470,6 +623,8 @@ pub(crate) async fn prepare_internal_worker_session(
let _ = cancel_sender.send(()).await;
let _ = (&mut run).await;
status.store(InternalWorkerSessionStatus::Stopped.encode(), std::sync::atomic::Ordering::Release);
let _ = event_tx.send(Event::Status { status: WorkerStatus::Paused });
let _ = event_tx.send(Event::Shutdown);
state_changed.notify_waiters();
let _ = done.send(());
return;
@@ -491,6 +646,10 @@ pub(crate) async fn prepare_internal_worker_session(
InternalWorkerSessionStatus::Stopped.encode(),
std::sync::atomic::Ordering::Release,
);
let _ = event_tx.send(Event::Status {
status: WorkerStatus::Paused,
});
let _ = event_tx.send(Event::Shutdown);
state_changed.notify_waiters();
let _ = done.send(());
return;
@@ -509,7 +668,14 @@ pub(crate) async fn spawn_prepared_internal_worker_session(
input: String,
on_turn_end: Option<Arc<dyn Fn(InternalWorkerSessionStatus) + Send + Sync>>,
) -> Result<InternalWorkerSessionHandle, InternalWorkerSessionError> {
let handle = prepare_internal_worker_session(worker, store, on_turn_end).await?;
let handle = prepare_internal_worker_session(
worker,
store,
InternalWorkerVisibility::ServicePrivate,
None,
on_turn_end,
)
.await?;
handle.send(input).await?;
Ok(handle)
}
@@ -709,6 +875,37 @@ impl session_store::WorkerMetadataStore for EphemeralSessionStore {
}
}
#[cfg(test)]
pub(crate) fn test_internal_worker_session(
visibility: InternalWorkerVisibility,
) -> (InternalWorkerSessionHandle, broadcast::Sender<Event>) {
let store = EphemeralSessionStore::default();
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 sink = SegmentLogSink::new();
spawn_internal_log_event_bridge(sink.clone(), event_tx.clone());
let handle = InternalWorkerSessionHandle {
command_tx,
status: Arc::new(std::sync::atomic::AtomicU8::new(
InternalWorkerSessionStatus::Idle.encode(),
)),
store,
session_id,
segment_id,
state_changed: Arc::new(tokio::sync::Notify::new()),
in_flight: InFlightEvents::new(event_tx.clone()),
event_tx: event_tx.clone(),
visibility,
last_error: Arc::new(Mutex::new(None)),
child_registry: None,
sink,
};
(handle, event_tx)
}
#[cfg(test)]
mod tests {
use std::pin::Pin;
crates/worker/src/spawn/comm_tools.rs
+1
View File
@@ -284,6 +284,7 @@ mod tests {
},
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
}
}
crates/worker/src/spawn/registry.rs
+240 -4
View File
@@ -10,17 +10,19 @@
use std::collections::HashSet;
use std::io;
use std::sync::{
Arc,
atomic::{AtomicBool, Ordering},
Arc, Mutex,
atomic::{AtomicBool, AtomicU64, Ordering},
};
use manifest::{Permission, ScopeRule, SharedScope};
use protocol::{Event, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot};
use session_store::{
WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerStoreError,
};
use tokio::sync::broadcast;
use tracing::warn;
use crate::internal_worker::InternalWorkerSessionHandle;
use crate::internal_worker::{InternalWorkerSessionHandle, InternalWorkerVisibility};
use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
use crate::runtime::worker_allocation;
@@ -30,6 +32,8 @@ pub(crate) struct InternalSpawnedWorkerRecord {
pub scope_delegated: Vec<ScopeRule>,
pub session: InternalWorkerSessionHandle,
scope_reclaimed: Arc<AtomicBool>,
protocol_revision: Arc<AtomicU64>,
forwarding_started: Arc<AtomicBool>,
}
impl InternalSpawnedWorkerRecord {
@@ -43,6 +47,8 @@ impl InternalSpawnedWorkerRecord {
scope_delegated,
session,
scope_reclaimed: Arc::new(AtomicBool::new(false)),
protocol_revision: Arc::new(AtomicU64::new(0)),
forwarding_started: Arc::new(AtomicBool::new(false)),
}
}
@@ -53,6 +59,19 @@ impl InternalSpawnedWorkerRecord {
fn restore_scope_reclaim(&self) {
self.scope_reclaimed.store(false, Ordering::Release);
}
fn protocol_ref(&self, parent_session_id: Option<String>) -> InternalWorkerRef {
InternalWorkerRef {
session_id: self.session.session_id_string(),
name: self.worker_name.clone(),
parent_session_id,
kind: InternalWorkerKind::SubWorker,
}
}
fn protocol_revision(&self) -> u64 {
self.protocol_revision.load(Ordering::Acquire)
}
}
pub(crate) struct InternalSpawnReservation {
@@ -73,7 +92,8 @@ impl InternalSpawnReservation {
.internal_records
.lock()
.map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))?
.push(record);
.push(record.clone());
self.registry.start_protocol_forwarding(record);
self.committed = true;
Ok(())
}
@@ -93,6 +113,7 @@ pub struct SpawnedWorkerRegistry {
internal_records: std::sync::Mutex<Vec<InternalSpawnedWorkerRecord>>,
internal_names: std::sync::Mutex<HashSet<String>>,
parent_scope: Option<SharedScope>,
parent_protocol: Mutex<Option<(broadcast::Sender<Event>, String)>>,
}
pub struct SpawnedWorkerRegistryLoad {
@@ -108,6 +129,7 @@ impl SpawnedWorkerRegistry {
internal_records: std::sync::Mutex::new(Vec::new()),
internal_names: std::sync::Mutex::new(HashSet::new()),
parent_scope: None,
parent_protocol: Mutex::new(None),
})
}
@@ -116,6 +138,7 @@ impl SpawnedWorkerRegistry {
internal_records: std::sync::Mutex::new(Vec::new()),
internal_names: std::sync::Mutex::new(HashSet::new()),
parent_scope: Some(parent_scope),
parent_protocol: Mutex::new(None),
})
}
@@ -193,6 +216,7 @@ impl SpawnedWorkerRegistry {
internal_records: std::sync::Mutex::new(Vec::new()),
internal_names: std::sync::Mutex::new(HashSet::new()),
parent_scope,
parent_protocol: Mutex::new(None),
}),
reclaimed_unreachable: !persisted_children.is_empty(),
})
@@ -220,6 +244,96 @@ impl SpawnedWorkerRegistry {
})
}
pub(crate) fn attach_parent_protocol(
&self,
event_tx: broadcast::Sender<Event>,
parent_session_id: String,
) {
*self.parent_protocol.lock().unwrap() = Some((event_tx, parent_session_id));
for record in self.internal_records.lock().unwrap().clone() {
self.start_protocol_forwarding(record);
}
}
fn start_protocol_forwarding(&self, record: InternalSpawnedWorkerRecord) {
if record.session.visibility() != InternalWorkerVisibility::ParentClient
|| record.forwarding_started.swap(true, Ordering::AcqRel)
{
return;
}
let Some((parent_tx, parent_session_id)) = self.parent_protocol.lock().unwrap().clone()
else {
record.forwarding_started.store(false, Ordering::Release);
return;
};
let worker = record.protocol_ref(Some(parent_session_id));
let protocol_revision = record.protocol_revision.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 revision = protocol_revision.fetch_add(1, Ordering::AcqRel) + 1;
let _ = parent_tx.send(Event::InternalWorker {
worker: worker.clone(),
revision,
event: Box::new(event),
});
if shutdown {
break;
}
}
Err(broadcast::error::RecvError::Lagged(skipped)) => {
let revision = protocol_revision.fetch_add(1, Ordering::AcqRel) + 1;
let _ = parent_tx.send(Event::InternalWorker {
worker: worker.clone(),
revision,
event: Box::new(Event::Error {
code: protocol::ErrorCode::Internal,
message: format!(
"internal Worker output lagged by {skipped} events; reconnect to resynchronize"
),
}),
});
}
Err(broadcast::error::RecvError::Closed) => break,
}
}
});
}
pub(crate) fn internal_worker_snapshots(&self) -> Vec<InternalWorkerSnapshot> {
let parent_session_id = self
.parent_protocol
.lock()
.unwrap()
.as_ref()
.map(|(_, id)| id.clone());
self.internal_records
.lock()
.unwrap()
.iter()
.filter(|record| record.session.visibility() == InternalWorkerVisibility::ParentClient)
.map(|record| {
let snapshot = record.session.protocol_snapshot();
InternalWorkerSnapshot {
worker: record.protocol_ref(parent_session_id.clone()),
revision: record.protocol_revision(),
entries: snapshot
.entries
.into_iter()
.filter_map(|entry| serde_json::to_value(entry).ok())
.collect(),
status: snapshot.status,
error: snapshot.error,
in_flight: snapshot.in_flight,
internal_workers: snapshot.internal_workers,
}
})
.collect()
}
pub(crate) fn get_internal(&self, worker_name: &str) -> Option<InternalSpawnedWorkerRecord> {
self.internal_records
.lock()
@@ -387,3 +501,125 @@ fn record_from_worker_state(child: &WorkerSpawnedChild) -> io::Result<SpawnedWor
fn store_error_to_io(error: WorkerStoreError) -> io::Error {
io::Error::other(error)
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use manifest::{Scope, ScopeConfig};
use session_store::LogEntry;
use super::*;
use crate::internal_worker::test_internal_worker_session;
fn registry() -> Arc<SpawnedWorkerRegistry> {
let scope = Scope::from_config(&ScopeConfig {
allow: vec![ScopeRule {
target: std::path::PathBuf::from("/tmp"),
permission: Permission::Read,
recursive: true,
}],
deny: Vec::new(),
})
.unwrap();
SpawnedWorkerRegistry::new_internal("parent".into(), SharedScope::new(scope))
}
fn record(
name: &str,
visibility: InternalWorkerVisibility,
) -> (InternalSpawnedWorkerRecord, broadcast::Sender<Event>) {
let (session, sender) = test_internal_worker_session(visibility);
(
InternalSpawnedWorkerRecord::new(name.into(), Vec::new(), session),
sender,
)
}
#[tokio::test]
async fn visible_internal_output_is_wrapped_after_registry_insertion() {
let registry = registry();
let (parent_tx, mut parent_rx) = broadcast::channel(16);
registry.attach_parent_protocol(parent_tx, "parent-session".into());
let (record, child_tx) = record("research", InternalWorkerVisibility::ParentClient);
registry
.internal_records
.lock()
.unwrap()
.push(record.clone());
registry.start_protocol_forwarding(record.clone());
child_tx
.send(Event::TextDone {
text: "answer".into(),
})
.unwrap();
let event = tokio::time::timeout(Duration::from_secs(1), parent_rx.recv())
.await
.unwrap()
.unwrap();
assert!(matches!(
event,
Event::InternalWorker { worker, revision: 1, event }
if worker.name == "research"
&& worker.parent_session_id.as_deref() == Some("parent-session")
&& matches!(*event, Event::TextDone { ref text } if text == "answer")
));
record.session.publish_test_entry(LogEntry::UserInput {
ts: 1,
segments: vec![protocol::Segment::text("question")],
extensions: Vec::new(),
});
let committed = tokio::time::timeout(Duration::from_secs(1), parent_rx.recv())
.await
.unwrap()
.unwrap();
assert!(matches!(
committed,
Event::InternalWorker { revision: 2, event, .. }
if matches!(*event, Event::UserMessage { .. })
));
let snapshots = registry.internal_worker_snapshots();
assert_eq!(snapshots.len(), 1);
assert_eq!(snapshots[0].revision, 2);
assert_eq!(snapshots[0].entries.len(), 1);
record.session.emit_test_text_delta("partial");
let streamed = tokio::time::timeout(Duration::from_secs(1), parent_rx.recv())
.await
.unwrap()
.unwrap();
assert!(matches!(
streamed,
Event::InternalWorker { revision: 3, event, .. }
if matches!(*event, Event::TextDelta { ref text } if text == "partial")
));
let snapshots = registry.internal_worker_snapshots();
assert_eq!(snapshots[0].revision, 3);
assert_eq!(snapshots[0].in_flight.blocks.len(), 1);
}
#[tokio::test]
async fn service_private_internal_output_is_never_disclosed() {
let registry = registry();
let (parent_tx, mut parent_rx) = broadcast::channel(16);
registry.attach_parent_protocol(parent_tx, "parent-session".into());
let (record, child_tx) = record("memory-helper", InternalWorkerVisibility::ServicePrivate);
registry
.internal_records
.lock()
.unwrap()
.push(record.clone());
registry.start_protocol_forwarding(record);
let _ = child_tx.send(Event::TextDone {
text: "private".into(),
});
assert!(
tokio::time::timeout(Duration::from_millis(50), parent_rx.recv())
.await
.is_err()
);
assert!(registry.internal_worker_snapshots().is_empty());
}
}
crates/worker/src/spawn/tool.rs
+7 -2
View File
@@ -24,7 +24,8 @@ use tokio::sync::mpsc;
use crate::PromptCatalogSource;
use crate::controller::register_worker_tools;
use crate::internal_worker::{
EphemeralSessionStore, InternalWorkerSessionStatus, prepare_internal_worker_session,
EphemeralSessionStore, InternalWorkerSessionStatus, InternalWorkerVisibility,
prepare_internal_worker_session,
};
use crate::prompt::catalog::PromptCatalog;
use crate::spawn::registry::SpawnedWorkerRegistry;
@@ -481,7 +482,7 @@ impl Tool for SubWorkerSpawnTool {
.join(&input.name)
.join("bash-output"),
self.runtime_base.clone(),
child_registry,
child_registry.clone(),
None,
)
.await
@@ -511,6 +512,8 @@ impl Tool for SubWorkerSpawnTool {
let session_result = prepare_internal_worker_session(
child,
store,
InternalWorkerVisibility::ParentClient,
Some(child_registry.clone()),
Some(Arc::new(move |status| {
if status == InternalWorkerSessionStatus::Failed {
if let Some(registry) = registry.upgrade() {
@@ -543,6 +546,8 @@ impl Tool for SubWorkerSpawnTool {
)));
}
};
child_registry
.attach_parent_protocol(session.protocol_sender(), session.session_id_string());
if let Some((ticket_id, capability_token)) = &reviewer_capability {
let workspace_id = self.workspace_context.workspace_id().ok_or_else(|| {