diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 75fcb079..bf1f3076 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -1160,7 +1160,10 @@ where match event { Ok(event) => { let _ = bridge_context.publish_protocol_event(event); - if bridge_handle.shared_state.get_status() == WorkerStatus::Idle { + if matches!( + bridge_handle.shared_state.get_status(), + WorkerStatus::Idle | WorkerStatus::Paused + ) { bridge_busy.store(false, Ordering::SeqCst); } } @@ -1240,6 +1243,13 @@ fn method_starts_turn(method: &Method) -> bool { ) } +fn method_can_start_turn_from_status(method: &Method, status: WorkerStatus) -> bool { + match method { + Method::Resume => matches!(status, WorkerStatus::Idle | WorkerStatus::Paused), + _ => status == WorkerStatus::Idle, + } +} + fn accepted_notify_run_state(status: WorkerStatus, auto_run: bool) -> WorkerExecutionRunState { match status { WorkerStatus::Running => WorkerExecutionRunState::Busy, @@ -1696,7 +1706,7 @@ where let starts_turn = method_starts_turn(&method); if starts_turn - && (worker.shared_state.get_status() != WorkerStatus::Idle + && (!method_can_start_turn_from_status(&method, worker.shared_state.get_status()) || busy .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) .is_err()) @@ -1842,7 +1852,7 @@ mod tests { use crate::observation::WorkerObservationCursor; use crate::working_directory::LocalGitWorktreeMaterializer; use async_trait::async_trait; - use futures::Stream; + use futures::{Stream, StreamExt}; use llm_engine::Engine; use llm_engine::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent}; use llm_engine::llm_client::{ClientError, LlmClient, Request}; @@ -1903,17 +1913,47 @@ mod tests { ); } + #[test] + fn resume_turn_claim_accepts_paused_and_idle_but_not_running_status() { + assert!(method_can_start_turn_from_status( + &Method::Resume, + WorkerStatus::Paused + )); + assert!(method_can_start_turn_from_status( + &Method::Resume, + WorkerStatus::Idle + )); + assert!(!method_can_start_turn_from_status( + &Method::Resume, + WorkerStatus::Running + )); + assert!(!method_can_start_turn_from_status( + &Method::Compact, + WorkerStatus::Paused + )); + } + + #[derive(Clone)] + enum MockResponse { + Complete(Vec), + Hang(Vec), + } + #[derive(Clone)] struct MockClient { - responses: Arc>>, + responses: Arc>, call_count: Arc, captured: Arc>>, } impl MockClient { fn new(events: Vec) -> Self { + Self::sequential(vec![MockResponse::Complete(events)]) + } + + fn sequential(responses: Vec) -> Self { Self { - responses: Arc::new(vec![events]), + responses: Arc::new(responses), call_count: Arc::new(AtomicUsize::new(0)), captured: Arc::new(Mutex::new(Vec::new())), } @@ -1933,8 +1973,20 @@ mod tests { { self.captured.lock().unwrap().push(request); let idx = self.call_count.fetch_add(1, Ordering::SeqCst); - let events = self.responses.get(idx).cloned().unwrap_or_default(); - Ok(Box::pin(futures::stream::iter(events.into_iter().map(Ok)))) + let response = self + .responses + .get(idx) + .cloned() + .unwrap_or_else(|| MockResponse::Complete(Vec::new())); + match response { + MockResponse::Complete(events) => { + Ok(Box::pin(futures::stream::iter(events.into_iter().map(Ok)))) + } + MockResponse::Hang(events) => Ok(Box::pin( + futures::stream::iter(events.into_iter().map(Ok)) + .chain(futures::stream::pending()), + )), + } } } @@ -2073,6 +2125,31 @@ mod tests { .collect() } + fn wait_for_adapter_state( + backend: &WorkerRuntimeExecutionBackend, + worker_ref: &WorkerRef, + expected_status: WorkerStatus, + expected_busy: bool, + ) { + let deadline = std::time::Instant::now() + Duration::from_secs(5); + loop { + let matches = { + let workers = backend.workers.lock().unwrap(); + let execution = workers.get(worker_ref).expect("live Worker execution"); + execution.handle.shared_state.get_status() == expected_status + && execution.busy.load(Ordering::SeqCst) == expected_busy + }; + if matches { + return; + } + assert!( + std::time::Instant::now() < deadline, + "timed out waiting for adapter state {expected_status:?}, busy={expected_busy}" + ); + std::thread::sleep(Duration::from_millis(10)); + } + } + fn simple_text_events() -> Vec { vec![ LlmEvent::text_block_start(0), @@ -2868,6 +2945,101 @@ mod tests { ); } + #[test] + #[cfg(feature = "ws-server")] + fn adapter_resumes_paused_turn_once_and_preserves_idle_not_paused_error() { + let hanging_events = || simple_text_events().into_iter().take(2).collect::>(); + let client = MockClient::sequential(vec![ + MockResponse::Hang(hanging_events()), + MockResponse::Hang(hanging_events()), + MockResponse::Complete(simple_text_events()), + ]); + let call_count = client.call_count.clone(); + let runtime_base = tempfile::tempdir().unwrap(); + let repo = create_clean_repo(); + let store = tempfile::tempdir().unwrap(); + let factory = MockFactory { + client, + runtime_base: runtime_base.path().to_path_buf(), + cwd: repo.path().to_path_buf(), + store_dir: store.path().join("sessions"), + worker_metadata_dir: store.path().join("workers"), + observed_cwds: Arc::new(Mutex::new(Vec::new())), + observed_workspace_clients: Arc::new(Mutex::new(Vec::new())), + }; + let backend = Arc::new(WorkerRuntimeExecutionBackend::new(factory).unwrap()); + let runtime = + EmbeddedRuntime::with_execution_backend(RuntimeOptions::default(), backend.clone()) + .unwrap(); + runtime.store_config_bundle(test_bundle()).unwrap(); + let detail = runtime + .create_worker(create_request("paused-resume")) + .unwrap(); + + runtime + .send_input(&detail.worker_ref, WorkerInput::user("pause and resume")) + .expect("start initial turn"); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running, true); + + let running_resume = runtime + .send_protocol_method(&detail.worker_ref, Method::Resume) + .expect_err("Resume while Running must be rejected"); + assert!( + running_resume + .to_string() + .contains("does not queue protocol methods"), + "unexpected Running Resume error: {running_resume}" + ); + + runtime + .send_protocol_method(&detail.worker_ref, Method::Pause) + .expect("pause initial turn"); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused, false); + + runtime + .send_protocol_method(&detail.worker_ref, Method::Resume) + .expect("resume paused turn"); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running, true); + + let duplicate_resume = runtime + .send_protocol_method(&detail.worker_ref, Method::Resume) + .expect_err("duplicate Resume must be rejected"); + assert!( + duplicate_resume + .to_string() + .contains("does not queue protocol methods"), + "unexpected duplicate Resume error: {duplicate_resume}" + ); + + runtime + .send_protocol_method(&detail.worker_ref, Method::Pause) + .expect("pause resumed turn"); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused, false); + runtime + .send_protocol_method(&detail.worker_ref, Method::Resume) + .expect("resume paused turn a second time"); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle, false); + assert_eq!(call_count.load(Ordering::SeqCst), 3); + + runtime + .send_protocol_method(&detail.worker_ref, Method::Resume) + .expect("Idle Resume preserves controller NotPaused semantics"); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle, false); + let events = runtime + .read_worker_observation_events(&detail.worker_ref, WorkerObservationCursor::zero()) + .expect("read protocol events"); + assert!(events.iter().any(|event| { + matches!( + &event.payload, + Event::Error { + code: protocol::ErrorCode::NotPaused, + .. + } + ) + })); + assert_eq!(call_count.load(Ordering::SeqCst), 3); + } + #[test] fn stopping_and_deleting_worker_preserves_bound_working_directory() { let client = MockClient::new(simple_text_events()); diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index fc11c320..933f9c89 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -12714,6 +12714,8 @@ mod tests { materializer: worker_runtime::working_directory::LocalGitWorktreeMaterializer, spawn_failure: std::sync::Mutex>, inputs: std::sync::Mutex>, + protocol_methods: + std::sync::Mutex>, } impl Default for DeterministicExecutionBackend { @@ -12733,6 +12735,7 @@ mod tests { ), spawn_failure: std::sync::Mutex::new(None), inputs: std::sync::Mutex::new(Vec::new()), + protocol_methods: std::sync::Mutex::new(Vec::new()), } } } @@ -12742,6 +12745,10 @@ mod tests { std::mem::take(&mut *self.inputs.lock().expect("inputs lock")) } + fn protocol_methods(&self) -> Vec<(worker_runtime::identity::WorkerRef, protocol::Method)> { + self.protocol_methods.lock().unwrap().clone() + } + fn fail_first_spawn(message: impl Into) -> Self { let backend = Self::default(); *backend.spawn_failure.lock().unwrap() = Some(message.into()); @@ -12835,6 +12842,21 @@ mod tests { } } + fn dispatch_method( + &self, + handle: &worker_runtime::execution::WorkerExecutionHandle, + method: protocol::Method, + ) -> worker_runtime::execution::WorkerExecutionResult { + self.protocol_methods + .lock() + .unwrap() + .push((handle.worker_ref().clone(), method)); + worker_runtime::execution::WorkerExecutionResult::accepted( + worker_runtime::execution::WorkerExecutionOperation::ProtocolMethod, + worker_runtime::execution::WorkerExecutionRunState::Idle, + ) + } + fn stop_worker( &self, _handle: &worker_runtime::execution::WorkerExecutionHandle, @@ -17109,7 +17131,7 @@ mod tests { #[tokio::test] async fn workspace_subscription_returns_workspace_snapshot() { let dir = tempfile::tempdir().unwrap(); - let api = test_api(dir.path()).await; + let (api, execution_backend) = test_api_with_recording_backend(dir.path()).await; let spawn_request = WorkerSpawnRequest { intent: WorkerSpawnIntent::WorkspaceCompanion, @@ -17310,6 +17332,42 @@ mod tests { break; } } + let resume = protocol::subscription::SubscriptionFrame::new( + protocol::subscription::SubscriptionFramePayload::WorkerProtocol( + protocol::subscription::SubscriptionWorkerProtocolMethod { + subscription_id: second_protocol_subscription_id, + method: protocol::Method::Resume, + }, + ), + ); + socket + .send(Message::Text( + serde_json::to_string(&resume).unwrap().into(), + )) + .await + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if execution_backend + .protocol_methods() + .iter() + .any(|(worker_ref, method)| { + worker_ref.worker_id.to_string() == worker_id + && matches!(method, protocol::Method::Resume) + }) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("Resume method should reach the Runtime execution backend"); + let protocol_methods = execution_backend.protocol_methods(); + assert!(protocol_methods.iter().any(|(worker_ref, method)| { + worker_ref.worker_id.to_string() == worker_id + && matches!(method, protocol::Method::Resume) + })); server.abort(); }