Merge ticket 00001M00PZS9R paused Worker Resume

This commit is contained in:
2026-08-15 03:55:20 +09:00
2 changed files with 238 additions and 8 deletions
+178 -6
View File
@@ -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<LlmEvent>),
Hang(Vec<LlmEvent>),
}
#[derive(Clone)]
struct MockClient {
responses: Arc<Vec<Vec<LlmEvent>>>,
responses: Arc<Vec<MockResponse>>,
call_count: Arc<AtomicUsize>,
captured: Arc<Mutex<Vec<Request>>>,
}
impl MockClient {
fn new(events: Vec<LlmEvent>) -> Self {
Self::sequential(vec![MockResponse::Complete(events)])
}
fn sequential(responses: Vec<MockResponse>) -> 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,9 +1973,21 @@ 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();
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()),
)),
}
}
}
#[cfg(feature = "ws-server")]
@@ -2073,6 +2125,31 @@ mod tests {
.collect()
}
fn wait_for_adapter_state(
backend: &WorkerRuntimeExecutionBackend<MockFactory>,
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<LlmEvent> {
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::<Vec<_>>();
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());
+59 -1
View File
@@ -12714,6 +12714,8 @@ mod tests {
materializer: worker_runtime::working_directory::LocalGitWorktreeMaterializer,
spawn_failure: std::sync::Mutex<Option<String>>,
inputs: std::sync::Mutex<Vec<(worker_runtime::identity::WorkerRef, String)>>,
protocol_methods:
std::sync::Mutex<Vec<(worker_runtime::identity::WorkerRef, protocol::Method)>>,
}
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<String>) -> 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();
}