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

fix: restrict worker send input to idle admission

Browse Source
This commit is contained in:
Hare
2026-10-02 02:39:10 +09:00
parent 4ac448d66c
commit d5f49cbf00
12 changed files with 503 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker-runtime/src/http_server.rs
+13
View File
@@ -1556,6 +1556,19 @@ fn authorize_runtime_protocol_method(
submission_request_id,
input,
},
protocol::Method::SubmitIfIdle {
submission_request_id,
input,
..
} => protocol::Method::SubmitIfIdle {
source: transport_source.cloned().unwrap_or_else(|| {
protocol::AuthenticatedInputSource::Backend {
operation_id: submission_request_id.clone(),
}
}),
submission_request_id,
input,
},
protocol::Method::NotifyTracked {
notification_request_id,
message,
crates/worker-runtime/src/interaction.rs
+2
View File
@@ -8,6 +8,8 @@ use serde::{Deserialize, Serialize};
#[serde(rename_all = "snake_case")]
pub enum WorkerInputKind {
User,
/// Agent input that must never enter the human Submit queue.
UserIfIdle,
Notify,
Compact,
ListRewindTargets,
crates/worker-runtime/src/runtime.rs
+21 -3
View File
@@ -2090,7 +2090,10 @@ impl Runtime {
mut input: WorkerInput,
) -> Result<WorkerInteractionAck, RuntimeError> {
validate_worker_input(&input)?;
if matches!(input.kind, WorkerInputKind::User | WorkerInputKind::Notify) {
if matches!(
input.kind,
WorkerInputKind::User | WorkerInputKind::UserIfIdle | WorkerInputKind::Notify
) {
let request_id = input
.submission_request_id
.clone()
@@ -2098,7 +2101,11 @@ impl Runtime {
.unwrap_or_else(|| Uuid::now_v7().to_string());
input.submission_request_id = Some(request_id);
}
let expected_submission_id = (input.kind == WorkerInputKind::User).then(|| {
let expected_submission_id = matches!(
input.kind,
WorkerInputKind::User | WorkerInputKind::UserIfIdle
)
.then(|| {
input
.submission_request_id
.clone()
@@ -2110,6 +2117,17 @@ impl Runtime {
.clone()
.expect("Notify request id must be assigned before dispatch")
});
if input.kind == WorkerInputKind::UserIfIdle {
let state = self.lock()?;
let worker = state.worker(worker_ref)?;
// Stopped/restoring Workers have no controller to enforce admission.
// Live busy state is checked atomically by the controller, not here.
if worker.status == WorkerStatus::Stopped || worker.has_pending_lifecycle_operation() {
return Err(RuntimeError::InvalidRequest(
"WorkerSendInput requires an Idle Worker and never queues input. Use WorkerNotify for information about work already in progress.".to_string(),
));
}
}
self.ensure_worker_execution(worker_ref)?;
let (backend, handle) = {
let state = self.lock()?;
@@ -5319,7 +5337,7 @@ fn input_protocol_event(input: &WorkerInput) -> Option<protocol::Event> {
match input.kind {
// Submit is projected only after the Worker commits UserInput. Queued
// payloads must never become model- or client-visible history early.
WorkerInputKind::User | WorkerInputKind::Notify => None,
WorkerInputKind::User | WorkerInputKind::UserIfIdle | WorkerInputKind::Notify => None,
WorkerInputKind::Compact
| WorkerInputKind::ListRewindTargets
| WorkerInputKind::RegisterPeer => Some(protocol::Event::SystemItem {
crates/worker-runtime/src/worker_backend.rs
+137 -11
View File
@@ -2982,7 +2982,7 @@ where
}
let (method, submission_request_id) = match input.kind {
WorkerInputKind::User => {
WorkerInputKind::User | WorkerInputKind::UserIfIdle => {
let Some(submission_id) = input
.submission_request_id
.filter(|submission_id| !submission_id.trim().is_empty())
@@ -2992,18 +2992,26 @@ where
"Runtime user input is missing its internal submission id",
);
};
(
let segments = input
.segments
.unwrap_or_else(|| vec![Segment::text(input.content.trim().to_string())]);
let source = protocol::AuthenticatedInputSource::Backend {
operation_id: submission_id.clone(),
};
let method = if input.kind == WorkerInputKind::UserIfIdle {
Method::SubmitIfIdle {
submission_request_id: submission_id.clone(),
input: segments,
source,
}
} else {
Method::SubmitTracked {
submission_request_id: submission_id.clone(),
input: input.segments.unwrap_or_else(|| {
vec![Segment::text(input.content.trim().to_string())]
}),
source: protocol::AuthenticatedInputSource::Backend {
operation_id: submission_id.clone(),
},
},
Some(submission_id),
)
input: segments,
source,
}
};
(method, Some(submission_id))
}
WorkerInputKind::Notify => {
unreachable!("Notify input is dispatched before ordinary input mapping")
@@ -5581,6 +5589,124 @@ mod tests {
assert_eq!(call_count.load(Ordering::SeqCst), 3);
}
#[test]
#[cfg(feature = "ws-server")]
#[serial_test::serial(worker_allocation)]
fn idle_only_input_acknowledges_idle_and_rejects_running_paused_and_stopped() {
let client = MockClient::sequential(vec![MockResponse::Hang(
simple_text_events().into_iter().take(2).collect(),
)]);
let runtime_base = tempfile::tempdir().unwrap();
let cwd = tempfile::tempdir().unwrap();
let store = tempfile::tempdir().unwrap();
let factory = MockFactory {
client,
runtime_base: runtime_base.path().to_path_buf(),
cwd: cwd.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("idle-only")).unwrap();
let input = |id: &str| {
let mut input = WorkerInput::user(id);
input.kind = WorkerInputKind::UserIfIdle;
input.submission_request_id = Some(id.to_string());
input
};
let ack = runtime
.send_input(&detail.worker_ref, input("idle-input"))
.unwrap();
assert_eq!(ack.submission.unwrap().submission_request_id, "idle-input");
wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running);
let error = runtime
.send_input(&detail.worker_ref, input("running-input"))
.unwrap_err();
assert!(error.to_string().contains("WorkerNotify"), "{error}");
runtime
.send_protocol_method(
&detail.worker_ref,
Method::Pause {
command: adapter_command(&backend, &detail.worker_ref),
},
)
.unwrap();
wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused);
let error = runtime
.send_input(&detail.worker_ref, input("paused-input"))
.unwrap_err();
assert!(error.to_string().contains("WorkerNotify"), "{error}");
let handle = backend
.workers
.lock()
.unwrap()
.get(&detail.worker_ref)
.unwrap()
.handle
.clone();
let pending = backend
.run_on_adapter_runtime(async move {
let mut events = handle.subscribe();
handle
.send(Method::ListPendingSubmissions)
.await
.map_err(|e| e.to_string())?;
tokio::time::timeout(Duration::from_secs(2), async {
loop {
if let Ok(Event::PendingSubmissionsChanged { pending }) =
events.recv().await
{
break pending;
}
}
})
.await
.map_err(|e| e.to_string())
})
.unwrap();
assert!(pending.submissions.is_empty());
// Human Submit still queues while paused, and Notify remains advisory.
let ack = runtime
.send_input(&detail.worker_ref, WorkerInput::user("human-input"))
.unwrap();
assert_eq!(
ack.submission.unwrap().disposition,
protocol::SubmissionDisposition::Queued
);
let mut notify = WorkerInput::user("advisory");
notify.kind = WorkerInputKind::Notify;
assert!(
runtime
.send_input(&detail.worker_ref, notify)
.unwrap()
.notification
.is_some()
);
wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused);
let handle = backend
.workers
.lock()
.unwrap()
.get(&detail.worker_ref)
.unwrap()
.handle
.clone();
let entries = serde_json::to_string(&handle.committed_entries()).unwrap();
assert!(!entries.contains("running-input"));
assert!(!entries.contains("paused-input"));
runtime.stop_worker(&detail.worker_ref, None).unwrap();
let error = runtime
.send_input(&detail.worker_ref, input("stopped-input"))
.unwrap_err();
assert!(error.to_string().contains("WorkerNotify"), "{error}");
}
#[test]
fn stopped_runtime_worker_can_restore_and_accept_input() {
let client = MockClient::new(simple_text_events());