feat: integrate worker submit queue

This commit is contained in:
2026-09-06 04:24:07 +09:00
51 changed files with 5022 additions and 929 deletions
+5 -2
View File
@@ -120,12 +120,15 @@ mod tests {
let mut client = Client::new(socket);
client
.send(&Method::run_text("hello"))
.send(&Method::submit_text(
protocol::new_submission_request_id(),
"hello",
))
.await
.expect("send method");
assert!(matches!(
decode_method(&client.socket.sent[0]),
Ok(Method::Run { .. })
Ok(Method::Submit { .. })
));
assert!(matches!(
client.next_event().await,
+5 -2
View File
@@ -89,12 +89,15 @@ mod tests {
let mut client = Client::new(socket);
client
.send(&Method::run_text("hello"))
.send(&Method::submit_text(
protocol::new_submission_request_id(),
"hello",
))
.await
.expect("send method");
assert!(matches!(
peer.next().await.as_deref().map(decode_method),
Some(Ok(Method::Run { .. }))
Some(Ok(Method::Submit { .. }))
));
peer.send(
+8 -2
View File
@@ -147,12 +147,18 @@ mod tests {
let mut client = Client::new(Socket::connect(&socket_path).await.unwrap());
client
.send(&Method::run_text("hello"))
.send(&Method::submit_text(
protocol::new_submission_request_id(),
"hello",
))
.await
.expect("send method");
let received = server.await.unwrap().expect("method message");
assert!(matches!(decode_method(&received), Ok(Method::Run { .. })));
assert!(matches!(
decode_method(&received),
Ok(Method::Submit { .. })
));
}
#[tokio::test]
+5 -2
View File
@@ -114,7 +114,7 @@ mod tests {
assert!(matches!(
message,
Message::Text(ref text)
if matches!(decode_method(text), Ok(Method::Run { .. }))
if matches!(decode_method(text), Ok(Method::Submit { .. }))
));
let event = encode_event(&Event::Status {
status: WorkerStatus::Idle,
@@ -126,7 +126,10 @@ mod tests {
let request = format!("ws://{address}").into_client_request().unwrap();
let mut client = Client::new(Socket::connect(request).await.unwrap());
client
.send(&Method::run_text("hello"))
.send(&Method::submit_text(
protocol::new_submission_request_id(),
"hello",
))
.await
.expect("send method");
assert!(matches!(
+211 -52
View File
@@ -11,6 +11,11 @@ use serde::{Deserialize, Serialize};
pub use identity::{WorkerId, WorkerIdParseError};
/// Allocate an opaque idempotency key for one client Submit request.
pub fn new_submission_request_id() -> String {
uuid::Uuid::now_v7().to_string()
}
fn default_true() -> bool {
true
}
@@ -27,21 +32,80 @@ fn is_false(value: &bool) -> bool {
// Method (Client → Worker via Unix Socket)
// ---------------------------------------------------------------------------
/// Trusted Server → Runtime transport header carrying the authenticated
/// browser Account identity for one Worker protocol connection.
///
/// Runtime accepts this only after its normal HTTP authentication succeeds;
/// serialized [`Method`] payloads cannot set authenticated source identity.
pub const AUTHENTICATED_ACCOUNT_ID_HEADER: &str = "x-yoi-authenticated-account-id";
/// Trusted source identity attached by an authenticated transport boundary.
///
/// Public clients cannot select this value directly. Runtime/Backend adapters
/// stamp it before forwarding an accepted Submit or Notify to a Worker.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum AuthenticatedInputSource {
/// Assigned whenever a serialized tracked method crosses an untrusted
/// protocol boundary. Receivers must handle it exactly like public input.
UntrustedWire,
Account {
account_id: String,
},
Worker {
runtime_id: String,
worker_id: String,
},
SubWorker {
session_id: String,
},
Backend {
operation_id: String,
},
}
impl Default for AuthenticatedInputSource {
fn default() -> Self {
Self::UntrustedWire
}
}
impl AuthenticatedInputSource {
pub fn namespace(&self) -> String {
match self {
Self::UntrustedWire => "untrusted-wire".into(),
Self::Account { account_id } => format!("account:{account_id}"),
Self::Worker {
runtime_id,
worker_id,
} => format!("worker:{runtime_id}:{worker_id}"),
Self::SubWorker { session_id } => format!("sub_worker:{session_id}"),
Self::Backend { operation_id } => format!("backend:{operation_id}"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(tag = "method", content = "params", rename_all = "snake_case")]
pub enum Method {
Run {
/// Durably accept typed input for immediate activation or the session FIFO.
///
/// `submission_request_id` is generated by the authenticated caller and is
/// used only for idempotent retry. Worker allocates the durable
/// `submission_id` returned by [`Event::SubmissionAccepted`].
Submit {
submission_request_id: String,
input: Vec<Segment>,
},
/// Runtime-internal Run carrying an opaque correlation id that is committed
/// with the resulting UserInput entry. This variant is not serializable on
/// the public Client → Worker protocol.
#[serde(skip)]
/// Authenticated transport form of Submit. Trusted adapters replace
/// public Submit before forwarding it to the Worker.
#[cfg_attr(feature = "typescript", ts(skip))]
RunTracked {
SubmitTracked {
submission_request_id: String,
input: Vec<Segment>,
submission_id: String,
#[serde(skip_deserializing, default)]
source: AuthenticatedInputSource,
},
/// Human-readable text injected into the target Worker's LLM context
/// as a non-blocking system message. `auto_run` controls whether an
@@ -50,25 +114,54 @@ pub enum Method {
/// No side effects beyond LLM context; use `WorkerEvent` for typed
/// lifecycle reports.
Notify {
notification_request_id: String,
message: String,
#[serde(default = "default_true", skip_serializing_if = "is_true")]
auto_run: bool,
},
/// Authenticated transport form of Notify.
#[cfg_attr(feature = "typescript", ts(skip))]
NotifyTracked {
notification_request_id: String,
message: String,
#[serde(default = "default_true", skip_serializing_if = "is_true")]
auto_run: bool,
#[serde(skip_deserializing, default)]
source: AuthenticatedInputSource,
},
/// Typed lifecycle report from a child Worker to its direct parent.
WorkerEvent(WorkerEvent),
/// Return the authoritative FIFO summary without exposing queued payloads.
ListPendingSubmissions,
/// Remove one queued submission. Running or already activated submissions
/// are immutable and therefore cannot be cancelled here.
CancelPendingSubmission {
submission_id: String,
expected_revision: u64,
},
/// Remove every queued submission while preserving the active run.
ClearPendingSubmissions {
expected_revision: u64,
},
/// Activate the next queued submission while the Worker is idle. This is an
/// explicit recovery operation and never resumes a paused run implicitly.
ContinuePending {
expected_revision: u64,
expected_head_id: String,
},
Resume,
Cancel,
/// Stop the in-flight turn and transition to `Paused`.
///
/// Unlike `Cancel` (which discards and returns to `Idle`), a paused
/// Worker can resume the interrupted work via `Resume`, or start a
/// fresh turn via `Run` (orphan `tool_use` items are closed with a
/// Worker can resume the interrupted work via `Resume`, or accept a
/// fresh `Submit` (orphan `tool_use` items are closed with a
/// synthetic tool result before the new user message is appended).
Pause,
/// Request an explicit compaction while the Worker is otherwise idle.
///
/// This is a typed control method: clients must not send `compact` as a
/// `Method::Run` user message.
/// `Method::Submit` user message.
Compact,
/// Ask the Worker to list valid rewind targets from its authoritative session log.
ListRewindTargets,
@@ -181,7 +274,7 @@ impl WorkerEvent {
/// One typed piece of a user submission.
///
/// `Method::Run` and `Event::UserMessage` carry `Vec<Segment>`. Dumb
/// `Method::Submit` and `Event::UserMessage` carry `Vec<Segment>`. Dumb
/// clients (CLI piping, scripts) only need to produce a single
/// `Segment::Text`; richer clients (TUI / GUI) construct typed atoms
/// (paste chips, file refs) and
@@ -404,12 +497,13 @@ impl Segment {
}
impl Method {
/// Convenience: a `Run` carrying a single `Segment::Text`.
/// Convenience: a `Submit` carrying a single `Segment::Text`.
/// Used by dumb clients, inter-Worker tools, and tests that only have
/// a string to forward.
pub fn run_text(s: impl Into<String>) -> Self {
Self::Run {
input: vec![Segment::text(s)],
pub fn submit_text(submission_request_id: impl Into<String>, text: impl Into<String>) -> Self {
Self::Submit {
submission_request_id: submission_request_id.into(),
input: vec![Segment::text(text)],
}
}
}
@@ -503,6 +597,39 @@ pub enum ToolResultDisposition {
OutcomeUnknown,
}
/// Durable acceptance result for one idempotent Submit request.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(rename_all = "snake_case")]
pub enum SubmissionDisposition {
Started,
Queued,
}
/// Bounded public projection of one pending submission. Payload segments and
/// provenance remain in the session log and are intentionally not exposed.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct PendingSubmissionSummary {
pub submission_id: String,
pub accepted_at_ms: u64,
pub segment_count: u32,
pub byte_len: u64,
}
/// Revisioned session-owned FIFO projection used by snapshots and live events.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct PendingSubmissionsSnapshot {
pub revision: u64,
#[serde(default)]
pub notification_count: u32,
#[serde(default)]
pub head_id: Option<String>,
#[serde(default)]
pub submissions: Vec<PendingSubmissionSummary>,
}
/// Canonical, storage-independent projection of committed session history.
///
/// Worker protocols expose this DTO instead of append-log records. New
@@ -511,6 +638,8 @@ pub enum ToolResultDisposition {
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct SessionSnapshot {
#[serde(default)]
pub pending_submissions: PendingSubmissionsSnapshot,
pub entries: Vec<SessionSnapshotEntry>,
}
@@ -609,16 +738,27 @@ pub struct SessionToolAttachment {
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(tag = "event", content = "data", rename_all = "snake_case")]
pub enum Event {
/// A user input message was accepted, persisted as
/// `LogEntry::AnnotatedUserInput`, and is about to start a new turn.
/// Broadcast to every subscribed client so TUI / GUI instances show
/// the same user line that reconnect snapshots would replay from
/// history; clients must not synthesize a separate pending/fake
/// message for accepted runs.
///
/// Fires exactly once per committed user input, after
/// `InvokeStart { kind: UserSend }` and before the first
/// `TurnStart`. Rejected runs (e.g. `AlreadyRunning`) do not emit.
/// Durable Submit acceptance. A `Started` receipt follows the atomic
/// UserInput commit; a `Queued` receipt follows the durable FIFO checkpoint.
/// Repeating the same request id and exact payload returns the same receipt
/// without appending or activating twice.
SubmissionAccepted {
submission_request_id: String,
submission_id: String,
disposition: SubmissionDisposition,
},
/// Correlated rejection before durable acceptance.
SubmissionRejected {
submission_request_id: String,
message: String,
},
/// Revisioned FIFO replacement following enqueue, activation, cancel, or clear.
PendingSubmissionsChanged {
pending: PendingSubmissionsSnapshot,
},
/// A user input message persisted as `LogEntry::AnnotatedUserInput` and
/// activated for a turn. Broadcast to every subscribed client so TUI / GUI
/// instances show the same user line that reconnect snapshots replay.
UserMessage {
segments: Vec<Segment>,
},
@@ -641,7 +781,7 @@ pub enum Event {
///
/// Marker event for the start of an Invoke range; the range extends
/// implicitly until the next `InvokeStart`. Fires for every accepted
/// `Method::Run` (kind=`UserSend`), `Method::Notify` (kind=`Notify`),
/// `Method::Submit` (kind=`UserSend`), `Method::Notify` (kind=`Notify`),
/// `Method::WorkerEvent` re-injection (kind=`WorkerEvent`), and any other
/// IDLE-breaking trigger. Mid-run interrupts (e.g. hook output,
/// typed system reminder insertion that doesn't break IDLE) do not
@@ -1193,7 +1333,7 @@ pub enum TurnResult {
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(rename_all = "snake_case")]
pub enum InvokeKind {
/// `Method::Run` — a user submission.
/// `Method::Submit` — a user submission.
UserSend,
/// `Method::Notify` — free-text notification injected into history.
Notify,
@@ -1216,7 +1356,7 @@ pub enum RunResult {
Finished,
Paused,
LimitReached,
/// The accepted Method::Run produced no assistant/tool output before
/// The accepted Method::Submit produced no assistant/tool output before
/// user interruption, so the Worker rolled the submit-time turn state back
/// to its pre-submit snapshot. Clients should treat the Worker as Idle and
/// restore the just-submitted input into the editable composer if desired.
@@ -1285,26 +1425,30 @@ mod tests {
use super::*;
#[test]
fn method_run_json_roundtrip() {
let json = r#"{"method":"run","params":{"input":[{"kind":"text","content":"Hello"}]}}"#;
fn method_submit_json_roundtrip_and_run_is_rejected() {
let json = r#"{"method":"submit","params":{"submission_request_id":"request-1","input":[{"kind":"text","content":"Hello"}]}}"#;
let method: Method = serde_json::from_str(json).unwrap();
match &method {
Method::Run { input } => {
Method::Submit { input, .. } => {
assert_eq!(input.len(), 1);
match &input[0] {
Segment::Text { content } => assert_eq!(content, "Hello"),
other => panic!("expected Text, got {other:?}"),
}
}
other => panic!("expected Run, got {other:?}"),
other => panic!("expected Submit, got {other:?}"),
}
let serialized = serde_json::to_string(&method).unwrap();
assert_eq!(serialized, json);
assert!(
serde_json::from_str::<Method>(r#"{"method":"run","params":{"input":[]}}"#).is_err()
);
}
#[test]
fn method_run_paste_segment_roundtrip() {
let method = Method::Run {
fn method_submit_paste_segment_roundtrip() {
let method = Method::Submit {
submission_request_id: "request-1".to_string(),
input: vec![
Segment::text("see "),
Segment::Paste {
@@ -1318,7 +1462,7 @@ mod tests {
let json = serde_json::to_string(&method).unwrap();
let decoded: Method = serde_json::from_str(&json).unwrap();
match decoded {
Method::Run { input } => {
Method::Submit { input, .. } => {
assert_eq!(input.len(), 2);
match &input[1] {
Segment::Paste {
@@ -1335,7 +1479,7 @@ mod tests {
other => panic!("expected Paste, got {other:?}"),
}
}
other => panic!("expected Run, got {other:?}"),
other => panic!("expected Submit, got {other:?}"),
}
}
@@ -1389,8 +1533,9 @@ mod tests {
}
#[test]
fn method_run_flow_segment_roundtrip() {
let method = Method::Run {
fn method_submit_flow_segment_roundtrip() {
let method = Method::Submit {
submission_request_id: "request-1".to_string(),
input: vec![
Segment::Flow {
selector: "builtin:coder-review".to_string(),
@@ -1404,7 +1549,7 @@ mod tests {
let decoded = serde_json::from_str::<Method>(&json).unwrap();
assert!(matches!(
decoded,
Method::Run { input }
Method::Submit { input, .. }
if matches!(
input.as_slice(),
[
@@ -1416,15 +1561,26 @@ mod tests {
}
#[test]
fn runtime_tracked_run_is_not_public_protocol_json() {
let method = Method::RunTracked {
fn authenticated_submit_replaces_wire_source_with_transport_identity() {
let method = Method::SubmitTracked {
input: vec![Segment::text("private")],
submission_id: "submission-1".to_string(),
submission_request_id: "request-1".to_string(),
source: AuthenticatedInputSource::Account {
account_id: "account-1".into(),
},
};
assert!(serde_json::to_string(&method).is_err());
let json = serde_json::to_string(&method).unwrap();
let decoded = serde_json::from_str::<Method>(&json).unwrap();
assert!(matches!(
decoded,
Method::SubmitTracked {
source: AuthenticatedInputSource::UntrustedWire,
..
}
));
assert!(
serde_json::from_str::<Method>(
r#"{"method":"run_tracked","input":[],"submission_id":"forged"}"#,
r#"{"method":"submit_tracked","input":[],"submission_request_id":"forged"}"#,
)
.is_err()
);
@@ -1442,16 +1598,16 @@ mod tests {
}
#[test]
fn method_run_with_unknown_segment_decodes() {
let json = r#"{"method":"run","params":{"input":[{"kind":"text","content":"hi"},{"kind":"future_thing","x":1}]}}"#;
fn method_submit_with_unknown_segment_decodes() {
let json = r#"{"method":"submit","params":{"submission_request_id":"request-1","input":[{"kind":"text","content":"hi"},{"kind":"future_thing","x":1}]}}"#;
let method: Method = serde_json::from_str(json).unwrap();
match method {
Method::Run { input } => {
Method::Submit { input, .. } => {
assert_eq!(input.len(), 2);
assert!(matches!(input[0], Segment::Text { .. }));
assert!(matches!(input[1], Segment::Unknown));
}
other => panic!("expected Run, got {other:?}"),
other => panic!("expected Submit, got {other:?}"),
}
}
@@ -1648,11 +1804,11 @@ mod tests {
#[test]
fn method_notify_json_roundtrip_defaults_to_auto_run() {
let json = r#"{"method":"notify","params":{"message":"turn done"}}"#;
let json = r#"{"method":"notify","params":{"notification_request_id":"notification-1","message":"turn done"}}"#;
let method: Method = serde_json::from_str(json).unwrap();
assert!(matches!(
method,
Method::Notify { ref message, auto_run: true } if message == "turn done"
Method::Notify { ref message, auto_run: true, .. } if message == "turn done"
));
let serialized = serde_json::to_string(&method).unwrap();
assert_eq!(serialized, json);
@@ -1660,11 +1816,11 @@ mod tests {
#[test]
fn method_notify_weak_json_roundtrip_serializes_auto_run_false() {
let json = r#"{"method":"notify","params":{"message":"progress","auto_run":false}}"#;
let json = r#"{"method":"notify","params":{"notification_request_id":"notification-1","message":"progress","auto_run":false}}"#;
let method: Method = serde_json::from_str(json).unwrap();
assert!(matches!(
method,
Method::Notify { ref message, auto_run: false } if message == "progress"
Method::Notify { ref message, auto_run: false, .. } if message == "progress"
));
assert_eq!(serde_json::to_string(&method).unwrap(), json);
}
@@ -1725,6 +1881,7 @@ mod tests {
fn event_snapshot_format() {
let event = Event::Snapshot {
session: SessionSnapshot {
pending_submissions: PendingSubmissionsSnapshot::default(),
entries: vec![SessionSnapshotEntry {
entry_id: "entry-1".into(),
timestamp: 1,
@@ -1776,6 +1933,7 @@ mod tests {
let event = Event::Snapshot {
session: SessionSnapshot {
pending_submissions: PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: Greeting {
@@ -1844,6 +2002,7 @@ mod tests {
fn event_segment_rotated_roundtrip() {
let event = Event::SegmentRotated {
session: SessionSnapshot {
pending_submissions: PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
};
+8 -5
View File
@@ -8,11 +8,11 @@ use crate::{
CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightSnapshot,
InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot,
InvokeKind, MemoryWorkerEvent, Method, PasteArtifactAvailability, PasteArtifactMediaType,
PasteArtifactRef, Permission, RewindSummary, RewindTarget, RewindTargetId, RunResult,
ScopeRule, Segment, SessionContentPart, SessionEntryProvenance, SessionMessageRole,
SessionSnapshot, SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment,
ToolResultDisposition, TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerEvent,
WorkerStatus,
PasteArtifactRef, PendingSubmissionSummary, PendingSubmissionsSnapshot, Permission,
RewindSummary, RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, SessionContentPart,
SessionEntryProvenance, SessionMessageRole, SessionSnapshot, SessionSnapshotEntry,
SessionSnapshotEntryData, SessionToolAttachment, SubmissionDisposition, ToolResultDisposition,
TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerEvent, WorkerStatus,
subscription::{
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
@@ -75,6 +75,9 @@ pub fn generated_protocol_types() -> String {
push_decl::<SessionToolAttachment>(&cfg, &mut output);
push_decl::<SessionSnapshotEntryData>(&cfg, &mut output);
push_decl::<SessionSnapshotEntry>(&cfg, &mut output);
push_decl::<PendingSubmissionSummary>(&cfg, &mut output);
push_decl::<PendingSubmissionsSnapshot>(&cfg, &mut output);
push_decl::<SubmissionDisposition>(&cfg, &mut output);
push_decl::<SessionSnapshot>(&cfg, &mut output);
push_decl::<InternalWorkerKind>(&cfg, &mut output);
push_decl::<InternalWorkerRef>(&cfg, &mut output);
+171 -9
View File
@@ -21,8 +21,10 @@ use crate::segment_log::LogEntry;
use crate::store::{Store, StoreError};
use crate::uploaded_file::{
bind_uploaded_file, clear_uploaded_file_binding, copy_committed_uploaded_files,
delete_uncommitted_uploaded_files, delete_uploaded_file, list_uploaded_file_refs,
read_uploaded_file, read_uploaded_file_by_id, write_uploaded_file,
delete_uncommitted_uploaded_files, delete_uploaded_file, finalize_uploaded_file_binding,
list_uploaded_file_refs, pin_uploaded_file, read_uploaded_file, read_uploaded_file_by_id,
reconcile_uploaded_file_pins, release_uploaded_file_pin, uploaded_file_has_pending_owner,
write_uploaded_file,
};
use crate::{
PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits, UploadedFileUploadContext,
@@ -518,6 +520,61 @@ impl Store for FsStore {
}
}
fn pin_uploaded_file(
&self,
session_id: SessionId,
reference: &UploadedFileRef,
owner_id: &str,
) -> Result<(), StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
pin_uploaded_file(&self.paste_artifact_dir(session_id), reference, owner_id)
}
fn release_uploaded_file_pin(
&self,
session_id: SessionId,
artifact_id: &str,
owner_id: &str,
) -> Result<(), StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
release_uploaded_file_pin(&self.paste_artifact_dir(session_id), artifact_id, owner_id)
}
fn finalize_uploaded_file_binding(
&self,
session_id: SessionId,
artifact_id: &str,
source_entry_id: &str,
) -> Result<(), StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
finalize_uploaded_file_binding(
&self.paste_artifact_dir(session_id),
artifact_id,
source_entry_id,
)
}
fn reconcile_uploaded_file_pins(
&self,
session_id: SessionId,
live_owner_ids: &[String],
) -> Result<u64, StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
reconcile_uploaded_file_pins(&self.paste_artifact_dir(session_id), live_owner_ids)
}
fn delete_uploaded_file(
&self,
session_id: SessionId,
@@ -541,13 +598,18 @@ impl Store for FsStore {
let Some(source_entry_id) = reference.source_entry_id.as_deref() else {
continue;
};
if !self.uploaded_file_is_referenced(session_id, &reference.artifact_id)? {
clear_uploaded_file_binding(&dir, &reference.artifact_id, source_entry_id)?;
if delete_uploaded_file(&dir, &reference.artifact_id)? {
removed = removed
.checked_add(1)
.ok_or(StoreError::ArtifactQuotaExceeded)?;
}
if self.uploaded_file_is_referenced(session_id, &reference.artifact_id)? {
finalize_uploaded_file_binding(&dir, &reference.artifact_id, source_entry_id)?;
continue;
}
if uploaded_file_has_pending_owner(&dir, &reference.artifact_id)? {
continue;
}
clear_uploaded_file_binding(&dir, &reference.artifact_id, source_entry_id)?;
if delete_uploaded_file(&dir, &reference.artifact_id)? {
removed = removed
.checked_add(1)
.ok_or(StoreError::ArtifactQuotaExceeded)?;
}
}
Ok(removed)
@@ -865,6 +927,106 @@ mod tests {
assert!(store.read_uploaded_file(owner, &reference).is_err());
}
#[test]
fn pending_upload_pin_survives_cleanup_until_release_or_history_binding() {
let tmp = tempfile::TempDir::new().unwrap();
let store = FsStore::new(tmp.path()).unwrap();
let session_id = new_session_id();
let limits = UploadedFileLimits {
max_file_bytes: 64,
max_session_bytes: 128,
};
let pending = store
.write_uploaded_file(session_id, "pending.txt", "text/plain", b"pending", limits)
.unwrap();
store
.pin_uploaded_file(session_id, &pending, "submission-1")
.unwrap();
assert!(matches!(
store.pin_uploaded_file(session_id, &pending, "submission-other"),
Err(StoreError::ArtifactAlreadyCommitted)
));
drop(store);
let store = FsStore::new(tmp.path()).unwrap();
assert_eq!(
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
0
);
assert_eq!(
store
.read_uploaded_file_by_id(session_id, &pending.artifact_id)
.unwrap()
.1,
b"pending"
);
let fork_session_id = new_session_id();
assert_eq!(
store
.copy_committed_uploaded_files(session_id, fork_session_id)
.unwrap(),
0
);
assert!(
store
.read_uploaded_file_by_id(fork_session_id, &pending.artifact_id)
.is_err()
);
let committed = store
.bind_uploaded_file(session_id, &pending, "entry-1")
.unwrap();
assert_eq!(
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
0
);
assert!(
store
.read_uploaded_file_by_id(session_id, &pending.artifact_id)
.is_ok()
);
store
.create_segment(
session_id,
new_segment_id(),
&[LogEntry::InputSegmentsCheckpoint {
ts: 1,
user_segments: vec![vec![protocol::Segment::UploadedFile {
file: committed.clone(),
}]],
}],
)
.unwrap();
assert_eq!(
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
0
);
assert!(
store
.release_uploaded_file_pin(session_id, &pending.artifact_id, "submission-1")
.is_err()
);
let releasable = store
.write_uploaded_file(session_id, "cancelled.txt", "text/plain", b"cancel", limits)
.unwrap();
store
.pin_uploaded_file(session_id, &releasable, "submission-2")
.unwrap();
store
.release_uploaded_file_pin(session_id, &releasable.artifact_id, "submission-2")
.unwrap();
assert_eq!(
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
1
);
assert!(
store
.read_uploaded_file_by_id(session_id, &releasable.artifact_id)
.is_err()
);
}
#[test]
fn uploaded_file_validation_and_shared_quota_fail_closed() {
let tmp = tempfile::TempDir::new().unwrap();
@@ -183,6 +183,7 @@ fn canonicalize_history_entry(
item,
metadata: legacy_metadata(segment_id, line_index, 0),
},
extensions: Vec::new(),
},
}
}
+5 -2
View File
@@ -71,7 +71,7 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess
entries.push(history_entry(entry, *ts, data));
}
}
LogEntry::AnnotatedSystemItem { ts, entry } => entries.push(system_entry(
LogEntry::AnnotatedSystemItem { ts, entry, .. } => entries.push(system_entry(
&entry.item,
entry.metadata.entry_id.0.clone(),
*ts,
@@ -100,7 +100,10 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess
}
}
SessionSnapshot { entries }
SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries,
}
}
fn extend_history(
+1
View File
@@ -287,6 +287,7 @@ pub fn append_system_item(
LogEntry::AnnotatedSystemItem {
ts: segment_log::now_millis(),
entry,
extensions: Vec::new(),
},
)
}
+10 -1
View File
@@ -112,6 +112,8 @@ pub enum LogEntry {
AnnotatedSystemItem {
ts: u64,
entry: LoggedSystemHistoryEntry,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
extensions: Vec<SessionExtension>,
},
/// Turn boundary. Records the turn count after increment.
@@ -312,12 +314,19 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState {
state.annotated_history.push(entry.clone());
state.history.push(Item::from(entry.item.clone()));
}
LogEntry::AnnotatedSystemItem { entry, .. } => {
LogEntry::AnnotatedSystemItem {
entry, extensions, ..
} => {
state.annotated_history.push(LoggedHistoryEntry {
item: LoggedItem::from(entry.item.to_history_item()),
metadata: entry.metadata.clone(),
});
state.history.push(entry.item.to_history_item());
state.extensions.extend(
extensions
.iter()
.map(|extension| (extension.domain.clone(), extension.payload.clone())),
);
}
LogEntry::TurnEnd { turn_count, .. } => {
if let Some(active_turn_count) = &mut state.active_run_turn_count {
+41
View File
@@ -226,6 +226,47 @@ pub trait Store: Send + Sync {
Err(StoreError::PasteArtifactUnsupported)
}
/// Retain an uploaded file while a durable pending operation owns it.
fn pin_uploaded_file(
&self,
_session_id: SessionId,
_reference: &UploadedFileRef,
_owner_id: &str,
) -> Result<(), StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
/// Release a pending-operation pin without changing committed ownership.
fn release_uploaded_file_pin(
&self,
_session_id: SessionId,
_artifact_id: &str,
_owner_id: &str,
) -> Result<(), StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
/// Complete the pending-to-history handoff after the history entry commits.
fn finalize_uploaded_file_binding(
&self,
_session_id: SessionId,
_artifact_id: &str,
_source_entry_id: &str,
) -> Result<(), StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
/// Clear pending-operation pins that have no owner in restored durable
/// Worker Session state. This repairs an interrupted pin-before-checkpoint
/// acceptance without disturbing live queue owners or committed history.
fn reconcile_uploaded_file_pins(
&self,
_session_id: SessionId,
_live_owner_ids: &[String],
) -> Result<u64, StoreError> {
Ok(0)
}
/// Delete an uncommitted uploaded file owned by `session_id`.
fn delete_uploaded_file(
&self,
+143 -2
View File
@@ -24,6 +24,12 @@ pub const DEFAULT_MAX_FILES_PER_SUBMISSION: usize = 8;
pub const DEFAULT_MAX_SESSION_UPLOADED_FILES: u64 = 256;
const MAX_FILE_NAME_CHARS: usize = 255;
const MAX_MEDIA_TYPE_BYTES: usize = 127;
fn validate_pending_owner_id(owner_id: &str) -> Result<()> {
if owner_id.is_empty() || owner_id.len() > 256 {
return Err(StoreError::ArtifactIntegrityMismatch);
}
Ok(())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct UploadedFileLimits {
@@ -59,6 +65,8 @@ struct StoredUploadedFile {
#[serde(default, skip_serializing_if = "Option::is_none")]
source_entry_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pending_owner_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
upload_context: Option<UploadedFileUploadContext>,
content_base64: String,
}
@@ -291,6 +299,7 @@ pub(crate) fn write_uploaded_file(
byte_len,
sha256: sha256.clone(),
source_entry_id: None,
pending_owner_id: None,
upload_context: context.cloned(),
content_base64: BASE64.encode(content),
};
@@ -338,6 +347,12 @@ pub(crate) fn read_uploaded_file_by_id(
Ok((reference, content))
}
pub(crate) fn uploaded_file_has_pending_owner(dir: &Path, artifact_id: &str) -> Result<bool> {
let path = record_path(dir, artifact_id)?;
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(path)?)?;
Ok(stored.pending_owner_id.is_some())
}
pub(crate) fn read_uploaded_file(dir: &Path, reference: &UploadedFileRef) -> Result<Vec<u8>> {
let (stored_reference, content) = read_uploaded_file_by_id(dir, &reference.artifact_id)?;
if stored_reference.file_name != reference.file_name
@@ -376,6 +391,98 @@ pub(crate) fn clear_uploaded_file_binding(
Ok(())
}
pub(crate) fn pin_uploaded_file(
dir: &Path,
reference: &UploadedFileRef,
owner_id: &str,
) -> Result<()> {
validate_pending_owner_id(owner_id)?;
if reference.source_entry_id.is_some() {
return Err(StoreError::ArtifactAlreadyCommitted);
}
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let path = record_path(dir, &reference.artifact_id)?;
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.file_name != reference.file_name
|| stored.media_type != reference.media_type
|| stored.created_at_ms != reference.created_at_ms
|| stored.byte_len != reference.byte_len
|| stored.sha256 != reference.sha256
{
return Err(StoreError::ArtifactIntegrityMismatch);
}
if stored.source_entry_id.is_some() {
return Err(StoreError::ArtifactAlreadyCommitted);
}
if let Some(existing_owner) = stored.pending_owner_id.as_deref() {
return if existing_owner == owner_id {
Ok(())
} else {
Err(StoreError::ArtifactAlreadyCommitted)
};
}
stored.pending_owner_id = Some(owner_id.to_owned());
let temp = dir.join(format!(".{}.file.pin.tmp", reference.artifact_id));
fs::write(&temp, serde_json::to_vec(&stored)?)?;
fs::rename(temp, path)?;
Ok(())
}
pub(crate) fn release_uploaded_file_pin(
dir: &Path,
artifact_id: &str,
owner_id: &str,
) -> Result<()> {
validate_pending_owner_id(owner_id)?;
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let path = record_path(dir, artifact_id)?;
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.pending_owner_id.as_deref() != Some(owner_id) {
return Err(StoreError::ArtifactIntegrityMismatch);
}
stored.pending_owner_id = None;
let temp = dir.join(format!(".{artifact_id}.file.unpin.tmp"));
fs::write(&temp, serde_json::to_vec(&stored)?)?;
fs::rename(temp, path)?;
Ok(())
}
pub(crate) fn finalize_uploaded_file_binding(
dir: &Path,
artifact_id: &str,
source_entry_id: &str,
) -> Result<()> {
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let path = record_path(dir, artifact_id)?;
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.source_entry_id.as_deref() != Some(source_entry_id) {
return Err(StoreError::ArtifactIntegrityMismatch);
}
if stored.pending_owner_id.is_none() {
return Ok(());
}
stored.pending_owner_id = None;
let temp = dir.join(format!(".{artifact_id}.file.finalize.tmp"));
fs::write(&temp, serde_json::to_vec(&stored)?)?;
fs::rename(temp, path)?;
Ok(())
}
pub(crate) fn bind_uploaded_file(
dir: &Path,
reference: &UploadedFileRef,
@@ -479,6 +586,40 @@ pub(crate) fn copy_committed_uploaded_files(source_dir: &Path, target_dir: &Path
Ok(copied)
}
pub(crate) fn reconcile_uploaded_file_pins(dir: &Path, live_owner_ids: &[String]) -> Result<u64> {
fs::create_dir_all(dir)?;
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let mut reconciled = 0_u64;
for entry in fs::read_dir(dir)? {
let entry = entry?;
let path = entry.path();
let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
continue;
};
let Some(artifact_id) = file_name.strip_suffix(".file.json") else {
continue;
};
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
let Some(owner_id) = stored.pending_owner_id.as_deref() else {
continue;
};
if live_owner_ids.iter().any(|live| live == owner_id) {
continue;
}
stored.pending_owner_id = None;
let temp = dir.join(format!(".{artifact_id}.file.reconcile.tmp"));
fs::write(&temp, serde_json::to_vec(&stored)?)?;
fs::rename(temp, path)?;
reconciled = reconciled.saturating_add(1);
}
Ok(reconciled)
}
pub(crate) fn delete_uncommitted_uploaded_files(dir: &Path) -> Result<u64> {
fs::create_dir_all(dir)?;
let aggregate_lock = fs::OpenOptions::new()
@@ -499,7 +640,7 @@ pub(crate) fn delete_uncommitted_uploaded_files(dir: &Path) -> Result<u64> {
continue;
}
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.source_entry_id.is_none() {
if stored.source_entry_id.is_none() && stored.pending_owner_id.is_none() {
fs::remove_file(path)?;
removed = removed
.checked_add(1)
@@ -523,7 +664,7 @@ pub(crate) fn delete_uploaded_file(dir: &Path, artifact_id: &str) -> Result<bool
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
Err(error) => return Err(error.into()),
};
if stored.source_entry_id.is_some() {
if stored.source_entry_id.is_some() || stored.pending_owner_id.is_some() {
return Err(StoreError::ArtifactAlreadyCommitted);
}
match fs::remove_file(path) {
+13 -3
View File
@@ -99,7 +99,10 @@ async fn in_process_host_runs_text_and_read_tool_then_shuts_down() {
let mut protocol_client = host.connect();
protocol_client
.send(&Method::run_text("read the probe"))
.send(&Method::submit_text(
protocol::new_submission_request_id(),
"read the probe",
))
.await
.expect("submit input");
@@ -336,11 +339,15 @@ async fn standalone_restore_preserves_history_tasks_notifications_and_cwd_scope(
let worker_id = host.worker_id();
let mut protocol_client = host.connect();
protocol_client
.send(&Method::run_text("first request"))
.send(&Method::submit_text(
protocol::new_submission_request_id(),
"first request",
))
.await?;
wait_for_run_end(&mut protocol_client).await?;
protocol_client
.send(&Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message: "persisted notification".to_string(),
auto_run: true,
})
@@ -394,7 +401,10 @@ async fn standalone_restore_preserves_history_tasks_notifications_and_cwd_scope(
assert!(snapshot.contains("persisted notification"), "{snapshot}");
protocol_client
.send(&Method::run_text("continue after restore"))
.send(&Method::submit_text(
protocol::new_submission_request_id(),
"continue after restore",
))
.await?;
wait_for_run_end(&mut protocol_client).await?;
let request = second_inspection
+125 -154
View File
@@ -102,23 +102,6 @@ struct RollbackSubmitState {
turn_before: usize,
}
#[derive(Clone)]
pub struct QueuedInput {
segments: Vec<Segment>,
preview: String,
}
impl QueuedInput {
fn new(segments: Vec<Segment>) -> Self {
let preview = Segment::flatten_to_text(&segments);
Self { segments, preview }
}
pub fn preview(&self) -> &str {
&self.preview
}
}
struct ComposerInputHistory {
entries: VecDeque<Vec<Segment>>,
browse: Option<ComposerInputHistoryBrowse>,
@@ -272,7 +255,7 @@ pub struct App {
/// Current transient actionbar notice. Notices are local UI state only:
/// they are never appended to transcript/session history or LLM context.
actionbar_notice: Option<ActionbarNotice>,
/// Normal composer input that is submitted as `Method::Run`.
/// Normal composer input that is submitted as `Method::Submit`.
pub input: InputBuffer,
/// Separate command-line input. It is never submitted as a user message.
pub command_input: InputBuffer,
@@ -333,9 +316,8 @@ pub struct App {
/// Top entry index of the task pane's visible window. Clamped on
/// render so it never points past the end of the list.
pub task_pane_scroll: usize,
/// TUI-local FIFO of user inputs submitted while the Worker is already running.
/// Entries have not been sent to the Worker yet, so they remain editable/cancellable locally.
queued_inputs: VecDeque<QueuedInput>,
/// Authoritative WorkerSession FIFO summary received from snapshot/live events.
pending_submissions: protocol::PendingSubmissionsSnapshot,
/// TUI-local readline-style composer input history. This is intentionally
/// client-side only: recalled entries are plain drafts until submitted again.
input_history: ComposerInputHistory,
@@ -395,7 +377,7 @@ impl App {
text_selection: TextSelectionState::default(),
task_pane_open: false,
task_pane_scroll: 0,
queued_inputs: VecDeque::new(),
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
input_history: ComposerInputHistory::new(),
input_history_store: None,
pending_submit_rollback: None,
@@ -768,18 +750,34 @@ impl App {
return None;
}
self.record_input_history(segments.clone());
if self.running {
self.queued_inputs.push_back(QueuedInput::new(segments));
self.input.clear();
self.completion = None;
return None;
}
self.input.clear();
Some(self.method_for_run(segments))
}
pub fn submit_notify_input(&mut self) -> Option<Method> {
let segments = self.input.submit_segments();
if segments_are_blank(&segments) {
return None;
}
if segments
.iter()
.any(|segment| matches!(segment, Segment::UploadedFile { .. }))
{
self.push_error("Notify accepts text only; remove attachments or queue a Submit.");
return None;
}
let message = Segment::flatten_to_text(&segments);
self.record_input_history(segments);
self.input.clear();
Some(Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message,
auto_run: true,
})
}
pub fn restore_unsent_run(&mut self, method: &Method) {
let Method::Run { input } = method else {
let Method::Submit { input, .. } = method else {
return;
};
self.pending_submit_rollback = None;
@@ -787,8 +785,9 @@ impl App {
self.input.replace_with_segments(input);
self.completion = None;
} else {
self.queued_inputs
.push_front(QueuedInput::new(input.clone()));
self.push_error(
"Submit transport failed; current Composer was preserved and the unsent input was not queued.",
);
}
}
@@ -804,7 +803,10 @@ impl App {
block_start: self.blocks.len(),
turn_before: self.turn_index,
});
Method::Run { input: segments }
Method::Submit {
submission_request_id: protocol::new_submission_request_id(),
input: segments,
}
}
fn record_input_history(&mut self, segments: Vec<Segment>) {
@@ -825,7 +827,7 @@ impl App {
}
pub fn queued_input_count(&self) -> usize {
self.queued_inputs.len()
self.pending_submissions.submissions.len()
}
#[cfg(test)]
@@ -910,36 +912,31 @@ impl App {
}
}
pub fn continue_pending_method(&self) -> Option<Method> {
Some(Method::ContinuePending {
expected_revision: self.pending_submissions.revision,
expected_head_id: self.pending_submissions.head_id.clone()?,
})
}
pub fn clear_pending_method(&self) -> Method {
Method::ClearPendingSubmissions {
expected_revision: self.pending_submissions.revision,
}
}
pub fn cancel_pending_method(&self, submission_id: String) -> Method {
Method::CancelPendingSubmission {
submission_id,
expected_revision: self.pending_submissions.revision,
}
}
pub fn next_queued_input_preview(&self) -> Option<&str> {
self.queued_inputs.front().map(QueuedInput::preview)
}
pub fn clear_queued_inputs(&mut self) -> usize {
let cleared = self.queued_inputs.len();
self.queued_inputs.clear();
cleared
}
pub fn restore_next_queued_input_to_composer(&mut self) -> bool {
if self.queued_inputs.is_empty() {
return false;
}
if !self.input.is_empty() {
self.push_error("Composer is not empty; clear it before editing queued input.");
return false;
}
let Some(queued) = self.queued_inputs.pop_front() else {
return false;
};
self.input_history.cancel_browse();
self.input.replace_with_segments(&queued.segments);
self.completion = None;
true
}
fn pop_next_queued_run(&mut self) -> Option<Method> {
let queued = self.queued_inputs.pop_front()?;
Some(self.method_for_run(queued.segments))
self.pending_submissions
.submissions
.first()
.map(|submission| submission.submission_id.as_str())
}
pub fn clear_actionbar_notice(&mut self) {
@@ -1123,6 +1120,11 @@ impl App {
}
match event {
Event::SubmissionAccepted { .. } => {}
Event::SubmissionRejected { message, .. } => self.push_error(message),
Event::PendingSubmissionsChanged { pending } => {
self.pending_submissions = pending;
}
Event::UserMessage { segments } => {
self.turn_index += 1;
self.blocks.push(Block::TurnHeader {
@@ -1372,9 +1374,6 @@ impl App {
WorkerStatus::Idle
}
});
if matches!(result, RunResult::Finished | RunResult::LimitReached) {
return self.pop_next_queued_run();
}
}
}
Event::CompactStart { .. } => {
@@ -1449,6 +1448,7 @@ impl App {
internal_workers,
} => {
self.rewind_refresh_fence = false;
self.pending_submissions = session.pending_submissions.clone();
self.restore_snapshot(&session, greeting, in_flight);
self.replace_internal_worker_snapshots(internal_workers);
self.set_worker_status(status);
@@ -2681,7 +2681,10 @@ mod rewind_refresh_tests {
});
app.handle_worker_event(Event::RewindApplied {
session: protocol::SessionSnapshot { entries: vec![] },
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: vec![],
},
input: vec![Segment::text("selected rewind input")],
summary: summary(3),
});
@@ -2700,7 +2703,10 @@ mod rewind_refresh_tests {
});
app.handle_worker_event(Event::RewindApplied {
session: protocol::SessionSnapshot { entries: vec![] },
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: vec![],
},
input: vec![Segment::text("rewound input")],
summary: summary(1),
});
@@ -2743,7 +2749,10 @@ mod rewind_refresh_tests {
});
app.handle_worker_event(Event::RewindApplied {
session: protocol::SessionSnapshot { entries: vec![] },
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: vec![],
},
input: vec![Segment::text("rewound input")],
summary: summary(2),
});
@@ -2877,7 +2886,7 @@ mod composer_history_persistence_tests {
path: "src/lib.rs".into(),
},
]);
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
let mut reloaded = App::new_with_input_history_store("test".into(), store);
assert!(reloaded.browse_input_history_older());
@@ -2958,7 +2967,7 @@ mod composer_history_persistence_tests {
app.insert_char(c);
}
match app.submit_input() {
Some(Method::Run { input }) => input,
Some(Method::Submit { input, .. }) => input,
other => panic!("expected Run, got {other:?}"),
}
}
@@ -3424,72 +3433,44 @@ mod completion_flow_tests {
}
#[test]
fn running_submit_is_queued_locally_and_clears_composer() {
fn running_submit_is_sent_to_the_worker_and_not_queued_locally() {
let mut app = App::new("test".into());
app.set_worker_status(WorkerStatus::Running);
insert_text(&mut app, "queued turn");
assert!(app.submit_input().is_none());
let method = app.submit_input();
assert_eq!(app.queued_input_count(), 1);
assert_eq!(app.next_queued_input_preview(), Some("queued turn"));
assert!(matches!(method, Some(Method::Submit { .. })));
assert_eq!(app.queued_input_count(), 0);
assert_eq!(input_text(&app), "");
}
#[test]
fn finished_run_auto_sends_next_queued_input() {
fn pending_submission_projection_is_worker_authoritative() {
let mut app = App::new("test".into());
app.set_worker_status(WorkerStatus::Running);
insert_text(&mut app, "next turn");
assert!(app.submit_input().is_none());
let method = app.handle_worker_event(Event::RunEnd {
result: RunResult::Finished,
app.handle_worker_event(Event::PendingSubmissionsChanged {
pending: protocol::PendingSubmissionsSnapshot {
revision: 3,
notification_count: 0,
head_id: Some("submission-1".into()),
submissions: vec![protocol::PendingSubmissionSummary {
submission_id: "submission-1".into(),
accepted_at_ms: 7,
segment_count: 2,
byte_len: 42,
}],
},
});
match method {
Some(Method::Run { input }) => {
assert_eq!(Segment::flatten_to_text(&input), "next turn");
}
other => panic!("expected queued Run, got {other:?}"),
}
assert_eq!(app.queued_input_count(), 0);
}
#[test]
fn limit_reached_run_auto_sends_next_queued_input() {
let mut app = App::new("test".into());
app.set_worker_status(WorkerStatus::Running);
insert_text(&mut app, "next after limit");
assert!(app.submit_input().is_none());
let method = app.handle_worker_event(Event::RunEnd {
result: RunResult::LimitReached,
});
match method {
Some(Method::Run { input }) => {
assert_eq!(Segment::flatten_to_text(&input), "next after limit");
}
other => panic!("expected queued Run, got {other:?}"),
}
assert_eq!(app.queued_input_count(), 0);
}
#[test]
fn paused_and_rolled_back_run_do_not_auto_send_queue() {
for result in [RunResult::Paused, RunResult::RolledBack] {
let mut app = App::new("test".into());
app.set_worker_status(WorkerStatus::Running);
insert_text(&mut app, "held turn");
assert!(app.submit_input().is_none());
let method = app.handle_worker_event(Event::RunEnd { result });
assert!(method.is_none());
assert_eq!(app.queued_input_count(), 1);
assert_eq!(app.next_queued_input_preview(), Some("held turn"));
}
assert_eq!(app.queued_input_count(), 1);
assert_eq!(app.next_queued_input_preview(), Some("submission-1"));
assert!(
app.handle_worker_event(Event::RunEnd {
result: RunResult::Finished,
})
.is_none()
);
assert_eq!(app.queued_input_count(), 1);
}
#[test]
@@ -3501,24 +3482,6 @@ mod completion_flow_tests {
assert_eq!(app.queued_input_count(), 0);
}
#[test]
fn queued_input_can_be_restored_to_composer_or_cleared() {
let mut app = App::new("test".into());
app.set_worker_status(WorkerStatus::Running);
insert_text(&mut app, "edit me");
assert!(app.submit_input().is_none());
assert!(app.restore_next_queued_input_to_composer());
assert_eq!(app.queued_input_count(), 0);
assert_eq!(input_text(&app), "edit me");
app.input.clear();
insert_text(&mut app, "clear me");
assert!(app.submit_input().is_none());
assert_eq!(app.clear_queued_inputs(), 1);
assert_eq!(app.queued_input_count(), 0);
}
fn insert_text(app: &mut App, text: &str) {
for c in text.chars() {
app.insert_char(c);
@@ -3530,7 +3493,7 @@ mod completion_flow_tests {
app.insert_char(c);
}
match app.submit_input() {
Some(Method::Run { input }) => input,
Some(Method::Submit { input, .. }) => input,
other => panic!("expected Run, got {other:?}"),
}
}
@@ -3675,6 +3638,7 @@ mod completion_flow_tests {
app.handle_worker_event(Event::Snapshot {
greeting: test_greeting(),
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
status: WorkerStatus::Running,
@@ -3783,6 +3747,7 @@ mod completion_flow_tests {
revision,
status: WorkerStatus::Idle,
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
in_flight: protocol::InFlightSnapshot::default(),
@@ -4000,6 +3965,7 @@ mod completion_flow_tests {
app.handle_worker_event(Event::Snapshot {
greeting: test_greeting(),
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
status: WorkerStatus::Idle,
@@ -4051,6 +4017,7 @@ mod completion_flow_tests {
app.handle_worker_event(Event::Snapshot {
greeting: test_greeting(),
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
status: WorkerStatus::Idle,
@@ -4064,6 +4031,7 @@ mod completion_flow_tests {
},
revision: 4,
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
status: WorkerStatus::Running,
@@ -4222,6 +4190,7 @@ mod completion_flow_tests {
app.handle_worker_event(Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting,
@@ -4437,23 +4406,23 @@ mod completion_flow_tests {
}
#[test]
fn input_history_records_queued_inputs_and_suppresses_consecutive_duplicates() {
fn input_history_records_running_submits_and_suppresses_consecutive_duplicates() {
let mut app = App::new("test".into());
app.running = true;
for c in "repeat".chars() {
app.insert_char(c);
}
assert!(app.submit_input().is_none());
assert!(app.submit_input().is_some());
assert_eq!(app.input_history_len(), 1);
assert_eq!(app.queued_input_count(), 1);
assert_eq!(app.queued_input_count(), 0);
for c in "repeat".chars() {
app.insert_char(c);
}
assert!(app.submit_input().is_none());
assert!(app.submit_input().is_some());
assert_eq!(app.input_history_len(), 1);
assert_eq!(app.queued_input_count(), 2);
assert_eq!(app.queued_input_count(), 0);
app.insert_char(' ');
assert!(app.submit_input().is_none());
@@ -4481,7 +4450,7 @@ mod completion_flow_tests {
},
];
app.input.replace_with_segments(&original);
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
assert!(app.browse_input_history_older());
assert_eq!(app.input.submit_segments(), original);
@@ -4493,7 +4462,7 @@ mod completion_flow_tests {
for c in "sent".chars() {
app.insert_char(c);
}
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
for c in "draft".chars() {
app.insert_char(c);
@@ -4511,7 +4480,7 @@ mod completion_flow_tests {
for c in "sent".chars() {
app.insert_char(c);
}
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
assert!(app.browse_input_history_older());
assert!(app.input_history_is_browsing());
@@ -4528,17 +4497,19 @@ mod completion_flow_tests {
for c in "first".chars() {
app.insert_char(c);
}
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
for c in "second".chars() {
app.insert_char(c);
}
assert!(matches!(app.submit_input(), Some(Method::Run { .. })));
assert!(matches!(app.submit_input(), Some(Method::Submit { .. })));
assert!(app.browse_input_history_older());
assert!(app.browse_input_history_older());
let method = app.submit_input();
match method {
Some(Method::Run { input }) => assert_eq!(Segment::flatten_to_text(&input), "first"),
Some(Method::Submit { input, .. }) => {
assert_eq!(Segment::flatten_to_text(&input), "first")
}
other => panic!("expected recalled run, got {other:?}"),
}
assert_eq!(app.input_history_len(), 3);
+156 -110
View File
@@ -270,8 +270,8 @@ impl<T: Socket> ConsoleConnection<T> {
async fn send(&mut self, method: &Method) -> Result<(), Box<dyn std::error::Error>> {
let mut prepared = method.clone();
let carries_attachments =
matches!(prepared, Method::Run { .. }) && !self.pending_attachments.is_empty();
if let Method::Run { input } = &mut prepared {
matches!(prepared, Method::Submit { .. }) && !self.pending_attachments.is_empty();
if let Method::Submit { input, .. } = &mut prepared {
input.extend(
self.pending_attachments
.iter()
@@ -569,6 +569,7 @@ async fn run_e2e_rewind_fixture(
app.connected = true;
app.handle_worker_event(Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
status: WorkerStatus::Idle,
@@ -697,6 +698,7 @@ async fn run_e2e_rewind_fixture(
if submitted_at.elapsed() >= apply_delay {
app.handle_worker_event(Event::RewindApplied {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
input: vec![Segment::text("rewind-live-refresh")],
@@ -916,7 +918,7 @@ async fn run_loop<T: Socket>(
}
fn attachment_command_path(method: &Method) -> Option<PathBuf> {
let Method::Run { input } = method else {
let Method::Submit { input, .. } = method else {
return None;
};
let [Segment::Text { content }] = input.as_slice() else {
@@ -927,7 +929,7 @@ fn attachment_command_path(method: &Method) -> Option<PathBuf> {
}
fn is_clear_attachments_command(method: &Method) -> bool {
let Method::Run { input } = method else {
let Method::Submit { input, .. } = method else {
return false;
};
matches!(
@@ -941,7 +943,7 @@ async fn send_console_method<T: Socket>(
client: &mut ConsoleConnection<T>,
method: &Method,
) -> Result<(), Box<dyn std::error::Error>> {
if matches!(method, Method::Run { .. }) && client.has_active_uploads() {
if matches!(method, Method::Submit { .. }) && client.has_active_uploads() {
app.restore_unsent_run(method);
app.flash_actionbar_notice(
"Attachment upload is still in progress; wait or use /clear-attachments.",
@@ -953,7 +955,7 @@ async fn send_console_method<T: Socket>(
}
let sends_attachments =
matches!(method, Method::Run { .. }) && !client.pending_attachments.is_empty();
matches!(method, Method::Submit { .. }) && !client.pending_attachments.is_empty();
if let Err(error) = client.send(method).await {
if sends_attachments {
app.restore_unsent_run(method);
@@ -1148,18 +1150,27 @@ fn handle_key(app: &mut App, key: KeyEvent) -> Option<Method> {
app.clear_command_input();
Some(None)
}
KeyCode::Char(c)
if c.eq_ignore_ascii_case(&'d') && alt && !ctrl && !app.is_command_mode() =>
{
Some(
app.next_queued_input_preview()
.map(str::to_owned)
.map(|submission_id| app.cancel_pending_method(submission_id)),
)
}
KeyCode::Char(c)
if c.eq_ignore_ascii_case(&'n') && alt && !ctrl && !app.is_command_mode() =>
{
Some(app.submit_notify_input())
}
KeyCode::Char(c)
if c.eq_ignore_ascii_case(&'q') && alt && !ctrl && !app.is_command_mode() =>
{
if app.restore_next_queued_input_to_composer() {
Some(app.refresh_completion())
} else {
Some(None)
}
Some(app.continue_pending_method())
}
KeyCode::Char(c) if c.eq_ignore_ascii_case(&'c') && alt && !ctrl => {
app.clear_queued_inputs();
Some(None)
Some(Some(app.clear_pending_method()))
}
KeyCode::Char('c') if ctrl => Some(handle_pause_or_quit(app)),
KeyCode::Char('x') if ctrl => Some(handle_cancel_or_shutdown(app)),
@@ -1427,7 +1438,6 @@ fn handle_cancel_or_shutdown(app: &mut App) -> Option<Method> {
WorkerStatus::Running | WorkerStatus::Paused
) {
app.shutdown_confirm = None;
app.clear_queued_inputs();
return Some(Method::Cancel);
}
if let Some(pressed_at) = app.shutdown_confirm
@@ -1450,7 +1460,6 @@ fn handle_cancel_or_shutdown(app: &mut App) -> Option<Method> {
/// Idle / Paused → 2-tap to quit the TUI (the Worker keeps running).
fn handle_pause_or_quit(app: &mut App) -> Option<Method> {
if app.worker_status == WorkerStatus::Running {
app.clear_queued_inputs();
return Some(Method::Pause);
}
if let Some(t) = app.quit_confirm
@@ -1476,8 +1485,8 @@ mod tests {
use crate::text_selection::{HistoryViewport, SelectionRow};
use async_trait::async_trait;
use protocol::{
Event, RewindTarget, RewindTargetId, RunResult, Segment, UploadedFileAvailability,
UploadedFileRef, WorkerStatus,
Event, RewindTarget, RewindTargetId, Segment, UploadedFileAvailability, UploadedFileRef,
WorkerStatus,
};
#[test]
@@ -1490,7 +1499,8 @@ mod tests {
#[test]
fn client_local_attachment_commands_are_typed_and_do_not_send_the_path() {
let attach = Method::Run {
let attach = Method::Submit {
submission_request_id: protocol::new_submission_request_id(),
input: vec![Segment::text("/attach /tmp/report.md")],
};
assert_eq!(
@@ -1499,7 +1509,8 @@ mod tests {
);
assert!(!is_clear_attachments_command(&attach));
let clear = Method::Run {
let clear = Method::Submit {
submission_request_id: protocol::new_submission_request_id(),
input: vec![Segment::text("/clear-attachments")],
};
assert!(is_clear_attachments_command(&clear));
@@ -1605,7 +1616,7 @@ mod tests {
}
#[tokio::test]
async fn queued_attachment_send_failure_restores_draft_without_exiting_console() {
async fn running_attachment_submit_failure_restores_draft_without_exiting_console() {
let file = UploadedFileRef {
artifact_id: "artifact-queued".into(),
file_name: "queued.txt".into(),
@@ -1631,13 +1642,10 @@ mod tests {
let mut app = App::new("worker".into());
app.set_worker_status(WorkerStatus::Running);
app.input.insert_str("queued inspect");
assert!(app.submit_input().is_none());
let method = app
.handle_worker_event(Event::RunEnd {
result: RunResult::Finished,
})
.expect("queued run must be released");
.submit_input()
.expect("running Submit is sent immediately");
send_console_method(&mut app, &mut connection, &method)
.await
.unwrap();
@@ -1960,7 +1968,7 @@ mod tests {
}
#[test]
fn running_enter_queues_instead_of_sending_run() {
fn running_enter_sends_submit_to_worker() {
let mut app = App::new("agent".to_string());
app.set_worker_status(WorkerStatus::Running);
for c in "queued".chars() {
@@ -1973,102 +1981,128 @@ mod tests {
);
}
assert!(handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)).is_none());
assert!(matches!(
handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)),
Some(Method::Submit { .. })
));
assert_eq!(app.queued_input_count(), 1);
assert_eq!(app.next_queued_input_preview(), Some("queued"));
assert_eq!(app.queued_input_count(), 0);
assert_eq!(input_text(&app), "");
}
#[test]
fn queued_input_keybindings_restore_and_clear() {
let mut app = App::new("agent".to_string());
fn running_alt_n_sends_explicit_notify_without_implicit_submit_conversion() {
let mut app = App::new("test".into());
app.set_worker_status(WorkerStatus::Running);
for c in "edit queued".chars() {
assert!(
handle_key(
&mut app,
KeyEvent::new(KeyCode::Char(c), KeyModifiers::NONE)
)
.is_none()
);
for character in "progress".chars() {
app.insert_char(character);
}
assert!(handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)).is_none());
assert!(
handle_key(
&mut app,
KeyEvent::new(KeyCode::Char('q'), KeyModifiers::ALT)
)
.is_none()
let method = handle_key(
&mut app,
KeyEvent::new(KeyCode::Char('n'), KeyModifiers::ALT),
);
assert_eq!(app.queued_input_count(), 0);
assert_eq!(input_text(&app), "edit queued");
app.input.clear();
for c in "clear queued".chars() {
assert!(
handle_key(
&mut app,
KeyEvent::new(KeyCode::Char(c), KeyModifiers::NONE)
)
.is_none()
);
}
assert!(handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)).is_none());
assert_eq!(app.queued_input_count(), 1);
assert!(
handle_key(
&mut app,
KeyEvent::new(KeyCode::Char('c'), KeyModifiers::ALT)
)
.is_none()
);
assert_eq!(app.queued_input_count(), 0);
assert!(matches!(
method,
Some(Method::Notify {
ref message,
auto_run: true,
..
}) if message == "progress"
));
assert_eq!(input_text(&app), "");
}
#[test]
fn pause_and_cancel_clear_queued_input() {
let mut app = App::new("agent".to_string());
app.set_worker_status(WorkerStatus::Running);
for c in "queued".chars() {
assert!(
handle_key(
&mut app,
KeyEvent::new(KeyCode::Char(c), KeyModifiers::NONE)
)
.is_none()
);
}
assert!(handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)).is_none());
assert_eq!(app.queued_input_count(), 1);
fn pending_queue_shortcuts_send_worker_operations() {
let mut app = App::new("test".into());
app.handle_worker_event(Event::PendingSubmissionsChanged {
pending: protocol::PendingSubmissionsSnapshot {
revision: 2,
notification_count: 0,
head_id: Some("submission-1".into()),
submissions: vec![protocol::PendingSubmissionSummary {
submission_id: "submission-1".into(),
accepted_at_ms: 1,
segment_count: 1,
byte_len: 6,
}],
},
});
let pause = handle_key(
let continue_next = handle_key(
&mut app,
KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL),
KeyEvent::new(KeyCode::Char('q'), KeyModifiers::ALT),
);
assert!(matches!(pause, Some(Method::Pause)));
assert_eq!(app.queued_input_count(), 0);
for c in "queued again".chars() {
assert!(
handle_key(
&mut app,
KeyEvent::new(KeyCode::Char(c), KeyModifiers::NONE)
)
.is_none()
);
}
assert!(handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)).is_none());
assert!(matches!(
continue_next,
Some(Method::ContinuePending {
expected_revision: 2,
ref expected_head_id,
}) if expected_head_id == "submission-1"
));
assert_eq!(app.queued_input_count(), 1);
let cancel = handle_key(
&mut app,
KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL),
KeyEvent::new(KeyCode::Char('d'), KeyModifiers::ALT),
);
assert!(matches!(cancel, Some(Method::Cancel)));
assert_eq!(app.queued_input_count(), 0);
assert!(matches!(
cancel,
Some(Method::CancelPendingSubmission {
expected_revision: 2,
ref submission_id,
}) if submission_id == "submission-1"
));
let clear = handle_key(
&mut app,
KeyEvent::new(KeyCode::Char('c'), KeyModifiers::ALT),
);
assert!(matches!(
clear,
Some(Method::ClearPendingSubmissions {
expected_revision: 2
})
));
assert_eq!(app.queued_input_count(), 1);
}
#[test]
fn pause_and_cancel_preserve_authoritative_pending_queue() {
let mut app = App::new("test".into());
app.handle_worker_event(Event::PendingSubmissionsChanged {
pending: protocol::PendingSubmissionsSnapshot {
revision: 2,
notification_count: 0,
head_id: Some("submission-1".into()),
submissions: vec![protocol::PendingSubmissionSummary {
submission_id: "submission-1".into(),
accepted_at_ms: 1,
segment_count: 1,
byte_len: 6,
}],
},
});
app.set_worker_status(WorkerStatus::Running);
assert!(matches!(
handle_key(
&mut app,
KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL),
),
Some(Method::Pause)
));
assert_eq!(app.queued_input_count(), 1);
app.set_worker_status(WorkerStatus::Running);
assert!(matches!(
handle_key(
&mut app,
KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL),
),
Some(Method::Cancel)
));
assert_eq!(app.queued_input_count(), 1);
}
#[test]
@@ -2535,13 +2569,19 @@ mod tests {
let mut app = App::new("agent".to_string());
app.handle_worker_event(Event::Snapshot {
greeting: test_greeting(),
session: protocol::SessionSnapshot { entries: vec![] },
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: vec![],
},
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
});
app.handle_worker_event(Event::RewindApplied {
session: protocol::SessionSnapshot { entries: vec![] },
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: vec![],
},
input: vec![Segment::Text {
content: "retry this".into(),
}],
@@ -2562,7 +2602,10 @@ mod tests {
let mut app = App::new("agent".to_string());
app.handle_worker_event(Event::Snapshot {
greeting: test_greeting(),
session: protocol::SessionSnapshot { entries: vec![] },
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: vec![],
},
status: WorkerStatus::Idle,
in_flight: Default::default(),
internal_workers: Vec::new(),
@@ -2570,7 +2613,10 @@ mod tests {
type_keys(&mut app, "draft");
app.handle_worker_event(Event::RewindApplied {
session: protocol::SessionSnapshot { entries: vec![] },
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: vec![],
},
input: vec![Segment::Text {
content: "retry this".into(),
}],
@@ -2918,12 +2964,12 @@ mod tests {
type_keys(&mut app, "first");
assert!(matches!(
handle_key(&mut app, key(KeyCode::Enter)),
Some(Method::Run { .. })
Some(Method::Submit { .. })
));
type_keys(&mut app, "second");
assert!(matches!(
handle_key(&mut app, key(KeyCode::Enter)),
Some(Method::Run { .. })
Some(Method::Submit { .. })
));
assert_eq!(input_text(&app), "");
@@ -2954,7 +3000,7 @@ mod tests {
type_keys(&mut app, "sent");
assert!(matches!(
handle_key(&mut app, key(KeyCode::Enter)),
Some(Method::Run { .. })
Some(Method::Submit { .. })
));
type_keys(&mut app, "draft\nbody");
app.move_cursor_start();
+22 -14
View File
@@ -1880,7 +1880,7 @@ fn actionbar_left_item(app: &App, now: Instant) -> Option<(String, Style)> {
}
if app.queued_input_count() > 0 {
return Some((
"Alt-q edit queued Alt-c clear queued".to_string(),
"Alt-n notify Alt-q continue Alt-d cancel queued Alt-c clear queued".to_string(),
Style::default().fg(Color::DarkGray),
));
}
@@ -2136,9 +2136,25 @@ mod tests {
use super::*;
use crate::app::{ActionbarNoticeLevel, ActionbarNoticeSource, App};
use crate::block::{ToolCallBlock, ToolCallState};
use protocol::WorkerStatus;
use protocol::Event;
use std::time::{Duration, Instant};
fn set_pending_submission(app: &mut App, id: &str) {
app.handle_worker_event(Event::PendingSubmissionsChanged {
pending: protocol::PendingSubmissionsSnapshot {
revision: 1,
notification_count: 0,
head_id: Some(id.into()),
submissions: vec![protocol::PendingSubmissionSummary {
submission_id: id.into(),
accepted_at_ms: 1,
segment_count: 1,
byte_len: 1,
}],
},
});
}
#[test]
fn run_status_line_matches_console_metrics_and_spinner_frame() {
let now = Instant::now();
@@ -2251,15 +2267,11 @@ mod tests {
#[test]
fn queue_status_text_includes_count_and_preview() {
let mut app = App::new("test".into());
app.set_worker_status(WorkerStatus::Running);
for c in "queued preview".chars() {
app.insert_char(c);
}
assert!(app.submit_input().is_none());
set_pending_submission(&mut app, "submission-1");
assert_eq!(
queue_status_text(&app),
Some("queued: 1 — queued preview".to_string())
Some("queued: 1 — submission-1".to_string())
);
}
@@ -2289,14 +2301,10 @@ mod tests {
Some("Worker keeps running. Press Ctrl-C again to exit TUI.".into())
);
app.set_worker_status(WorkerStatus::Running);
for c in "queued turn".chars() {
app.insert_char(c);
}
assert!(app.submit_input().is_none());
set_pending_submission(&mut app, "submission-1");
assert_eq!(
actionbar_left_item(&app, now).map(|(text, _)| text),
Some("Alt-q edit queued Alt-c clear queued".into())
Some("Alt-n notify Alt-q continue Alt-d cancel queued Alt-c clear queued".into())
);
app.enter_command_mode();
+21 -16
View File
@@ -41,14 +41,12 @@ pub enum WorkerExecutionOperation {
Cancel,
}
/// Evidence that a user input reached the durable Worker session boundary.
///
/// This is intentionally distinct from accepting a method on the Worker's
/// in-memory channel. For Flow submissions, the committed UserInput entry also
/// carries the initial Flow runtime-state extension.
/// Evidence that a Submit request reached the durable Worker session boundary.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkerInputCommitAck {
pub struct WorkerSubmissionAck {
pub submission_request_id: String,
pub submission_id: String,
pub disposition: protocol::SubmissionDisposition,
}
/// Typed execution result class. Results are transient operation outcomes and
@@ -61,7 +59,7 @@ pub struct WorkerExecutionResult {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub message: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input_commit: Option<WorkerInputCommitAck>,
pub submission: Option<WorkerSubmissionAck>,
}
/// Backend result class for a Worker execution operation.
@@ -85,22 +83,26 @@ impl WorkerExecutionResult {
outcome: WorkerExecutionOutcome::Accepted,
run_state,
message: None,
input_commit: None,
submission: None,
}
}
pub fn accepted_input_committed(
pub fn accepted_submission(
operation: WorkerExecutionOperation,
run_state: WorkerExecutionRunState,
submission_request_id: impl Into<String>,
submission_id: impl Into<String>,
disposition: protocol::SubmissionDisposition,
) -> Self {
Self {
operation,
outcome: WorkerExecutionOutcome::Accepted,
run_state,
message: None,
input_commit: Some(WorkerInputCommitAck {
submission: Some(WorkerSubmissionAck {
submission_request_id: submission_request_id.into(),
submission_id: submission_id.into(),
disposition,
}),
}
}
@@ -111,7 +113,7 @@ impl WorkerExecutionResult {
outcome: WorkerExecutionOutcome::Busy,
run_state: WorkerExecutionRunState::Busy,
message: Some(message.into()),
input_commit: None,
submission: None,
}
}
@@ -121,7 +123,7 @@ impl WorkerExecutionResult {
outcome: WorkerExecutionOutcome::Rejected,
run_state: WorkerExecutionRunState::Stopped,
message: Some(message.into()),
input_commit: None,
submission: None,
}
}
@@ -131,7 +133,7 @@ impl WorkerExecutionResult {
outcome: WorkerExecutionOutcome::Errored,
run_state: WorkerExecutionRunState::Errored,
message: Some(message.into()),
input_commit: None,
submission: None,
}
}
@@ -141,7 +143,7 @@ impl WorkerExecutionResult {
outcome: WorkerExecutionOutcome::Unsupported,
run_state: WorkerExecutionRunState::Stopped,
message: Some(message.into()),
input_commit: None,
submission: None,
}
}
@@ -618,14 +620,17 @@ mod tests {
use super::*;
#[test]
fn input_commit_ack_survives_json_round_trip() {
let result = WorkerExecutionResult::accepted_input_committed(
fn submission_ack_survives_json_round_trip() {
let result = WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
"request-1",
"submission-1",
protocol::SubmissionDisposition::Started,
);
let json = serde_json::to_string(&result).unwrap();
assert!(json.contains("\"submission_request_id\":\"request-1\""));
assert!(json.contains("\"submission_id\":\"submission-1\""));
assert_eq!(
serde_json::from_str::<WorkerExecutionResult>(&json).unwrap(),
+142 -5
View File
@@ -1240,10 +1240,12 @@ async fn worker_protocol_ws(
auth: Option<Extension<RuntimeAuthContext>>,
Path(worker_id): Path<String>,
Query(query): Query<RuntimeWorkerEventsWsQuery>,
headers: HeaderMap,
ws: WebSocketUpgrade,
) -> Result<Response, RuntimeHttpRestError> {
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
let scope = auth_workspace_scope(&state, auth.as_ref())?;
let input_source = authenticated_protocol_input_source(&headers)?;
match scope.as_ref() {
Some(scope) => state
.runtime
@@ -1254,17 +1256,89 @@ async fn worker_protocol_ws(
.map_err(RuntimeHttpRestError::runtime)?;
Ok(ws
.on_upgrade(move |socket| {
worker_protocol_ws_session(state.runtime, scope, worker_ref, query, socket)
worker_protocol_ws_session(
state.runtime,
scope,
worker_ref,
query,
input_source,
socket,
)
})
.into_response())
}
#[cfg(feature = "ws-server")]
fn authenticated_protocol_input_source(
headers: &HeaderMap,
) -> Result<Option<protocol::AuthenticatedInputSource>, RuntimeHttpRestError> {
let Some(value) = headers.get(protocol::AUTHENTICATED_ACCOUNT_ID_HEADER) else {
return Ok(None);
};
let account_id = value.to_str().map_err(|_| {
RuntimeHttpRestError::new(
StatusCode::BAD_REQUEST,
"authenticated_input_source_invalid",
"authenticated Worker input source is invalid",
)
})?;
if account_id.trim().is_empty() || account_id.len() > 128 {
return Err(RuntimeHttpRestError::new(
StatusCode::BAD_REQUEST,
"authenticated_input_source_invalid",
"authenticated Worker input source is invalid",
));
}
Ok(Some(protocol::AuthenticatedInputSource::Account {
account_id: account_id.to_owned(),
}))
}
#[cfg(feature = "ws-server")]
fn authorize_runtime_protocol_method(
method: protocol::Method,
transport_source: Option<&protocol::AuthenticatedInputSource>,
) -> protocol::Method {
match method {
protocol::Method::SubmitTracked {
submission_request_id,
input,
..
} => protocol::Method::SubmitTracked {
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,
auto_run,
..
} => protocol::Method::NotifyTracked {
source: transport_source.cloned().unwrap_or_else(|| {
protocol::AuthenticatedInputSource::Backend {
operation_id: notification_request_id.clone(),
}
}),
notification_request_id,
message,
auto_run,
},
other => other,
}
}
#[cfg(feature = "ws-server")]
async fn worker_protocol_ws_session(
runtime: Runtime,
scope: Option<RuntimeWorkspaceScope>,
worker_ref: WorkerRef,
query: RuntimeWorkerEventsWsQuery,
input_source: Option<protocol::AuthenticatedInputSource>,
mut socket: WebSocket,
) {
let mut cursor = match query.cursor.as_deref() {
@@ -1347,6 +1421,8 @@ async fn worker_protocol_ws_session(
match inbound {
Some(Ok(WsMessage::Text(text))) => match decode_method(&text) {
Ok(method) => {
let method =
authorize_runtime_protocol_method(method, input_source.as_ref());
let result = match scope.as_ref() {
Some(scope) => {
runtime.send_protocol_method_scoped(scope, &worker_ref, method)
@@ -2219,6 +2295,63 @@ mod tests {
);
}
#[test]
fn runtime_protocol_replaces_serialized_tracked_source() {
let wire = serde_json::to_string(&protocol::Method::SubmitTracked {
submission_request_id: "request-1".into(),
input: vec![protocol::Segment::text("hello")],
source: protocol::AuthenticatedInputSource::Account {
account_id: "forged".into(),
},
})
.unwrap();
let decoded: protocol::Method = serde_json::from_str(&wire).unwrap();
assert!(matches!(
decoded,
protocol::Method::SubmitTracked {
source: protocol::AuthenticatedInputSource::UntrustedWire,
..
}
));
assert!(matches!(
authorize_runtime_protocol_method(decoded, None),
protocol::Method::SubmitTracked {
source: protocol::AuthenticatedInputSource::Backend { operation_id },
..
} if operation_id == "request-1"
));
}
#[test]
fn runtime_protocol_uses_transport_authenticated_account_source() {
let mut headers = HeaderMap::new();
headers.insert(
protocol::AUTHENTICATED_ACCOUNT_ID_HEADER,
"account-1".parse().unwrap(),
);
let source = authenticated_protocol_input_source(&headers)
.unwrap()
.expect("account source header must resolve");
let wire = serde_json::to_string(&protocol::Method::NotifyTracked {
notification_request_id: "notification-1".into(),
message: "hello".into(),
auto_run: true,
source: protocol::AuthenticatedInputSource::Account {
account_id: "forged".into(),
},
})
.unwrap();
let decoded: protocol::Method = serde_json::from_str(&wire).unwrap();
assert!(matches!(
authorize_runtime_protocol_method(decoded, Some(&source)),
protocol::Method::NotifyTracked {
source: protocol::AuthenticatedInputSource::Account { account_id },
..
} if account_id == "account-1"
));
}
#[test]
fn attachment_routes_require_worker_input_permission() {
assert_eq!(
@@ -2870,11 +3003,13 @@ mod tests {
_handle: &WorkerExecutionHandle,
input: WorkerInput,
) -> WorkerExecutionResult {
if let Some(submission_id) = input.submission_id {
WorkerExecutionResult::accepted_input_committed(
if let Some(submission_id) = input.submission_request_id {
WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
submission_id.clone(),
submission_id,
protocol::SubmissionDisposition::Started,
)
} else {
WorkerExecutionResult::accepted(
@@ -3194,11 +3329,13 @@ mod ws_tests {
_handle: &WorkerExecutionHandle,
input: WorkerInput,
) -> WorkerExecutionResult {
if let Some(submission_id) = input.submission_id {
WorkerExecutionResult::accepted_input_committed(
if let Some(submission_id) = input.submission_request_id {
WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
submission_id.clone(),
submission_id,
protocol::SubmissionDisposition::Started,
)
} else {
WorkerExecutionResult::accepted(
+24 -6
View File
@@ -25,10 +25,10 @@ impl WorkerInputKind {
pub struct WorkerInput {
pub kind: WorkerInputKind,
pub content: String,
/// Runtime-generated correlation id. This is never accepted from public
/// JSON input and is consumed only by the execution backend.
#[serde(skip)]
pub submission_id: Option<String>,
/// Authenticated client-generated idempotency key. Runtime generates one
/// only for trusted internal callers that omit it.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub submission_request_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub segments: Option<Vec<Segment>>,
}
@@ -38,7 +38,7 @@ impl WorkerInput {
Self {
kind: WorkerInputKind::User,
content: content.into(),
submission_id: None,
submission_request_id: None,
segments: None,
}
}
@@ -47,7 +47,7 @@ impl WorkerInput {
Self {
kind: WorkerInputKind::Notify,
content: content.into(),
submission_id: None,
submission_request_id: None,
segments: None,
}
}
@@ -57,6 +57,21 @@ impl WorkerInput {
mod tests {
use super::WorkerInput;
#[test]
fn submission_request_id_round_trips_for_authenticated_client_retry() {
let input: WorkerInput = serde_json::from_value(serde_json::json!({
"kind": "user",
"content": "message",
"submission_request_id": "request-1"
}))
.unwrap();
assert_eq!(input.submission_request_id.as_deref(), Some("request-1"));
assert_eq!(
serde_json::to_value(input).unwrap()["submission_request_id"],
"request-1"
);
}
#[test]
fn notify_is_an_operation_and_legacy_system_kind_is_rejected() {
assert_eq!(
@@ -78,4 +93,7 @@ mod tests {
pub struct WorkerInteractionAck {
pub worker_ref: WorkerRef,
pub status: WorkerStatus,
/// Present for User Submit and absent for non-Submit interactions.
#[serde(skip_serializing_if = "Option::is_none")]
pub submission: Option<crate::execution::WorkerSubmissionAck>,
}
+62 -50
View File
@@ -748,8 +748,12 @@ impl Runtime {
let state = self.lock()?;
state.worker(&worker_ref)?.request.initial_input.clone()
} {
let expected_submission_id = Uuid::now_v7().to_string();
initial_input.submission_id = Some(expected_submission_id.clone());
let expected_submission_id = initial_input
.submission_request_id
.clone()
.filter(|request_id| !request_id.trim().is_empty())
.unwrap_or_else(|| Uuid::now_v7().to_string());
initial_input.submission_request_id = Some(expected_submission_id.clone());
let dispatch_result = backend.dispatch_input(&handle, initial_input.clone());
if !dispatch_result.is_accepted() {
let _ = backend.stop_worker(&handle);
@@ -763,9 +767,9 @@ impl Runtime {
});
}
let has_commit_ack = dispatch_result
.input_commit
.submission
.as_ref()
.is_some_and(|ack| ack.submission_id == expected_submission_id);
.is_some_and(|ack| ack.submission_request_id == expected_submission_id);
if !has_commit_ack {
let _ = backend.stop_worker(&handle);
self.rollback_failed_create(&worker_ref)?;
@@ -1146,13 +1150,18 @@ impl Runtime {
mut input: WorkerInput,
) -> Result<WorkerInteractionAck, RuntimeError> {
validate_worker_input(&input)?;
let expected_submission_id = if input.kind == WorkerInputKind::User {
let submission_id = Uuid::now_v7().to_string();
input.submission_id = Some(submission_id.clone());
Some(submission_id)
} else {
None
};
let expected_submission_id =
if matches!(input.kind, WorkerInputKind::User | WorkerInputKind::Notify) {
let submission_id = input
.submission_request_id
.clone()
.filter(|request_id| !request_id.trim().is_empty())
.unwrap_or_else(|| Uuid::now_v7().to_string());
input.submission_request_id = Some(submission_id.clone());
Some(submission_id)
} else {
None
};
self.ensure_worker_execution(worker_ref)?;
let (backend, handle) = {
let state = self.lock()?;
@@ -1191,13 +1200,13 @@ impl Runtime {
}
if let Some(expected_submission_id) = expected_submission_id
&& dispatch_result
.input_commit
.submission
.as_ref()
.is_none_or(|ack| ack.submission_id != expected_submission_id)
.is_none_or(|ack| ack.submission_request_id != expected_submission_id)
{
let result = WorkerExecutionResult::rejected(
WorkerExecutionOperation::Input,
"execution backend did not acknowledge the committed Runtime submission id",
"execution backend did not acknowledge the committed Runtime submission request id",
);
self.record_execution_result(worker_ref, result.clone())?;
return Err(RuntimeError::WorkerExecutionRejected {
@@ -1209,6 +1218,7 @@ impl Runtime {
});
}
let submission = dispatch_result.submission.clone();
let mut state = self.lock()?;
state.ensure_running()?;
let worker = state.worker_mut(worker_ref)?;
@@ -1225,6 +1235,7 @@ impl Runtime {
Ok(WorkerInteractionAck {
worker_ref: worker_ref.clone(),
status,
submission,
})
}
@@ -1706,6 +1717,7 @@ impl Runtime {
}
Ok(protocol::Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -3250,17 +3262,9 @@ fn validate_worker_input(input: &WorkerInput) -> Result<(), RuntimeError> {
#[cfg(feature = "ws-server")]
fn input_protocol_event(input: &WorkerInput) -> Option<protocol::Event> {
match input.kind {
WorkerInputKind::User => Some(protocol::Event::UserMessage {
segments: input.segments.clone().unwrap_or_else(|| {
vec![protocol::Segment::Text {
content: input.content.clone(),
}]
}),
}),
// The committed `SystemItem::Notification` is the sole agent-visible
// and Console-visible authority for Notify. A synthetic observation
// here would display the same notification twice.
WorkerInputKind::Notify => None,
// 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::Compact
| WorkerInputKind::ListRewindTargets
| WorkerInputKind::RegisterPeer => Some(protocol::Event::SystemItem {
@@ -3435,6 +3439,7 @@ mod tests {
);
let snapshot = protocol::Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -3475,7 +3480,7 @@ mod tests {
let input = WorkerInput {
kind: WorkerInputKind::User,
content: String::new(),
submission_id: None,
submission_request_id: None,
segments: Some(vec![protocol::Segment::Flow {
selector: "builtin:coder-review".to_string(),
}]),
@@ -3488,7 +3493,7 @@ mod tests {
let input = WorkerInput {
kind: WorkerInputKind::User,
content: String::new(),
submission_id: None,
submission_request_id: None,
segments: Some(Vec::new()),
};
assert!(matches!(
@@ -3503,7 +3508,7 @@ mod tests {
request.initial_input = Some(WorkerInput {
kind: WorkerInputKind::User,
content: String::new(),
submission_id: None,
submission_request_id: None,
segments: Some(vec![protocol::Segment::Flow {
selector: "builtin:coder-review".to_string(),
}]),
@@ -4005,7 +4010,7 @@ mod tests {
_handle: &WorkerExecutionHandle,
input: WorkerInput,
) -> WorkerExecutionResult {
let submission_id = input.submission_id.clone();
let submission_id = input.submission_request_id.clone();
self.dispatched_inputs.lock().unwrap().push(input);
let mut result = self
.dispatch_result
@@ -4013,19 +4018,21 @@ mod tests {
.unwrap()
.clone()
.unwrap_or_else(|| {
WorkerExecutionResult::accepted_input_committed(
WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
"request-test",
"test-submission",
protocol::SubmissionDisposition::Started,
)
});
if !self
.preserve_commit_ack_submission_id
.load(Ordering::SeqCst)
&& let (Some(ack), Some(submission_id)) =
(result.input_commit.as_mut(), submission_id)
(result.submission.as_mut(), submission_id)
{
ack.submission_id = submission_id;
ack.submission_request_id = submission_id;
}
result
}
@@ -4717,10 +4724,12 @@ mod tests {
#[test]
fn create_worker_uses_committed_input_ack_run_state() {
let (runtime, backend) = runtime_and_backend();
backend.set_dispatch_result(WorkerExecutionResult::accepted_input_committed(
backend.set_dispatch_result(WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
"request-test",
"test-submission",
protocol::SubmissionDisposition::Started,
));
let mut request = task_request("committed initial input is already idle");
request.initial_input = Some(WorkerInput::user("start the ticket"));
@@ -4731,13 +4740,15 @@ mod tests {
}
#[test]
fn create_worker_rejects_mismatched_input_commit_acknowledgement() {
fn create_worker_rejects_mismatched_submission_acknowledgement() {
let (runtime, backend) = runtime_and_backend();
backend.preserve_commit_ack_submission_id();
backend.set_dispatch_result(WorkerExecutionResult::accepted_input_committed(
backend.set_dispatch_result(WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
"request-test",
"forged-submission",
protocol::SubmissionDisposition::Started,
));
let mut request = task_request("mismatched initial input commit ack");
request.initial_input = Some(WorkerInput::user("start the ticket"));
@@ -4866,6 +4877,7 @@ mod tests {
&detail.worker_ref,
protocol::Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: vec![protocol::SessionSnapshotEntry {
entry_id: "restored-log-entry".to_owned(),
timestamp: 1,
@@ -4937,10 +4949,14 @@ mod tests {
_handle: &WorkerExecutionHandle,
input: WorkerInput,
) -> WorkerExecutionResult {
WorkerExecutionResult::accepted_input_committed(
WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
input.submission_id.expect("Runtime submission id"),
"request-test",
input
.submission_request_id
.expect("Runtime submission request id"),
protocol::SubmissionDisposition::Started,
)
}
}
@@ -5030,7 +5046,7 @@ mod tests {
request.initial_input = Some(WorkerInput {
kind: WorkerInputKind::User,
content: String::new(),
submission_id: None,
submission_request_id: None,
segments: Some(vec![
protocol::Segment::Flow {
selector: "builtin:coder-review".to_string(),
@@ -5068,7 +5084,7 @@ mod tests {
let input = WorkerInput {
kind: WorkerInputKind::User,
content: String::new(),
submission_id: None,
submission_request_id: None,
segments: Some(vec![protocol::Segment::Flow {
selector: "builtin:coder-review".to_string(),
}]),
@@ -5083,11 +5099,11 @@ mod tests {
assert_eq!(dispatched[0].kind, input.kind);
assert_eq!(dispatched[0].content, input.content);
assert_eq!(dispatched[0].segments, input.segments);
let submission_id = dispatched[0]
.submission_id
let submission_request_id = dispatched[0]
.submission_request_id
.as_deref()
.expect("Runtime submission id");
Uuid::parse_str(submission_id).expect("submission id UUID");
.expect("Runtime submission request id");
Uuid::parse_str(submission_request_id).expect("submission request id UUID");
}
#[cfg(feature = "ws-server")]
@@ -5113,11 +5129,7 @@ mod tests {
let observations = runtime
.read_worker_observation_events(&detail.worker_ref, WorkerObservationCursor::zero())
.unwrap();
assert_eq!(observations.len(), 1);
assert!(matches!(
observations[0].payload,
protocol::Event::UserMessage { .. }
));
assert!(observations.is_empty());
runtime
.observe_worker_event(
@@ -5135,8 +5147,8 @@ mod tests {
let observations = runtime
.read_worker_observation_events(&detail.worker_ref, WorkerObservationCursor::zero())
.unwrap();
assert_eq!(observations.len(), 2);
let protocol::Event::SystemItem { item } = &observations[1].payload else {
assert_eq!(observations.len(), 1);
let protocol::Event::SystemItem { item } = &observations[0].payload else {
panic!("committed notification observation must be a system item");
};
assert_eq!(item["kind"], "notification");
+154 -140
View File
@@ -39,7 +39,7 @@ use crate::working_directory::{
};
use async_trait::async_trait;
use protocol::{ErrorCode, Event, Method, Segment, WorkerStatus};
use session_store::{CombinedStore, LogEntry, WorkerAggregateStore, WorkerSessionStore};
use session_store::{CombinedStore, WorkerAggregateStore, WorkerSessionStore};
#[cfg(test)]
use session_store::{FsStore, FsWorkerStore};
use tokio::runtime::Runtime;
@@ -57,11 +57,10 @@ use worker::feature::builtin::{
#[cfg(feature = "ws-server")]
use worker::ipc::protocol_session::{live_log_entry_event, subscribe_worker_protocol_session};
use worker::{
PreparedWorker, PromptCatalogSource, SegmentLogSink, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN,
Worker, WorkerBootstrap, WorkerBootstrapError, WorkerBootstrapLayout,
WorkerControllerTransport, WorkerError, WorkerFilesystemAuthority, WorkerHandle,
WorkerSharedState, WorkerWorkspaceContext, WorkspaceClient, WorkspaceId,
bash_output_dir_for_worker_id,
PreparedWorker, PromptCatalogSource, SegmentLogSink, Worker, WorkerBootstrap,
WorkerBootstrapError, WorkerBootstrapLayout, WorkerControllerTransport, WorkerError,
WorkerFilesystemAuthority, WorkerHandle, WorkerSharedState, WorkerWorkspaceContext,
WorkspaceClient, WorkspaceId, bash_output_dir_for_worker_id,
};
const DEFAULT_BACKEND_ID: &str = "worker-crate";
@@ -70,17 +69,6 @@ const RUNTIME_TASK_TIMEOUT: Duration = Duration::from_secs(10);
// returns a typed execution error instead of leaving the outer waiter to time out.
const USER_INPUT_COMMIT_TIMEOUT: Duration = Duration::from_secs(9);
fn user_input_has_submission(entry: &LogEntry, submission_id: &str) -> bool {
let extensions = match entry {
LogEntry::AnnotatedUserInput { extensions, .. } => extensions,
_ => return false,
};
extensions.iter().any(|extension| {
extension.domain == WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN
&& extension.payload["submission_id"].as_str() == Some(submission_id)
})
}
pub struct RuntimeWorkerController {
pub handle: WorkerHandle,
pub shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
@@ -1342,126 +1330,75 @@ where
.unwrap_or_else(|message| WorkerExecutionResult::errored(operation, message))
}
fn send_user_input_and_wait_for_commit(
fn send_submit_and_wait_for_acceptance(
&self,
operation: WorkerExecutionOperation,
worker: WorkerHandle,
method: Method,
submission_id: String,
submission_request_id: String,
accepted_run_state: WorkerExecutionRunState,
) -> WorkerExecutionResult {
let acknowledged_submission_id = submission_id.clone();
let request_id = submission_request_id.clone();
self.run_on_adapter_runtime(async move {
// Subscribe before enqueueing the input so the acknowledgement cannot
// race with a fast Worker commit. The opaque submission id is stored in
// the same UserInput entry as the transformed Flow input and its state.
let (_, mut committed_entries) = worker.sink.subscribe_with_snapshot();
let committed_probe = worker.clone();
// Subscribe before enqueueing so a fast durable acceptance cannot
// race the Runtime acknowledgement.
let mut events = worker.subscribe();
worker
.send(method)
.await
.map_err(|err| format!("failed to send Worker method: {err}"))?;
let timeout_probe = committed_probe.clone();
let timeout_submission_id = submission_id.clone();
let acknowledgement = tokio::time::timeout(USER_INPUT_COMMIT_TIMEOUT, async move {
let input_was_committed = || {
committed_probe
.committed_entries()
.iter()
.any(|entry| user_input_has_submission(entry, &submission_id))
};
tokio::time::timeout(USER_INPUT_COMMIT_TIMEOUT, async move {
loop {
tokio::select! {
entry = committed_entries.recv() => {
match entry {
Ok(entry) if user_input_has_submission(&entry, &submission_id) => {
return Ok(());
}
Ok(_) => {}
Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
if input_was_committed() {
return Ok(());
}
return Err(format!(
"worker input commit acknowledgement lagged by {skipped} entry event(s)"
));
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
if input_was_committed() {
return Ok(());
}
return Err(
"worker entry stream closed before user input was committed"
.to_string(),
);
}
}
match events.recv().await {
Ok(Event::SubmissionAccepted {
submission_request_id,
submission_id,
disposition,
}) if submission_request_id == request_id => {
return Ok((submission_id, disposition));
}
event = events.recv() => {
match event {
Ok(Event::Error { message, .. }) => {
if input_was_committed() {
return Ok(());
}
return Err(format!(
"worker rejected user input before session commit: {message}"
));
}
Ok(Event::Shutdown) => {
if input_was_committed() {
return Ok(());
}
return Err(
"worker shut down before user input was committed".to_string()
);
}
Ok(_) => {}
Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
if input_was_committed() {
return Ok(());
}
return Err(format!(
"worker input commit acknowledgement lagged by {skipped} protocol event(s)"
));
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
if input_was_committed() {
return Ok(());
}
return Err(
"worker event stream closed before user input was committed"
.to_string(),
);
}
}
Ok(Event::SubmissionRejected {
submission_request_id,
message,
}) if submission_request_id == request_id => {
return Err(format!("worker rejected Submit: {message}"));
}
Ok(Event::Error { message, .. }) => {
return Err(format!(
"worker rejected Submit before durable acceptance: {message}"
));
}
Ok(Event::Shutdown) => {
return Err(
"worker shut down before Submit was durably accepted".to_string()
);
}
Ok(_) => {}
Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
return Err(format!(
"worker Submit acknowledgement lagged by {skipped} protocol event(s)"
));
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
return Err(
"worker event stream closed before Submit was durably accepted"
.to_string(),
);
}
}
}
})
.await;
match acknowledgement {
Ok(result) => result,
Err(_) => {
if timeout_probe
.committed_entries()
.iter()
.any(|entry| user_input_has_submission(entry, &timeout_submission_id))
{
Ok(())
} else {
Err("timed out waiting for worker user input commit".to_string())
}
}
}
.await
.map_err(|_| "timed out waiting for durable Worker Submit acceptance".to_string())?
})
.map(|_| {
WorkerExecutionResult::accepted_input_committed(
.map(|(submission_id, disposition)| {
WorkerExecutionResult::accepted_submission(
operation,
accepted_run_state,
acknowledged_submission_id,
submission_request_id,
submission_id,
disposition,
)
})
.unwrap_or_else(|message| WorkerExecutionResult::errored(operation, message))
@@ -1582,9 +1519,10 @@ impl<F> Drop for WorkerRuntimeExecutionBackend<F> {
fn method_starts_turn(method: &Method) -> bool {
matches!(
method,
Method::Run { .. }
| Method::RunTracked { .. }
Method::Submit { .. }
| Method::SubmitTracked { .. }
| Method::Notify { auto_run: true, .. }
| Method::NotifyTracked { auto_run: true, .. }
| Method::Resume
| Method::Compact
)
@@ -1609,9 +1547,10 @@ fn accepted_notify_run_state(status: WorkerStatus, auto_run: bool) -> WorkerExec
fn accepted_run_state_for_method(method: &Method) -> WorkerExecutionRunState {
match method {
Method::Run { .. }
| Method::RunTracked { .. }
Method::Submit { .. }
| Method::SubmitTracked { .. }
| Method::Notify { auto_run: true, .. }
| Method::NotifyTracked { auto_run: true, .. }
| Method::Resume
| Method::Compact => WorkerExecutionRunState::Busy,
Method::Shutdown => WorkerExecutionRunState::Stopped,
@@ -1959,12 +1898,19 @@ where
&& busy
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_ok();
let notification_request_id = input
.submission_request_id
.unwrap_or_else(protocol::new_submission_request_id);
let result = self.send_method(
WorkerExecutionOperation::Input,
worker,
Method::Notify {
Method::NotifyTracked {
notification_request_id: notification_request_id.clone(),
message: input.content,
auto_run: true,
source: protocol::AuthenticatedInputSource::Backend {
operation_id: notification_request_id,
},
},
accepted_run_state,
);
@@ -1975,21 +1921,23 @@ where
return result;
}
if worker.shared_state.get_status() != WorkerStatus::Idle
|| busy
let is_user_submit = input.kind == WorkerInputKind::User;
let status = worker.shared_state.get_status();
let claimed_here = status == WorkerStatus::Idle
&& busy
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_err()
{
.is_ok();
if !is_user_submit && !claimed_here {
return WorkerExecutionResult::busy(
WorkerExecutionOperation::Input,
"Worker is already running; runtime adapter v0 does not queue input",
"Worker is already running",
);
}
let (method, submission_id) = match input.kind {
let (method, submission_request_id) = match input.kind {
WorkerInputKind::User => {
let Some(submission_id) = input
.submission_id
.submission_request_id
.filter(|submission_id| !submission_id.trim().is_empty())
else {
busy.store(false, Ordering::SeqCst);
@@ -1999,11 +1947,14 @@ where
);
};
(
Method::RunTracked {
Method::SubmitTracked {
submission_request_id: submission_id.clone(),
input: input.segments.unwrap_or_else(|| {
vec![Segment::text(input.content.trim().to_string())]
}),
submission_id: submission_id.clone(),
source: protocol::AuthenticatedInputSource::Backend {
operation_id: submission_id.clone(),
},
},
Some(submission_id),
)
@@ -2021,21 +1972,22 @@ where
),
};
let accepted_run_state = match method {
Method::Run { .. }
| Method::RunTracked { .. }
Method::Submit { .. }
| Method::SubmitTracked { .. }
| Method::Notify { .. }
| Method::NotifyTracked { .. }
| Method::Compact => WorkerExecutionRunState::Busy,
_ => WorkerExecutionRunState::Idle,
};
let accepted_is_idle = accepted_run_state == WorkerExecutionRunState::Idle;
let waits_for_user_input_commit = submission_id.is_some();
let waits_for_submission_acceptance = submission_request_id.is_some();
let result = if waits_for_user_input_commit {
self.send_user_input_and_wait_for_commit(
let result = if waits_for_submission_acceptance {
self.send_submit_and_wait_for_acceptance(
WorkerExecutionOperation::Input,
worker,
method,
submission_id.expect("tracked Run has submission id"),
submission_request_id.expect("Submit must have a submission request id"),
accepted_run_state,
)
} else {
@@ -2046,7 +1998,9 @@ where
accepted_run_state,
)
};
if accepted_is_idle || result.outcome != crate::execution::WorkerExecutionOutcome::Accepted
if accepted_is_idle
|| (claimed_here
&& result.outcome != crate::execution::WorkerExecutionOutcome::Accepted)
{
busy.store(false, Ordering::SeqCst);
}
@@ -2116,8 +2070,12 @@ where
}
};
if let Method::Notify { auto_run, .. } = &method {
let auto_run = *auto_run;
if let Some(auto_run) = match &method {
Method::Notify { auto_run, .. } | Method::NotifyTracked { auto_run, .. } => {
Some(*auto_run)
}
_ => None,
} {
let status = worker.shared_state.get_status();
let accepted_run_state = accepted_notify_run_state(status, auto_run);
let claimed_here = status == WorkerStatus::Idle
@@ -3400,6 +3358,60 @@ mod tests {
);
}
#[test]
fn running_worker_accepts_a_second_submit_as_queued() {
let client = MockClient::sequential(vec![MockResponse::Hang(vec![])]);
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).unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let detail = runtime
.create_worker(create_request("queued-submit"))
.unwrap();
let mut first_input = WorkerInput::user("first");
first_input.submission_request_id = Some("request-first".into());
let first = runtime
.send_input(&detail.worker_ref, first_input.clone())
.unwrap();
assert_eq!(
first.submission.as_ref().map(|ack| ack.disposition),
Some(protocol::SubmissionDisposition::Started)
);
let retry = runtime.send_input(&detail.worker_ref, first_input).unwrap();
assert_eq!(retry.submission, first.submission);
let mut conflicting_retry = WorkerInput::user("different");
conflicting_retry.submission_request_id = Some("request-first".into());
assert!(
runtime
.send_input(&detail.worker_ref, conflicting_retry)
.is_err(),
"same request id with a different payload must fail"
);
let mut second_input = WorkerInput::user("second");
second_input.submission_request_id = Some("request-second".into());
let second = runtime
.send_input(&detail.worker_ref, second_input)
.unwrap();
assert_eq!(
second.submission.as_ref().map(|ack| ack.disposition),
Some(protocol::SubmissionDisposition::Queued)
);
}
#[test]
fn create_with_initial_input_returns_after_session_commit() {
let client = MockClient::new(simple_text_events());
@@ -3449,8 +3461,10 @@ mod tests {
};
extensions
.iter()
.find(|extension| extension.domain == WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN)
.and_then(|extension| extension.payload["submission_id"].as_str())
.find(|extension| extension.domain == "worker.pending_activations.v1")
.and_then(|extension| {
extension.payload["receipts"][0]["submission_id"].as_str()
})
})
.expect("committed input submission id");
uuid::Uuid::parse_str(submission_id).expect("opaque submission id is a UUID");
+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
File diff suppressed because it is too large Load Diff
+28 -3
View File
@@ -1012,7 +1012,19 @@ 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
let notification_request_id = protocol::new_submission_request_id();
connect_and_send(
socket_path,
&Method::NotifyTracked {
notification_request_id: notification_request_id.clone(),
message,
auto_run,
source: protocol::AuthenticatedInputSource::Backend {
operation_id: notification_request_id,
},
},
)
.await
}
fn json_content<T: Serialize>(value: &T) -> Result<String, ToolError> {
@@ -1482,6 +1494,7 @@ mod tests {
writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -1517,6 +1530,7 @@ mod tests {
writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -1536,7 +1550,10 @@ mod tests {
.await
.unwrap();
let method = reader.next::<Method>().await.unwrap().unwrap();
if let Method::Notify { message, auto_run } = method {
if let Method::NotifyTracked {
message, auto_run, ..
} = method
{
assert!(auto_run);
tx.send(message).await.unwrap();
} else {
@@ -1608,6 +1625,7 @@ mod tests {
writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -1634,6 +1652,7 @@ mod tests {
writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -1653,7 +1672,10 @@ mod tests {
.await
.unwrap();
let method = reader.next::<Method>().await.unwrap().unwrap();
if let Method::Notify { message, auto_run } = method {
if let Method::NotifyTracked {
message, auto_run, ..
} = method
{
assert!(!auto_run);
tx.send(message).await.unwrap();
} else {
@@ -1738,6 +1760,7 @@ mod tests {
writer
.write(&Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: protocol::Greeting {
@@ -1790,6 +1813,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,
},
})
}
}
+31 -6
View File
@@ -176,12 +176,23 @@ 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>,
Option<session_store::LoggedSessionHistoryOrigin>,
)],
) -> 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, history_provenance) in items {
let entry = writer.commit_system_item_with_extensions(
item.clone(),
extensions.clone(),
history_provenance.clone(),
)?;
self.pending_committed_history
.lock()
.expect("pending committed history poisoned")
@@ -190,6 +201,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(), None))
.collect::<Vec<_>>(),
)
}
fn current_turn_index(&self) -> usize {
self.next_turn_index
.load(Ordering::Relaxed)
@@ -327,7 +348,11 @@ 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>,
Option<session_store::LoggedSessionHistoryOrigin>,
)> = 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 +370,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(), entry.history_provenance()));
}
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,
+70 -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::{LoggedSessionHistoryOrigin, SessionExtension, SystemItem};
use tracing::warn;
use crate::prompt::catalog::{CatalogError, PromptCatalog};
@@ -41,8 +41,33 @@ 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>,
history_provenance: Option<LoggedSessionHistoryOrigin>,
},
WorkerEvent {
event: WorkerEvent,
},
}
impl PendingNotify {
pub(crate) fn extensions(&self) -> Vec<SessionExtension> {
match self {
PendingNotify::Notify { extensions, .. } => extensions.clone(),
PendingNotify::WorkerEvent { .. } => Vec::new(),
}
}
pub(crate) fn history_provenance(&self) -> Option<LoggedSessionHistoryOrigin> {
match self {
PendingNotify::Notify {
history_provenance, ..
} => history_provenance.clone(),
PendingNotify::WorkerEvent { .. } => None,
}
}
}
/// Shared, mutex-guarded buffer of pending entries.
@@ -62,7 +87,46 @@ 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(),
history_provenance: None,
});
}
pub fn push_durable_notify(
&self,
message: String,
auto_run: bool,
history_provenance: LoggedSessionHistoryOrigin,
extension: SessionExtension,
) {
self.push_entry(PendingNotify::Notify {
message,
auto_run,
extensions: vec![extension],
history_provenance: Some(history_provenance),
});
}
pub(crate) fn replace_durable_notification_extension(
&self,
extension: SessionExtension,
) -> bool {
let mut queue = self.inner.lock().expect("notify buffer poisoned");
let Some(extensions) = queue.iter_mut().rev().find_map(|pending| match pending {
PendingNotify::Notify {
auto_run: false,
extensions,
..
} if !extensions.is_empty() => Some(extensions),
_ => None,
}) else {
return false;
};
*extensions = vec![extension];
true
}
/// Push a typed worker-event entry onto the queue.
@@ -202,6 +266,8 @@ mod tests {
let entry = PendingNotify::Notify {
message: "hello".into(),
auto_run: false,
extensions: Vec::new(),
history_provenance: None,
};
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()
+78 -23
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
@@ -219,33 +219,50 @@ fn parse_spawn_profile_selector(raw: Option<&str>) -> Result<SpawnProfileSelecto
#[derive(Clone)]
pub(crate) enum ParentNotificationTarget {
Controller(mpsc::WeakSender<Method>),
Buffer(crate::ipc::notify_buffer::NotifyBuffer),
Controller {
sender: mpsc::WeakSender<Method>,
fallback: Arc<dyn Fn(Method) + Send + Sync>,
},
Durable(Arc<dyn Fn(Method) + Send + Sync>),
}
impl ParentNotificationTarget {
fn notify(&self, message: String, auto_run: bool) {
pub(crate) fn with_controller_fallback(
sender: mpsc::WeakSender<Method>,
fallback: ParentNotificationTarget,
) -> Self {
let ParentNotificationTarget::Durable(fallback) = fallback else {
unreachable!("controller fallback must use durable pending authority");
};
Self::Controller { sender, fallback }
}
pub(crate) fn notify(&self, child_session_id: String, message: String, auto_run: bool) {
let method = Method::NotifyTracked {
notification_request_id: protocol::new_submission_request_id(),
message,
auto_run,
source: protocol::AuthenticatedInputSource::SubWorker {
session_id: child_session_id,
},
};
match self {
Self::Controller(parent_method_tx) => {
let Some(parent_method_tx) = parent_method_tx.upgrade() else {
tracing::warn!(
"parent Worker controller closed before Internal SubWorker completion notification"
);
Self::Controller { sender, fallback } => {
let Some(parent_method_tx) = sender.upgrade() else {
fallback(method);
return;
};
let fallback = fallback.clone();
tokio::spawn(async move {
if let Err(error) = parent_method_tx
.send(Method::Notify { message, auto_run })
.await
{
if let Err(error) = parent_method_tx.send(method).await {
tracing::warn!(
%error,
"failed to notify parent Worker about Internal SubWorker completion"
"failed to notify parent Controller; using durable pending authority"
);
fallback(error.0);
}
});
}
Self::Buffer(parent_notifies) => parent_notifies.push_notify(message, auto_run),
Self::Durable(notify) => notify(method),
}
}
}
@@ -550,7 +567,7 @@ impl Tool for SubWorkerSpawnTool {
let message = format!(
"SubWorker `{child_name}` turn ended with status {status:?}. Inspect its committed session with worker-observation tools before making completion decisions."
);
parent_notifications.notify(message, true);
parent_notifications.notify(child_name.clone(), message, true);
})),
)
.await;
@@ -1134,12 +1151,41 @@ enabled = false
#[tokio::test]
async fn parent_controller_notification_target_does_not_keep_channel_open() {
let (parent_method_tx, mut parent_method_rx) = mpsc::channel(1);
let target = ParentNotificationTarget::Controller(parent_method_tx.downgrade());
let captured = Arc::new(std::sync::Mutex::new(false));
let captured_for_fallback = captured.clone();
let target = ParentNotificationTarget::with_controller_fallback(
parent_method_tx.downgrade(),
ParentNotificationTarget::Durable(Arc::new(move |_| {
*captured_for_fallback.lock().unwrap() = true;
})),
);
drop(parent_method_tx);
assert!(parent_method_rx.recv().await.is_none());
target.notify("late completion".to_string(), true);
target.notify("child-session".into(), "late completion".to_string(), true);
assert!(*captured.lock().unwrap());
}
#[test]
fn durable_parent_notification_target_preserves_child_source() {
let captured = Arc::new(std::sync::Mutex::new(None));
let captured_for_target = captured.clone();
let target = ParentNotificationTarget::Durable(Arc::new(move |method| {
*captured_for_target.lock().unwrap() = Some(method);
}));
target.notify("child-session".into(), "completed".into(), true);
assert!(matches!(
captured.lock().unwrap().take(),
Some(Method::NotifyTracked {
message,
auto_run: true,
source: protocol::AuthenticatedInputSource::SubWorker { session_id },
..
}) if session_id == "child-session" && message == "completed"
));
}
#[tokio::test]
@@ -1185,7 +1231,10 @@ enabled = false
let tool = SubWorkerSpawnTool::new(
"parent".into(),
workspace_context,
ParentNotificationTarget::Controller(parent_method_tx.downgrade()),
ParentNotificationTarget::with_controller_fallback(
parent_method_tx.downgrade(),
ParentNotificationTarget::Durable(Arc::new(|_| {})),
),
runtime.path().to_path_buf(),
bash_output_dir.clone(),
workspace_root.clone(),
@@ -1282,10 +1331,13 @@ enabled = false
.expect("parent method channel remains open");
assert!(matches!(
completion,
Method::Notify {
Method::NotifyTracked {
message,
auto_run: true,
} if message.contains("SubWorker `reviewer-child` turn ended with status Idle")
source: protocol::AuthenticatedInputSource::SubWorker { session_id },
..
} if session_id == "reviewer-child"
&& message.contains("SubWorker `reviewer-child` turn ended with status Idle")
));
assert!(!runtime.path().join("reviewer-child/sock").exists());
@@ -1434,7 +1486,10 @@ enabled = false
let tool = SubWorkerSpawnTool::new(
"parent".into(),
workspace_context,
ParentNotificationTarget::Controller(parent_method_tx.downgrade()),
ParentNotificationTarget::with_controller_fallback(
parent_method_tx.downgrade(),
ParentNotificationTarget::Durable(Arc::new(|_| {})),
),
runtime.path().to_path_buf(),
bash_output_dir.clone(),
workspace_root.clone(),
File diff suppressed because it is too large Load Diff
+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 {
+328 -64
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;
@@ -826,7 +844,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);
@@ -875,7 +899,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;
@@ -928,7 +958,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);
@@ -975,7 +1011,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;
@@ -1017,7 +1059,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!(
@@ -1066,7 +1114,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 {
@@ -1131,7 +1182,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())
@@ -1169,7 +1226,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;
@@ -1183,7 +1246,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;
@@ -1201,7 +1270,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;
@@ -1236,10 +1311,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..."),
@@ -1249,35 +1322,67 @@ 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_snapshot = 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_snapshot = Some(pending)
}
_ = 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_snapshot.is_some() {
break;
}
}
assert!(saw_already_running, "should see already_running error");
assert_eq!(accepted, Some(protocol::SubmissionDisposition::Queued));
let pending_snapshot = pending_snapshot.expect("pending snapshot");
assert_eq!(pending_snapshot.submissions.len(), 1);
handle.send(Method::Pause).await.unwrap();
wait_for_status(&handle, WorkerStatus::Paused).await;
handle
.send(Method::ContinuePending {
expected_revision: pending_snapshot.revision,
expected_head_id: pending_snapshot.head_id.expect("pending head"),
})
.await
.unwrap();
let rejection = tokio::time::timeout(std::time::Duration::from_secs(1), async {
loop {
if let Ok(Event::Error { code, message }) = rx.recv().await
&& code == worker::ErrorCode::InvalidRequest
&& message.contains("requires an idle Worker")
{
break message;
}
}
})
.await
.expect("paused ContinuePending rejection");
assert!(rejection.contains("Resume or Cancel"));
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Paused);
}
#[tokio::test]
@@ -1365,7 +1470,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
@@ -1437,7 +1543,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();
@@ -1485,7 +1597,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(),
}],
@@ -1538,7 +1651,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;
@@ -1586,6 +1705,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,
})
@@ -1626,6 +1746,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
@@ -1671,14 +1804,18 @@ async fn notify_while_idle_with_auto_run_false_waits_for_explicit_run() {
let client_for_assert = client.clone();
let worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
let notification_request_id = protocol::new_submission_request_id();
handle
.send(Method::Notify {
message: "progress snapshot".into(),
auto_run: false,
})
.await
.unwrap();
for _ in 0..2 {
handle
.send(Method::Notify {
notification_request_id: notification_request_id.clone(),
message: "progress snapshot".into(),
auto_run: false,
})
.await
.unwrap();
}
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Idle);
@@ -1687,7 +1824,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() {
@@ -1867,9 +2010,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,
})
@@ -1902,6 +2052,66 @@ async fn notify_while_running_does_not_emit_already_running_error() {
wait_for_status(&handle, WorkerStatus::Idle).await;
}
#[tokio::test]
async fn weak_notify_while_running_is_deduped_and_survives_until_next_submit() {
let client = MockClient::sequential(vec![
MockResponse::Hang(Vec::new()),
MockResponse::Complete(simple_text_events()),
]);
let client_for_assert = client.clone();
let worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"first",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Running).await;
let notification_request_id = protocol::new_submission_request_id();
for _ in 0..2 {
handle
.send(Method::Notify {
notification_request_id: notification_request_id.clone(),
message: "durable weak notice".into(),
auto_run: false,
})
.await
.unwrap();
}
handle.send(Method::Cancel).await.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let mut rx = handle.subscribe();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"second",
))
.await
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if matches!(rx.recv().await, Ok(Event::TurnEnd { .. })) {
break;
}
}
})
.await
.expect("second submit completes");
let requests = client_for_assert.captured_requests();
let notice_count = requests[1]
.items
.iter()
.filter_map(|item| item.as_text())
.filter(|text| text.contains("durable weak notice"))
.count();
assert_eq!(notice_count, 1);
}
#[tokio::test]
async fn status_json_reflects_worker_name() {
let client = MockClient::new(simple_text_events());
@@ -1936,7 +2146,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;
@@ -2243,7 +2459,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.
@@ -2332,7 +2554,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.
@@ -2369,7 +2591,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.
@@ -2400,7 +2628,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,
@@ -2531,7 +2765,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,
@@ -2599,7 +2839,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!(
@@ -2688,7 +2931,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();
@@ -2721,7 +2970,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;
@@ -2755,7 +3007,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,
@@ -2769,7 +3027,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;
@@ -2816,7 +3077,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!(
+9 -7
View File
@@ -538,7 +538,7 @@ fn initial_worker_input(segments: &[Segment]) -> Option<EmbeddedWorkerInput> {
Some(EmbeddedWorkerInput {
kind: EmbeddedWorkerInputKind::User,
content: Segment::flatten_to_text(segments),
submission_id: None,
submission_request_id: None,
segments: Some(segments.to_vec()),
})
}
@@ -2706,7 +2706,7 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
WorkerInputKind::RegisterPeer => EmbeddedWorkerInputKind::RegisterPeer,
},
content: request.content,
submission_id: None,
submission_request_id: None,
segments: request.segments,
};
match self.runtime.send_input(&worker_ref, input) {
@@ -3934,7 +3934,7 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
WorkerInputKind::RegisterPeer => EmbeddedWorkerInputKind::RegisterPeer,
},
content: request.content,
submission_id: None,
submission_request_id: None,
segments: request.segments,
};
match self.post_json::<_, RuntimeHttpWorkerInputResponse>(
@@ -5195,7 +5195,7 @@ mod tests {
"missing test context",
);
};
let submission_id = input.submission_id.clone();
let submission_request_id = input.submission_request_id.clone();
let content = input.content;
std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(10));
@@ -5212,11 +5212,13 @@ mod tests {
status: protocol::WorkerStatus::Idle,
});
});
if let Some(submission_id) = submission_id {
worker_runtime::execution::WorkerExecutionResult::accepted_input_committed(
if let Some(submission_request_id) = submission_request_id {
worker_runtime::execution::WorkerExecutionResult::accepted_submission(
worker_runtime::execution::WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
submission_id,
submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
)
} else {
worker_runtime::execution::WorkerExecutionResult::accepted(
@@ -32,11 +32,13 @@ impl WorkerExecutionBackend for TestExecutionBackend {
_handle: &WorkerExecutionHandle,
input: worker_runtime::interaction::WorkerInput,
) -> WorkerExecutionResult {
if let Some(submission_id) = input.submission_id {
WorkerExecutionResult::accepted_input_committed(
if let Some(submission_request_id) = input.submission_request_id {
WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
submission_id,
submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
)
} else {
WorkerExecutionResult::accepted(
+196 -33
View File
@@ -8574,13 +8574,15 @@ async fn scoped_list_runtimes(
async fn scoped_workspace_protocol_ws(
State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath(workspace_id): AxumPath<String>,
ws: axum::extract::ws::WebSocketUpgrade,
) -> std::result::Result<Response, Response> {
validate_workspace_scope(&api, &workspace_id).map_err(|error| error.into_response())?;
let input_source = authenticated_browser_input_source(&actor);
Ok(ws
.on_upgrade(move |socket| {
crate::workspace_subscription::serve_workspace_subscription(api, socket)
crate::workspace_subscription::serve_workspace_subscription(api, socket, input_source)
})
.into_response())
}
@@ -9286,7 +9288,7 @@ async fn scoped_capture_worker_observation_session(
return Err(ApiError::from(Error::UnknownWorker { worker: target }));
}
let mut connection = connect_workspace_worker_protocol(&api, &target).await?;
let mut connection = connect_workspace_worker_protocol(&api, &target, None).await?;
let event = tokio::time::timeout(std::time::Duration::from_secs(10), connection.events.recv())
.await
.map_err(|_| {
@@ -11495,6 +11497,7 @@ async fn scoped_cancel_runtime_worker(
async fn scoped_worker_protocol_ws(
ws: WebSocketUpgrade,
State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
) -> Response {
if let Err(err) = validate_workspace_scope(&api, &path.workspace_id) {
@@ -11502,6 +11505,7 @@ async fn scoped_worker_protocol_ws(
}
worker_protocol_ws(
State(api),
Extension(actor),
AxumPath((path.worker.runtime_id, path.worker.worker_id)),
ws,
)
@@ -13896,8 +13900,45 @@ async fn cancel_runtime_worker(
Ok(Json(result))
}
fn authenticated_browser_input_source(actor: &RequestActor) -> protocol::AuthenticatedInputSource {
protocol::AuthenticatedInputSource::Account {
account_id: actor.account_id.clone(),
}
}
pub(crate) fn authorize_browser_worker_method(
method: protocol::Method,
source: &protocol::AuthenticatedInputSource,
) -> std::result::Result<protocol::Method, &'static str> {
match method {
protocol::Method::Submit {
submission_request_id,
input,
} => Ok(protocol::Method::SubmitTracked {
submission_request_id,
input,
source: source.clone(),
}),
protocol::Method::Notify {
notification_request_id,
message,
auto_run,
} => Ok(protocol::Method::NotifyTracked {
notification_request_id,
message,
auto_run,
source: source.clone(),
}),
protocol::Method::SubmitTracked { .. } | protocol::Method::NotifyTracked { .. } => {
Err("authenticated Worker input source is server-owned")
}
other => Ok(other),
}
}
async fn worker_protocol_ws(
State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>,
ws: WebSocketUpgrade,
) -> impl IntoResponse {
@@ -13921,7 +13962,8 @@ async fn worker_protocol_ws(
.into_response();
}
};
ws.on_upgrade(move |socket| worker_protocol_ws_session(source, socket))
let input_source = authenticated_browser_input_source(&actor);
ws.on_upgrade(move |socket| worker_protocol_ws_session(source, socket, input_source))
}
pub(crate) struct WorkspaceWorkerProtocolConnection {
@@ -13932,6 +13974,7 @@ pub(crate) struct WorkspaceWorkerProtocolConnection {
pub(crate) async fn connect_workspace_worker_protocol(
api: &WorkspaceApi,
worker: &RuntimeWorkerRef,
input_source: Option<&protocol::AuthenticatedInputSource>,
) -> Result<WorkspaceWorkerProtocolConnection> {
let source = match api.observation_proxy.source(worker) {
Ok(source) => source,
@@ -13948,15 +13991,39 @@ pub(crate) async fn connect_workspace_worker_protocol(
}
};
match source {
RuntimeObservationSource::RemoteWs(config) => connect_remote_worker_protocol(config).await,
RuntimeObservationSource::RemoteWs(config) => {
connect_remote_worker_protocol(config, input_source).await
}
RuntimeObservationSource::Embedded(source) => {
connect_embedded_worker_protocol(source).await
}
}
}
fn insert_authenticated_input_source_header(
headers: &mut HeaderMap,
input_source: Option<&protocol::AuthenticatedInputSource>,
) -> Result<()> {
let Some(input_source) = input_source else {
return Ok(());
};
let protocol::AuthenticatedInputSource::Account { account_id } = input_source else {
return Err(Error::Config(
"remote Worker protocol transport supports only Account input source".into(),
));
};
headers.insert(
protocol::AUTHENTICATED_ACCOUNT_ID_HEADER,
account_id.parse().map_err(|error| {
Error::Config(format!("invalid authenticated Account identity: {error}"))
})?,
);
Ok(())
}
async fn connect_remote_worker_protocol(
config: RuntimeObservationSourceConfig,
input_source: Option<&protocol::AuthenticatedInputSource>,
) -> Result<WorkspaceWorkerProtocolConnection> {
let mut request = config
.endpoint
@@ -13971,6 +14038,7 @@ async fn connect_remote_worker_protocol(
})?,
);
}
insert_authenticated_input_source_header(request.headers_mut(), input_source)?;
let (socket, _) =
connect_async(request)
.await
@@ -14048,13 +14116,17 @@ async fn connect_embedded_worker_protocol(
Ok(WorkspaceWorkerProtocolConnection { methods, events })
}
async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: WebSocket) {
async fn worker_protocol_ws_session(
source: RuntimeObservationSource,
socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
match source {
RuntimeObservationSource::RemoteWs(config) => {
remote_worker_protocol_ws_session(config, socket).await;
remote_worker_protocol_ws_session(config, socket, input_source).await;
}
RuntimeObservationSource::Embedded(source) => {
embedded_worker_protocol_ws_session(source, socket).await;
embedded_worker_protocol_ws_session(source, socket, input_source).await;
}
}
}
@@ -14062,6 +14134,7 @@ async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: We
async fn remote_worker_protocol_ws_session(
config: RuntimeObservationSourceConfig,
socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
let mut request = match config.endpoint.clone().into_client_request() {
Ok(request) => request,
@@ -14089,6 +14162,16 @@ async fn remote_worker_protocol_ws_session(
}
}
}
if let Err(error) =
insert_authenticated_input_source_header(request.headers_mut(), Some(&input_source))
{
let mut socket = socket;
let event = protocol_error_event(format!(
"failed to build authenticated Account identity header: {error}"
));
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
let (upstream, _) = match connect_async(request).await {
Ok(connection) => connection,
@@ -14110,14 +14193,33 @@ async fn remote_worker_protocol_ws_session(
inbound = client_stream.next() => {
match inbound {
Some(Ok(WsMessage::Text(text))) => {
if upstream_sink.send(TungsteniteMessage::Text(text.to_string().into())).await.is_err() {
let method = match protocol::stream::decode_method(text.as_ref()) {
Ok(method) => match authorize_browser_worker_method(method, &input_source) {
Ok(method) => method,
Err(message) => {
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event(message)) {
let _ = client_sink.send(WsMessage::Text(event.into())).await;
}
break;
}
},
Err(error) => {
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event(error.to_string())) {
let _ = client_sink.send(WsMessage::Text(event.into())).await;
}
break;
}
};
let Ok(method) = protocol::stream::encode_method(&method) else { break };
if upstream_sink.send(TungsteniteMessage::Text(method.into())).await.is_err() {
break;
}
}
Some(Ok(WsMessage::Binary(binary))) => {
if upstream_sink.send(TungsteniteMessage::Binary(binary.to_vec().into())).await.is_err() {
break;
Some(Ok(WsMessage::Binary(_))) => {
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event("binary Worker methods are not accepted")) {
let _ = client_sink.send(WsMessage::Text(event.into())).await;
}
break;
}
Some(Ok(WsMessage::Close(_))) | None => {
let _ = upstream_sink.send(TungsteniteMessage::Close(None)).await;
@@ -14173,6 +14275,7 @@ async fn remote_worker_protocol_ws_session(
async fn embedded_worker_protocol_ws_session(
source: crate::observation::EmbeddedRuntimeObservationSource,
mut socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
let mut upstream = match RuntimeObservationClient::connect(&RuntimeObservationSource::Embedded(
source.clone(),
@@ -14192,24 +14295,32 @@ async fn embedded_worker_protocol_ws_session(
inbound = socket.next() => {
match inbound {
Some(Ok(WsMessage::Text(text))) => match decode_method(&text) {
Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) {
Ok(events) => {
for event in events {
Ok(method) => match authorize_browser_worker_method(method, &input_source) {
Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) {
Ok(events) => {
for event in events {
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
}
Err(error) => {
let event = protocol_error_event(error.to_string());
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
}
Err(error) => {
let event = protocol_error_event(error.to_string());
if !send_protocol_event(&mut socket, &event).await {
return;
}
},
Err(message) => {
let event = protocol_error_event(message);
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
},
Err(error) => {
let event =
protocol_error_event(format!("malformed protocol method frame: {error}"));
let event = protocol_error_event(format!(
"malformed protocol method frame: {error}"
));
if !send_protocol_event(&mut socket, &event).await {
return;
}
@@ -16219,6 +16330,48 @@ mod tests {
&tail[..end]
}
#[test]
fn browser_worker_methods_receive_server_owned_account_source() {
let source = protocol::AuthenticatedInputSource::Account {
account_id: "account-1".into(),
};
let method = authorize_browser_worker_method(
protocol::Method::Submit {
submission_request_id: "request-1".into(),
input: vec![protocol::Segment::text("hello")],
},
&source,
)
.unwrap();
assert!(matches!(
method,
protocol::Method::SubmitTracked {
source: protocol::AuthenticatedInputSource::Account { ref account_id },
..
} if account_id == "account-1"
));
assert!(authorize_browser_worker_method(method, &source).is_err());
}
#[test]
fn remote_worker_protocol_header_preserves_authenticated_account_source() {
let mut headers = HeaderMap::new();
insert_authenticated_input_source_header(
&mut headers,
Some(&protocol::AuthenticatedInputSource::Account {
account_id: "account-1".into(),
}),
)
.unwrap();
assert_eq!(
headers
.get(protocol::AUTHENTICATED_ACCOUNT_ID_HEADER)
.unwrap(),
"account-1"
);
}
#[test]
fn merge_request_http_paths_observe_refs_through_runtime_provider_authority() {
let source = include_str!("server.rs");
@@ -18468,7 +18621,7 @@ mod tests {
.get(handle.worker_ref())
.cloned()
.expect("execution context");
let submission_id = input.submission_id.clone();
let submission_request_id = input.submission_request_id.clone();
let content = input.content.clone();
std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(25));
@@ -18476,11 +18629,13 @@ mod tests {
text: format!("server companion echoed: {content}"),
});
});
if let Some(submission_id) = submission_id {
worker_runtime::execution::WorkerExecutionResult::accepted_input_committed(
if let Some(submission_request_id) = submission_request_id {
worker_runtime::execution::WorkerExecutionResult::accepted_submission(
worker_runtime::execution::WorkerExecutionOperation::Input,
worker_runtime::execution::WorkerExecutionRunState::Idle,
submission_id,
submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
)
} else {
worker_runtime::execution::WorkerExecutionResult::accepted(
@@ -27069,6 +27224,16 @@ mod tests {
(runtime, worker_ref, endpoint)
}
fn test_browser_request_actor() -> RequestActor {
RequestActor {
user_id: "test-user".into(),
account_id: format!("account-{TEST_WORKSPACE_ID}"),
handle: "test".into(),
display_name: "Test".into(),
auth_method: ActorAuthMethod::BrowserSession,
}
}
async fn spawn_workspace_proxy(
source: RuntimeObservationSourceConfig,
) -> (String, tempfile::TempDir) {
@@ -27087,11 +27252,8 @@ mod tests {
.unwrap();
let app_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let app_addr = app_listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(app_listener, build_inner_router(api))
.await
.unwrap()
});
let app = build_inner_router(api).layer(Extension(test_browser_request_actor()));
tokio::spawn(async move { axum::serve(app_listener, app).await.unwrap() });
(
format!("ws://{app_addr}/api/runtimes/{runtime_id}/workers/{worker_id}/protocol/ws"),
dir,
@@ -27103,7 +27265,8 @@ mod tests {
let dir = tempfile::tempdir().unwrap();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = build_inner_router(test_api(dir.path()).await);
let app = build_inner_router(test_api(dir.path()).await)
.layer(Extension(test_browser_request_actor()));
let server = tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
});
@@ -27151,7 +27314,7 @@ mod tests {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = build_inner_router(api);
let app = build_inner_router(api).layer(Extension(test_browser_request_actor()));
let server = tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
});
@@ -11,7 +11,9 @@ use tokio::sync::mpsc;
use worker_runtime::identity::RuntimeWorkerRef;
use crate::runtime_subscription::{BrokerSubscriptionEvent, RuntimeSubscriptionBroker};
use crate::server::{WorkspaceApi, connect_workspace_worker_protocol};
use crate::server::{
WorkspaceApi, authorize_browser_worker_method, connect_workspace_worker_protocol,
};
use crate::store::WorkspaceResourceKind;
const OUTBOUND_CAPACITY: usize = 256;
@@ -21,7 +23,11 @@ struct ActiveSubscription {
methods: Option<mpsc::Sender<protocol::Method>>,
}
pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebSocket) {
pub(crate) async fn serve_workspace_subscription(
api: WorkspaceApi,
socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
let broker = api.runtime_subscription_broker().clone();
let (mut socket_sender, mut socket_receiver) = socket.split();
let (control_outbound, mut control_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
@@ -85,7 +91,13 @@ pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebS
runtime_id: Some(runtime_id),
} => {
let worker = RuntimeWorkerRef::new(&runtime_id, worker_id.as_str());
match connect_workspace_worker_protocol(&api, &worker).await {
match connect_workspace_worker_protocol(
&api,
&worker,
Some(&input_source),
)
.await
{
Ok(connection) => {
let methods = connection.methods.clone();
let task = tokio::spawn(run_worker_protocol(
@@ -153,7 +165,12 @@ pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebS
else {
break;
};
if methods.send(message.method).await.is_err() {
let Ok(method) =
authorize_browser_worker_method(message.method, &input_source)
else {
break;
};
if methods.send(method).await.is_err() {
break;
}
}