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/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(|| {