feat: persist worker submit activation queue

This commit is contained in:
2026-09-05 22:06:59 +09:00
parent aa96bbedbc
commit bb56283063
41 changed files with 2336 additions and 811 deletions
+1
View File
@@ -20,6 +20,7 @@ protocol = { workspace = true, features = ["json-schema"] }
client = { workspace = true }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
sha2 = { workspace = true }
reqwest = { version = "0.13", default-features = false, features = ["blocking", "native-tls"] }
thiserror = { workspace = true }
tokio = { workspace = true, features = ["fs", "io-util", "macros", "net", "process", "rt-multi-thread", "signal", "sync", "time"] }
+4 -1
View File
@@ -101,7 +101,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
// Send a run method
handle
.send(Method::run_text("What is the capital of France?"))
.send(Method::submit_text(
protocol::new_submission_request_id(),
"What is the capital of France?",
))
.await?;
// Wait for completion
+364 -152
View File
@@ -5,7 +5,7 @@ use std::sync::atomic::Ordering;
use agen::EngineError;
use agen::llm_client::client::LlmClient;
use session_store::WorkerMetadataStore;
use session_store::{LogEntry, SessionExtension, Store};
use session_store::{LogEntry, Store};
use tokio::sync::{broadcast, mpsc, oneshot};
use crate::discovery::WorkerDiscovery;
@@ -23,16 +23,12 @@ use crate::shutdown_after_idle::{
};
use crate::spawn::registry::SpawnedWorkerRegistry;
use crate::spawn::tool::sub_worker_spawn_tool;
use crate::worker::{
SystemItemCommitter, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError,
WorkerRunResult,
};
use crate::worker::{SystemItemCommitter, Worker, WorkerError, WorkerRunResult};
use protocol::{
AlertLevel, AlertSource, CommandEvent as ProtocolCommandEvent,
CommandSnapshot as ProtocolCommandSnapshot, CommandStatus as ProtocolCommandStatus,
CommandStream as ProtocolCommandStream, CommandStreamSlice as ProtocolCommandStreamSlice,
ErrorCode, Event, Method, RewindTargetId, RunResult, Segment, TurnResult, UploadedFileRef,
WorkerStatus,
ErrorCode, Event, Method, RewindTargetId, RunResult, TurnResult, UploadedFileRef, WorkerStatus,
};
use workdir::{
CommandEvent as WorkdirCommandEvent, CommandSnapshot as WorkdirCommandSnapshot,
@@ -58,6 +54,7 @@ pub struct WorkerHandle {
spawned_registry: Arc<SpawnedWorkerRegistry>,
artifact_store: Arc<dyn Store>,
session_id: session_store::SessionId,
pending_activations: Arc<std::sync::Mutex<crate::worker::PendingActivationState>>,
}
impl WorkerHandle {
@@ -131,8 +128,15 @@ impl WorkerHandle {
let in_flight = snapshot_from_guard(&in_flight_guard);
(entries, entry_rx, in_flight)
};
let mut session =
session_store::public_snapshot::project_current_session_snapshot(&entries);
session.pending_submissions = self
.pending_activations
.lock()
.expect("pending activation state poisoned")
.snapshot();
let event = Event::Snapshot {
session: session_store::public_snapshot::project_current_session_snapshot(&entries),
session,
greeting: self.shared_state.greeting.clone(),
status: self.shared_state.get_status(),
in_flight,
@@ -213,21 +217,41 @@ async fn finish_controller_run<C, St>(
/// `Worker::*` entry point — `RunForNotification` carries none because
/// `worker.run_for_notification()` drains the NotifyBuffer on its own.
enum PendingRun {
Run(Vec<Segment>),
RunTracked {
input: Vec<Segment>,
extension: SessionExtension,
},
Submit(crate::worker::PendingSubmission),
/// Self-initiated turn kicked from the notify buffer. The carried
/// `InvokeKind` is the trigger that flipped the Worker from IDLE
/// (Notify or WorkerEvent) and is recorded by the Invoke marker
/// committed at the start of `worker.run_for_notification`.
RunForNotification(protocol::InvokeKind),
RunForNotification {
invoke_kind: protocol::InvokeKind,
notification_request_id: Option<String>,
},
Resume,
}
fn prepare_pending_run<St: Store + Clone>(
pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
notify_buffer: &NotifyBuffer,
) -> Result<Option<PendingRun>, crate::worker::PendingSubmissionError> {
Ok(match pending_submissions.prepare_next_activation()? {
Some(crate::worker::PendingActivation::Submission(submission)) => {
Some(PendingRun::Submit(submission))
}
Some(crate::worker::PendingActivation::Notification(notification)) => {
let extension = pending_submissions.notification_activation_extension();
let notification_request_id = notification.notification_request_id.clone();
notify_buffer.push_durable_notify(notification.message, extension);
Some(PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::Notify,
notification_request_id: Some(notification_request_id),
})
}
None => None,
})
}
impl PendingRun {
/// Whether this turn was kicked off by the parent (via `Method::Run`
/// Whether this turn was kicked off by the parent (via `Method::Submit`
/// or `Method::Resume`). Used by [`drive_turn`] to gate upward
/// `WorkerEvent::TurnEnded` / `WorkerEvent::Errored` reports so the parent
/// only sees completion signals for work it actually delegated.
@@ -235,16 +259,12 @@ impl PendingRun {
/// notify buffer (Notify / inbound WorkerEvent) and stays silent.
fn is_parent_originated(&self) -> bool {
match self {
PendingRun::Run(_) | PendingRun::RunTracked { .. } | PendingRun::Resume => true,
PendingRun::RunForNotification(_) => false,
PendingRun::Submit(_) | PendingRun::Resume => true,
PendingRun::RunForNotification { .. } => false,
}
}
}
fn should_auto_run_notification(status: WorkerStatus, auto_run: bool) -> bool {
auto_run && status == WorkerStatus::Idle
}
// ---------------------------------------------------------------------------
// WorkerController — actor that owns a Worker
// ---------------------------------------------------------------------------
@@ -552,6 +572,7 @@ impl WorkerController {
let artifact_store: Arc<dyn Store> = Arc::new(worker.store().clone());
let session_id = worker.session_id();
let pending_activations = worker.pending_activation_state();
let handle = WorkerHandle {
method_tx,
working_event_tx: working_event_tx.clone(),
@@ -563,6 +584,7 @@ impl WorkerController {
spawned_registry: spawned_registry.clone(),
artifact_store,
session_id,
pending_activations,
};
let socket_server = match transport {
@@ -1291,6 +1313,7 @@ async fn controller_loop<C, St>(
spawned_registry.clone(),
);
let mut pending: Option<PendingRun> = None;
let pending_submissions = worker.pending_submission_handle();
loop {
// Top-of-iteration: if an event handler staged a run, fire it
@@ -1307,8 +1330,8 @@ async fn controller_loop<C, St>(
// interrupted/error turn from being carried into the next snapshot.
worker.clear_in_flight_events();
let parent_originated = run.is_parent_originated();
let user_input_run = matches!(&run, PendingRun::Run(_) | PendingRun::RunTracked { .. });
if !user_input_run {
let user_input_submit = matches!(&run, PendingRun::Submit(_));
if !user_input_submit {
set_controller_status(
&shared_state,
&runtime_dir,
@@ -1317,37 +1340,21 @@ async fn controller_loop<C, St>(
)
.await;
}
let (mut new_status, shutdown) = match run {
PendingRun::Run(input) => {
let notification_request_id = match &run {
PendingRun::RunForNotification {
notification_request_id,
..
} => notification_request_id.clone(),
_ => None,
};
let (mut new_status, shutdown, may_drain_pending) = match run {
PendingRun::Submit(submission) => {
let (input_commit_tx, input_commit_rx) = oneshot::channel();
let committed_submission = submission.clone();
let extension = pending_submissions.activation_extension();
drive_turn(
worker.run_with_input_extensions_and_commit_hook(
input,
Vec::new(),
move || {
let _ = input_commit_tx.send(());
},
),
&mut method_rx,
&working_event_tx,
&cancel_tx,
&pause_tx,
&shared_state,
&runtime_dir,
Some(input_commit_rx),
&notify_buffer,
self_parent_socket.as_ref(),
&spawner_name,
&spawned_registry,
parent_originated,
)
.await
}
PendingRun::RunTracked { input, extension } => {
let (input_commit_tx, input_commit_rx) = oneshot::channel();
drive_turn(
worker.run_with_input_extensions_and_commit_hook(
input,
submission.input,
vec![extension],
move || {
let _ = input_commit_tx.send(());
@@ -1359,8 +1366,9 @@ async fn controller_loop<C, St>(
&pause_tx,
&shared_state,
&runtime_dir,
Some(input_commit_rx),
Some((input_commit_rx, committed_submission)),
&notify_buffer,
&pending_submissions,
self_parent_socket.as_ref(),
&spawner_name,
&spawned_registry,
@@ -1368,9 +1376,9 @@ async fn controller_loop<C, St>(
)
.await
}
PendingRun::RunForNotification(kind) => {
PendingRun::RunForNotification { invoke_kind, .. } => {
drive_turn(
worker.run_for_notification(kind),
worker.run_for_notification(invoke_kind),
&mut method_rx,
&working_event_tx,
&cancel_tx,
@@ -1379,6 +1387,7 @@ async fn controller_loop<C, St>(
&runtime_dir,
None,
&notify_buffer,
&pending_submissions,
self_parent_socket.as_ref(),
&spawner_name,
&spawned_registry,
@@ -1397,6 +1406,7 @@ async fn controller_loop<C, St>(
&runtime_dir,
None,
&notify_buffer,
&pending_submissions,
self_parent_socket.as_ref(),
&spawner_name,
&spawned_registry,
@@ -1405,10 +1415,32 @@ async fn controller_loop<C, St>(
.await
}
};
if !shutdown && new_status == WorkerStatus::Idle && notify_buffer.has_auto_run_pending()
{
pending = Some(PendingRun::RunForNotification(protocol::InvokeKind::Notify));
new_status = WorkerStatus::Running;
if let Some(notification_request_id) = notification_request_id {
pending_submissions.finish_notification_activation(&notification_request_id);
}
if !shutdown && may_drain_pending && new_status == WorkerStatus::Idle {
match prepare_pending_run(&pending_submissions, &notify_buffer) {
Ok(Some(next)) => {
pending = Some(next);
new_status = WorkerStatus::Running;
}
Ok(None) => {
if notify_buffer.has_auto_run_pending() {
pending = Some(PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::Notify,
notification_request_id: None,
});
new_status = WorkerStatus::Running;
}
}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: error.to_string(),
});
}
}
}
finish_controller_run(
&mut worker,
@@ -1435,61 +1467,118 @@ async fn controller_loop<C, St>(
};
match method {
Method::Run { input } => {
if shared_state.get_status() == WorkerStatus::Running {
// Defensive: the inner select! inside drive_turn
// already rejects `Run` while a turn is live, so
// this branch is only reachable across a race window
// around status flips.
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn".into(),
});
continue;
}
// Stage the run without a speculative user-message echo.
// `Worker::run` validates the input, commits
// `LogEntry::AnnotatedUserInput`, and the session-log sink turns that
// committed entry into the live `Event::UserMessage`. That
// keeps every client ordered against `SegmentStart` replay and
// makes persisted history the single source of visible user
// input. Paused→Run cleanup (orphan tool_result closure +
// interrupt system note) is applied inside `Worker::run` itself
// when the worker's `last_run_interrupted` flag is set.
pending = Some(PendingRun::Run(input));
}
Method::RunTracked {
Method::Submit {
submission_request_id,
input,
submission_id,
} => {
// Runtime-correlated submissions retain their opaque id in the
// same durable UserInput record used for Flow state.
let extension = SessionExtension::new(
WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN,
serde_json::json!({ "submission_id": submission_id }),
);
pending = Some(PendingRun::RunTracked { input, extension });
}
Method::Notify { message, auto_run } => {
// Client-side live echo is delivered as `Event::SystemItem`
// once the interceptor commits the corresponding
// `LogEntry::AnnotatedSystemItem` entry — drained out of the
// notify buffer + broadcast through the sink. No
// separate echo here.
worker.push_notify(message, auto_run);
// RUNNING: the in-flight turn drains the buffer at its next
// pending_history_appends; if an auto-run notification remains
// at turn end, the Controller stages a follow-up notification
// turn. Paused notifications remain queued until Resume/Run.
// IDLE: `auto_run` notifications stage RunForNotification;
// weak progress notices stay queued until an explicit run.
if should_auto_run_notification(shared_state.get_status(), auto_run) {
pending = Some(PendingRun::RunForNotification(protocol::InvokeKind::Notify));
| Method::SubmitTracked {
submission_request_id,
input,
} => {
let request_id = submission_request_id.clone();
match pending_submissions.accept(submission_request_id, input, true) {
Ok(acceptance) => {
if let Some(activation) = acceptance.activation {
pending = Some(PendingRun::Submit(activation));
} else {
let _ = working_event_tx.send(Event::SubmissionAccepted {
submission_request_id: acceptance.submission_request_id,
submission_id: acceptance.submission_id,
disposition: acceptance.disposition,
});
}
}
Err(error) => {
let _ = working_event_tx.send(Event::SubmissionRejected {
submission_request_id: request_id,
message: error.to_string(),
});
}
}
}
Method::Notify {
notification_request_id,
message,
auto_run,
} => {
if auto_run {
match pending_submissions.accept_notification(notification_request_id, message)
{
Ok(true) => match prepare_pending_run(&pending_submissions, &notify_buffer)
{
Ok(Some(next)) => pending = Some(next),
Ok(None) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: error.to_string(),
});
}
},
Ok(false) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
}
}
} else {
worker.push_notify(message, false);
}
}
Method::ListPendingSubmissions => {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_submissions.snapshot(),
});
}
Method::CancelPendingSubmission { submission_id } => {
match pending_submissions.cancel(&submission_id) {
Ok(pending_snapshot) => {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_snapshot,
});
}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
}
}
}
Method::ClearPendingSubmissions => match pending_submissions.clear() {
Ok(pending_snapshot) => {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_snapshot,
});
}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: error.to_string(),
});
}
},
Method::ContinuePending => {
match prepare_pending_run(&pending_submissions, &notify_buffer) {
Ok(Some(next)) => pending = Some(next),
Ok(None) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: "pending activation queue is empty".into(),
});
}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: error.to_string(),
});
}
}
}
Method::Resume => {
if shared_state.get_status() != WorkerStatus::Paused {
let _ = working_event_tx.send(Event::Error {
@@ -1703,9 +1792,10 @@ async fn controller_loop<C, St>(
// notification is not stranded. Matches the
// `Method::Notify` idle path.
if shared_state.get_status() == WorkerStatus::Idle {
pending = Some(PendingRun::RunForNotification(
protocol::InvokeKind::WorkerEvent,
));
pending = Some(PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::WorkerEvent,
notification_request_id: None,
});
}
}
}
@@ -1788,12 +1878,12 @@ async fn handle_inbound_worker_event(
/// as `Errored` — only the worker-execution `Err` branch below fires.
///
/// `parent_originated` further restricts both upward reports to turns
/// the parent actually delegated (`Method::Run` / `Method::Resume`).
/// the parent actually delegated (`Method::Submit` / `Method::Resume`).
/// `Method::Notify` / inbound `WorkerEvent` auto-kicks complete silently
/// so the parent's history does not get flooded with child-internal
/// turn boundaries.
#[allow(clippy::too_many_arguments)]
async fn drive_turn<F>(
async fn drive_turn<F, St>(
worker_future: F,
method_rx: &mut mpsc::Receiver<Method>,
working_event_tx: &broadcast::Sender<Event>,
@@ -1801,15 +1891,17 @@ async fn drive_turn<F>(
pause_tx: &mpsc::Sender<()>,
shared_state: &Arc<WorkerSharedState>,
runtime_dir: &RuntimeDir,
mut input_commit_rx: Option<oneshot::Receiver<()>>,
mut input_commit: Option<(oneshot::Receiver<()>, crate::worker::PendingSubmission)>,
notify_buffer: &NotifyBuffer,
pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
parent_socket: Option<&PathBuf>,
self_name: &str,
spawned_registry: &Arc<SpawnedWorkerRegistry>,
parent_originated: bool,
) -> (WorkerStatus, bool)
) -> (WorkerStatus, bool, bool)
where
F: std::future::Future<Output = Result<WorkerRunResult, WorkerError>>,
St: Store + Clone,
{
tokio::pin!(worker_future);
let mut shutdown_requested = false;
@@ -1822,13 +1914,25 @@ where
// Running snapshot contract deterministic even for immediate clients.
biased;
committed = async {
input_commit_rx
input_commit
.as_mut()
.map(|(receiver, _)| receiver)
.expect("input commit receiver guarded by select condition")
.await
}, if input_commit_rx.is_some() => {
input_commit_rx = None;
}, if input_commit.is_some() => {
let submission = input_commit.take().map(|(_, submission)| submission);
if committed.is_ok() {
if let Some(submission) = submission {
pending_submissions.finish_activation(&submission.submission_id);
let _ = working_event_tx.send(Event::SubmissionAccepted {
submission_request_id: submission.submission_request_id,
submission_id: submission.submission_id,
disposition: protocol::SubmissionDisposition::Started,
});
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_submissions.snapshot(),
});
}
set_controller_status(
shared_state,
runtime_dir,
@@ -1836,11 +1940,33 @@ where
WorkerStatus::Running,
)
.await;
} else if let Some(submission) = submission {
pending_submissions.abort_activation(submission);
}
}
result = &mut worker_future => {
if let Some((mut receiver, submission)) = input_commit.take() {
match receiver.try_recv() {
Ok(()) => {
pending_submissions.finish_activation(&submission.submission_id);
let _ = working_event_tx.send(Event::SubmissionAccepted {
submission_request_id: submission.submission_request_id,
submission_id: submission.submission_id,
disposition: protocol::SubmissionDisposition::Started,
});
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_submissions.snapshot(),
});
}
Err(_) => pending_submissions.abort_activation(submission),
}
}
return match result {
Ok(r) => {
let may_drain_pending = matches!(
&r,
WorkerRunResult::Finished | WorkerRunResult::LimitReached
);
let (status, run_result) = match r {
WorkerRunResult::Finished if pause_requested => {
(WorkerStatus::Paused, RunResult::Paused)
@@ -1851,7 +1977,7 @@ where
WorkerRunResult::RolledBack => (WorkerStatus::Idle, RunResult::RolledBack),
WorkerRunResult::Interrupted { .. } if pause_requested => {
let _ = working_event_tx.send(Event::RunEnd { result: RunResult::Paused });
return (WorkerStatus::Paused, shutdown_requested);
return (WorkerStatus::Paused, shutdown_requested, false);
}
WorkerRunResult::Interrupted { code, message } => {
let _ = working_event_tx.send(Event::Error {
@@ -1867,7 +1993,7 @@ where
},
);
}
return (WorkerStatus::Idle, shutdown_requested);
return (WorkerStatus::Idle, shutdown_requested, false);
}
};
let _ = working_event_tx.send(Event::RunEnd { result: run_result });
@@ -1879,7 +2005,7 @@ where
},
);
}
(status, shutdown_requested)
(status, shutdown_requested, may_drain_pending)
}
Err(WorkerError::Engine(EngineError::Cancelled)) if pause_requested => {
// User-initiated Pause. Report the transition to
@@ -1888,7 +2014,7 @@ where
// that channel is reserved for worker runtime
// failures, not deliberate interruptions.
let _ = working_event_tx.send(Event::RunEnd { result: RunResult::Paused });
(WorkerStatus::Paused, shutdown_requested)
(WorkerStatus::Paused, shutdown_requested, false)
}
Err(e) => {
let code = worker_error_code(&e);
@@ -1906,11 +2032,11 @@ where
},
);
}
(WorkerStatus::Idle, shutdown_requested)
(WorkerStatus::Idle, shutdown_requested, false)
}
};
}
method = method_rx.recv() => {
method = method_rx.recv(), if input_commit.is_none() => {
match method {
Some(Method::Cancel) => {
let _ = cancel_tx.try_send(());
@@ -1923,12 +2049,71 @@ where
shutdown_requested = true;
let _ = cancel_tx.try_send(());
}
Some(Method::Run { .. } | Method::RunTracked { .. } | Method::Resume) => {
Some(Method::Submit {
submission_request_id,
input,
}
| Method::SubmitTracked {
submission_request_id,
input,
}) => {
let request_id = submission_request_id.clone();
match pending_submissions.accept(submission_request_id, input, false) {
Ok(acceptance) => {
let _ = working_event_tx.send(Event::SubmissionAccepted {
submission_request_id: acceptance.submission_request_id,
submission_id: acceptance.submission_id,
disposition: acceptance.disposition,
});
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_submissions.snapshot(),
});
}
Err(error) => {
let _ = working_event_tx.send(Event::SubmissionRejected {
submission_request_id: request_id,
message: error.to_string(),
});
}
}
}
Some(Method::Resume | Method::ContinuePending) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn".into(),
});
}
Some(Method::ListPendingSubmissions) => {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_submissions.snapshot(),
});
}
Some(Method::CancelPendingSubmission { submission_id }) => {
match pending_submissions.cancel(&submission_id) {
Ok(pending) => {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged { pending });
}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
}
}
}
Some(Method::ClearPendingSubmissions) => {
match pending_submissions.clear() {
Ok(pending) => {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged { pending });
}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: error.to_string(),
});
}
}
}
Some(Method::Compact | Method::ListRewindTargets | Method::RewindTo { .. }) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
@@ -1936,11 +2121,28 @@ where
.into(),
});
}
Some(Method::Notify { message, auto_run }) => {
// Live echo arrives via `Event::SystemItem` once
// the in-flight turn's next `pending_history_appends`
// drains this entry through the interceptor.
notify_buffer.push_notify(message, auto_run);
Some(Method::Notify {
notification_request_id,
message,
auto_run,
}) => {
if auto_run {
if let Err(error) = pending_submissions.accept_notification(
notification_request_id,
message,
) {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
} else {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_submissions.snapshot(),
});
}
} else {
notify_buffer.push_notify(message, false);
}
}
Some(Method::ListCompletions { .. }) => {}
Some(Method::ListWorkers | Method::RestoreWorker { .. } | Method::RegisterPeer { .. }) => {
@@ -1969,7 +2171,7 @@ where
None => {
let _ = cancel_tx.try_send(());
shared_state.set_status(WorkerStatus::Idle);
return (WorkerStatus::Idle, false);
return (WorkerStatus::Idle, false, false);
}
}
}
@@ -2134,21 +2336,16 @@ mod tests {
#[test]
fn pending_run_parent_origin_table() {
assert!(PendingRun::Run(Vec::new()).is_parent_originated());
assert!(PendingRun::Resume.is_parent_originated());
assert!(
!PendingRun::RunForNotification(protocol::InvokeKind::Notify).is_parent_originated()
!PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::Notify,
notification_request_id: None,
}
.is_parent_originated()
);
}
#[test]
fn notification_auto_run_gate_only_allows_idle_auto_run() {
assert!(should_auto_run_notification(WorkerStatus::Idle, true));
assert!(!should_auto_run_notification(WorkerStatus::Idle, false));
assert!(!should_auto_run_notification(WorkerStatus::Running, true));
assert!(!should_auto_run_notification(WorkerStatus::Paused, true));
}
struct DriveTurnEnv {
// Held to keep the channel alive; without this `method_rx.recv()`
// would observe channel-closed and confuse the select! arm.
@@ -2161,6 +2358,7 @@ mod tests {
_pause_rx: mpsc::Receiver<()>,
shared_state: Arc<WorkerSharedState>,
notify_buffer: NotifyBuffer,
pending_submissions: crate::worker::PendingSubmissionHandle<session_store::FsStore>,
spawned_registry: Arc<SpawnedWorkerRegistry>,
parent_socket_path: PathBuf,
runtime_dir: Arc<RuntimeDir>,
@@ -2194,6 +2392,8 @@ mod tests {
},
));
let notify_buffer = NotifyBuffer::new();
let pending_submissions =
crate::worker::PendingSubmissionHandle::for_test(&temp.path().join("pending-sessions"));
let spawned_registry = SpawnedWorkerRegistry::new(runtime_dir.clone());
let parent_socket_path = temp.path().join("parent.sock");
@@ -2207,6 +2407,7 @@ mod tests {
_pause_rx: pause_rx,
shared_state,
notify_buffer,
pending_submissions,
spawned_registry,
parent_socket_path,
runtime_dir,
@@ -2225,6 +2426,7 @@ mod tests {
writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -2259,7 +2461,7 @@ mod tests {
let recv = tokio::spawn(recv_worker_event(listener, Duration::from_secs(2)));
let worker_future = async { Ok::<_, WorkerError>(WorkerRunResult::Finished) };
let (status, shutdown) = drive_turn(
let (status, shutdown, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.working_event_tx,
@@ -2269,6 +2471,7 @@ mod tests {
&env.runtime_dir,
None,
&env.notify_buffer,
&env.pending_submissions,
Some(&env.parent_socket_path),
"child-worker",
&env.spawned_registry,
@@ -2302,7 +2505,7 @@ mod tests {
Ok::<_, WorkerError>(WorkerRunResult::Finished)
};
let started_at = std::time::Instant::now();
let (status, shutdown) = drive_turn(
let (status, shutdown, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.working_event_tx,
@@ -2312,6 +2515,7 @@ mod tests {
&env.runtime_dir,
None,
&env.notify_buffer,
&env.pending_submissions,
None,
"child-worker",
&env.spawned_registry,
@@ -2332,7 +2536,7 @@ mod tests {
let listener = UnixListener::bind(&env.parent_socket_path).expect("bind listener");
let worker_future = async { Ok::<_, WorkerError>(WorkerRunResult::Finished) };
let (status, _) = drive_turn(
let (status, _, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.working_event_tx,
@@ -2342,6 +2546,7 @@ mod tests {
&env.runtime_dir,
None,
&env.notify_buffer,
&env.pending_submissions,
Some(&env.parent_socket_path),
"child-worker",
&env.spawned_registry,
@@ -2370,7 +2575,7 @@ mod tests {
"boom from test".into(),
)))
};
let (status, _) = drive_turn(
let (status, _, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.working_event_tx,
@@ -2380,6 +2585,7 @@ mod tests {
&env.runtime_dir,
None,
&env.notify_buffer,
&env.pending_submissions,
Some(&env.parent_socket_path),
"child-worker",
&env.spawned_registry,
@@ -2414,7 +2620,7 @@ mod tests {
"boom from notify".into(),
)))
};
let (status, _) = drive_turn(
let (status, _, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.working_event_tx,
@@ -2424,6 +2630,7 @@ mod tests {
&env.runtime_dir,
None,
&env.notify_buffer,
&env.pending_submissions,
Some(&env.parent_socket_path),
"child-worker",
&env.spawned_registry,
@@ -2456,7 +2663,7 @@ mod tests {
tokio::time::sleep(Duration::from_millis(50)).await;
Ok::<_, WorkerError>(WorkerRunResult::Finished)
};
let (status, shutdown) = drive_turn(
let (status, shutdown, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.working_event_tx,
@@ -2466,6 +2673,7 @@ mod tests {
&env.runtime_dir,
None,
&env.notify_buffer,
&env.pending_submissions,
Some(&env.parent_socket_path),
"parent",
&env.spawned_registry,
@@ -2495,7 +2703,7 @@ mod tests {
tokio::time::sleep(Duration::from_millis(50)).await;
Ok::<_, WorkerError>(WorkerRunResult::Finished)
};
let (status, shutdown) = drive_turn(
let (status, shutdown, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.working_event_tx,
@@ -2505,6 +2713,7 @@ mod tests {
&env.runtime_dir,
None,
&env.notify_buffer,
&env.pending_submissions,
Some(&env.parent_socket_path),
"parent",
&env.spawned_registry,
@@ -2522,6 +2731,7 @@ mod tests {
let mut env = make_env().await;
env._method_tx
.send(Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message: "continue".into(),
auto_run: true,
})
@@ -2532,7 +2742,7 @@ mod tests {
tokio::time::sleep(Duration::from_millis(50)).await;
Ok::<_, WorkerError>(WorkerRunResult::Finished)
};
let (status, shutdown) = drive_turn(
let (status, shutdown, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.working_event_tx,
@@ -2542,6 +2752,7 @@ mod tests {
&env.runtime_dir,
None,
&env.notify_buffer,
&env.pending_submissions,
Some(&env.parent_socket_path),
"parent",
&env.spawned_registry,
@@ -2551,8 +2762,8 @@ mod tests {
assert_eq!(status, WorkerStatus::Idle);
assert!(!shutdown);
assert_eq!(env.notify_buffer.len(), 1);
assert!(env.notify_buffer.has_auto_run_pending());
assert_eq!(env.notify_buffer.len(), 0);
assert_eq!(env.pending_submissions.snapshot().notification_count, 1);
}
#[tokio::test]
@@ -2568,7 +2779,7 @@ mod tests {
tokio::time::sleep(Duration::from_millis(50)).await;
Ok::<_, WorkerError>(WorkerRunResult::Finished)
};
let (status, shutdown) = drive_turn(
let (status, shutdown, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.working_event_tx,
@@ -2578,6 +2789,7 @@ mod tests {
&env.runtime_dir,
None,
&env.notify_buffer,
&env.pending_submissions,
Some(&env.parent_socket_path),
"child-worker",
&env.spawned_registry,
+24 -3
View File
@@ -1012,7 +1012,15 @@ async fn send_peer_notify(socket_path: &Path, message: String) -> io::Result<()>
}
async fn send_notify(socket_path: &Path, message: String, auto_run: bool) -> io::Result<()> {
connect_and_send(socket_path, &Method::Notify { message, auto_run }).await
connect_and_send(
socket_path,
&Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message,
auto_run,
},
)
.await
}
fn json_content<T: Serialize>(value: &T) -> Result<String, ToolError> {
@@ -1482,6 +1490,7 @@ mod tests {
writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -1517,6 +1526,7 @@ mod tests {
writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -1536,7 +1546,10 @@ mod tests {
.await
.unwrap();
let method = reader.next::<Method>().await.unwrap().unwrap();
if let Method::Notify { message, auto_run } = method {
if let Method::Notify {
message, auto_run, ..
} = method
{
assert!(auto_run);
tx.send(message).await.unwrap();
} else {
@@ -1608,6 +1621,7 @@ mod tests {
writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -1634,6 +1648,7 @@ mod tests {
writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -1653,7 +1668,10 @@ mod tests {
.await
.unwrap();
let method = reader.next::<Method>().await.unwrap().unwrap();
if let Method::Notify { message, auto_run } = method {
if let Method::Notify {
message, auto_run, ..
} = method
{
assert!(!auto_run);
tx.send(message).await.unwrap();
} else {
@@ -1738,6 +1756,7 @@ mod tests {
writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -1790,6 +1809,8 @@ mod tests {
let _ = writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(
),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -803,7 +803,10 @@ mod tests {
.collect();
Ok(WorkerSessionCapture {
segment_id: "segment".to_string(),
session: protocol::SessionSnapshot { entries },
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries,
},
})
}
}
+21 -6
View File
@@ -176,12 +176,16 @@ impl WorkerInterceptor {
/// `Item::system_message`s reach the worker via
/// `ContinueWith` / `pending_history_appends`, so on-disk order
/// matches worker-history order.
fn commit_system_items(&self, items: &[SystemItem]) -> Result<(), session_store::StoreError> {
fn commit_system_items_with_extensions(
&self,
items: &[(SystemItem, Vec<session_store::SessionExtension>)],
) -> Result<(), session_store::StoreError> {
let Some(writer) = self.log_writer.as_ref() else {
return Ok(());
};
for item in items {
let entry = writer.commit_system_item(item.clone())?;
for (item, extensions) in items {
let entry =
writer.commit_system_item_with_extensions(item.clone(), extensions.clone())?;
self.pending_committed_history
.lock()
.expect("pending committed history poisoned")
@@ -190,6 +194,16 @@ impl WorkerInterceptor {
Ok(())
}
fn commit_system_items(&self, items: &[SystemItem]) -> Result<(), session_store::StoreError> {
self.commit_system_items_with_extensions(
&items
.iter()
.cloned()
.map(|item| (item, Vec::new()))
.collect::<Vec<_>>(),
)
}
fn current_turn_index(&self) -> usize {
self.next_turn_index
.load(Ordering::Relaxed)
@@ -327,7 +341,8 @@ impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
projection_digest: projection.catalog_digest.clone(),
logical_name: "internal.notify_wrapper".to_string(),
};
let mut system_items: Vec<SystemItem> = Vec::with_capacity(drained.len());
let mut system_items: Vec<(SystemItem, Vec<session_store::SessionExtension>)> =
Vec::with_capacity(drained.len());
let mut items: Vec<Item> = Vec::with_capacity(drained.len());
for entry in &drained {
let system_item = match build_system_item_with_provenance(
@@ -345,9 +360,9 @@ impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
}
};
items.push(system_item.to_history_item());
system_items.push(system_item);
system_items.push((system_item, entry.extensions()));
}
if let Err(error) = self.commit_system_items(&system_items) {
if let Err(error) = self.commit_system_items_with_extensions(&system_items) {
self.pending_notifies.requeue_front(drained);
return Err(InterceptorError::new(
InterceptorErrorCategory::Dependency,
+32 -4
View File
@@ -25,7 +25,7 @@ use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use protocol::WorkerEvent;
use session_store::SystemItem;
use session_store::{SessionExtension, SystemItem};
use tracing::warn;
use crate::prompt::catalog::{CatalogError, PromptCatalog};
@@ -41,8 +41,23 @@ const CAPACITY: usize = 128;
/// is available.
#[derive(Debug, Clone)]
pub enum PendingNotify {
Notify { message: String, auto_run: bool },
WorkerEvent { event: WorkerEvent },
Notify {
message: String,
auto_run: bool,
extensions: Vec<SessionExtension>,
},
WorkerEvent {
event: WorkerEvent,
},
}
impl PendingNotify {
pub(crate) fn extensions(&self) -> Vec<SessionExtension> {
match self {
PendingNotify::Notify { extensions, .. } => extensions.clone(),
PendingNotify::WorkerEvent { .. } => Vec::new(),
}
}
}
/// Shared, mutex-guarded buffer of pending entries.
@@ -62,7 +77,19 @@ impl NotifyBuffer {
/// oldest entry is dropped and a `tracing::warn` is emitted — the
/// caller should never hit this in normal operation.
pub fn push_notify(&self, message: String, auto_run: bool) {
self.push_entry(PendingNotify::Notify { message, auto_run });
self.push_entry(PendingNotify::Notify {
message,
auto_run,
extensions: Vec::new(),
});
}
pub fn push_durable_notify(&self, message: String, extension: SessionExtension) {
self.push_entry(PendingNotify::Notify {
message,
auto_run: true,
extensions: vec![extension],
});
}
/// Push a typed worker-event entry onto the queue.
@@ -202,6 +229,7 @@ mod tests {
let entry = PendingNotify::Notify {
message: "hello".into(),
auto_run: false,
extensions: Vec::new(),
};
let catalog = PromptCatalog::builtins_only().unwrap();
let item = build_system_item(&entry, &catalog).unwrap();
+5 -5
View File
@@ -57,9 +57,9 @@ pub use session_history::{
};
pub use shared_state::WorkerSharedState;
pub use worker::{
LocalWorkingDirectory, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError,
WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceClient,
WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspacePromptCatalogResolution,
WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, WorkspaceWorkerDiscoveryRequest,
apply_worker_manifest, marker_workspace_client, unavailable_workspace_client,
LocalWorkingDirectory, Worker, WorkerError, WorkerFilesystemAuthority, WorkerRunResult,
WorkerWorkspaceContext, WorkspaceClient, WorkspaceClientError, WorkspaceId, WorkspaceIdError,
WorkspacePromptCatalogResolution, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse,
WorkspaceWorkerDiscoveryRequest, apply_worker_manifest, marker_workspace_client,
unavailable_workspace_client,
};
+1
View File
@@ -291,6 +291,7 @@ mod tests {
prompt_provenance: None,
},
),
extensions: Vec::new(),
}
}
+1
View File
@@ -72,6 +72,7 @@ mod tests {
fn snapshot(entries: Vec<serde_json::Value>) -> Event {
Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: entries
.into_iter()
.enumerate()
+7 -2
View File
@@ -58,7 +58,7 @@ struct SubWorkerSpawnInput {
/// a host path and grants no authority. When omitted, the Workdir root is used.
#[serde(default)]
cwd: Option<String>,
/// First message sent to the spawned SubWorker via `Method::Run`.
/// First message sent to the spawned SubWorker via `Method::Submit`.
task: String,
/// Allow rules delegated to the spawned SubWorker. Must be a subset of the
/// spawner's explicit delegation authority; direct tool scope alone is not
@@ -235,7 +235,11 @@ impl ParentNotificationTarget {
};
tokio::spawn(async move {
if let Err(error) = parent_method_tx
.send(Method::Notify { message, auto_run })
.send(Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message,
auto_run,
})
.await
{
tracing::warn!(
@@ -1267,6 +1271,7 @@ enabled = false
Method::Notify {
message,
auto_run: true,
..
} if message.contains("SubWorker `reviewer-child` turn ended with status Idle")
));
assert!(!runtime.path().join("reviewer-child/sock").exists());
+851 -6
View File
@@ -1,3 +1,4 @@
use std::collections::VecDeque;
#[cfg(test)]
use std::path::Path;
use std::path::PathBuf;
@@ -68,6 +69,130 @@ const LARGE_PASTE_INLINE_MAX_BYTES: usize = 32 * 1024;
const WORKER_ORCHESTRATION_INSTRUCTION_ID: &str = "worker.orchestration";
const WORKER_ORCHESTRATION_PROMPT_REF: &str = "common.worker_orchestration";
const FEATURE_HOOK_CHAIN_TIMEOUT: Duration = Duration::from_secs(30);
const SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN: &str = "worker.pending_activations.v1";
const MAX_PENDING_SUBMISSIONS: usize = 32;
const MAX_PENDING_SUBMISSION_BYTES: u64 = 1024 * 1024;
const MAX_PENDING_ARTIFACT_REFS: usize = 64;
const MAX_ACTIVATION_REQUEST_ID_BYTES: usize = 128;
const MAX_SUBMISSION_RECEIPTS: usize = 128;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub(crate) struct PendingSubmission {
pub(crate) submission_request_id: String,
pub(crate) submission_id: String,
payload_digest: String,
accepted_at_ms: u64,
activation_sequence: u64,
provenance: WorkerHistoryProvenance,
#[serde(default)]
was_queued: bool,
pub(crate) input: Vec<Segment>,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
struct SubmissionReceipt {
submission_request_id: String,
submission_id: String,
payload_digest: String,
disposition: protocol::SubmissionDisposition,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub(crate) struct PendingNotification {
pub(crate) notification_request_id: String,
pub(crate) message: String,
payload_digest: String,
accepted_at_ms: u64,
activation_sequence: u64,
provenance: WorkerHistoryProvenance,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
struct NotificationReceipt {
notification_request_id: String,
payload_digest: String,
}
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub(crate) struct PendingActivationState {
revision: u64,
next_activation_sequence: u64,
/// A prepared activation remains in checkpoints until the same atomic
/// UserInput record commits the clearing checkpoint. Restore puts it back
/// at the FIFO head.
activating: Option<PendingSubmission>,
activating_notification: Option<PendingNotification>,
pending: VecDeque<PendingSubmission>,
pending_notifications: VecDeque<PendingNotification>,
receipts: VecDeque<SubmissionReceipt>,
notification_receipts: VecDeque<NotificationReceipt>,
}
impl PendingActivationState {
pub(crate) fn snapshot(&self) -> protocol::PendingSubmissionsSnapshot {
protocol::PendingSubmissionsSnapshot {
revision: self.revision,
notification_count: u32::try_from(self.pending_notifications.len()).unwrap_or(u32::MAX),
submissions: self
.pending
.iter()
.map(|pending| protocol::PendingSubmissionSummary {
submission_id: pending.submission_id.clone(),
accepted_at_ms: pending.accepted_at_ms,
segment_count: u32::try_from(pending.input.len()).unwrap_or(u32::MAX),
byte_len: submission_payload_len(&pending.input),
})
.collect(),
}
}
fn remember_notification_receipt(&mut self, receipt: NotificationReceipt) {
self.notification_receipts.push_back(receipt);
while self.notification_receipts.len() > MAX_SUBMISSION_RECEIPTS {
self.notification_receipts.pop_front();
}
}
fn remember_receipt(&mut self, receipt: SubmissionReceipt) {
self.receipts.push_back(receipt);
while self.receipts.len() > MAX_SUBMISSION_RECEIPTS {
self.receipts.pop_front();
}
}
}
fn submission_payload_len(input: &[Segment]) -> u64 {
serde_json::to_vec(input)
.map(|bytes| u64::try_from(bytes.len()).unwrap_or(u64::MAX))
.unwrap_or(u64::MAX)
}
fn submission_payload_digest(input: &[Segment]) -> String {
use sha2::Digest as _;
sha2::Sha256::digest(serde_json::to_vec(input).unwrap_or_default())
.iter()
.map(|byte| format!("{byte:02x}"))
.collect()
}
fn submission_artifact_ref_count(input: &[Segment]) -> usize {
input
.iter()
.filter(|segment| {
matches!(
segment,
Segment::PasteArtifact { .. } | Segment::UploadedFile { .. }
)
})
.count()
}
fn pending_activation_extension(state: &PendingActivationState) -> SessionExtension {
SessionExtension {
domain: SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.to_owned(),
payload: serde_json::to_value(state).expect("pending activation state must serialize"),
}
}
fn hook_run_exit(exit: &EngineRunExit) -> RunCommittedExit {
match exit {
@@ -970,15 +1095,489 @@ where
}
}
/// Type-erased commit handle for the interceptor. Lets the
/// interceptor commit `SystemItem`s without being generic over the
#[derive(Debug, Clone)]
pub(crate) enum PendingActivation {
Submission(PendingSubmission),
Notification(PendingNotification),
}
#[derive(Debug, Clone)]
pub(crate) struct SubmissionAcceptance {
pub(crate) submission_request_id: String,
pub(crate) submission_id: String,
pub(crate) disposition: protocol::SubmissionDisposition,
pub(crate) activation: Option<PendingSubmission>,
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum PendingSubmissionError {
#[error("submission_request_id must not be empty")]
EmptyRequestId,
#[error("activation request id exceeds {MAX_ACTIVATION_REQUEST_ID_BYTES} bytes")]
RequestIdLimit,
#[error("submission input must contain at least one typed segment")]
EmptyInput,
#[error("submission request id was already used with a different payload")]
IdempotencyConflict,
#[error("pending submission queue is full (maximum {MAX_PENDING_SUBMISSIONS})")]
CountLimit,
#[error("pending submission bytes exceed {MAX_PENDING_SUBMISSION_BYTES}")]
ByteLimit,
#[error("pending submission artifact references exceed {MAX_PENDING_ARTIFACT_REFS}")]
ArtifactLimit,
#[error("pending submission not found: {0}")]
NotFound(String),
#[error("pending submission state persistence failed: {0}")]
Store(#[from] StoreError),
}
#[derive(Clone)]
pub(crate) struct PendingSubmissionHandle<St: Clone> {
state: Arc<Mutex<PendingActivationState>>,
writer: LogWriterHandle<St>,
}
impl<St> PendingSubmissionHandle<St>
where
St: Store + Clone,
{
fn persist_locked(&self, state: &PendingActivationState) -> Result<(), PendingSubmissionError> {
self.writer.append_entry_locked(LogEntry::Extension {
ts: segment_log::now_millis(),
domain: SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.to_owned(),
payload: serde_json::to_value(state).expect("pending activation state must serialize"),
})?;
Ok(())
}
pub(crate) fn accept(
&self,
submission_request_id: String,
input: Vec<Segment>,
activate_now: bool,
) -> Result<SubmissionAcceptance, PendingSubmissionError> {
if submission_request_id.trim().is_empty() {
return Err(PendingSubmissionError::EmptyRequestId);
}
if submission_request_id.len() > MAX_ACTIVATION_REQUEST_ID_BYTES {
return Err(PendingSubmissionError::RequestIdLimit);
}
if input.is_empty() {
return Err(PendingSubmissionError::EmptyInput);
}
let payload_digest = submission_payload_digest(&input);
let _append_guard = self
.writer
.state
.append_lock
.lock()
.expect("segment append lock poisoned");
let mut current = self
.state
.lock()
.expect("pending activation state poisoned");
let original = current.clone();
if let Some(receipt) = current
.receipts
.iter()
.find(|receipt| receipt.submission_request_id == submission_request_id)
{
if receipt.payload_digest != payload_digest {
return Err(PendingSubmissionError::IdempotencyConflict);
}
return Ok(SubmissionAcceptance {
submission_request_id,
submission_id: receipt.submission_id.clone(),
disposition: receipt.disposition,
activation: None,
});
}
let submission_id = uuid::Uuid::now_v7().to_string();
let pending = PendingSubmission {
submission_request_id: submission_request_id.clone(),
submission_id: submission_id.clone(),
payload_digest: payload_digest.clone(),
accepted_at_ms: segment_log::now_millis(),
activation_sequence: current.next_activation_sequence,
provenance: WorkerHistoryProvenance::LegacyUnknown,
was_queued: !activate_now,
input,
};
current.next_activation_sequence = current.next_activation_sequence.saturating_add(1);
let disposition = if activate_now {
protocol::SubmissionDisposition::Started
} else {
protocol::SubmissionDisposition::Queued
};
current.remember_receipt(SubmissionReceipt {
submission_request_id: submission_request_id.clone(),
submission_id: submission_id.clone(),
payload_digest,
disposition,
});
current.revision = current.revision.saturating_add(1);
if activate_now {
current.activating = Some(pending.clone());
} else {
let count = current
.pending
.len()
.saturating_add(current.pending_notifications.len())
.saturating_add(1);
if count > MAX_PENDING_SUBMISSIONS {
*current = original;
return Err(PendingSubmissionError::CountLimit);
}
let bytes = current
.pending
.iter()
.map(|pending| submission_payload_len(&pending.input))
.sum::<u64>()
.saturating_add(
current
.pending_notifications
.iter()
.map(|pending| u64::try_from(pending.message.len()).unwrap_or(u64::MAX))
.sum::<u64>(),
)
.saturating_add(submission_payload_len(&pending.input));
if bytes > MAX_PENDING_SUBMISSION_BYTES {
*current = original;
return Err(PendingSubmissionError::ByteLimit);
}
let artifact_refs = current
.pending
.iter()
.map(|pending| submission_artifact_ref_count(&pending.input))
.sum::<usize>()
.saturating_add(submission_artifact_ref_count(&pending.input));
if artifact_refs > MAX_PENDING_ARTIFACT_REFS {
*current = original;
return Err(PendingSubmissionError::ArtifactLimit);
}
current.pending.push_back(pending.clone());
if let Err(error) = self.persist_locked(&current) {
*current = original;
return Err(error);
}
}
Ok(SubmissionAcceptance {
submission_request_id,
submission_id,
disposition,
activation: activate_now.then_some(pending),
})
}
pub(crate) fn accept_notification(
&self,
notification_request_id: String,
message: String,
) -> Result<bool, PendingSubmissionError> {
if notification_request_id.trim().is_empty() {
return Err(PendingSubmissionError::EmptyRequestId);
}
if notification_request_id.len() > MAX_ACTIVATION_REQUEST_ID_BYTES {
return Err(PendingSubmissionError::RequestIdLimit);
}
let payload_digest = submission_payload_digest(&[Segment::text(message.clone())]);
let _append_guard = self
.writer
.state
.append_lock
.lock()
.expect("segment append lock poisoned");
let mut state = self
.state
.lock()
.expect("pending activation state poisoned");
if let Some(receipt) = state
.notification_receipts
.iter()
.find(|receipt| receipt.notification_request_id == notification_request_id)
{
if receipt.payload_digest != payload_digest {
return Err(PendingSubmissionError::IdempotencyConflict);
}
return Ok(false);
}
if state
.pending
.len()
.saturating_add(state.pending_notifications.len())
>= MAX_PENDING_SUBMISSIONS
{
return Err(PendingSubmissionError::CountLimit);
}
let queued_bytes = state
.pending
.iter()
.map(|pending| submission_payload_len(&pending.input))
.sum::<u64>()
.saturating_add(
state
.pending_notifications
.iter()
.map(|pending| u64::try_from(pending.message.len()).unwrap_or(u64::MAX))
.sum::<u64>(),
)
.saturating_add(u64::try_from(message.len()).unwrap_or(u64::MAX));
if queued_bytes > MAX_PENDING_SUBMISSION_BYTES {
return Err(PendingSubmissionError::ByteLimit);
}
let original = state.clone();
let activation_sequence = state.next_activation_sequence;
state.next_activation_sequence = state.next_activation_sequence.saturating_add(1);
state.pending_notifications.push_back(PendingNotification {
notification_request_id: notification_request_id.clone(),
message,
payload_digest: payload_digest.clone(),
accepted_at_ms: segment_log::now_millis(),
activation_sequence,
provenance: WorkerHistoryProvenance::BackendInstruction {
operation_id: Some(notification_request_id.clone()),
},
});
state.remember_notification_receipt(NotificationReceipt {
notification_request_id,
payload_digest,
});
state.revision = state.revision.saturating_add(1);
if let Err(error) = self.persist_locked(&state) {
*state = original;
return Err(error);
}
Ok(true)
}
pub(crate) fn prepare_next_activation(
&self,
) -> Result<Option<PendingActivation>, PendingSubmissionError> {
let _append_guard = self
.writer
.state
.append_lock
.lock()
.expect("segment append lock poisoned");
let mut state = self
.state
.lock()
.expect("pending activation state poisoned");
if state.activating.is_some() || state.activating_notification.is_some() {
return Ok(None);
}
let submission_sequence = state.pending.front().map(|item| item.activation_sequence);
let notification_sequence = state
.pending_notifications
.front()
.map(|item| item.activation_sequence);
if notification_sequence.is_some()
&& (submission_sequence.is_none() || notification_sequence < submission_sequence)
{
let notification = state
.pending_notifications
.pop_front()
.expect("notification sequence came from queue head");
state.activating_notification = Some(notification.clone());
state.revision = state.revision.saturating_add(1);
return Ok(Some(PendingActivation::Notification(notification)));
}
if submission_sequence.is_some() {
let pending = state
.pending
.pop_front()
.expect("submission sequence came from queue head");
state.activating = Some(pending.clone());
state.revision = state.revision.saturating_add(1);
return Ok(Some(PendingActivation::Submission(pending)));
}
Ok(None)
}
pub(crate) fn abort_activation(&self, pending: PendingSubmission) {
let mut state = self
.state
.lock()
.expect("pending activation state poisoned");
if state.activating.as_ref().map(|item| &item.submission_id) != Some(&pending.submission_id)
{
return;
}
state.activating = None;
if pending.was_queued {
state.pending.push_front(pending);
} else {
state
.receipts
.retain(|receipt| receipt.submission_id != pending.submission_id);
}
state.revision = state.revision.saturating_add(1);
}
pub(crate) fn activation_extension(&self) -> SessionExtension {
let state = self
.state
.lock()
.expect("pending activation state poisoned");
let mut committed = state.clone();
if let Some(activating) = &committed.activating
&& let Some(receipt) = committed
.receipts
.iter_mut()
.find(|receipt| receipt.submission_id == activating.submission_id)
{
receipt.disposition = protocol::SubmissionDisposition::Started;
}
committed.activating = None;
committed.revision = committed.revision.saturating_add(1);
pending_activation_extension(&committed)
}
pub(crate) fn finish_activation(&self, submission_id: &str) {
let mut state = self
.state
.lock()
.expect("pending activation state poisoned");
if state
.activating
.as_ref()
.map(|item| item.submission_id.as_str())
== Some(submission_id)
{
if let Some(receipt) = state
.receipts
.iter_mut()
.find(|receipt| receipt.submission_id == submission_id)
{
receipt.disposition = protocol::SubmissionDisposition::Started;
}
state.activating = None;
state.revision = state.revision.saturating_add(1);
}
}
pub(crate) fn notification_activation_extension(&self) -> SessionExtension {
let state = self
.state
.lock()
.expect("pending activation state poisoned");
let mut committed = state.clone();
committed.activating_notification = None;
committed.revision = committed.revision.saturating_add(1);
pending_activation_extension(&committed)
}
pub(crate) fn finish_notification_activation(&self, notification_request_id: &str) {
let mut state = self
.state
.lock()
.expect("pending activation state poisoned");
if state
.activating_notification
.as_ref()
.map(|item| item.notification_request_id.as_str())
== Some(notification_request_id)
{
state.activating_notification = None;
state.revision = state.revision.saturating_add(1);
}
}
pub(crate) fn snapshot(&self) -> protocol::PendingSubmissionsSnapshot {
self.state
.lock()
.expect("pending activation state poisoned")
.snapshot()
}
pub(crate) fn cancel(
&self,
submission_id: &str,
) -> Result<protocol::PendingSubmissionsSnapshot, PendingSubmissionError> {
let _append_guard = self
.writer
.state
.append_lock
.lock()
.expect("segment append lock poisoned");
let mut state = self
.state
.lock()
.expect("pending activation state poisoned");
let original = state.clone();
let Some(index) = state
.pending
.iter()
.position(|pending| pending.submission_id == submission_id)
else {
return Err(PendingSubmissionError::NotFound(submission_id.to_owned()));
};
state.pending.remove(index);
state.revision = state.revision.saturating_add(1);
if let Err(error) = self.persist_locked(&state) {
*state = original;
return Err(error);
}
Ok(state.snapshot())
}
pub(crate) fn clear(
&self,
) -> Result<protocol::PendingSubmissionsSnapshot, PendingSubmissionError> {
let _append_guard = self
.writer
.state
.append_lock
.lock()
.expect("segment append lock poisoned");
let mut state = self
.state
.lock()
.expect("pending activation state poisoned");
let original = state.clone();
state.pending.clear();
state.pending_notifications.clear();
state.revision = state.revision.saturating_add(1);
if let Err(error) = self.persist_locked(&state) {
*state = original;
return Err(error);
}
Ok(state.snapshot())
}
}
impl PendingSubmissionHandle<session_store::FsStore> {
#[cfg(test)]
pub(crate) fn for_test(root: &std::path::Path) -> Self {
let store = session_store::FsStore::new(root).expect("test session store");
let session_id = session_store::new_session_id();
let segment_id = session_store::new_segment_id();
store
.create_segment(session_id, segment_id, &[])
.expect("test session segment");
Self {
state: Arc::new(Mutex::new(PendingActivationState::default())),
writer: LogWriterHandle {
store,
state: SegmentState::new(session_id, segment_id, 0),
sink: SegmentLogSink::new(),
in_flight: None,
},
}
}
}
/// Type-erased commit handle for the interceptor. Lets the interceptor commit `SystemItem`s without being generic over the
/// concrete `Store` type.
pub trait SystemItemCommitter: Send + Sync {
fn commit_log_entry(&self, entry: LogEntry) -> Result<(), StoreError>;
fn commit_system_item(
fn commit_system_item_with_extensions(
&self,
item: SystemItem,
extensions: Vec<SessionExtension>,
) -> Result<HistoryEntry<SessionHistoryMetadata>, StoreError> {
let metadata = new_history_metadata(
WorkerHistoryProvenance::BackendInstruction { operation_id: None },
@@ -991,6 +1590,7 @@ pub trait SystemItemCommitter: Send + Sync {
item,
metadata: metadata.clone(),
},
extensions,
})?;
Ok(HistoryEntry::new(history_item, metadata))
}
@@ -1027,8 +1627,6 @@ where
}
}
pub const WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN: &str = "worker.input-submission.v1";
#[derive(Clone)]
struct PreparedFlowProjection {
selector: String,
@@ -1049,6 +1647,7 @@ pub struct WorkerSession {
session_id: SessionId,
revision: u64,
history: History<SessionHistoryMetadata>,
pending_activations: Arc<Mutex<PendingActivationState>>,
}
impl WorkerSession {
@@ -1058,9 +1657,39 @@ impl WorkerSession {
session_id,
revision,
history: History::from_entries(entries),
pending_activations: Arc::new(Mutex::new(PendingActivationState::default())),
}
}
fn restore_pending_activations(&mut self, extensions: &[(String, serde_json::Value)]) {
let Some(payload) = extensions.iter().rev().find_map(|(domain, payload)| {
(domain == SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN).then_some(payload)
}) else {
return;
};
if let Ok(mut state) = serde_json::from_value::<PendingActivationState>(payload.clone()) {
if let Some(activating) = state.activating.take() {
state.pending.push_front(activating);
state.revision = state.revision.saturating_add(1);
}
if let Some(activating) = state.activating_notification.take() {
state.pending_notifications.push_front(activating);
state.revision = state.revision.saturating_add(1);
}
*self
.pending_activations
.lock()
.expect("pending activation state poisoned") = state;
}
}
pub fn pending_submissions(&self) -> protocol::PendingSubmissionsSnapshot {
self.pending_activations
.lock()
.expect("pending activation state poisoned")
.snapshot()
}
pub fn session_id(&self) -> SessionId {
self.session_id
}
@@ -1309,6 +1938,21 @@ impl<C: LlmClient + 'static, St: Store + Clone + 'static> Worker<C, St> {
}
}
pub(crate) fn pending_activation_state(&self) -> Arc<Mutex<PendingActivationState>> {
self.session.pending_activations.clone()
}
pub(crate) fn pending_submission_handle(&self) -> PendingSubmissionHandle<St> {
PendingSubmissionHandle {
state: self.session.pending_activations.clone(),
writer: self.log_writer_handle(),
}
}
pub fn pending_submissions(&self) -> protocol::PendingSubmissionsSnapshot {
self.session.pending_submissions()
}
/// Attach a type-erased system-item commit handle. The controller
/// calls this once during spawn so the interceptor can commit
/// `SystemItem`s directly without owning a generic store handle.
@@ -1670,6 +2314,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
},
metadata: skill_metadata.clone(),
},
extensions: Vec::new(),
})?;
let history_entry = HistoryEntry::new(agen::Item::system_message(body), skill_metadata);
let mut annotate = history_annotator(
@@ -1960,6 +2605,28 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.truncate(loc.session_id, loc.segment_id, truncate_entries)?;
self.segment_state.set_entries_written(truncate_entries);
self.sink.truncate_silent(truncate_entries);
let pending_state = self
.session
.pending_activations
.lock()
.expect("pending activation state poisoned")
.clone();
if !pending_state.pending.is_empty()
|| pending_state.activating.is_some()
|| pending_state.activating_notification.is_some()
|| !pending_state.receipts.is_empty()
{
let checkpoint = LogEntry::Extension {
ts: segment_log::now_millis(),
domain: SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.to_owned(),
payload: serde_json::to_value(&pending_state).map_err(|error| {
RewindError::Invalid(format!(
"serialize pending submissions during rewind: {error}"
))
})?,
};
self.commit_entry(checkpoint)?;
}
let history_entries = restore_history_entries(loc.session_id, loc.segment_id, &retained)
.map_err(|error| RewindError::Invalid(error.into()))?;
@@ -2525,7 +3192,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// Convenience: run with a single `Segment::Text`.
///
/// Equivalent to `run(vec![Segment::text(s)])`. The dumb-client
/// counterpart of [`protocol::Method::run_text`]; primarily for
/// counterpart of [`protocol::Method::submit_text`]; primarily for
/// tests and tools that have only a string in hand.
pub async fn run_text(&mut self, s: impl Into<String>) -> Result<WorkerRunResult, WorkerError>
where
@@ -3042,6 +3709,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
},
metadata: interrupt_metadata.clone(),
},
extensions: Vec::new(),
})?;
let interrupt_entry =
HistoryEntry::new(agen::Item::system_message(system_note), interrupt_metadata);
@@ -4428,6 +5096,22 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
{
initial_entries.push(checkpoint);
}
initial_entries.push(LogEntry::Extension {
ts: segment_log::now_millis(),
domain: SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.to_owned(),
payload: serde_json::to_value(
&*self
.session
.pending_activations
.lock()
.expect("pending activation state poisoned"),
)
.map_err(|error| {
WorkerError::InvalidState(format!(
"serialize pending submissions during compaction: {error}"
))
})?,
});
if let Some(flow_state) = self
.flow_runtime_state
.lock()
@@ -5248,6 +5932,9 @@ where
history_persistence_wired: false,
log_writer: None,
};
worker
.session
.restore_pending_activations(&state.extensions);
worker.apply_permissions_from_manifest();
worker.apply_prune_from_manifest();
worker.write_worker_metadata_active(SegmentLocation {
@@ -8453,6 +9140,164 @@ mod build_summary_prompt_tests {
);
}
#[test]
fn pending_submission_queue_is_durable_idempotent_and_bounded() {
let temp = tempfile::tempdir().unwrap();
let handle = PendingSubmissionHandle::for_test(temp.path());
let input = vec![Segment::text("queued")];
let accepted = handle
.accept("request-1".into(), input.clone(), false)
.unwrap();
assert_eq!(
accepted.disposition,
protocol::SubmissionDisposition::Queued
);
assert_eq!(handle.snapshot().submissions.len(), 1);
let replay = handle
.accept("request-1".into(), input.clone(), false)
.unwrap();
assert_eq!(replay.submission_id, accepted.submission_id);
assert!(replay.activation.is_none());
assert_eq!(handle.snapshot().submissions.len(), 1);
assert!(matches!(
handle.accept("request-1".into(), vec![Segment::text("different")], false),
Err(PendingSubmissionError::IdempotencyConflict)
));
let entries = handle
.writer
.store
.read_all(
handle.writer.state.session_id(),
handle.writer.state.segment_id(),
)
.unwrap();
let payload = entries
.iter()
.rev()
.find_map(|entry| match entry {
LogEntry::Extension {
domain, payload, ..
} if domain == SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN => {
Some(payload.clone())
}
_ => None,
})
.unwrap();
let restored: PendingActivationState = serde_json::from_value(payload).unwrap();
assert_eq!(restored.pending.len(), 1);
assert_eq!(restored.pending[0].submission_id, accepted.submission_id);
let snapshot = handle.cancel(&accepted.submission_id).unwrap();
assert!(snapshot.submissions.is_empty());
assert!(matches!(
handle.cancel(&accepted.submission_id),
Err(PendingSubmissionError::NotFound(_))
));
for index in 0..MAX_PENDING_SUBMISSIONS {
handle
.accept(
format!("limit-{index}"),
vec![Segment::text(format!("value-{index}"))],
false,
)
.unwrap();
}
assert!(matches!(
handle.accept("over-limit".into(), vec![Segment::text("too much")], false),
Err(PendingSubmissionError::CountLimit)
));
assert_eq!(handle.snapshot().submissions.len(), MAX_PENDING_SUBMISSIONS);
let cleared = handle.clear().unwrap();
assert!(cleared.submissions.is_empty());
assert_eq!(cleared.notification_count, 0);
}
#[test]
fn notification_and_submit_share_activation_order_and_notification_dedupes() {
let temp = tempfile::tempdir().unwrap();
let handle = PendingSubmissionHandle::for_test(temp.path());
assert!(
handle
.accept_notification("notification-1".into(), "notice".into())
.unwrap()
);
assert!(
!handle
.accept_notification("notification-1".into(), "notice".into())
.unwrap()
);
assert!(matches!(
handle.accept_notification("notification-1".into(), "different".into()),
Err(PendingSubmissionError::IdempotencyConflict)
));
handle
.accept("request-1".into(), vec![Segment::text("submit")], false)
.unwrap();
let first = handle.prepare_next_activation().unwrap().unwrap();
assert!(matches!(
first,
PendingActivation::Notification(PendingNotification { ref message, .. })
if message == "notice"
));
let committed = handle.notification_activation_extension();
let committed_state: PendingActivationState =
serde_json::from_value(committed.payload).unwrap();
assert!(committed_state.pending_notifications.is_empty());
assert!(committed_state.activating_notification.is_none());
handle.finish_notification_activation("notification-1");
let second = handle.prepare_next_activation().unwrap().unwrap();
assert!(matches!(second, PendingActivation::Submission(_)));
}
#[test]
fn restoring_an_in_flight_activation_requeues_it_at_the_fifo_head() {
let mut session = WorkerSession::new(session_store::new_session_id(), Vec::new());
let state = PendingActivationState {
revision: 4,
next_activation_sequence: 2,
activating: Some(PendingSubmission {
submission_request_id: "request-1".into(),
submission_id: "submission-1".into(),
payload_digest: submission_payload_digest(&[Segment::text("first")]),
accepted_at_ms: 1,
activation_sequence: 0,
provenance: WorkerHistoryProvenance::LegacyUnknown,
was_queued: false,
input: vec![Segment::text("first")],
}),
activating_notification: None,
pending: VecDeque::from([PendingSubmission {
submission_request_id: "request-2".into(),
submission_id: "submission-2".into(),
payload_digest: submission_payload_digest(&[Segment::text("second")]),
accepted_at_ms: 2,
activation_sequence: 1,
provenance: WorkerHistoryProvenance::LegacyUnknown,
was_queued: true,
input: vec![Segment::text("second")],
}]),
pending_notifications: VecDeque::new(),
receipts: VecDeque::new(),
notification_receipts: VecDeque::new(),
};
session.restore_pending_activations(&[(
SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.into(),
serde_json::to_value(state).unwrap(),
)]);
let state = session
.pending_activations
.lock()
.expect("pending activation state poisoned");
assert!(state.activating.is_none());
assert_eq!(state.pending.len(), 2);
assert_eq!(state.pending[0].submission_id, "submission-1");
assert_eq!(state.pending[1].submission_id, "submission-2");
}
fn minimal_manifest() -> WorkerManifest {
let toml_str = r#"
[worker]
+4 -1
View File
@@ -630,7 +630,10 @@ async fn controller_compact_method_emits_start_and_done() {
let mut rx = handle.subscribe();
handle
.send(Method::run_text("seed history"))
.send(Method::submit_text(
protocol::new_submission_request_id(),
"seed history",
))
.await
.expect("send run");
loop {
+235 -57
View File
@@ -617,7 +617,13 @@ async fn feature_flags_default_to_core_tool_surface_only() {
let worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
handle.send(Method::run_text("Hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let request = wait_for_captured_request(&client_for_assert).await;
@@ -672,7 +678,13 @@ permission = "write"
let worker = make_worker_with_pwd_and_manifest(client, manifest).await.0;
let handle = spawn_controller(worker).await;
handle.send(Method::run_text("Hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let request = wait_for_captured_request(&client_for_assert).await;
@@ -758,7 +770,13 @@ permission = "write"
let worker = make_worker_with_pwd_and_manifest(client, &manifest).await.0;
let handle = spawn_controller(worker).await;
handle.send(Method::run_text("Hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let request = wait_for_captured_request(&client_for_assert).await;
@@ -814,7 +832,13 @@ async fn builtin_orchestrator_exposes_worker_remove_and_workdir_delete() {
.await;
let handle = spawn_controller(worker).await;
handle.send(Method::run_text("Hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let request = wait_for_captured_request(&client_for_assert).await;
let installed = request_tool_names(&request);
@@ -863,7 +887,13 @@ permission = "write"
.0;
let handle = spawn_controller(worker).await;
handle.send(Method::run_text("Hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let request = wait_for_captured_request(&client_for_assert).await;
@@ -916,7 +946,13 @@ permission = "write"
)
.await;
let handle = spawn_controller(worker).await;
handle.send(Method::run_text("Hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let request = wait_for_captured_request(&client_for_assert).await;
let names = request_tool_names(&request);
@@ -963,7 +999,13 @@ async fn run_end_returns_to_idle_without_busy_status() {
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
handle.send(Method::run_text("Hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
let mut saw_run_end = false;
let mut saw_idle_status = false;
@@ -1005,7 +1047,13 @@ async fn provider_stream_error_records_run_errored() {
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
handle.send(Method::run_text("ping")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"ping",
))
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
@@ -1054,7 +1102,10 @@ async fn snapshot_includes_user_input_for_in_flight_turn() {
let mut events = handle.subscribe();
handle
.send(Method::run_text("hello in-flight"))
.send(Method::submit_text(
protocol::new_submission_request_id(),
"hello in-flight",
))
.await
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(2), async {
@@ -1119,7 +1170,13 @@ async fn attach_snapshot_includes_current_status() {
let worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
handle.send(Method::run_text("Hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Running).await;
let stream = tokio::net::UnixStream::connect(handle.runtime_dir.socket_path())
@@ -1157,7 +1214,13 @@ async fn run_updates_shared_state_to_idle_after_completion() {
let worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
handle.send(Method::run_text("Hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
// Wait for the run to complete
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
@@ -1171,7 +1234,13 @@ async fn run_populates_history() {
let worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
handle.send(Method::run_text("Hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
@@ -1189,7 +1258,13 @@ async fn events_are_broadcast() {
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
handle.send(Method::run_text("Hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
let mut saw_turn_start = false;
let mut saw_text_delta = false;
@@ -1224,10 +1299,8 @@ async fn events_are_broadcast() {
}
#[tokio::test]
async fn double_run_returns_error() {
// Keep the first turn in-flight until the test drops the handle. A
// finite stream can finish before the second Method reaches the
// controller in the full test suite, making this assertion racy.
async fn submit_while_running_is_durably_queued() {
// Keep the first turn in-flight until the second Submit is accepted.
let events = vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, "slow..."),
@@ -1237,35 +1310,44 @@ async fn double_run_returns_error() {
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
// Send first run and wait until the controller has entered Running.
handle.send(Method::run_text("first")).await.unwrap();
handle
.send(Method::submit_text("request-first", "first"))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Running).await;
handle
.send(Method::submit_text("request-second", "second"))
.await
.unwrap();
// Now the second run must be rejected by drive_turn's live Method arm.
handle.send(Method::run_text("second")).await.unwrap();
// Look for the error event
let mut saw_already_running = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
tokio::select! {
event = rx.recv() => {
match event {
Ok(Event::Error { code, .. }) => {
if code == worker::ErrorCode::AlreadyRunning {
saw_already_running = true;
break;
}
}
Err(_) => break,
_ => {}
}
let mut accepted = None;
let mut pending_count = None;
while tokio::time::Instant::now() < deadline {
match tokio::time::timeout(std::time::Duration::from_millis(100), rx.recv()).await {
Ok(Ok(Event::SubmissionAccepted {
submission_request_id,
disposition,
..
})) if submission_request_id == "request-second" => accepted = Some(disposition),
Ok(Ok(Event::PendingSubmissionsChanged { pending }))
if pending.submissions.len() == 1 =>
{
pending_count = Some(1)
}
_ = tokio::time::sleep_until(deadline) => break,
Ok(Ok(Event::Error { code, message })) if code == worker::ErrorCode::AlreadyRunning => {
panic!("Submit was busy-rejected: {message}")
}
_ => {}
}
if accepted.is_some() && pending_count.is_some() {
break;
}
}
assert!(saw_already_running, "should see already_running error");
assert_eq!(accepted, Some(protocol::SubmissionDisposition::Queued));
assert_eq!(pending_count, Some(1));
handle.send(Method::Pause).await.unwrap();
}
#[tokio::test]
@@ -1353,7 +1435,8 @@ async fn run_with_paste_segment_inlines_content_and_emits_typed_user_message() {
protocol::Segment::text(" thanks"),
];
handle
.send(Method::Run {
.send(Method::Submit {
submission_request_id: protocol::new_submission_request_id(),
input: segments.clone(),
})
.await
@@ -1425,7 +1508,13 @@ async fn run_with_resolvable_file_ref_attaches_system_message_after_user() {
path: "notes.md".into(),
},
];
handle.send(Method::Run { input: segments }).await.unwrap();
handle
.send(Method::Submit {
submission_request_id: protocol::new_submission_request_id(),
input: segments,
})
.await
.unwrap();
// Wait for the turn to complete.
let mut rx = handle.subscribe();
@@ -1473,7 +1562,8 @@ async fn run_with_file_ref_uses_manifest_file_upload_limit() {
let handle = spawn_controller(worker).await;
handle
.send(Method::Run {
.send(Method::Submit {
submission_request_id: protocol::new_submission_request_id(),
input: vec![protocol::Segment::FileRef {
path: "long.txt".into(),
}],
@@ -1526,7 +1616,13 @@ async fn run_with_unresolved_segment_emits_alert_and_placeholder() {
path: "src/lib.rs".into(),
},
];
handle.send(Method::Run { input: segments }).await.unwrap();
handle
.send(Method::Submit {
submission_request_id: protocol::new_submission_request_id(),
input: segments,
})
.await
.unwrap();
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
let mut saw_alert_for_file_ref = false;
@@ -1574,6 +1670,7 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
handle
.send(Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message: "turn finished".into(),
auto_run: true,
})
@@ -1614,6 +1711,19 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
saw_notify_in_mirror,
"Method::Notify should commit a SystemItem::Notification entry; mirror = {entries:?}"
);
let queue_checkpoint_is_atomic = entries.iter().any(|entry| match entry {
LogEntry::AnnotatedSystemItem { extensions, .. } => extensions.iter().any(|extension| {
extension.domain == "worker.pending_activations.v1"
&& extension.payload["pending_notifications"]
.as_array()
.is_some_and(Vec::is_empty)
}),
_ => false,
});
assert!(
queue_checkpoint_is_atomic,
"notification history and queue claim must share one log entry"
);
// Exactly one request was made; it must contain the formatted
// notification as one of the items (committed to history by
@@ -1662,6 +1772,7 @@ async fn notify_while_idle_with_auto_run_false_waits_for_explicit_run() {
handle
.send(Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message: "progress snapshot".into(),
auto_run: false,
})
@@ -1675,7 +1786,13 @@ async fn notify_while_idle_with_auto_run_false_waits_for_explicit_run() {
"weak Notify must not stage RunForNotification while idle"
);
handle.send(Method::run_text("continue")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"continue",
))
.await
.unwrap();
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
if !client_for_assert.captured_requests().is_empty() {
@@ -1855,9 +1972,16 @@ async fn notify_while_running_does_not_emit_already_running_error() {
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
handle.send(Method::run_text("start")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"start",
))
.await
.unwrap();
handle
.send(Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message: "ping".into(),
auto_run: true,
})
@@ -1924,7 +2048,13 @@ async fn socket_run_receives_events() {
let mut writer = JsonLineWriter::new(writer);
// Send run method via socket
writer.write(&Method::run_text("Hello")).await.unwrap();
writer
.write(&Method::submit_text(
protocol::new_submission_request_id(),
"Hello",
))
.await
.unwrap();
// Collect events
let mut saw_turn_start = false;
@@ -2231,7 +2361,13 @@ async fn pause_then_resume_transitions_and_preserves_history_consistency() {
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
handle.send(Method::run_text("hello")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"hello",
))
.await
.unwrap();
// Wait for the partial text_delta to confirm the first stream is
// live before we pause.
@@ -2320,7 +2456,7 @@ async fn pause_then_resume_transitions_and_preserves_history_consistency() {
assert!(!has_tool_call, "no orphan tool_call in history");
}
/// Paused with an orphan `tool_use` in history + a fresh `Method::Run`
/// Paused with an orphan `tool_use` in history + a fresh `Method::Submit`
/// must produce a wire-valid next LLM request: the orphan is closed
/// with a synthetic `tool_result`, a system note is inserted, and the
/// new user input is appended.
@@ -2357,7 +2493,13 @@ async fn paused_then_run_closes_orphan_tool_use_for_next_request() {
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
handle.send(Method::run_text("first")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"first",
))
.await
.unwrap();
// Wait for ToolCallDone — the ToolCall is committed to history
// right before the Engine enters tool execution and pends.
@@ -2388,7 +2530,13 @@ async fn paused_then_run_closes_orphan_tool_use_for_next_request() {
// `last_run_interrupted` and runs its interrupt-prep step, which
// closes the orphan + injects a system note before the fresh user
// message.
handle.send(Method::run_text("new request")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"new request",
))
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
e,
@@ -2519,7 +2667,13 @@ async fn paused_cancel_abandons_resume_and_next_input_is_fresh_run() {
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
handle.send(Method::run_text("first")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"first",
))
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
e,
@@ -2587,7 +2741,10 @@ async fn paused_cancel_abandons_resume_and_next_input_is_fresh_run() {
);
handle
.send(Method::run_text("fresh request"))
.send(Method::submit_text(
protocol::new_submission_request_id(),
"fresh request",
))
.await
.unwrap();
assert!(
@@ -2676,7 +2833,13 @@ async fn empty_turn_cancel_rolls_back_submit_entries_and_emits_signal() {
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
handle.send(Method::run_text("rollback me")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"rollback me",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Running).await;
handle.send(Method::Cancel).await.unwrap();
@@ -2709,7 +2872,10 @@ async fn empty_turn_pause_rolls_back_and_snapshot_does_not_restore_input() {
let mut rx = handle.subscribe();
handle
.send(Method::run_text("pause rollback"))
.send(Method::submit_text(
protocol::new_submission_request_id(),
"pause rollback",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Running).await;
@@ -2743,7 +2909,13 @@ async fn empty_turn_rollback_removes_only_the_most_recent_turn() {
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
handle.send(Method::run_text("first kept")).await.unwrap();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"first kept",
))
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
e,
@@ -2757,7 +2929,10 @@ async fn empty_turn_rollback_removes_only_the_most_recent_turn() {
wait_for_status(&handle, WorkerStatus::Idle).await;
handle
.send(Method::run_text("second rolled back"))
.send(Method::submit_text(
protocol::new_submission_request_id(),
"second rolled back",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Running).await;
@@ -2804,7 +2979,10 @@ async fn pause_after_assistant_token_does_not_rollback() {
let mut rx = handle.subscribe();
handle
.send(Method::run_text("keep this turn"))
.send(Method::submit_text(
protocol::new_submission_request_id(),
"keep this turn",
))
.await
.unwrap();
assert!(