diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 82ab1a08..7ad5f772 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1169,22 +1169,34 @@ async fn controller_loop( // 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), ¬ify_buffer, self_parent_socket.as_ref(), &spawner_name, @@ -1194,12 +1206,21 @@ async fn controller_loop( .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), ¬ify_buffer, self_parent_socket.as_ref(), &spawner_name, @@ -1215,6 +1236,8 @@ async fn controller_loop( &event_tx, &cancel_tx, &shared_state, + &runtime_dir, + None, ¬ify_buffer, self_parent_socket.as_ref(), &spawner_name, @@ -1230,6 +1253,8 @@ async fn controller_loop( &event_tx, &cancel_tx, &shared_state, + &runtime_dir, + None, ¬ify_buffer, self_parent_socket.as_ref(), &spawner_name, @@ -1627,6 +1652,8 @@ async fn drive_turn( event_tx: &broadcast::Sender, cancel_tx: &mpsc::Sender<()>, shared_state: &Arc, + runtime_dir: &RuntimeDir, + mut input_commit_rx: Option>, 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, parent_socket_path: PathBuf, - _runtime_dir: Arc, + runtime_dir: Arc, _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", diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 38e2195a..8bd5cd25 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -2752,10 +2752,28 @@ impl Worker { pub(crate) async fn run_with_input_extensions( &mut self, input: Vec, - mut input_extensions: Vec, + input_extensions: Vec, ) -> Result 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( + &mut self, + input: Vec, + mut input_extensions: Vec, + on_input_committed: F, + ) -> Result + 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 Worker { .expect("flow_runtime_state poisoned") = Some(state); } self.user_segments.push(input.clone()); + on_input_committed(); // Resolve `@` file refs to system messages stashed for the // WorkerInterceptor to attach right after the user message. Resolution diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 04836dbe..d58cfd04 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -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();