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

fix(console): keep live context metadata race-free

Browse Source
This commit is contained in:
Hare
2026-10-02 16:25:43 +09:00
parent 3d37f056ad
commit 07cbff9a5f
5 changed files with 307 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/tui/src/app.rs
+48
View File
@@ -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());
crates/worker/src/internal_worker.rs
+9 -1
View File
@@ -1220,6 +1220,14 @@ impl session_store::WorkerMetadataStore for EphemeralSessionStore {
#[cfg(test)]
pub(crate) fn test_internal_worker_session(
visibility: InternalWorkerVisibility,
) -> (InternalWorkerSessionHandle, broadcast::Sender<Event>) {
test_internal_worker_session_with_greeting(visibility, None)
}
#[cfg(test)]
pub(crate) fn test_internal_worker_session_with_greeting(
visibility: InternalWorkerVisibility,
greeting: Option<protocol::Greeting>,
) -> (InternalWorkerSessionHandle, broadcast::Sender<Event>) {
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)),
crates/worker/src/ipc/protocol_session.rs
+46 -2
View File
@@ -16,9 +16,23 @@ pub struct WorkerProtocolSessionStreams {
pub events: broadcast::Receiver<Event>,
}
fn subscribe_before_snapshot<Subscription, 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();
crates/worker/src/spawn/registry.rs
+156 -3
View File
@@ -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<Event>,
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<Event> {
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<SpawnedWorkerRegistry> {
let scope = Scope::from_config(&ScopeConfig {
@@ -1142,7 +1198,15 @@ mod tests {
name: &str,
visibility: InternalWorkerVisibility,
) -> (InternalSpawnedWorkerRecord, broadcast::Sender<Event>) {
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<protocol::Greeting>,
) -> (InternalSpawnedWorkerRecord, broadcast::Sender<Event>) {
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();
web/workspace/src/lib/workspace/console/model.test.ts
+48
View File
@@ -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",