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

fix: fence running snapshots on input commit

Browse Source
This commit is contained in:
Hare
2026-08-27 15:42:20 +09:00
parent 7aa06afc45
commit 183c37446e
3 changed files with 113 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+75 -11
View File
@@ -1169,22 +1169,34 @@ async fn controller_loop<C, St>(
// clear at run start prevents stale partial output left by an older
// interrupted/error turn from being carried into the next snapshot.
worker.clear_in_flight_events();
set_controller_status(
&shared_state,
&runtime_dir,
&event_tx,
WorkerStatus::Running,
)
.await;
let parent_originated = run.is_parent_originated();
let user_input_run = matches!(&run, PendingRun::Run(_) | PendingRun::RunTracked { .. });
if !user_input_run {
set_controller_status(
&shared_state,
&runtime_dir,
&event_tx,
WorkerStatus::Running,
)
.await;
}
let (mut new_status, shutdown) = match run {
PendingRun::Run(input) => {
let (input_commit_tx, input_commit_rx) = oneshot::channel();
drive_turn(
worker.run(input),
worker.run_with_input_extensions_and_commit_hook(
input,
Vec::new(),
move || {
let _ = input_commit_tx.send(());
},
),
&mut method_rx,
&event_tx,
&cancel_tx,
&shared_state,
&runtime_dir,
Some(input_commit_rx),
&notify_buffer,
self_parent_socket.as_ref(),
&spawner_name,
@@ -1194,12 +1206,21 @@ async fn controller_loop<C, St>(
.await
}
PendingRun::RunTracked { input, extension } => {
let (input_commit_tx, input_commit_rx) = oneshot::channel();
drive_turn(
worker.run_with_input_extensions(input, vec![extension]),
worker.run_with_input_extensions_and_commit_hook(
input,
vec![extension],
move || {
let _ = input_commit_tx.send(());
},
),
&mut method_rx,
&event_tx,
&cancel_tx,
&shared_state,
&runtime_dir,
Some(input_commit_rx),
&notify_buffer,
self_parent_socket.as_ref(),
&spawner_name,
@@ -1215,6 +1236,8 @@ async fn controller_loop<C, St>(
&event_tx,
&cancel_tx,
&shared_state,
&runtime_dir,
None,
&notify_buffer,
self_parent_socket.as_ref(),
&spawner_name,
@@ -1230,6 +1253,8 @@ async fn controller_loop<C, St>(
&event_tx,
&cancel_tx,
&shared_state,
&runtime_dir,
None,
&notify_buffer,
self_parent_socket.as_ref(),
&spawner_name,
@@ -1627,6 +1652,8 @@ async fn drive_turn<F>(
event_tx: &broadcast::Sender<Event>,
cancel_tx: &mpsc::Sender<()>,
shared_state: &Arc<WorkerSharedState>,
runtime_dir: &RuntimeDir,
mut input_commit_rx: Option<oneshot::Receiver<()>>,
notify_buffer: &NotifyBuffer,
parent_socket: Option<&PathBuf>,
self_name: &str,
@@ -1642,6 +1669,27 @@ where
loop {
tokio::select! {
// If input commit and provider completion become ready together, expose
// Running only after processing the commit fence. This makes the
// Running snapshot contract deterministic even for immediate clients.
biased;
committed = async {
input_commit_rx
.as_mut()
.expect("input commit receiver guarded by select condition")
.await
}, if input_commit_rx.is_some() => {
input_commit_rx = None;
if committed.is_ok() {
set_controller_status(
shared_state,
runtime_dir,
event_tx,
WorkerStatus::Running,
)
.await;
}
}
result = &mut worker_future => {
return match result {
Ok(r) => {
@@ -1974,7 +2022,7 @@ mod tests {
notify_buffer: NotifyBuffer,
spawned_registry: Arc<SpawnedWorkerRegistry>,
parent_socket_path: PathBuf,
_runtime_dir: Arc<RuntimeDir>,
runtime_dir: Arc<RuntimeDir>,
_temp: TempDir,
}
@@ -2017,7 +2065,7 @@ mod tests {
notify_buffer,
spawned_registry,
parent_socket_path,
_runtime_dir: runtime_dir,
runtime_dir,
_temp: temp,
}
}
@@ -2071,6 +2119,8 @@ mod tests {
&env.event_tx,
&env.cancel_tx,
&env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer,
Some(&env.parent_socket_path),
"child-worker",
@@ -2103,6 +2153,8 @@ mod tests {
&env.event_tx,
&env.cancel_tx,
&env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer,
Some(&env.parent_socket_path),
"child-worker",
@@ -2138,6 +2190,8 @@ mod tests {
&env.event_tx,
&env.cancel_tx,
&env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer,
Some(&env.parent_socket_path),
"child-worker",
@@ -2179,6 +2233,8 @@ mod tests {
&env.event_tx,
&env.cancel_tx,
&env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer,
Some(&env.parent_socket_path),
"child-worker",
@@ -2218,6 +2274,8 @@ mod tests {
&env.event_tx,
&env.cancel_tx,
&env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer,
Some(&env.parent_socket_path),
"parent",
@@ -2254,6 +2312,8 @@ mod tests {
&env.event_tx,
&env.cancel_tx,
&env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer,
Some(&env.parent_socket_path),
"parent",
@@ -2288,6 +2348,8 @@ mod tests {
&env.event_tx,
&env.cancel_tx,
&env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer,
Some(&env.parent_socket_path),
"parent",
@@ -2321,6 +2383,8 @@ mod tests {
&env.event_tx,
&env.cancel_tx,
&env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer,
Some(&env.parent_socket_path),
"child-worker",
crates/worker/src/worker.rs
+20 -1
View File
@@ -2752,10 +2752,28 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
pub(crate) async fn run_with_input_extensions(
&mut self,
input: Vec<Segment>,
mut input_extensions: Vec<SessionExtension>,
input_extensions: Vec<SessionExtension>,
) -> Result<WorkerRunResult, WorkerError>
where
St: Clone + 'static,
{
self.run_with_input_extensions_and_commit_hook(input, input_extensions, || {})
.await
}
/// Run user input and invoke `on_input_committed` only after the annotated
/// input has crossed both the durable Store and live SegmentLogSink commit
/// boundaries. The Controller uses this fence before exposing `Running`, so
/// every in-flight snapshot for a user turn includes its committed input.
pub(crate) async fn run_with_input_extensions_and_commit_hook<F>(
&mut self,
input: Vec<Segment>,
mut input_extensions: Vec<SessionExtension>,
on_input_committed: F,
) -> Result<WorkerRunResult, WorkerError>
where
St: Clone + 'static,
F: FnOnce(),
{
let (input, pending_flow_state, flow_projection) = self.prepare_flow_input(input)?;
if let Some(state) = pending_flow_state.as_ref() {
@@ -2810,6 +2828,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.expect("flow_runtime_state poisoned") = Some(state);
}
self.user_segments.push(input.clone());
on_input_committed();
// Resolve `@<path>` file refs to system messages stashed for the
// WorkerInterceptor to attach right after the user message. Resolution
crates/worker/tests/controller_test.rs
+18 -1
View File
@@ -806,13 +806,30 @@ async fn snapshot_includes_user_input_for_in_flight_turn() {
let client = MockClient::sequential(vec![MockResponse::Hang(simple_text_events())]);
let worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
let mut events = handle.subscribe();
handle
.send(Method::run_text("hello in-flight"))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Running).await;
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if matches!(
events.recv().await,
Ok(Event::Status {
status: WorkerStatus::Running,
})
) {
break;
}
}
})
.await
.expect("running status event");
// The Running event is the in-flight visibility fence: the committed
// annotated input must already be available to an immediately attaching
// subscriber rather than racing behind this status transition.
let stream = tokio::net::UnixStream::connect(handle.runtime_dir.socket_path())
.await
.unwrap();