diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index 0f45861d9..eeff41061 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -4092,6 +4092,54 @@ mod completion_flow_tests { assert_eq!(child_app.session_context_tokens, 12_000); } + #[test] + fn live_internal_worker_snapshot_installs_child_metadata_before_context_updates() { + let mut app = App::new("parent".into()); + let worker = test_internal_worker_snapshot("child", "research", 1).worker; + let mut greeting = test_greeting(); + greeting.worker_name = "research".into(); + greeting.model = "child-model".into(); + greeting.context_window = 64_000; + greeting.context_usage = Some(protocol::ContextUsage { + tokens: 12_000, + source: protocol::ContextTokenSource::Measured, + }); + + app.handle_worker_event(Event::InternalWorker { + worker: worker.clone(), + revision: 1, + event: Box::new(Event::Snapshot { + session: protocol::SessionSnapshot { + pending_submissions: protocol::PendingSubmissionsSnapshot::default(), + entries: Vec::new(), + }, + greeting, + state: WorkerStatus::Running.into(), + in_flight: Default::default(), + internal_workers: Vec::new(), + }), + }); + app.handle_worker_event(Event::InternalWorker { + worker, + revision: 2, + event: Box::new(Event::ContextUsage { + usage: Some(protocol::ContextUsage { + tokens: 14_000, + source: protocol::ContextTokenSource::Estimated, + }), + }), + }); + + let child = &app.internal_workers[0].app; + assert_eq!(child.greeting.as_ref().unwrap().model, "child-model"); + assert_eq!(child.context_window, 64_000); + assert_eq!(child.session_context_tokens, 14_000); + assert_eq!( + child.session_context_source, + Some(protocol::ContextTokenSource::Estimated) + ); + } + #[test] fn worker_view_cycle_uses_stable_session_identity_and_wraps_to_main() { let mut app = App::new("parent".into()); diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index 76c8b6b91..b431f2bf9 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -1220,6 +1220,14 @@ impl session_store::WorkerMetadataStore for EphemeralSessionStore { #[cfg(test)] pub(crate) fn test_internal_worker_session( visibility: InternalWorkerVisibility, +) -> (InternalWorkerSessionHandle, broadcast::Sender) { + test_internal_worker_session_with_greeting(visibility, None) +} + +#[cfg(test)] +pub(crate) fn test_internal_worker_session_with_greeting( + visibility: InternalWorkerVisibility, + greeting: Option, ) -> (InternalWorkerSessionHandle, broadcast::Sender) { let store = EphemeralSessionStore::default(); let session_id = session_store::new_session_id(); @@ -1252,7 +1260,7 @@ pub(crate) fn test_internal_worker_session( visibility, last_error: Arc::new(Mutex::new(None)), last_outcome: Arc::new(Mutex::new(None)), - greeting: None, + greeting: greeting.map(|greeting| Arc::new(RwLock::new(greeting))), child_registry: None, sink, fail_stop: Arc::new(std::sync::atomic::AtomicBool::new(false)), diff --git a/crates/worker/src/ipc/protocol_session.rs b/crates/worker/src/ipc/protocol_session.rs index cf83db223..cdfcdf878 100644 --- a/crates/worker/src/ipc/protocol_session.rs +++ b/crates/worker/src/ipc/protocol_session.rs @@ -16,9 +16,23 @@ pub struct WorkerProtocolSessionStreams { pub events: broadcast::Receiver, } +fn subscribe_before_snapshot( + subscribe: impl FnOnce() -> Subscription, + snapshot: impl FnOnce() -> Snapshot, +) -> (Snapshot, Subscription) { + // Subscription must exist before any snapshot field is read. A concurrent + // state update can then appear in both lanes, but it can never be absent + // from both the snapshot and the queued live stream. + let subscription = subscribe(); + let snapshot = snapshot(); + (snapshot, subscription) +} + pub fn subscribe_worker_protocol_session(handle: &WorkerHandle) -> WorkerProtocolSessionStreams { - let (snapshot_event, log_entries) = handle.snapshot_event_with_entry_subscription(); - let (alert_snapshot, events) = handle.alerter.subscribe_with_snapshot(); + let ((snapshot_event, log_entries), (alert_snapshot, events)) = subscribe_before_snapshot( + || handle.alerter.subscribe_with_snapshot(), + || handle.snapshot_event_with_entry_subscription(), + ); WorkerProtocolSessionStreams { snapshot_event, alert_snapshot, @@ -129,6 +143,36 @@ mod tests { } } + #[test] + fn context_publication_between_subscription_and_snapshot_cannot_be_missed() { + let (events, _) = broadcast::channel(4); + let context_tokens = std::cell::Cell::new(20_u64); + + let (snapshot_tokens, mut receiver) = subscribe_before_snapshot( + || events.subscribe(), + || { + context_tokens.set(25); + events + .send(Event::ContextUsage { + usage: Some(protocol::ContextUsage { + tokens: 25, + source: protocol::ContextTokenSource::Estimated, + }), + }) + .unwrap(); + context_tokens.get() + }, + ); + + assert_eq!(snapshot_tokens, 25); + assert!(matches!( + receiver.try_recv(), + Ok(Event::ContextUsage { + usage: Some(protocol::ContextUsage { tokens: 25, .. }) + }) + )); + } + #[test] fn persisted_run_transitions_keep_identity_across_live_and_snapshot_projection() { let temp = tempfile::tempdir().unwrap(); diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index bfd280e47..5b706f792 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -29,7 +29,9 @@ use tokio::sync::{Notify, broadcast}; use tracing::warn; use workdir::WorkdirScopeLeaseSet; -use crate::internal_worker::{InternalWorkerSessionHandle, InternalWorkerVisibility}; +use crate::internal_worker::{ + InternalWorkerSessionHandle, InternalWorkerSessionSnapshot, InternalWorkerVisibility, +}; use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord}; use crate::runtime::worker_allocation; @@ -648,6 +650,14 @@ impl SpawnedWorkerRegistry { let protocol_emit_lock = record.protocol_emit_lock.clone(); let protocol_terminal = record.protocol_terminal.clone(); let mut child_rx = record.session.subscribe_events(); + publish_initial_internal_snapshot( + &parent_tx, + &worker, + &protocol_revision, + &protocol_emit_lock, + &protocol_terminal, + &record.session, + ); tokio::spawn(async move { loop { match child_rx.recv().await { @@ -710,6 +720,14 @@ impl SpawnedWorkerRegistry { let protocol_emit_lock = record.protocol_emit_lock.clone(); let protocol_terminal = record.protocol_terminal.clone(); let mut child_rx = record.session.subscribe_events(); + publish_initial_internal_snapshot( + &parent_tx, + &worker, + &protocol_revision, + &protocol_emit_lock, + &protocol_terminal, + &record.session, + ); tokio::spawn(async move { loop { match child_rx.recv().await { @@ -985,6 +1003,41 @@ impl SpawnedWorkerRegistry { } } +fn publish_initial_internal_snapshot( + parent_tx: &broadcast::Sender, + worker: &InternalWorkerRef, + protocol_revision: &AtomicU64, + protocol_emit_lock: &Mutex<()>, + protocol_terminal: &AtomicBool, + session: &InternalWorkerSessionHandle, +) { + let Some(event) = internal_worker_live_snapshot(session.protocol_snapshot()) else { + return; + }; + let _emit_guard = protocol_emit_lock + .lock() + .unwrap_or_else(|error| error.into_inner()); + if protocol_terminal.load(Ordering::Acquire) { + return; + } + let revision = protocol_revision.fetch_add(1, Ordering::AcqRel) + 1; + let _ = parent_tx.send(Event::InternalWorker { + worker: worker.clone(), + revision, + event: Box::new(event), + }); +} + +fn internal_worker_live_snapshot(snapshot: InternalWorkerSessionSnapshot) -> Option { + Some(Event::Snapshot { + session: snapshot.session, + greeting: snapshot.greeting?, + state: snapshot.status.into(), + in_flight: snapshot.in_flight, + internal_workers: snapshot.internal_workers, + }) +} + fn internal_worker_snapshot( worker: InternalWorkerRef, revision: u64, @@ -1122,7 +1175,10 @@ mod tests { use session_store::LogEntry; use super::*; - use crate::internal_worker::{InternalWorkerSessionStatus, test_internal_worker_session}; + use crate::internal_worker::{ + InternalWorkerSessionStatus, test_internal_worker_session, + test_internal_worker_session_with_greeting, + }; fn registry() -> Arc { let scope = Scope::from_config(&ScopeConfig { @@ -1142,7 +1198,15 @@ mod tests { name: &str, visibility: InternalWorkerVisibility, ) -> (InternalSpawnedWorkerRecord, broadcast::Sender) { - let (session, sender) = test_internal_worker_session(visibility); + record_with_greeting(name, visibility, None).await + } + + async fn record_with_greeting( + name: &str, + visibility: InternalWorkerVisibility, + greeting: Option, + ) -> (InternalSpawnedWorkerRecord, broadcast::Sender) { + let (session, sender) = test_internal_worker_session_with_greeting(visibility, greeting); let root = std::path::PathBuf::from("/tmp"); let scope = Scope::from_config(&ScopeConfig { allow: vec![ScopeRule { @@ -1194,6 +1258,95 @@ mod tests { ) } + #[tokio::test] + async fn visible_internal_worker_starts_with_own_live_snapshot_metadata() { + let registry = registry(); + let (parent_tx, mut parent_rx) = broadcast::channel(16); + registry.attach_parent_protocol(parent_tx, "parent-session".into()); + let greeting = protocol::Greeting { + worker_name: "research".into(), + cwd: "/child".into(), + provider: "provider".into(), + model: "child-model".into(), + reasoning: Some(protocol::ReasoningConfig::Effort { + effort: "high".into(), + }), + scope_summary: "child scope".into(), + tools: Vec::new(), + context_window: 64_000, + context_tokens: 12_000, + context_usage: Some(protocol::ContextUsage { + tokens: 12_000, + source: protocol::ContextTokenSource::Measured, + }), + }; + let (record, child_tx) = record_with_greeting( + "research", + InternalWorkerVisibility::ParentClient, + Some(greeting), + ) + .await; + registry + .internal_records + .lock() + .unwrap() + .push(record.clone()); + + registry.start_protocol_forwarding(record); + + let initial = tokio::time::timeout(Duration::from_secs(1), parent_rx.recv()) + .await + .unwrap() + .unwrap(); + assert!(matches!( + initial, + Event::InternalWorker { + revision: 1, + event, + .. + } if matches!( + *event, + Event::Snapshot { + greeting: protocol::Greeting { + ref model, + context_window: 64_000, + .. + }, + .. + } if model == "child-model" + ) + )); + + child_tx + .send(Event::ContextUsage { + usage: Some(protocol::ContextUsage { + tokens: 14_000, + source: protocol::ContextTokenSource::Estimated, + }), + }) + .unwrap(); + let update = tokio::time::timeout(Duration::from_secs(1), parent_rx.recv()) + .await + .unwrap() + .unwrap(); + assert!(matches!( + update, + Event::InternalWorker { + revision: 2, + event, + .. + } if matches!( + *event, + Event::ContextUsage { + usage: Some(protocol::ContextUsage { + tokens: 14_000, + source: protocol::ContextTokenSource::Estimated, + }) + } + ) + )); + } + #[tokio::test] async fn visible_internal_output_is_wrapped_after_registry_insertion() { let registry = registry(); diff --git a/web/workspace/src/lib/workspace/console/model.test.ts b/web/workspace/src/lib/workspace/console/model.test.ts index f5b43718b..3f71f8d80 100644 --- a/web/workspace/src/lib/workspace/console/model.test.ts +++ b/web/workspace/src/lib/workspace/console/model.test.ts @@ -3050,6 +3050,54 @@ Deno.test("Internal Worker output stays separate and revision-fenced", () => { assertEquals(resolveConsoleWorkerView(projection, "missing").sessionId, null); }); +Deno.test("live-created Internal Worker receives own metadata without reconnect", () => { + const worker = { + session_id: "live-child", + name: "research", + parent_session_id: "parent-session", + kind: "sub_worker" as const, + }; + const childSnapshot = snapshotEvent("/child") as SnapshotEvent; + childSnapshot.data.greeting.worker_name = "research"; + childSnapshot.data.greeting.model = "child-model"; + childSnapshot.data.greeting.reasoning = { kind: "effort", effort: "high" }; + childSnapshot.data.greeting.context_window = 64_000; + childSnapshot.data.greeting.context_usage = { tokens: 12_000, source: "measured" }; + + const projection = projectConsole([ + { + eventId: "child-snapshot", + event: { + event: "internal_worker", + data: { worker, revision: 1, event: childSnapshot }, + }, + }, + { + eventId: "child-context", + event: { + event: "internal_worker", + data: { + worker, + revision: 2, + event: { + event: "context_usage", + data: { usage: { tokens: 14_000, source: "estimated" } }, + }, + }, + }, + }, + ]); + + assertEquals(projection.workerMetadata, null); + assertEquals(projection.internalWorkers[0].console.workerMetadata, { + model: "child-model", + reasoning: { kind: "effort", effort: "high" }, + contextWindow: 64_000, + contextTokens: 14_000, + contextSource: "estimated", + }); +}); + Deno.test("Internal Worker snapshot output continues without entering the parent view", () => { const worker = { session_id: "child-reconnect",