feat: stream internal subworker output through parent

This commit is contained in:
2026-08-19 09:24:14 +09:00
parent fe74d7c4b8
commit ce62e09919
18 changed files with 990 additions and 19 deletions
+110
View File
@@ -278,6 +278,47 @@ impl Method {
// Event (Worker → Client via Unix Socket broadcast) // 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)] #[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(tag = "event", content = "data", rename_all = "snake_case")] #[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. /// run but is not yet represented by committed snapshot entries.
#[serde(default, skip_serializing_if = "InFlightSnapshot::is_empty")] #[serde(default, skip_serializing_if = "InFlightSnapshot::is_empty")]
in_flight: InFlightSnapshot, 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`. /// Server-side segment log rotated to a fresh `SegmentStart`.
/// ///
@@ -1269,6 +1322,7 @@ mod tests {
}, },
status: WorkerStatus::Paused, status: WorkerStatus::Paused,
in_flight: InFlightSnapshot::default(), in_flight: InFlightSnapshot::default(),
internal_workers: Vec::new(),
}; };
let json = serde_json::to_string(&event).unwrap(); let json = serde_json::to_string(&event).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).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 json = serde_json::to_string(&event).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).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] #[test]
fn worker_discovery_events_roundtrip() { fn worker_discovery_events_roundtrip() {
let events = [ let events = [
+7 -3
View File
@@ -4,9 +4,10 @@ use ts_rs::{Config, TS};
use crate::{ use crate::{
Alert, AlertLevel, AlertSource, CompletionEntry, CompletionKind, ErrorCode, Event, Greeting, Alert, AlertLevel, AlertSource, CompletionEntry, CompletionKind, ErrorCode, Event, Greeting,
InFlightBlock, InFlightSnapshot, InFlightToolCallState, InvokeKind, MemoryWorkerEvent, Method, InFlightBlock, InFlightSnapshot, InFlightToolCallState, InternalWorkerKind, InternalWorkerRef,
Permission, RewindSummary, RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, InternalWorkerSnapshot, InvokeKind, MemoryWorkerEvent, Method, Permission, RewindSummary,
TurnResult, WorkerEvent, WorkerStatus, RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, TurnResult, WorkerEvent,
WorkerStatus,
subscription::{ subscription::{
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame, EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest, SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
@@ -53,6 +54,9 @@ pub fn generated_protocol_types() -> String {
push_decl::<RewindSummary>(&cfg, &mut output); push_decl::<RewindSummary>(&cfg, &mut output);
push_decl::<InFlightBlock>(&cfg, &mut output); push_decl::<InFlightBlock>(&cfg, &mut output);
push_decl::<InFlightSnapshot>(&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::<Greeting>(&cfg, &mut output);
push_decl::<Alert>(&cfg, &mut output); push_decl::<Alert>(&cfg, &mut output);
push_decl::<MemoryWorkerEvent>(&cfg, &mut output); push_decl::<MemoryWorkerEvent>(&cfg, &mut output);
+155 -2
View File
@@ -4,8 +4,8 @@ use std::time::{Duration, Instant};
use protocol::{ use protocol::{
AlertLevel, AlertSource, CompletionEntry, CompletionKind, ErrorCode, Event, InFlightBlock, AlertLevel, AlertSource, CompletionEntry, CompletionKind, ErrorCode, Event, InFlightBlock,
InFlightSnapshot, InFlightToolCallState, Method, RewindTarget, RunResult, Segment, InFlightSnapshot, InFlightToolCallState, InternalWorkerRef, InternalWorkerSnapshot, Method,
WorkerStatus, RewindTarget, RunResult, Segment, WorkerStatus,
}; };
use crate::block::{ 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 struct App {
pub worker_name: String, pub worker_name: String,
pub connected: bool, pub connected: bool,
@@ -274,6 +280,9 @@ pub struct App {
/// Turn/protocol errors retained when a real `SegmentStart` replaces the /// Turn/protocol errors retained when a real `SegmentStart` replaces the
/// replayable conversation rows during segment rotation. /// replayable conversation rows during segment rotation.
run_error_messages: Vec<String>, 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 scroll: Scroll,
pub mode: Mode, pub mode: Mode,
pub cache: FileCache, pub cache: FileCache,
@@ -351,6 +360,7 @@ impl App {
quit_confirm: None, quit_confirm: None,
blocks: Vec::new(), blocks: Vec::new(),
run_error_messages: Vec::new(), run_error_messages: Vec::new(),
internal_workers: Vec::new(),
scroll: Scroll::default(), scroll: Scroll::default(),
mode: Mode::Normal, mode: Mode::Normal,
cache: FileCache::new(), cache: FileCache::new(),
@@ -1296,11 +1306,18 @@ impl App {
greeting, greeting,
status, status,
in_flight, in_flight,
internal_workers,
} => { } => {
self.rewind_refresh_fence = false; self.rewind_refresh_fence = false;
self.restore_snapshot(&entries, greeting, in_flight); self.restore_snapshot(&entries, greeting, in_flight);
self.replace_internal_worker_snapshots(internal_workers);
self.set_worker_status(status); self.set_worker_status(status);
} }
Event::InternalWorker {
worker,
revision,
event,
} => self.apply_internal_worker_event(worker, revision, *event),
Event::Status { status } => { Event::Status { status } => {
self.rewind_refresh_fence = false; self.rewind_refresh_fence = false;
self.set_worker_status(status); self.set_worker_status(status);
@@ -1980,6 +1997,60 @@ impl App {
/// LogEntry variant into the same blocks live events would have /// LogEntry variant into the same blocks live events would have
/// produced. Followed by `Event::Entry` updates for anything /// produced. Followed by `Event::Entry` updates for anything
/// committed after the snapshot. /// 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( fn restore_snapshot(
&mut self, &mut self,
entries: &[serde_json::Value], entries: &[serde_json::Value],
@@ -3276,6 +3347,7 @@ mod completion_flow_tests {
entries: vec![session_start_value], entries: vec![session_start_value],
status: WorkerStatus::Running, status: WorkerStatus::Running,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}); });
assert!(matches!(app.worker_status, WorkerStatus::Running)); assert!(matches!(app.worker_status, WorkerStatus::Running));
@@ -3321,6 +3393,7 @@ mod completion_flow_tests {
entries: vec![serde_json::to_value(run_errored).unwrap()], entries: vec![serde_json::to_value(run_errored).unwrap()],
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}); });
let errors = app let errors = app
@@ -3398,6 +3471,7 @@ mod completion_flow_tests {
}, },
], ],
}, },
internal_workers: Vec::new(),
}); });
app.handle_worker_event(Event::TextDelta { text: "lo".into() }); app.handle_worker_event(Event::TextDelta { text: "lo".into() });
@@ -3434,6 +3508,83 @@ mod completion_flow_tests {
assert!(app.blocks.is_empty()); 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] #[test]
fn live_system_item_notification_appends_notify_block() { fn live_system_item_notification_appends_notify_block() {
let mut app = App::new("test".into()); let mut app = App::new("test".into());
@@ -3552,6 +3703,7 @@ mod completion_flow_tests {
greeting, greeting,
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}); });
assert_eq!(app.context_window, 123_000); assert_eq!(app.context_window, 123_000);
@@ -3749,6 +3901,7 @@ mod completion_flow_tests {
entries: assistant_item_entries, entries: assistant_item_entries,
status: WorkerStatus::Running, status: WorkerStatus::Running,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}); });
let tasks = app.task_store.tasks(); let tasks = app.task_store.tasks();
+2
View File
@@ -2019,6 +2019,7 @@ mod tests {
entries: vec![], entries: vec![],
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}); });
app.handle_worker_event(Event::RewindApplied { app.handle_worker_event(Event::RewindApplied {
entries: vec![], entries: vec![],
@@ -2045,6 +2046,7 @@ mod tests {
entries: vec![], entries: vec![],
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}); });
type_keys(&mut app, "draft"); type_keys(&mut app, "draft");
+2
View File
@@ -859,6 +859,7 @@ async fn ticket_queue_notification_sends_notify_when_socket_available() {
}, },
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}) })
.await .await
.unwrap(); .unwrap();
@@ -900,6 +901,7 @@ async fn send_notify_only_can_deliver_weak_notification_without_auto_run() {
}, },
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}) })
.await .await
.unwrap(); .unwrap();
+22
View File
@@ -381,6 +381,28 @@ pub fn compute_history(app: &App, width: u16) -> HistoryLayout {
i += 1; 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 // Step 2: pre-wrap every logical line to char-based terminal rows so
// scroll math is exact. Track the logical → wrapped mapping so // scroll math is exact. Track the logical → wrapped mapping so
// turn-start indices get translated into wrapped-row coordinates. // turn-start indices get translated into wrapped-row coordinates.
+1
View File
@@ -914,6 +914,7 @@ mod tests {
greeting: test_greeting(), greeting: test_greeting(),
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}, },
]; ];
+2
View File
@@ -1435,6 +1435,7 @@ impl Runtime {
}, },
status: protocol::WorkerStatus::Idle, status: protocol::WorkerStatus::Idle,
in_flight: protocol::InFlightSnapshot { blocks: Vec::new() }, in_flight: protocol::InFlightSnapshot { blocks: Vec::new() },
internal_workers: Vec::new(),
}) })
} }
@@ -3770,6 +3771,7 @@ mod tests {
}, },
status: protocol::WorkerStatus::Running, status: protocol::WorkerStatus::Running,
in_flight: protocol::InFlightSnapshot { blocks: Vec::new() }, in_flight: protocol::InFlightSnapshot { blocks: Vec::new() },
internal_workers: Vec::new(),
}, },
); );
+6 -1
View File
@@ -48,6 +48,7 @@ pub struct WorkerHandle {
/// it on every new connection (Event::Snapshot) and forwards /// it on every new connection (Event::Snapshot) and forwards
/// subsequent commits (Event::Entry) on the receiver. /// subsequent commits (Event::Entry) on the receiver.
pub sink: SegmentLogSink, pub sink: SegmentLogSink,
spawned_registry: Arc<SpawnedWorkerRegistry>,
} }
impl WorkerHandle { impl WorkerHandle {
@@ -84,6 +85,7 @@ impl WorkerHandle {
greeting: self.shared_state.greeting.clone(), greeting: self.shared_state.greeting.clone(),
status: self.shared_state.get_status(), status: self.shared_state.get_status(),
in_flight, in_flight,
internal_workers: self.spawned_registry.internal_worker_snapshots(),
}; };
(event, entry_rx) (event, entry_rx)
} }
@@ -413,6 +415,7 @@ impl WorkerController {
wire_event_bridges_on_engine(&mut worker, &event_tx, &alerter, &in_flight); wire_event_bridges_on_engine(&mut worker, &event_tx, &alerter, &in_flight);
// === 3. Tool registration (builtin / memory / spawn-orchestration) === // === 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( let fs_for_view = register_worker_tools(
&mut worker, &mut worker,
bash_output_dir, bash_output_dir,
@@ -460,6 +463,7 @@ impl WorkerController {
alerter: alerter.clone(), alerter: alerter.clone(),
in_flight: in_flight.clone(), in_flight: in_flight.clone(),
sink: worker.sink(), sink: worker.sink(),
spawned_registry: spawned_registry.clone(),
}; };
let socket_server = match transport { let socket_server = match transport {
@@ -502,7 +506,7 @@ impl WorkerController {
/// per-item history commit callback so every assistant / tool item /// per-item history commit callback so every assistant / tool item
/// landing in `worker.history` becomes a singular `LogEntry::AssistantItem` /// landing in `worker.history` becomes a singular `LogEntry::AssistantItem`
/// / `ToolResult` commit through the sync writer. /// / `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>, worker: &mut Worker<C, St>,
event_tx: &broadcast::Sender<Event>, event_tx: &broadcast::Sender<Event>,
alerter: &Alerter, alerter: &Alerter,
@@ -1888,6 +1892,7 @@ mod tests {
}, },
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}) })
.await .await
.ok()?; .ok()?;
+6
View File
@@ -1494,6 +1494,7 @@ mod tests {
}, },
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}) })
.await .await
.unwrap(); .unwrap();
@@ -1526,6 +1527,7 @@ mod tests {
}, },
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}) })
.await .await
.unwrap(); .unwrap();
@@ -1614,6 +1616,7 @@ mod tests {
}, },
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}) })
.await .await
.unwrap(); .unwrap();
@@ -1637,6 +1640,7 @@ mod tests {
}, },
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}) })
.await .await
.unwrap(); .unwrap();
@@ -1738,6 +1742,7 @@ mod tests {
}, },
status: WorkerStatus::Paused, status: WorkerStatus::Paused,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}) })
.await .await
.unwrap(); .unwrap();
@@ -1787,6 +1792,7 @@ mod tests {
}, },
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
}) })
.await; .await;
}); });
+203 -6
View File
@@ -12,10 +12,18 @@ use std::sync::{Arc, Mutex};
use llm_engine::timeline::event::UsageEvent; use llm_engine::timeline::event::UsageEvent;
use llm_engine::{Engine, llm_client::LlmClient}; use llm_engine::{Engine, llm_client::LlmClient};
use manifest::{Scope, WorkerManifest}; use manifest::{Scope, WorkerManifest};
use protocol::{Event, InFlightSnapshot, WorkerStatus};
use session_store::{LogEntry, SegmentId, SessionId, Store, StoreError, TraceEntry}; use session_store::{LogEntry, SegmentId, SessionId, Store, StoreError, TraceEntry};
use tokio::sync::broadcast;
use uuid::Uuid; use uuid::Uuid;
use crate::controller::wire_event_bridges_on_engine;
use crate::feature::FeatureRegistryBuilder; 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::{ use crate::worker::{
Worker, WorkerError, WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, 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)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum InternalWorkerSessionStatus { pub(crate) enum InternalWorkerSessionStatus {
Idle, Idle,
@@ -246,8 +268,18 @@ enum InternalWorkerSessionCommand {
/// Parent-owned handle for a long-lived Internal Worker session. /// Parent-owned handle for a long-lived Internal Worker session.
/// ///
/// The handle exposes only typed turn, history, status, and stop operations. The underlying Worker, /// The handle exposes typed turn, history, status, presentation snapshot, event subscription, and
/// Engine, ephemeral Store, and cancellation sender remain inside the actor task. /// 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)] #[derive(Clone)]
pub(crate) struct InternalWorkerSessionHandle { pub(crate) struct InternalWorkerSessionHandle {
command_tx: tokio::sync::mpsc::Sender<InternalWorkerSessionCommand>, command_tx: tokio::sync::mpsc::Sender<InternalWorkerSessionCommand>,
@@ -256,6 +288,12 @@ pub(crate) struct InternalWorkerSessionHandle {
session_id: SessionId, session_id: SessionId,
segment_id: SegmentId, segment_id: SegmentId,
state_changed: Arc<tokio::sync::Notify>, 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 { impl InternalWorkerSessionHandle {
@@ -267,6 +305,54 @@ impl InternalWorkerSessionHandle {
InternalWorkerSessionStatus::decode(self.status.load(std::sync::atomic::Ordering::Acquire)) 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> { pub(crate) fn entries(&self) -> Vec<LogEntry> {
self.store self.store
.read_all(self.session_id, self.segment_id) .read_all(self.session_id, self.segment_id)
@@ -305,8 +391,17 @@ impl InternalWorkerSessionHandle {
std::sync::atomic::Ordering::Release, std::sync::atomic::Ordering::Release,
); );
self.state_changed.notify_waiters(); 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); return Err(InternalWorkerSessionError::Unavailable);
} }
let _ = self.event_tx.send(Event::Status {
status: WorkerStatus::Running,
});
Ok(()) Ok(())
} }
@@ -423,11 +518,48 @@ pub(crate) async fn spawn_internal_worker_session(
spawn_prepared_internal_worker_session(worker, store, input, None).await 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( pub(crate) async fn prepare_internal_worker_session(
mut worker: Worker<Box<dyn LlmClient>, EphemeralSessionStore>, mut worker: Worker<Box<dyn LlmClient>, EphemeralSessionStore>,
store: EphemeralSessionStore, store: EphemeralSessionStore,
visibility: InternalWorkerVisibility,
child_registry: Option<Arc<SpawnedWorkerRegistry>>,
on_turn_end: Option<Arc<dyn Fn(InternalWorkerSessionStatus) + Send + Sync>>, on_turn_end: Option<Arc<dyn Fn(InternalWorkerSessionStatus) + Send + Sync>>,
) -> Result<InternalWorkerSessionHandle, InternalWorkerSessionError> { ) -> 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 session_id = worker.session_id();
let segment_id = worker.segment_id(); let segment_id = worker.segment_id();
let (command_tx, mut command_rx) = tokio::sync::mpsc::channel(8); 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(), InternalWorkerSessionStatus::Idle.encode(),
)); ));
let state_changed = Arc::new(tokio::sync::Notify::new()); let state_changed = Arc::new(tokio::sync::Notify::new());
let last_error = Arc::new(Mutex::new(None));
let handle = InternalWorkerSessionHandle { let handle = InternalWorkerSessionHandle {
command_tx, command_tx,
status: status.clone(), status: status.clone(),
@@ -442,6 +575,12 @@ pub(crate) async fn prepare_internal_worker_session(
session_id, session_id,
segment_id, segment_id,
state_changed: state_changed.clone(), 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 { tokio::spawn(async move {
@@ -453,11 +592,25 @@ pub(crate) async fn prepare_internal_worker_session(
loop { loop {
tokio::select! { tokio::select! {
result = &mut run => { result = &mut run => {
let turn_status = match result { let (turn_status, error) = match result {
Ok(_) => InternalWorkerSessionStatus::Idle, Ok(_) => (InternalWorkerSessionStatus::Idle, None),
Err(_) => InternalWorkerSessionStatus::Failed, Err(error) => (
InternalWorkerSessionStatus::Failed,
Some(error.to_string()),
),
}; };
status.store(turn_status.encode(), std::sync::atomic::Ordering::Release); 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 { if let Some(callback) = &on_turn_end {
callback(turn_status); callback(turn_status);
} }
@@ -470,6 +623,8 @@ pub(crate) async fn prepare_internal_worker_session(
let _ = cancel_sender.send(()).await; let _ = cancel_sender.send(()).await;
let _ = (&mut run).await; let _ = (&mut run).await;
status.store(InternalWorkerSessionStatus::Stopped.encode(), std::sync::atomic::Ordering::Release); 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(); state_changed.notify_waiters();
let _ = done.send(()); let _ = done.send(());
return; return;
@@ -491,6 +646,10 @@ pub(crate) async fn prepare_internal_worker_session(
InternalWorkerSessionStatus::Stopped.encode(), InternalWorkerSessionStatus::Stopped.encode(),
std::sync::atomic::Ordering::Release, std::sync::atomic::Ordering::Release,
); );
let _ = event_tx.send(Event::Status {
status: WorkerStatus::Paused,
});
let _ = event_tx.send(Event::Shutdown);
state_changed.notify_waiters(); state_changed.notify_waiters();
let _ = done.send(()); let _ = done.send(());
return; return;
@@ -509,7 +668,14 @@ pub(crate) async fn spawn_prepared_internal_worker_session(
input: String, input: String,
on_turn_end: Option<Arc<dyn Fn(InternalWorkerSessionStatus) + Send + Sync>>, on_turn_end: Option<Arc<dyn Fn(InternalWorkerSessionStatus) + Send + Sync>>,
) -> Result<InternalWorkerSessionHandle, InternalWorkerSessionError> { ) -> 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?; handle.send(input).await?;
Ok(handle) 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)] #[cfg(test)]
mod tests { mod tests {
use std::pin::Pin; use std::pin::Pin;
+1
View File
@@ -284,6 +284,7 @@ mod tests {
}, },
status: WorkerStatus::Idle, status: WorkerStatus::Idle,
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(),
} }
} }
+240 -4
View File
@@ -10,17 +10,19 @@
use std::collections::HashSet; use std::collections::HashSet;
use std::io; use std::io;
use std::sync::{ use std::sync::{
Arc, Arc, Mutex,
atomic::{AtomicBool, Ordering}, atomic::{AtomicBool, AtomicU64, Ordering},
}; };
use manifest::{Permission, ScopeRule, SharedScope}; use manifest::{Permission, ScopeRule, SharedScope};
use protocol::{Event, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot};
use session_store::{ use session_store::{
WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerStoreError, WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerStoreError,
}; };
use tokio::sync::broadcast;
use tracing::warn; use tracing::warn;
use crate::internal_worker::InternalWorkerSessionHandle; use crate::internal_worker::{InternalWorkerSessionHandle, InternalWorkerVisibility};
use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord}; use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
use crate::runtime::worker_allocation; use crate::runtime::worker_allocation;
@@ -30,6 +32,8 @@ pub(crate) struct InternalSpawnedWorkerRecord {
pub scope_delegated: Vec<ScopeRule>, pub scope_delegated: Vec<ScopeRule>,
pub session: InternalWorkerSessionHandle, pub session: InternalWorkerSessionHandle,
scope_reclaimed: Arc<AtomicBool>, scope_reclaimed: Arc<AtomicBool>,
protocol_revision: Arc<AtomicU64>,
forwarding_started: Arc<AtomicBool>,
} }
impl InternalSpawnedWorkerRecord { impl InternalSpawnedWorkerRecord {
@@ -43,6 +47,8 @@ impl InternalSpawnedWorkerRecord {
scope_delegated, scope_delegated,
session, session,
scope_reclaimed: Arc::new(AtomicBool::new(false)), 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) { fn restore_scope_reclaim(&self) {
self.scope_reclaimed.store(false, Ordering::Release); 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 { pub(crate) struct InternalSpawnReservation {
@@ -73,7 +92,8 @@ impl InternalSpawnReservation {
.internal_records .internal_records
.lock() .lock()
.map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))? .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; self.committed = true;
Ok(()) Ok(())
} }
@@ -93,6 +113,7 @@ pub struct SpawnedWorkerRegistry {
internal_records: std::sync::Mutex<Vec<InternalSpawnedWorkerRecord>>, internal_records: std::sync::Mutex<Vec<InternalSpawnedWorkerRecord>>,
internal_names: std::sync::Mutex<HashSet<String>>, internal_names: std::sync::Mutex<HashSet<String>>,
parent_scope: Option<SharedScope>, parent_scope: Option<SharedScope>,
parent_protocol: Mutex<Option<(broadcast::Sender<Event>, String)>>,
} }
pub struct SpawnedWorkerRegistryLoad { pub struct SpawnedWorkerRegistryLoad {
@@ -108,6 +129,7 @@ impl SpawnedWorkerRegistry {
internal_records: std::sync::Mutex::new(Vec::new()), internal_records: std::sync::Mutex::new(Vec::new()),
internal_names: std::sync::Mutex::new(HashSet::new()), internal_names: std::sync::Mutex::new(HashSet::new()),
parent_scope: None, parent_scope: None,
parent_protocol: Mutex::new(None),
}) })
} }
@@ -116,6 +138,7 @@ impl SpawnedWorkerRegistry {
internal_records: std::sync::Mutex::new(Vec::new()), internal_records: std::sync::Mutex::new(Vec::new()),
internal_names: std::sync::Mutex::new(HashSet::new()), internal_names: std::sync::Mutex::new(HashSet::new()),
parent_scope: Some(parent_scope), 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_records: std::sync::Mutex::new(Vec::new()),
internal_names: std::sync::Mutex::new(HashSet::new()), internal_names: std::sync::Mutex::new(HashSet::new()),
parent_scope, parent_scope,
parent_protocol: Mutex::new(None),
}), }),
reclaimed_unreachable: !persisted_children.is_empty(), 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> { pub(crate) fn get_internal(&self, worker_name: &str) -> Option<InternalSpawnedWorkerRecord> {
self.internal_records self.internal_records
.lock() .lock()
@@ -387,3 +501,125 @@ fn record_from_worker_state(child: &WorkerSpawnedChild) -> io::Result<SpawnedWor
fn store_error_to_io(error: WorkerStoreError) -> io::Error { fn store_error_to_io(error: WorkerStoreError) -> io::Error {
io::Error::other(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());
}
}
+7 -2
View File
@@ -24,7 +24,8 @@ use tokio::sync::mpsc;
use crate::PromptCatalogSource; use crate::PromptCatalogSource;
use crate::controller::register_worker_tools; use crate::controller::register_worker_tools;
use crate::internal_worker::{ use crate::internal_worker::{
EphemeralSessionStore, InternalWorkerSessionStatus, prepare_internal_worker_session, EphemeralSessionStore, InternalWorkerSessionStatus, InternalWorkerVisibility,
prepare_internal_worker_session,
}; };
use crate::prompt::catalog::PromptCatalog; use crate::prompt::catalog::PromptCatalog;
use crate::spawn::registry::SpawnedWorkerRegistry; use crate::spawn::registry::SpawnedWorkerRegistry;
@@ -481,7 +482,7 @@ impl Tool for SubWorkerSpawnTool {
.join(&input.name) .join(&input.name)
.join("bash-output"), .join("bash-output"),
self.runtime_base.clone(), self.runtime_base.clone(),
child_registry, child_registry.clone(),
None, None,
) )
.await .await
@@ -511,6 +512,8 @@ impl Tool for SubWorkerSpawnTool {
let session_result = prepare_internal_worker_session( let session_result = prepare_internal_worker_session(
child, child,
store, store,
InternalWorkerVisibility::ParentClient,
Some(child_registry.clone()),
Some(Arc::new(move |status| { Some(Arc::new(move |status| {
if status == InternalWorkerSessionStatus::Failed { if status == InternalWorkerSessionStatus::Failed {
if let Some(registry) = registry.upgrade() { 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 { if let Some((ticket_id, capability_token)) = &reviewer_capability {
let workspace_id = self.workspace_context.workspace_id().ok_or_else(|| { let workspace_id = self.workspace_context.workspace_id().ok_or_else(|| {
+12 -1
View File
@@ -53,6 +53,12 @@ export type InFlightBlock = { "kind": "text", text: string, finished?: boolean,
export type InFlightSnapshot = { blocks?: Array<InFlightBlock>, }; export type InFlightSnapshot = { blocks?: Array<InFlightBlock>, };
export type InternalWorkerKind = "sub_worker";
export type InternalWorkerRef = { session_id: string, name: string, parent_session_id?: string | null, kind: InternalWorkerKind, };
export type InternalWorkerSnapshot = { worker: InternalWorkerRef, revision: number, entries: Array<unknown>, status: WorkerStatus, error?: string | null, in_flight?: InFlightSnapshot, internal_workers?: Array<InternalWorkerSnapshot>, };
export type Greeting = { worker_name: string, cwd: string, provider: string, model: string, scope_summary: string, tools: Array<string>, export type Greeting = { worker_name: string, cwd: string, provider: string, model: string, scope_summary: string, tools: Array<string>,
/** /**
* Model context window in tokens. Always filled by the Worker greeting. * Model context window in tokens. Always filled by the Worker greeting.
@@ -167,4 +173,9 @@ output?: string | null, is_error: boolean, } } | { "event": "usage", "data": { i
* Unfinished model output that has already streamed in the current * Unfinished model output that has already streamed in the current
* run but is not yet represented by committed snapshot entries. * run but is not yet represented by committed snapshot entries.
*/ */
in_flight?: InFlightSnapshot, } } | { "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" }; 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" };
@@ -1266,6 +1266,82 @@ Deno.test("projectConsole mirrors live TaskCreate and TaskUpdate calls", () => {
}]); }]);
}); });
Deno.test("Internal Worker output stays separate and revision-fenced", () => {
const worker = {
session_id: "child-session",
name: "research",
parent_session_id: "parent-session",
kind: "sub_worker" as const,
};
const projector = createConsoleProjector();
let projection = projector.append([{
eventId: "1",
event: {
event: "internal_worker",
data: {
worker,
revision: 2,
event: { event: "text_done", data: { text: "child output" } },
},
},
}]);
assertEquals(projection.lines, []);
assertEquals(projection.internalWorkers.length, 1);
assertEquals(projection.internalWorkers[0].console.lines[0].body, "child output");
projection = projector.append([{
eventId: "2",
event: {
event: "internal_worker",
data: {
worker,
revision: 1,
event: { event: "text_done", data: { text: "stale" } },
},
},
}]);
assertEquals(projection.internalWorkers[0].console.lines.length, 1);
});
Deno.test("parent snapshot authoritatively replaces Internal Worker projections", () => {
const event = snapshotEvent("/repo");
if (event.event !== "snapshot") throw new Error("snapshot fixture expected");
event.data.internal_workers = [{
worker: {
session_id: "replacement",
name: "replacement",
parent_session_id: "parent-session",
kind: "sub_worker",
},
revision: 4,
entries: [],
status: "idle",
in_flight: { blocks: [] },
internal_workers: [],
}];
const projector = createConsoleProjector();
projector.append([{
eventId: "old",
event: {
event: "internal_worker",
data: {
worker: {
session_id: "old",
name: "old",
parent_session_id: "parent-session",
kind: "sub_worker",
},
revision: 1,
event: { event: "status", data: { status: "running" } },
},
},
}]);
const projection = projector.append([{ eventId: "snapshot", event }]);
assertEquals(projection.internalWorkers.map((worker) => worker.worker.session_id), [
"replacement",
]);
});
Deno.test("snapshot restores TaskStore state from system history", () => { Deno.test("snapshot restores TaskStore state from system history", () => {
const taskSnapshot = const taskSnapshot =
`[Session TaskStore snapshot]\n\n\`\`\`json\n{\n "tasks": [{"taskid": 3, "status": "pending", "subject": "Restored", "description": "From compaction"}]\n}\n\`\`\``; `[Session TaskStore snapshot]\n\n\`\`\`json\n{\n "tasks": [{"taskid": 3, "status": "pending", "subject": "Restored", "description": "From compaction"}]\n}\n\`\`\``;
@@ -3,6 +3,8 @@ import type {
Event as ProtocolEvent, Event as ProtocolEvent,
InFlightBlock, InFlightBlock,
InFlightToolCallState, InFlightToolCallState,
InternalWorkerRef,
InternalWorkerSnapshot,
Segment, Segment,
} from "$lib/generated/protocol"; } from "$lib/generated/protocol";
import { workspaceRoute } from "$lib/workspace/api/http"; import { workspaceRoute } from "$lib/workspace/api/http";
@@ -63,6 +65,26 @@ export type ConsoleLine = {
toolCall?: ToolCallView; toolCall?: ToolCallView;
}; };
export type InternalWorkerProjection = {
worker: InternalWorkerRef;
revision: number;
console: ConsoleProjection;
};
export type FlattenedInternalWorkerProjection = InternalWorkerProjection & {
depth: number;
};
export function flattenInternalWorkers(
workers: InternalWorkerProjection[],
depth = 0,
): FlattenedInternalWorkerProjection[] {
return workers.flatMap((worker) => [
{ ...worker, depth },
...flattenInternalWorkers(worker.console.internalWorkers, depth + 1),
]);
}
export type ConsoleProjection = { export type ConsoleProjection = {
lines: ConsoleLine[]; lines: ConsoleLine[];
tasks: ConsoleTask[]; tasks: ConsoleTask[];
@@ -71,6 +93,7 @@ export type ConsoleProjection = {
usage: string | null; usage: string | null;
cwd: string | null; cwd: string | null;
lastEventId: string | null; lastEventId: string | null;
internalWorkers: InternalWorkerProjection[];
}; };
export type ConsoleTimelineLineSelection = { export type ConsoleTimelineLineSelection = {
@@ -155,6 +178,7 @@ export function emptyConsoleProjection(): ConsoleProjection {
usage: null, usage: null,
cwd: null, cwd: null,
lastEventId: null, lastEventId: null,
internalWorkers: [],
}; };
} }
@@ -193,9 +217,50 @@ function projectVisibleConsole(
return { return {
...projection, ...projection,
lines: aggregateReadToolLines(projection.lines), lines: aggregateReadToolLines(projection.lines),
internalWorkers: projection.internalWorkers.map((worker) => ({
...worker,
console: projectVisibleConsole(worker.console),
})),
}; };
} }
function projectInternalWorkerSnapshot(
snapshot: InternalWorkerSnapshot,
eventId: string,
cwd: string | null,
): InternalWorkerProjection {
const console = snapshotProjectionFromEntries(
`${eventId}:internal:${snapshot.worker.session_id}:snapshot`,
snapshot.entries,
cwd,
);
console.status = snapshot.status;
for (const block of snapshot.in_flight?.blocks ?? []) {
console.lines.push(
inFlightLine(
`${eventId}:internal:${snapshot.worker.session_id}:in-flight`,
block,
cwd,
),
);
}
if (snapshot.error) {
console.lines.push({
id: `${eventId}:internal:${snapshot.worker.session_id}:error`,
kind: "error",
title: "Error",
body: snapshot.error,
eventId,
source: "event",
error: true,
});
}
console.internalWorkers = (snapshot.internal_workers ?? []).map((child) =>
projectInternalWorkerSnapshot(child, eventId, cwd)
);
return { worker: snapshot.worker, revision: snapshot.revision, console };
}
export function applyProtocolEvent( export function applyProtocolEvent(
projection: ConsoleProjection, projection: ConsoleProjection,
envelope: { eventId: string; event: ProtocolEvent }, envelope: { eventId: string; event: ProtocolEvent },
@@ -208,6 +273,7 @@ export function applyProtocolEvent(
usage: projection.usage, usage: projection.usage,
cwd: projection.cwd, cwd: projection.cwd,
lastEventId: envelope.eventId, lastEventId: envelope.eventId,
internalWorkers: [...projection.internalWorkers],
}; };
const event = envelope.event; const event = envelope.event;
@@ -322,6 +388,34 @@ export function applyProtocolEvent(
for (const block of event.data.in_flight?.blocks ?? []) { for (const block of event.data.in_flight?.blocks ?? []) {
next.lines.push(inFlightLine(envelope.eventId, block, next.cwd)); next.lines.push(inFlightLine(envelope.eventId, block, next.cwd));
} }
next.internalWorkers = (event.data.internal_workers ?? []).map((worker) =>
projectInternalWorkerSnapshot(worker, envelope.eventId, next.cwd)
);
break;
}
case "internal_worker": {
const existingIndex = next.internalWorkers.findIndex((worker) =>
worker.worker.session_id === event.data.worker.session_id
);
const existing = existingIndex >= 0
? next.internalWorkers[existingIndex]
: {
worker: event.data.worker,
revision: 0,
console: emptyConsoleProjection(),
};
if (event.data.revision <= existing.revision) break;
const updated: InternalWorkerProjection = {
worker: event.data.worker,
revision: event.data.revision,
console: applyProtocolEvent(existing.console, {
eventId:
`${envelope.eventId}:internal:${event.data.worker.session_id}:${event.data.revision}`,
event: event.data.event,
}),
};
if (existingIndex >= 0) next.internalWorkers[existingIndex] = updated;
else next.internalWorkers.push(updated);
break; break;
} }
case "status": case "status":
@@ -1184,6 +1278,7 @@ function snapshotProjectionFromEntries(
usage: null, usage: null,
cwd, cwd,
lastEventId: eventId, lastEventId: eventId,
internalWorkers: [],
}; };
entries.forEach((entry, index) => entries.forEach((entry, index) =>
applyLogEntry(projection, `${eventId}-snapshot-${index}`, entry) applyLogEntry(projection, `${eventId}-snapshot-${index}`, entry)
@@ -18,6 +18,7 @@
import { fitTextarea } from "$lib/workspace/console/textarea-fit"; import { fitTextarea } from "$lib/workspace/console/textarea-fit";
import { import {
createConsoleProjector, createConsoleProjector,
flattenInternalWorkers,
isConsoleProjectionEvent, isConsoleProjectionEvent,
selectConsoleTimelineLines, selectConsoleTimelineLines,
type ConsoleEventInput, type ConsoleEventInput,
@@ -152,6 +153,9 @@
const lines = $derived(consoleProjection.lines); const lines = $derived(consoleProjection.lines);
const tasks = $derived(consoleProjection.tasks); const tasks = $derived(consoleProjection.tasks);
const internalWorkers = $derived(
flattenInternalWorkers(consoleProjection.internalWorkers),
);
const timelineLayout = $derived( const timelineLayout = $derived(
buildTimelineLayout(lines, eventObservedAtVersion, consoleScroll), buildTimelineLayout(lines, eventObservedAtVersion, consoleScroll),
); );
@@ -1245,6 +1249,28 @@
</ol> </ol>
{/if} {/if}
</article> </article>
{#each internalWorkers as internal (internal.worker.session_id)}
<section
class="card internal-worker-pane"
style={`--internal-worker-depth: ${internal.depth}`}
aria-label={`SubWorker ${internal.worker.name}`}
>
<header class="internal-worker-header">
<strong>{internal.worker.name}</strong>
<span>{internal.console.status ?? "unknown"}</span>
</header>
{#if internal.console.lines.length === 0}
<p>No output yet.</p>
{:else}
<ol class="console-log">
{#each internal.console.lines as item (item.id)}
<ConsoleLineItem {item} />
{/each}
</ol>
{/if}
</section>
{/each}
</div> </div>
<ConsoleTimeline <ConsoleTimeline
@@ -1771,6 +1797,23 @@
margin-right: auto; margin-right: auto;
} }
.internal-worker-pane {
margin: 0.75rem 0 0 calc((var(--internal-worker-depth) + 1) * 1rem);
border-left: 3px solid var(--color-border-strong, currentColor);
}
.internal-worker-header {
display: flex;
justify-content: space-between;
gap: 1rem;
margin-bottom: 0.5rem;
font-family: var(--font-mono);
}
.internal-worker-header span {
color: var(--color-text-muted);
}
@media (max-width: 960px) { @media (max-width: 960px) {
.console-history.with-task-pane { .console-history.with-task-pane {
grid-template-columns: minmax(0, 1fr); grid-template-columns: minmax(0, 1fr);