diff --git a/Cargo.lock b/Cargo.lock index e6038be2..8c6a126e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6650,6 +6650,7 @@ dependencies = [ "serial_test", "session-metrics", "session-store", + "sha2 0.11.0", "tempfile", "thiserror 2.0.18", "ticket", diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs index 517d6a7c..b760ce00 100644 --- a/crates/client/src/client.rs +++ b/crates/client/src/client.rs @@ -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, diff --git a/crates/client/src/transport/in_process.rs b/crates/client/src/transport/in_process.rs index b1db8e6b..3111808a 100644 --- a/crates/client/src/transport/in_process.rs +++ b/crates/client/src/transport/in_process.rs @@ -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( diff --git a/crates/client/src/transport/unix_socket.rs b/crates/client/src/transport/unix_socket.rs index 0262bff2..089ed83e 100644 --- a/crates/client/src/transport/unix_socket.rs +++ b/crates/client/src/transport/unix_socket.rs @@ -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] diff --git a/crates/client/src/transport/websocket.rs b/crates/client/src/transport/websocket.rs index e8640573..b4c1ed84 100644 --- a/crates/client/src/transport/websocket.rs +++ b/crates/client/src/transport/websocket.rs @@ -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!( diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index 9f0ab12c..fd5190c7 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -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, }, - /// 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, - 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`. Dumb +/// `Method::Submit` and `Event::UserMessage` carry `Vec`. 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) -> Self { - Self::Run { - input: vec![Segment::text(s)], + pub fn submit_text(submission_request_id: impl Into, text: impl Into) -> 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, + #[serde(default)] + pub submissions: Vec, +} + /// 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, } @@ -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, }, @@ -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::(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::(&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::(&json).unwrap(); + assert!(matches!( + decoded, + Method::SubmitTracked { + source: AuthenticatedInputSource::UntrustedWire, + .. + } + )); assert!( serde_json::from_str::( - 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(), }, }; diff --git a/crates/protocol/src/typescript.rs b/crates/protocol/src/typescript.rs index bdac40ca..bf63c8d6 100644 --- a/crates/protocol/src/typescript.rs +++ b/crates/protocol/src/typescript.rs @@ -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::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); diff --git a/crates/session-store/src/fs_store.rs b/crates/session-store/src/fs_store.rs index a0e80a75..19a2c6d0 100644 --- a/crates/session-store/src/fs_store.rs +++ b/crates/session-store/src/fs_store.rs @@ -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 { + 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(); diff --git a/crates/session-store/src/legacy_session_log.rs b/crates/session-store/src/legacy_session_log.rs index 40dd0fd5..bb9dd1dd 100644 --- a/crates/session-store/src/legacy_session_log.rs +++ b/crates/session-store/src/legacy_session_log.rs @@ -183,6 +183,7 @@ fn canonicalize_history_entry( item, metadata: legacy_metadata(segment_id, line_index, 0), }, + extensions: Vec::new(), }, } } diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs index f3b61754..b1793584 100644 --- a/crates/session-store/src/public_snapshot.rs +++ b/crates/session-store/src/public_snapshot.rs @@ -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( diff --git a/crates/session-store/src/segment.rs b/crates/session-store/src/segment.rs index 163b2317..75d997c2 100644 --- a/crates/session-store/src/segment.rs +++ b/crates/session-store/src/segment.rs @@ -287,6 +287,7 @@ pub fn append_system_item( LogEntry::AnnotatedSystemItem { ts: segment_log::now_millis(), entry, + extensions: Vec::new(), }, ) } diff --git a/crates/session-store/src/segment_log.rs b/crates/session-store/src/segment_log.rs index 9daf85c2..ee509cbb 100644 --- a/crates/session-store/src/segment_log.rs +++ b/crates/session-store/src/segment_log.rs @@ -112,6 +112,8 @@ pub enum LogEntry { AnnotatedSystemItem { ts: u64, entry: LoggedSystemHistoryEntry, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + extensions: Vec, }, /// 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 { diff --git a/crates/session-store/src/store.rs b/crates/session-store/src/store.rs index a22fd6c7..33daae39 100644 --- a/crates/session-store/src/store.rs +++ b/crates/session-store/src/store.rs @@ -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 { + Ok(0) + } + /// Delete an uncommitted uploaded file owned by `session_id`. fn delete_uploaded_file( &self, diff --git a/crates/session-store/src/uploaded_file.rs b/crates/session-store/src/uploaded_file.rs index 4ce37588..6583dd25 100644 --- a/crates/session-store/src/uploaded_file.rs +++ b/crates/session-store/src/uploaded_file.rs @@ -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, #[serde(default, skip_serializing_if = "Option::is_none")] + pending_owner_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] upload_context: Option, 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 { + 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> { 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 { + 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 { 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 { 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 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) { diff --git a/crates/standalone/tests/host.rs b/crates/standalone/tests/host.rs index 510cc68c..45c8f620 100644 --- a/crates/standalone/tests/host.rs +++ b/crates/standalone/tests/host.rs @@ -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 diff --git a/crates/tools/src/bash.rs b/crates/tools/src/bash.rs index fb68fcbc..433dec0c 100644 --- a/crates/tools/src/bash.rs +++ b/crates/tools/src/bash.rs @@ -118,6 +118,7 @@ impl Tool for BashTool { command: params.command, timeout_secs, output_limit: INLINE_BYTE_BUDGET, + cwd: None, spill_dir: Some(self.output_dir.clone()), tool_call_id: Some(call_id.clone()), }) diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index b1f4fd04..d23430c2 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -102,23 +102,6 @@ struct RollbackSubmitState { turn_before: usize, } -#[derive(Clone)] -pub struct QueuedInput { - segments: Vec, - preview: String, -} - -impl QueuedInput { - fn new(segments: Vec) -> Self { - let preview = Segment::flatten_to_text(&segments); - Self { segments, preview } - } - - pub fn preview(&self) -> &str { - &self.preview - } -} - struct ComposerInputHistory { entries: VecDeque>, browse: Option, @@ -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, - /// 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, + /// 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 { + 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) { @@ -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 { + 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 { - 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); diff --git a/crates/tui/src/console/mod.rs b/crates/tui/src/console/mod.rs index 0106b403..a7ff703e 100644 --- a/crates/tui/src/console/mod.rs +++ b/crates/tui/src/console/mod.rs @@ -270,8 +270,8 @@ impl ConsoleConnection { async fn send(&mut self, method: &Method) -> Result<(), Box> { 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( } fn attachment_command_path(method: &Method) -> Option { - 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 { } 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( client: &mut ConsoleConnection, method: &Method, ) -> Result<(), Box> { - 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( } 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 { 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 { 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 { /// Idle / Paused → 2-tap to quit the TUI (the Worker keeps running). fn handle_pause_or_quit(app: &mut App) -> Option { 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(); diff --git a/crates/tui/src/ui.rs b/crates/tui/src/ui.rs index 35b4fcbc..58b790f1 100644 --- a/crates/tui/src/ui.rs +++ b/crates/tui/src/ui.rs @@ -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(); diff --git a/crates/workdir/src/delegation.rs b/crates/workdir/src/delegation.rs deleted file mode 100644 index 17598dab..00000000 --- a/crates/workdir/src/delegation.rs +++ /dev/null @@ -1,1189 +0,0 @@ -use std::collections::HashMap; -use std::path::Path; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::sync::{Arc, Mutex, Weak}; - -use async_trait::async_trait; -use fs_operation::{ - EditRequest, EditResult, FsPath, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, - ListResult, ReadRequest, ReadResult, StatRequest, StatResult, WriteRequest, WriteResult, -}; -use tokio::sync::broadcast; - -use crate::{ - CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, - CommandSnapshot, CommandStatus, Workdir, WorkdirError, WorkdirSession, - WorkdirSessionCapabilities, WorkdirSessionCapability, WorkdirSessionHandle, -}; - -#[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WorkdirDelegationPermission { - Read, - Write, -} - -#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WorkdirDelegationRule { - pub target: FsPath, - pub permission: WorkdirDelegationPermission, - pub recursive: bool, -} - -#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WorkdirDelegationRequest { - pub rules: Vec, - pub cwd: FsPath, -} - -pub struct WorkdirDelegation { - pub scoped_session: WorkdirSessionHandle, - pub capabilities: WorkdirSessionCapabilities, - validity: Arc, -} - -impl std::fmt::Debug for WorkdirDelegation { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("WorkdirDelegation") - .field("workdir", &self.scoped_session.workdir()) - .field("capabilities", &self.capabilities) - .field("active", &self.is_active()) - .finish() - } -} - -impl WorkdirDelegation { - pub fn is_active(&self) -> bool { - self.validity.is_active() - } - - pub fn release(&self) { - self.validity.active.store(false, Ordering::Release); - } -} - -impl Drop for WorkdirDelegation { - fn drop(&mut self) { - self.release(); - } -} - -pub struct AppliedWorkdirDelegation { - pub scoped_session: WorkdirSessionHandle, - _leases: Vec, -} - -impl std::fmt::Debug for AppliedWorkdirDelegation { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AppliedWorkdirDelegation") - .field("workdir", self.scoped_session.workdir()) - .field("lease_count", &self._leases.len()) - .finish() - } -} - -pub async fn apply_delegation_chain( - source: WorkdirSessionHandle, - requests: impl IntoIterator, -) -> Result { - let mut current = source; - let mut leases = Vec::new(); - for request in requests { - let authority = if current.is_delegation_capable() { - current.clone() - } else { - delegation_capable_session(current.clone()) - }; - let lease = authority.delegate(request).await?; - current = lease.scoped_session.clone(); - leases.push(lease); - } - Ok(AppliedWorkdirDelegation { - scoped_session: current, - _leases: leases, - }) -} - -#[derive(Debug)] -struct SessionValidity { - active: AtomicBool, - parent: Option>, -} - -impl SessionValidity { - fn root() -> Arc { - Arc::new(Self { - active: AtomicBool::new(true), - parent: None, - }) - } - - fn child(parent: Arc) -> Arc { - Arc::new(Self { - active: AtomicBool::new(true), - parent: Some(parent), - }) - } - - fn is_active(&self) -> bool { - self.active.load(Ordering::Acquire) - && self.parent.as_ref().is_none_or(|parent| parent.is_active()) - } -} - -#[derive(Clone, Debug)] -struct ActiveWriteLease { - validity: Weak, - rules: Vec, -} - -struct DelegatingWorkdirSession { - source: WorkdirSessionHandle, - cwd: FsPath, - scope: Option>, - capabilities: WorkdirSessionCapabilities, - validity: Arc, - child_write_leases: Mutex>, - next_lease_id: AtomicU64, - closes_source: bool, -} - -impl std::fmt::Debug for DelegatingWorkdirSession { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DelegatingWorkdirSession") - .field("workdir", &self.source.workdir()) - .field("scope", &self.scope) - .field("capabilities", &self.capabilities) - .field("active", &self.validity.is_active()) - .finish_non_exhaustive() - } -} - -/// Wrap a provider session with logical-path delegation and parent write gates. -pub fn delegation_capable_session(source: WorkdirSessionHandle) -> WorkdirSessionHandle { - let capabilities = source.capabilities(); - Arc::new(DelegatingWorkdirSession { - source, - cwd: FsPath::new("").expect("empty Workdir path is valid"), - scope: None, - capabilities, - validity: SessionValidity::root(), - child_write_leases: Mutex::new(HashMap::new()), - next_lease_id: AtomicU64::new(1), - closes_source: true, - }) -} - -impl DelegatingWorkdirSession { - fn ensure_active(&self) -> Result<(), WorkdirError> { - if self.validity.is_active() { - Ok(()) - } else { - Err(WorkdirError::SessionClosed) - } - } - - fn ensure_capability( - &self, - required: WorkdirSessionCapability, - operation: &'static str, - ) -> Result<(), WorkdirError> { - self.ensure_active()?; - if self.capabilities.supports(required) { - Ok(()) - } else { - Err(WorkdirError::Denied(format!( - "delegated workdir session does not permit {operation}" - ))) - } - } - - fn ensure_path( - &self, - path: &FsPath, - permission: WorkdirDelegationPermission, - ) -> Result<(), WorkdirError> { - self.ensure_active()?; - if let Some(scope) = &self.scope { - if !scope - .iter() - .any(|rule| rule_allows_path(rule, path, permission)) - { - return Err(WorkdirError::Denied(format!( - "logical workdir path `{path}` is outside the delegated {permission:?} scope" - ))); - } - } - if permission == WorkdirDelegationPermission::Write { - self.ensure_parent_write_available(path)?; - } - Ok(()) - } - - fn resolve_path(&self, path: &FsPath) -> Result { - if self.cwd.as_str().is_empty() { - return Ok(path.clone()); - } - let joined = Path::new(self.cwd.as_str()).join(path.as_str()); - let joined = joined.to_str().ok_or_else(|| { - WorkdirError::Denied("logical Workdir path is not valid UTF-8".into()) - })?; - FsPath::new(joined).map_err(|error| WorkdirError::Denied(error.to_string())) - } - - fn ensure_read( - &self, - path: &FsPath, - capability: WorkdirSessionCapability, - ) -> Result<(), WorkdirError> { - self.ensure_capability(capability, "read operations")?; - self.ensure_path(path, WorkdirDelegationPermission::Read) - } - - fn ensure_write( - &self, - path: &FsPath, - capability: WorkdirSessionCapability, - ) -> Result<(), WorkdirError> { - self.ensure_capability(capability, "write operations")?; - self.ensure_path(path, WorkdirDelegationPermission::Write) - } - - fn ensure_command(&self) -> Result<(), WorkdirError> { - self.ensure_capability(WorkdirSessionCapability::Command, "command execution") - } - - fn ensure_parent_write_available(&self, path: &FsPath) -> Result<(), WorkdirError> { - let mut leases = self - .child_write_leases - .lock() - .expect("workdir delegation lease mutex poisoned"); - leases.retain(|_, lease| lease.validity.upgrade().is_some_and(|v| v.is_active())); - if leases.values().any(|lease| { - lease.rules.iter().any(|rule| { - rule.permission == WorkdirDelegationPermission::Write - && rule_allows_path(rule, path, WorkdirDelegationPermission::Write) - }) - }) { - Err(WorkdirError::Denied(format!( - "logical workdir path `{path}` is leased to a child session" - ))) - } else { - Ok(()) - } - } - - fn validate_delegation_rules( - &self, - rules: &[WorkdirDelegationRule], - ) -> Result { - self.ensure_active()?; - if rules.is_empty() { - return Err(WorkdirError::Denied( - "workdir delegation requires at least one logical scope rule".into(), - )); - } - let writable = rules - .iter() - .any(|rule| rule.permission == WorkdirDelegationPermission::Write); - if !self.capabilities.supports(WorkdirSessionCapability::Read) - || (writable - && (!self.capabilities.supports(WorkdirSessionCapability::Write) - || !self.capabilities.supports(WorkdirSessionCapability::Edit) - || !self - .capabilities - .supports(WorkdirSessionCapability::Command))) - { - return Err(WorkdirError::Denied( - "parent workdir session cannot delegate the requested capabilities".into(), - )); - } - for requested in rules { - if let Some(scope) = &self.scope { - if !scope - .iter() - .any(|parent| rule_contains_rule(parent, requested)) - { - return Err(WorkdirError::Denied(format!( - "logical workdir scope `{}` exceeds the parent delegation", - requested.target - ))); - } - } - } - let mut delegated = vec![WorkdirSessionCapability::Read]; - for capability in [ - WorkdirSessionCapability::Glob, - WorkdirSessionCapability::Grep, - ] { - if self.capabilities.supports(capability) { - delegated.push(capability); - } - } - if writable { - delegated.push(WorkdirSessionCapability::Write); - delegated.push(WorkdirSessionCapability::Edit); - delegated.push(WorkdirSessionCapability::Command); - } - Ok(WorkdirSessionCapabilities::from_capabilities(delegated)) - } -} - -#[async_trait] -impl WorkdirSession for DelegatingWorkdirSession { - fn workdir(&self) -> &Workdir { - self.source.workdir() - } - - fn capabilities(&self) -> WorkdirSessionCapabilities { - self.capabilities - } - - fn is_delegation_capable(&self) -> bool { - true - } - - fn transports_delegation_context(&self) -> bool { - self.source.transports_delegation_context() - } - - async fn capture_delegation_source( - &self, - request: &WorkdirDelegationRequest, - ) -> Result { - self.ensure_active()?; - if self.scope.is_some() { - return Err(WorkdirError::Denied( - "scoped Workdir sessions cannot expose their provider source".into(), - )); - } - self.source.capture_delegation_source(request).await - } - - async fn delegate( - &self, - request: WorkdirDelegationRequest, - ) -> Result { - let capabilities = self.validate_delegation_rules(&request.rules)?; - if !request - .rules - .iter() - .any(|rule| rule_allows_path(rule, &request.cwd, WorkdirDelegationPermission::Read)) - { - return Err(WorkdirError::Denied(format!( - "delegated cwd `{}` is outside the delegated readable scope", - request.cwd - ))); - } - let source = self.source.capture_delegation_source(&request).await?; - let validity = SessionValidity::child(self.validity.clone()); - let id = self.next_lease_id.fetch_add(1, Ordering::Relaxed); - if request - .rules - .iter() - .any(|rule| rule.permission == WorkdirDelegationPermission::Write) - { - self.child_write_leases - .lock() - .expect("workdir delegation lease mutex poisoned") - .insert( - id, - ActiveWriteLease { - validity: Arc::downgrade(&validity), - rules: request.rules.clone(), - }, - ); - } - let child: WorkdirSessionHandle = Arc::new(DelegatingWorkdirSession { - source, - cwd: request.cwd, - scope: Some(request.rules), - capabilities, - validity: validity.clone(), - child_write_leases: Mutex::new(HashMap::new()), - next_lease_id: AtomicU64::new(1), - closes_source: false, - }); - let scoped_session: WorkdirSessionHandle = - if capabilities == WorkdirSessionCapabilities::READ_ONLY { - Arc::new(ReadOnlyWorkdirSession::new(child)) - } else { - child - }; - Ok(WorkdirDelegation { - scoped_session, - capabilities, - validity, - }) - } - - async fn stat(&self, mut request: StatRequest) -> Result { - let path = self.resolve_path(&request.path)?; - self.ensure_read(&path, WorkdirSessionCapability::Read)?; - if !self.source.transports_delegation_context() { - request.path = path; - } - self.source.stat(request).await - } - - async fn read(&self, mut request: ReadRequest) -> Result { - let path = self.resolve_path(&request.path)?; - self.ensure_read(&path, WorkdirSessionCapability::Read)?; - if !self.source.transports_delegation_context() { - request.path = path; - } - self.source.read(request).await - } - - async fn write(&self, mut request: WriteRequest) -> Result { - let path = self.resolve_path(&request.path)?; - self.ensure_write(&path, WorkdirSessionCapability::Write)?; - if !self.source.transports_delegation_context() { - request.path = path; - } - self.source.write(request).await - } - - async fn edit(&self, mut request: EditRequest) -> Result { - let path = self.resolve_path(&request.path)?; - self.ensure_write(&path, WorkdirSessionCapability::Edit)?; - if !self.source.transports_delegation_context() { - request.path = path; - } - self.source.edit(request).await - } - - async fn list(&self, mut request: ListRequest) -> Result { - let path = self.resolve_path(&request.path)?; - self.ensure_read(&path, WorkdirSessionCapability::Read)?; - if !self.source.transports_delegation_context() { - request.path = path; - } - self.source.list(request).await - } - - async fn glob(&self, mut request: GlobRequest) -> Result { - let path = self.resolve_path(&request.path)?; - self.ensure_read(&path, WorkdirSessionCapability::Glob)?; - if !self.source.transports_delegation_context() { - request.path = path; - } - self.source.glob(request).await - } - - async fn grep(&self, mut request: GrepRequest) -> Result { - let path = self.resolve_path(&request.path)?; - self.ensure_read(&path, WorkdirSessionCapability::Grep)?; - if !self.source.transports_delegation_context() { - request.path = path; - } - self.source.grep(request).await - } - - async fn start_command(&self, request: CommandRequest) -> Result { - self.ensure_command()?; - self.source.start_command(request).await - } - - async fn command_status(&self, handle: CommandHandle) -> Result { - self.ensure_command()?; - self.source.command_status(handle).await - } - - async fn command_output( - &self, - request: CommandOutputRequest, - ) -> Result { - self.ensure_command()?; - self.source.command_output(request).await - } - - async fn cancel_command(&self, handle: CommandHandle) -> Result<(), WorkdirError> { - self.ensure_command()?; - self.source.cancel_command(handle).await - } - - fn subscribe_command_events(&self) -> Option> { - self.ensure_capability(WorkdirSessionCapability::Command, "command observation") - .ok()?; - self.source.subscribe_command_events() - } - - fn command_snapshot(&self) -> Vec { - if self - .ensure_capability(WorkdirSessionCapability::Command, "command observation") - .is_err() - { - return Vec::new(); - } - self.source.command_snapshot() - } - - async fn close(&self) -> Result<(), WorkdirError> { - self.validity.active.store(false, Ordering::Release); - if self.closes_source { - self.source.close().await - } else { - Ok(()) - } - } -} - -/// A fail-closed read-only view over an already scoped delegated session. -#[derive(Debug)] -pub struct ReadOnlyWorkdirSession { - inner: WorkdirSessionHandle, -} - -impl ReadOnlyWorkdirSession { - pub fn new(inner: WorkdirSessionHandle) -> Self { - Self { inner } - } -} - -#[async_trait] -impl WorkdirSession for ReadOnlyWorkdirSession { - fn workdir(&self) -> &Workdir { - self.inner.workdir() - } - - fn capabilities(&self) -> WorkdirSessionCapabilities { - WorkdirSessionCapabilities::READ_ONLY - } - - fn is_delegation_capable(&self) -> bool { - true - } - - fn transports_delegation_context(&self) -> bool { - self.inner.transports_delegation_context() - } - - async fn delegate( - &self, - request: WorkdirDelegationRequest, - ) -> Result { - if request - .rules - .iter() - .any(|rule| rule.permission == WorkdirDelegationPermission::Write) - { - return Err(WorkdirError::Denied( - "read-only workdir session cannot delegate write access".into(), - )); - } - self.inner.delegate(request).await - } - - async fn stat(&self, request: StatRequest) -> Result { - self.inner.stat(request).await - } - - async fn read(&self, request: ReadRequest) -> Result { - self.inner.read(request).await - } - - async fn write(&self, _request: WriteRequest) -> Result { - Err(WorkdirError::Denied("read-only workdir session".into())) - } - - async fn edit(&self, _request: EditRequest) -> Result { - Err(WorkdirError::Denied("read-only workdir session".into())) - } - - async fn list(&self, request: ListRequest) -> Result { - self.inner.list(request).await - } - - async fn glob(&self, request: GlobRequest) -> Result { - self.inner.glob(request).await - } - - async fn grep(&self, request: GrepRequest) -> Result { - self.inner.grep(request).await - } - - async fn start_command(&self, _request: CommandRequest) -> Result { - Err(WorkdirError::Denied("read-only workdir session".into())) - } - - async fn command_status(&self, _handle: CommandHandle) -> Result { - Err(WorkdirError::Denied("read-only workdir session".into())) - } - - async fn command_output( - &self, - _request: CommandOutputRequest, - ) -> Result { - Err(WorkdirError::Denied("read-only workdir session".into())) - } - - async fn cancel_command(&self, _handle: CommandHandle) -> Result<(), WorkdirError> { - Err(WorkdirError::Denied("read-only workdir session".into())) - } - - async fn close(&self) -> Result<(), WorkdirError> { - self.inner.close().await - } -} - -fn rule_allows_path( - rule: &WorkdirDelegationRule, - path: &FsPath, - required: WorkdirDelegationPermission, -) -> bool { - if required == WorkdirDelegationPermission::Write - && rule.permission != WorkdirDelegationPermission::Write - { - return false; - } - path_in_rule(rule, path) -} - -fn path_in_rule(rule: &WorkdirDelegationRule, path: &FsPath) -> bool { - let target = Path::new(rule.target.as_str()); - let path = Path::new(path.as_str()); - if path == target { - return true; - } - let Ok(suffix) = path.strip_prefix(target) else { - return false; - }; - let depth = suffix.components().count(); - rule.recursive || depth <= 1 -} - -fn rule_contains_rule(parent: &WorkdirDelegationRule, child: &WorkdirDelegationRule) -> bool { - if child.permission == WorkdirDelegationPermission::Write - && parent.permission != WorkdirDelegationPermission::Write - { - return false; - } - if !path_in_rule(parent, &child.target) { - return false; - } - if parent.recursive { - return true; - } - !child.recursive && parent.target == child.target -} - -#[cfg(test)] -mod tests { - use std::fs; - - use manifest::{Permission, Scope, ScopeConfig, ScopeRule, SharedScope}; - use tempfile::TempDir; - - use super::*; - use crate::LocalWorkdirSession; - - fn fs_path(path: &str) -> FsPath { - FsPath::new(path).unwrap() - } - - fn session(root: &Path) -> WorkdirSessionHandle { - let scope = SharedScope::new( - Scope::from_config(&ScopeConfig { - allow: vec![ScopeRule { - target: root.to_path_buf(), - permission: Permission::Write, - recursive: true, - }], - deny: Vec::new(), - }) - .unwrap(), - ); - delegation_capable_session(Arc::new(LocalWorkdirSession::materialized_bound( - Workdir::new("delegation-test"), - root.to_path_buf(), - root.to_path_buf(), - scope, - WorkdirSessionCapabilities::ALL, - ))) - } - - fn request(path: &str, permission: WorkdirDelegationPermission) -> WorkdirDelegationRequest { - WorkdirDelegationRequest { - rules: vec![WorkdirDelegationRule { - target: fs_path(path), - permission, - recursive: true, - }], - cwd: fs_path(path), - } - } - - fn read(path: &str) -> ReadRequest { - ReadRequest { - path: fs_path(path), - offset: 0, - limit: 20, - max_bytes: 1024, - } - } - - fn write(path: &str, content: &str) -> WriteRequest { - WriteRequest { - path: fs_path(path), - content: content.as_bytes().to_vec(), - expected_hash: None, - } - } - - async fn run_command( - session: &WorkdirSessionHandle, - command: impl Into, - tool_call_id: impl Into, - ) -> CommandOutput { - let handle = session - .start_command(CommandRequest { - command: command.into(), - timeout_secs: 5, - output_limit: 1024, - spill_dir: None, - tool_call_id: Some(tool_call_id.into()), - }) - .await - .unwrap(); - session - .command_output(CommandOutputRequest { - handle, - cursor: 0, - limit: 1024, - wait: true, - }) - .await - .unwrap() - } - - #[tokio::test] - async fn delegation_capable_session_forwards_command_telemetry() { - let root = TempDir::new().unwrap(); - let parent = session(root.path()); - let mut events = parent - .subscribe_command_events() - .expect("delegation wrapper must preserve command observation"); - let handle = parent - .start_command(CommandRequest { - command: "printf ready; sleep 0.2; printf done".into(), - timeout_secs: 5, - output_limit: 1024, - spill_dir: None, - tool_call_id: Some("tool-delegated".into()), - }) - .await - .unwrap(); - - let first_output = loop { - let event = tokio::time::timeout(std::time::Duration::from_secs(2), events.recv()) - .await - .expect("delegated command telemetry should not stall") - .unwrap(); - if let CommandEvent::Output { content, .. } = event { - break content; - } - }; - assert_eq!(first_output, "ready"); - let snapshots = parent.command_snapshot(); - assert_eq!(snapshots.len(), 1); - assert_eq!(snapshots[0].command_id, handle.0); - assert_eq!(snapshots[0].status, CommandStatus::Running); - assert_eq!(snapshots[0].stdout.content, "ready"); - - let output = parent - .command_output(CommandOutputRequest { - handle, - cursor: 0, - limit: 1024, - wait: true, - }) - .await - .unwrap(); - assert_eq!(output.status, CommandStatus::Completed); - assert_eq!(output.content, "readydone"); - assert!(parent.command_snapshot().is_empty()); - } - - #[test] - fn non_recursive_rule_covers_target_and_direct_children_only() { - let rule = WorkdirDelegationRule { - target: fs_path("docs"), - permission: WorkdirDelegationPermission::Read, - recursive: false, - }; - assert!(path_in_rule(&rule, &fs_path("docs"))); - assert!(path_in_rule(&rule, &fs_path("docs/readme.md"))); - assert!(!path_in_rule(&rule, &fs_path("docs/guides/start.md"))); - } - - #[tokio::test] - async fn read_only_delegation_allows_prefix_and_denies_mutation() { - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("docs")).unwrap(); - fs::create_dir_all(root.path().join("secret")).unwrap(); - fs::write(root.path().join("docs/readme.md"), "visible").unwrap(); - fs::write(root.path().join("secret/key"), "hidden").unwrap(); - let parent = session(root.path()); - - let child = parent - .delegate(request("docs", WorkdirDelegationPermission::Read)) - .await - .unwrap(); - assert_eq!(child.capabilities, WorkdirSessionCapabilities::READ_ONLY); - assert_eq!( - child - .scoped_session - .read(read("readme.md")) - .await - .unwrap() - .bytes, - b"visible" - ); - assert!(matches!( - child.scoped_session.write(write("new.md", "no")).await, - Err(WorkdirError::Denied(_)) - )); - assert!( - !child - .capabilities - .supports(WorkdirSessionCapability::Command) - ); - assert!(child.scoped_session.subscribe_command_events().is_none()); - assert!(child.scoped_session.command_snapshot().is_empty()); - assert!(matches!( - child - .scoped_session - .start_command(CommandRequest { - command: "printf denied".into(), - timeout_secs: 5, - output_limit: 1024, - spill_dir: None, - tool_call_id: Some("read-only-command".into()), - }) - .await, - Err(WorkdirError::Denied(_)) - )); - } - - #[cfg(unix)] - #[tokio::test] - async fn provider_scope_denies_read_through_symlink_outside_grant() { - use std::os::unix::fs::symlink; - - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("granted")).unwrap(); - fs::create_dir_all(root.path().join("secret")).unwrap(); - fs::write(root.path().join("secret/key"), "hidden").unwrap(); - symlink("../secret/key", root.path().join("granted/link")).unwrap(); - let parent = session(root.path()); - let child = parent - .delegate(request("granted", WorkdirDelegationPermission::Read)) - .await - .unwrap(); - - let result = child.scoped_session.read(read("link")).await; - assert!( - result.is_err(), - "symlink read escaped provider scope: {result:?}" - ); - } - - #[cfg(unix)] - #[tokio::test] - async fn provider_scope_denies_write_through_symlink_outside_grant() { - use std::os::unix::fs::symlink; - - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("granted")).unwrap(); - fs::create_dir_all(root.path().join("secret")).unwrap(); - symlink("../secret", root.path().join("granted/outside")).unwrap(); - let parent = session(root.path()); - let child = parent - .delegate(request("granted", WorkdirDelegationPermission::Write)) - .await - .unwrap(); - - let result = child - .scoped_session - .write(write("outside/new", "forbidden")) - .await; - assert!( - result.is_err(), - "symlink write escaped provider scope: {result:?}" - ); - assert!(!root.path().join("secret/new").exists()); - } - - #[cfg(unix)] - #[tokio::test] - async fn write_delegation_rejects_symlink_target_before_lease() { - use std::os::unix::fs::symlink; - - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("granted")).unwrap(); - fs::create_dir_all(root.path().join("secret")).unwrap(); - symlink("../secret", root.path().join("granted/outside")).unwrap(); - let parent = session(root.path()); - - assert!(matches!( - parent - .delegate(request( - "granted/outside", - WorkdirDelegationPermission::Write - )) - .await, - Err(WorkdirError::Denied(_)) - )); - parent - .write(write("secret/parent", "still-authoritative")) - .await - .unwrap(); - } - - #[tokio::test] - async fn write_lease_keeps_typed_parent_writes_exclusive_without_blocking_commands() { - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("leased")).unwrap(); - fs::create_dir_all(root.path().join("other")).unwrap(); - let parent = session(root.path()); - let child = parent - .delegate(request("leased", WorkdirDelegationPermission::Write)) - .await - .unwrap(); - assert!( - child - .capabilities - .supports(WorkdirSessionCapability::Command) - ); - let child_output = run_command( - &child.scoped_session, - "printf child-command", - "delegated-child-command", - ) - .await; - assert_eq!(child_output.content, "child-command"); - let parent_output = run_command( - &parent, - "printf parent-write > leased/from-command; printf parent-command", - "parent-command-during-child-write", - ) - .await; - assert_eq!(parent_output.status, CommandStatus::Completed); - assert_eq!(parent_output.content, "parent-command"); - assert_eq!( - fs::read_to_string(root.path().join("leased/from-command")).unwrap(), - "parent-write" - ); - - assert!(matches!( - parent.write(write("leased/file", "parent")).await, - Err(WorkdirError::Denied(_)) - )); - parent.write(write("other/file", "parent")).await.unwrap(); - child - .scoped_session - .write(write("file", "child")) - .await - .unwrap(); - child.release(); - assert!(matches!( - child - .scoped_session - .start_command(CommandRequest { - command: "printf revoked".into(), - timeout_secs: 5, - output_limit: 1024, - spill_dir: None, - tool_call_id: Some("revoked-child-command".into()), - }) - .await, - Err(WorkdirError::SessionClosed) - )); - parent - .write(write("leased/parent", "parent")) - .await - .unwrap(); - assert!(matches!( - child.scoped_session.read(read("file")).await, - Err(WorkdirError::SessionClosed) - )); - } - - #[tokio::test] - async fn nested_delegation_is_attenuated_and_parent_revocation_cascades() { - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("docs/sub")).unwrap(); - fs::create_dir_all(root.path().join("docs/peer")).unwrap(); - fs::write(root.path().join("docs/sub/a"), "a").unwrap(); - fs::write(root.path().join("docs/peer/b"), "b").unwrap(); - let root_session = session(root.path()); - let child = root_session - .delegate(request("docs", WorkdirDelegationPermission::Read)) - .await - .unwrap(); - let nested = child - .scoped_session - .delegate(request("docs/sub", WorkdirDelegationPermission::Read)) - .await - .unwrap(); - - nested.scoped_session.read(read("a")).await.unwrap(); - assert!( - child - .scoped_session - .delegate(request("other", WorkdirDelegationPermission::Read)) - .await - .is_err() - ); - assert!( - child - .scoped_session - .delegate(request("docs/sub", WorkdirDelegationPermission::Write)) - .await - .is_err() - ); - - child.release(); - assert!(matches!( - nested.scoped_session.read(read("a")).await, - Err(WorkdirError::SessionClosed) - )); - } - - #[tokio::test] - async fn nested_write_leases_do_not_block_command_capable_ancestors() { - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("docs/sub")).unwrap(); - let root_session = session(root.path()); - let child = root_session - .delegate(request("docs", WorkdirDelegationPermission::Write)) - .await - .unwrap(); - let nested = child - .scoped_session - .delegate(request("docs/sub", WorkdirDelegationPermission::Write)) - .await - .unwrap(); - - for (session, label) in [ - (&root_session, "root"), - (&child.scoped_session, "child"), - (&nested.scoped_session, "nested"), - ] { - let output = run_command( - session, - format!("printf {label}"), - format!("{label}-command-during-nested-write"), - ) - .await; - assert_eq!(output.status, CommandStatus::Completed); - assert_eq!(output.content, label); - } - - assert!(matches!( - root_session.write(write("docs/root", "blocked")).await, - Err(WorkdirError::Denied(_)) - )); - assert!(matches!( - child - .scoped_session - .write(write("sub/child", "blocked")) - .await, - Err(WorkdirError::Denied(_)) - )); - nested - .scoped_session - .write(write("nested", "allowed")) - .await - .unwrap(); - - nested.release(); - child.release(); - } - - #[tokio::test] - async fn reapplied_write_delegation_chain_forwards_command_lifecycle() { - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("delegated")).unwrap(); - let applied = apply_delegation_chain( - session(root.path()), - [request("delegated", WorkdirDelegationPermission::Write)], - ) - .await - .unwrap(); - - let output = run_command( - &applied.scoped_session, - "printf reapplied", - "reapplied-command", - ) - .await; - assert_eq!(output.status, CommandStatus::Completed); - assert_eq!(output.content, "reapplied"); - } - - #[tokio::test] - async fn applied_chain_cannot_replace_outer_provider_attenuation() { - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("outer")).unwrap(); - fs::create_dir_all(root.path().join("outside")).unwrap(); - let result = apply_delegation_chain( - Arc::new(LocalWorkdirSession::materialized_bound( - Workdir::new("delegation-chain-test"), - root.path().to_path_buf(), - root.path().to_path_buf(), - SharedScope::new( - Scope::from_config(&ScopeConfig { - allow: vec![ScopeRule { - target: root.path().to_path_buf(), - permission: Permission::Write, - recursive: true, - }], - deny: Vec::new(), - }) - .unwrap(), - ), - WorkdirSessionCapabilities::ALL, - )), - [ - request("outer", WorkdirDelegationPermission::Read), - request("outside", WorkdirDelegationPermission::Read), - ], - ) - .await; - assert!(matches!(result, Err(WorkdirError::Denied(_)))); - } - - #[tokio::test] - async fn closing_parent_invalidates_delegated_sessions() { - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("docs")).unwrap(); - fs::write(root.path().join("docs/a"), "a").unwrap(); - let parent = session(root.path()); - let child = parent - .delegate(request("docs", WorkdirDelegationPermission::Read)) - .await - .unwrap(); - - parent.close().await.unwrap(); - assert!(matches!( - parent - .start_command(CommandRequest { - command: "printf closed".into(), - timeout_secs: 5, - output_limit: 1024, - spill_dir: None, - tool_call_id: Some("closed-parent-command".into()), - }) - .await, - Err(WorkdirError::SessionClosed) - )); - assert!(matches!( - child.scoped_session.read(read("a")).await, - Err(WorkdirError::SessionClosed) - )); - } -} diff --git a/crates/workdir/src/http.rs b/crates/workdir/src/http.rs index 72733e8b..f89b608b 100644 --- a/crates/workdir/src/http.rs +++ b/crates/workdir/src/http.rs @@ -68,12 +68,10 @@ pub enum WorkdirSessionOperation { CommandCancel(CommandHandle), } -/// Wire envelope for an operation and its optional provider-enforced child scope. +/// Wire envelope for one provider operation. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct WorkdirSessionOperationRequest { - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub delegations: Vec, pub operation: WorkdirSessionOperation, } @@ -289,7 +287,7 @@ mod client { use reqwest::{Client, StatusCode, Url}; use super::*; - use crate::{Workdir, WorkdirSession, WorkdirSessionHandle}; + use crate::{Workdir, WorkdirSession}; /// Provides a fresh bearer token for each Runtime request. Backend /// implementations can mint short-lived capability tokens without making a @@ -324,7 +322,6 @@ mod client { workdir: Workdir, session_id: WorkdirSessionId, capabilities: WorkdirSessionCapabilities, - delegations: Vec, closed: AtomicBool, } @@ -377,7 +374,6 @@ mod client { workdir: Workdir::new(opened.workdir_id.as_str()), session_id: opened.session_id, capabilities: opened.capabilities, - delegations: Vec::new(), closed: AtomicBool::new(false), }) } @@ -404,10 +400,7 @@ mod client { "operations", ], )?; - let operation = WorkdirSessionOperationRequest { - delegations: self.delegations.clone(), - operation, - }; + let operation = WorkdirSessionOperationRequest { operation }; let response = self .client .post(url) @@ -436,37 +429,6 @@ mod client { self.capabilities } - fn transports_delegation_context(&self) -> bool { - true - } - - async fn capture_delegation_source( - &self, - request: &crate::WorkdirDelegationRequest, - ) -> Result { - if self.closed.load(Ordering::Acquire) { - return Err(WorkdirError::SessionClosed); - } - let mut delegations = self.delegations.clone(); - delegations.push(request.clone()); - let candidate = Arc::new(Self { - client: self.client.clone(), - base_url: self.base_url.clone(), - authorization: self.authorization.clone(), - workdir: self.workdir.clone(), - session_id: self.session_id.clone(), - capabilities: self.capabilities, - delegations, - closed: AtomicBool::new(false), - }); - candidate - .stat(StatRequest { - path: fs_operation::FsPath::new("").expect("empty Workdir path is valid"), - }) - .await?; - Ok(candidate) - } - async fn stat(&self, request: StatRequest) -> Result { match self.operate(WorkdirSessionOperation::Stat(request)).await? { WorkdirSessionOperationResult::Stat(result) => Ok(result), diff --git a/crates/workdir/src/lib.rs b/crates/workdir/src/lib.rs index 5d58da26..4cfa91c1 100644 --- a/crates/workdir/src/lib.rs +++ b/crates/workdir/src/lib.rs @@ -5,10 +5,10 @@ //! bound to one Worker. Tools consume sessions; they do not own Workdir //! materialization or cleanup. -mod delegation; pub mod http; mod local; mod operation; +mod scope; pub mod workspace; use std::path::{Path, PathBuf}; @@ -18,11 +18,6 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use tokio::sync::broadcast; -pub use delegation::{ - AppliedWorkdirDelegation, ReadOnlyWorkdirSession, WorkdirDelegation, - WorkdirDelegationPermission, WorkdirDelegationRequest, WorkdirDelegationRule, - apply_delegation_chain, delegation_capable_session, -}; pub use fs_operation::{ ContentHash, EditRequest, EditResult, EntryKind, FsPath as WorkdirPath, GlobRequest, GlobResult, GrepOutputMode, GrepRequest, GrepResult, ListEntry, ListRequest, ListResult, @@ -32,6 +27,10 @@ pub use local::{ LocalWorkdirSession, SymlinkInfo, WorkdirSessionResource, direct_symlink, first_symlink, }; pub use operation::*; +pub use scope::{ + ReadOnlyWorkdirSession, WorkdirScopeLease, WorkdirToolBroker, WorkdirToolScope, + WorkdirToolScopePermission, WorkdirToolScopeRule, +}; /// Persistent, opaque identity of one materialized Workdir. #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -148,39 +147,6 @@ pub trait WorkdirSession: std::fmt::Debug + Send + Sync { fn workdir(&self) -> &Workdir; fn capabilities(&self) -> WorkdirSessionCapabilities; - fn is_delegation_capable(&self) -> bool { - false - } - - /// Whether this session transports the delegation chain to another - /// provider boundary that will apply logical cwd/path resolution there. - fn transports_delegation_context(&self) -> bool { - false - } - - /// Capture a provider-specific source for a delegated child session. - /// Remote providers use this boundary to pin attachment identity without - /// exposing transport handles or host paths. - async fn capture_delegation_source( - &self, - _request: &WorkdirDelegationRequest, - ) -> Result { - Err(WorkdirError::Denied( - "workdir provider does not support delegated sessions".into(), - )) - } - - /// Attenuate this session into a revocable child lease. Only sessions - /// created with [`delegation_capable_session`] implement this operation. - async fn delegate( - &self, - _request: WorkdirDelegationRequest, - ) -> Result { - Err(WorkdirError::Denied( - "workdir session is not delegation-capable".into(), - )) - } - async fn stat(&self, request: StatRequest) -> Result; async fn read(&self, request: ReadRequest) -> Result; async fn write(&self, request: WriteRequest) -> Result; diff --git a/crates/workdir/src/local.rs b/crates/workdir/src/local.rs index 8f9556dd..94efeceb 100644 --- a/crates/workdir/src/local.rs +++ b/crates/workdir/src/local.rs @@ -18,7 +18,7 @@ use std::sync::{Arc, Mutex as StdMutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use async_trait::async_trait; -use manifest::{Permission, Scope, ScopeConfig, ScopeRule, SharedScope}; +use manifest::{Scope, SharedScope}; use sha2::{Digest, Sha256}; use tokio::process::Command; use tokio::sync::{Mutex, broadcast, watch}; @@ -28,10 +28,8 @@ use crate::{ CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, ReadRequest, - ReadResult, StatRequest, StatResult, Workdir, WorkdirDelegationPermission, - WorkdirDelegationRequest, WorkdirError, WorkdirPath, WorkdirSession, - WorkdirSessionCapabilities, WorkdirSessionCapability, WorkdirSessionHandle, WriteRequest, - WriteResult, + ReadResult, StatRequest, StatResult, Workdir, WorkdirError, WorkdirPath, WorkdirSession, + WorkdirSessionCapabilities, WorkdirSessionCapability, WriteRequest, WriteResult, }; #[cfg(test)] use crate::{EntryKind, WriteOutcome}; @@ -558,69 +556,6 @@ impl WorkdirSession for LocalWorkdirSession { self.inner.capabilities } - async fn capture_delegation_source( - &self, - request: &WorkdirDelegationRequest, - ) -> Result { - let host_rules = request - .rules - .iter() - .map(|rule| ScopeRule { - target: self.inner.root.join(rule.target.as_str()), - permission: match rule.permission { - WorkdirDelegationPermission::Read => Permission::Read, - WorkdirDelegationPermission::Write => Permission::Write, - }, - recursive: rule.recursive, - }) - .collect::>(); - for (logical, host) in request.rules.iter().zip(&host_rules) { - if logical.permission == WorkdirDelegationPermission::Write { - let resolved = Scope::resolved_target(host) - .map_err(|error| WorkdirError::Denied(error.to_string()))?; - if resolved != host.target { - return Err(WorkdirError::Denied(format!( - "write delegation target `{}` traverses a symlink", - logical.target - ))); - } - } - } - let parent_scope = self.inner.scope.snapshot(); - for rule in &host_rules { - if !parent_scope - .allows_rule(rule) - .map_err(|error| WorkdirError::Denied(error.to_string()))? - { - return Err(WorkdirError::Denied(format!( - "delegated provider scope `{}` exceeds the parent session", - rule.target.display() - ))); - } - } - let child_scope = Scope::from_config(&ScopeConfig { - allow: host_rules, - deny: Vec::new(), - }) - .map_err(|error| WorkdirError::Denied(error.to_string()))?; - let child_cwd = self.inner.root.join(request.cwd.as_str()); - if !child_scope.is_readable(&child_cwd) - || !std::fs::metadata(&child_cwd).is_ok_and(|metadata| metadata.is_dir()) - { - return Err(WorkdirError::Denied(format!( - "delegated cwd `{}` is not a readable Workdir directory", - request.cwd - ))); - } - Ok(Arc::new(LocalWorkdirSession::materialized_bound( - self.inner.workdir.clone(), - self.inner.root.clone(), - self.inner.root.clone(), - SharedScope::new(child_scope), - self.inner.capabilities, - ))) - } - async fn stat(&self, request: StatRequest) -> Result { self.ensure_capability(WorkdirSessionCapability::Read)?; let logical = request.path.clone(); @@ -694,9 +629,20 @@ impl WorkdirSession for LocalWorkdirSession { { return Err(WorkdirError::OutOfScope(spill_dir.to_path_buf())); } + let cwd = if let Some(logical_cwd) = request.cwd.as_ref() { + let cwd = self.resolve(logical_cwd); + let scope = self.inner.scope.snapshot(); + if !scope.is_readable(&cwd) + || !std::fs::metadata(&cwd).is_ok_and(|metadata| metadata.is_dir()) + { + return Err(WorkdirError::OutOfScope(cwd)); + } + cwd + } else { + self.inner.cwd.clone() + }; let id = self.inner.next_command_id.fetch_add(1, Ordering::Relaxed); let handle = CommandHandle(format!("command-{id}")); - let cwd = self.inner.cwd.clone(); let (completion_tx, completion) = watch::channel(false); let command_id = handle.0.clone(); let telemetry = self.inner.command_telemetry.clone(); @@ -1516,6 +1462,7 @@ mod tests { command: "sleep 30".to_owned(), timeout_secs: 60, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: None, }, @@ -2043,6 +1990,7 @@ mod tests { command: "pwd && printf provider-command".into(), timeout_secs: 5, output_limit: 4096, + cwd: None, spill_dir: None, tool_call_id: None, }, @@ -2141,6 +2089,7 @@ mod tests { command: "printf hidden".into(), timeout_secs: 5, output_limit: 1, + cwd: None, spill_dir: Some(spill.path().to_path_buf()), tool_call_id: None, }, @@ -2178,6 +2127,7 @@ mod tests { command: "i=0; while [ $i -lt 200 ]; do printf 'line-%03d\\n' \"$i\"; i=$((i+1)); done; printf 'FINAL-NEEDLE\\n'".into(), timeout_secs: 5, output_limit: 64, + cwd: None, spill_dir: Some(spill.path().to_path_buf()), tool_call_id: None, }, @@ -2224,6 +2174,7 @@ mod tests { command: "printf 'aéz'".into(), timeout_secs: 5, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: None, }, @@ -2449,6 +2400,7 @@ mod tests { command: "printf ready; printf warning >&2; sleep 0.2; printf done".into(), timeout_secs: 5, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: Some("tool-7".into()), }, @@ -2553,6 +2505,7 @@ mod tests { command: "sleep 30".into(), timeout_secs: 1, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: None, }, @@ -2623,6 +2576,7 @@ mod tests { command: "sleep 30".into(), timeout_secs: 60, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: None, }, diff --git a/crates/workdir/src/operation.rs b/crates/workdir/src/operation.rs index 67858685..5af8b54a 100644 --- a/crates/workdir/src/operation.rs +++ b/crates/workdir/src/operation.rs @@ -11,6 +11,10 @@ pub struct CommandRequest { pub command: String, pub timeout_secs: u64, pub output_limit: usize, + /// Workdir-relative command directory. Providers validate it against the + /// active session before process start. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cwd: Option, /// Provider-local directory where complete output is retained when the /// inline result exceeds `output_limit`. pub spill_dir: Option, diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs new file mode 100644 index 00000000..131ae6cf --- /dev/null +++ b/crates/workdir/src/scope.rs @@ -0,0 +1,1944 @@ +use std::collections::{HashMap, HashSet}; +use std::path::Path; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, Weak}; + +use async_trait::async_trait; +use fs_operation::{ + EditRequest, EditResult, FsPath, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, + ListResult, ReadRequest, ReadResult, StatRequest, StatResult, WriteRequest, WriteResult, +}; +use tokio::sync::broadcast; + +const MAX_SCOPED_COMMANDS: usize = 16; + +use crate::{ + CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, + CommandSnapshot, CommandStatus, CommandStream, Workdir, WorkdirError, WorkdirSession, + WorkdirSessionCapabilities, WorkdirSessionCapability, WorkdirSessionHandle, +}; + +#[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WorkdirToolScopePermission { + Read, + Write, +} + +#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WorkdirToolScopeRule { + pub target: FsPath, + pub permission: WorkdirToolScopePermission, + pub recursive: bool, +} + +#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WorkdirToolScope { + pub rules: Vec, + pub cwd: FsPath, + pub command: bool, +} + +#[derive(Clone)] +pub struct WorkdirToolBroker { + authority: Arc, + session: WorkdirSessionHandle, + event_forwarder: Option>>>>, +} + +impl std::fmt::Debug for WorkdirToolBroker { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkdirToolBroker") + .field("workdir", self.session.workdir()) + .field("capabilities", &self.session.capabilities()) + .finish_non_exhaustive() + } +} + +impl WorkdirToolBroker { + /// Own the parent Worker's active session and mediate every scoped child operation. + pub fn new(source: WorkdirSessionHandle) -> Self { + let capabilities = source.capabilities(); + let (command_events, _) = broadcast::channel(64); + let authority = Arc::new(ScopedWorkdirSession { + source, + cwd: FsPath::new("").expect("empty Workdir path is valid"), + scope: None, + capabilities, + validity: SessionValidity::root(), + child_write_leases: Mutex::new(HashMap::new()), + next_lease_id: AtomicU64::new(1), + close_lock: Arc::new(tokio::sync::Mutex::new(())), + owned_commands: Arc::new(Mutex::new(HashSet::new())), + pending_command_events: Arc::new(Mutex::new(HashMap::new())), + starting_tool_calls: Arc::new(Mutex::new(HashSet::new())), + forwarded_starts: Arc::new(Mutex::new(HashSet::new())), + forwarded_terminals: Arc::new(Mutex::new(HashSet::new())), + command_events, + closes_source: true, + #[cfg(test)] + command_start_gate: Mutex::new(None), + }); + Self { + session: authority.clone(), + authority, + event_forwarder: None, + } + } + + /// Session used only by tools registered by the owning Worker. + pub fn tool_session(&self) -> WorkdirSessionHandle { + self.session.clone() + } + + /// Create a revocable, attenuated tool route without delegating a provider session. + pub async fn scope( + &self, + request: WorkdirToolScope, + ) -> Result { + self.authority.scope(request).await + } +} + +impl std::ops::Deref for WorkdirToolBroker { + type Target = WorkdirSessionHandle; + + fn deref(&self) -> &Self::Target { + &self.session + } +} + +pub struct WorkdirScopeLease { + broker: WorkdirToolBroker, + pub capabilities: WorkdirSessionCapabilities, + validity: Arc, + cleanup_pending: Arc, + close_lock: Arc>, +} + +impl std::fmt::Debug for WorkdirScopeLease { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WorkdirScopeLease") + .field("workdir", self.broker.session.workdir()) + .field("capabilities", &self.capabilities) + .field("active", &self.is_active()) + .finish() + } +} + +impl WorkdirScopeLease { + pub fn broker(&self) -> WorkdirToolBroker { + self.broker.clone() + } + + pub fn tool_session(&self) -> WorkdirSessionHandle { + self.broker.tool_session() + } + + pub async fn scope( + &self, + request: WorkdirToolScope, + ) -> Result { + self.broker.scope(request).await + } + + pub async fn close(&self) -> Result<(), WorkdirError> { + let _close_guard = self.close_lock.lock().await; + if !self.cleanup_pending.load(Ordering::Acquire) { + return Ok(()); + } + self.validity.active.store(false, Ordering::Release); + let command_ids = self + .broker + .authority + .owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .iter() + .cloned() + .collect::>(); + let mut first_error = None; + for command_id in command_ids { + let handle = CommandHandle(command_id.clone()); + let cancel = self + .broker + .authority + .source + .cancel_command(handle.clone()) + .await; + let terminal = self + .broker + .authority + .source + .command_output(CommandOutputRequest { + handle, + cursor: 0, + limit: 1, + wait: true, + }) + .await; + match (cancel, terminal) { + (_, Ok(output)) => { + self.broker.authority.publish_terminal_if_missing( + &command_id, + output.status, + output.exit_code, + output.next_cursor.unwrap_or(output.content.len()) as u64, + &output.content, + ); + self.broker + .authority + .owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .remove(&command_id); + } + (Ok(()), Err(WorkdirError::UnknownCommand(_))) + | (Err(WorkdirError::UnknownCommand(_)), Err(WorkdirError::UnknownCommand(_))) => { + self.broker.authority.publish_terminal_if_missing( + &command_id, + CommandStatus::Cancelled, + None, + 0, + "", + ); + self.broker + .authority + .owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .remove(&command_id); + } + (Err(error), _) | (_, Err(error)) => { + first_error.get_or_insert(error); + } + } + } + if let Some(error) = first_error { + return Err(error); + } + tokio::task::yield_now().await; + self.finish_release(); + Ok(()) + } + + pub fn is_active(&self) -> bool { + self.validity.is_active() + } + + /// Revoke a scope whose owner has already terminalized every tool call. + /// Use [`Self::close`] when commands may still be live. + pub fn revoke(&self) { + self.finish_release(); + } + + fn finish_release(&self) { + self.validity.active.store(false, Ordering::Release); + self.cleanup_pending.store(false, Ordering::Release); + if let Some(forwarder) = &self.broker.event_forwarder + && let Some(handle) = forwarder + .lock() + .expect("scoped command forwarder mutex poisoned") + .take() + { + handle.abort(); + } + } +} + +impl std::ops::Deref for WorkdirScopeLease { + type Target = WorkdirSessionHandle; + + fn deref(&self) -> &Self::Target { + &self.broker.session + } +} + +impl Drop for WorkdirScopeLease { + fn drop(&mut self) { + self.finish_release(); + } +} + +#[derive(Debug)] +struct SessionValidity { + active: AtomicBool, + parent: Option>, +} + +impl SessionValidity { + fn root() -> Arc { + Arc::new(Self { + active: AtomicBool::new(true), + parent: None, + }) + } + + fn child(parent: Arc) -> Arc { + Arc::new(Self { + active: AtomicBool::new(true), + parent: Some(parent), + }) + } + + fn is_active(&self) -> bool { + self.active.load(Ordering::Acquire) + && self.parent.as_ref().is_none_or(|parent| parent.is_active()) + } +} + +#[derive(Clone, Debug)] +struct ActiveWriteLease { + validity: Weak, + cleanup_pending: Weak, + rules: Vec, +} + +#[cfg(test)] +struct TestCommandStartGate { + entered: tokio::sync::Notify, + release: tokio::sync::Notify, +} + +struct ScopedWorkdirSession { + source: WorkdirSessionHandle, + cwd: FsPath, + scope: Option>, + capabilities: WorkdirSessionCapabilities, + validity: Arc, + child_write_leases: Mutex>, + next_lease_id: AtomicU64, + close_lock: Arc>, + owned_commands: Arc>>, + pending_command_events: Arc>>>, + starting_tool_calls: Arc>>, + forwarded_starts: Arc>>, + forwarded_terminals: Arc>>, + command_events: broadcast::Sender, + closes_source: bool, + #[cfg(test)] + command_start_gate: Mutex>>, +} + +impl std::fmt::Debug for ScopedWorkdirSession { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ScopedWorkdirSession") + .field("workdir", &self.source.workdir()) + .field("scope", &self.scope) + .field("capabilities", &self.capabilities) + .field("active", &self.validity.is_active()) + .finish_non_exhaustive() + } +} + +impl ScopedWorkdirSession { + fn ensure_active(&self) -> Result<(), WorkdirError> { + if self.validity.is_active() { + Ok(()) + } else { + Err(WorkdirError::SessionClosed) + } + } + + fn ensure_capability( + &self, + required: WorkdirSessionCapability, + operation: &'static str, + ) -> Result<(), WorkdirError> { + self.ensure_active()?; + if self.capabilities.supports(required) { + Ok(()) + } else { + Err(WorkdirError::Denied(format!( + "scoped Workdir tools do not permit {operation}" + ))) + } + } + + fn ensure_path( + &self, + path: &FsPath, + permission: WorkdirToolScopePermission, + ) -> Result<(), WorkdirError> { + self.ensure_active()?; + if let Some(scope) = &self.scope { + if !scope + .iter() + .any(|rule| rule_allows_path(rule, path, permission)) + { + return Err(WorkdirError::Denied(format!( + "logical workdir path `{path}` is outside the scoped {permission:?} scope" + ))); + } + } + if permission == WorkdirToolScopePermission::Write { + self.ensure_parent_write_available(path)?; + } + Ok(()) + } + + fn resolve_path(&self, path: &FsPath) -> Result { + if self.cwd.as_str().is_empty() { + return Ok(path.clone()); + } + let joined = Path::new(self.cwd.as_str()).join(path.as_str()); + let joined = joined.to_str().ok_or_else(|| { + WorkdirError::Denied("logical Workdir path is not valid UTF-8".into()) + })?; + FsPath::new(joined).map_err(|error| WorkdirError::Denied(error.to_string())) + } + + fn ensure_read( + &self, + path: &FsPath, + capability: WorkdirSessionCapability, + ) -> Result<(), WorkdirError> { + self.ensure_capability(capability, "read operations")?; + self.ensure_path(path, WorkdirToolScopePermission::Read) + } + + fn ensure_write( + &self, + path: &FsPath, + capability: WorkdirSessionCapability, + ) -> Result<(), WorkdirError> { + self.ensure_capability(capability, "write operations")?; + self.ensure_path(path, WorkdirToolScopePermission::Write) + } + + fn ensure_command(&self) -> Result<(), WorkdirError> { + self.ensure_capability(WorkdirSessionCapability::Command, "command execution") + } + + fn ensure_owned_command(&self, handle: &CommandHandle) -> Result<(), WorkdirError> { + self.ensure_command()?; + if self.scope.is_none() + || self + .owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .contains(&handle.0) + { + Ok(()) + } else { + Err(WorkdirError::UnknownCommand(handle.0.clone())) + } + } + + fn publish_terminal_if_missing( + &self, + command_id: &str, + status: CommandStatus, + exit_code: Option, + offset: u64, + fallback_output: &str, + ) { + let mut terminals = self + .forwarded_terminals + .lock() + .expect("forwarded terminal command mutex poisoned"); + if !terminals.insert(command_id.to_string()) { + return; + } + if !fallback_output.is_empty() { + let _ = self.command_events.send(CommandEvent::Output { + command_id: command_id.to_string(), + stream: CommandStream::Stdout, + start_offset: 0, + end_offset: fallback_output.len() as u64, + content: fallback_output.to_string(), + observed_at_ms: unix_timestamp_ms(), + }); + } + let _ = self.command_events.send(CommandEvent::Terminal { + command_id: command_id.to_string(), + status, + exit_code, + stdout_end_offset: offset, + stderr_end_offset: 0, + observed_at_ms: unix_timestamp_ms(), + }); + } + + fn ensure_parent_write_available(&self, path: &FsPath) -> Result<(), WorkdirError> { + let mut leases = self + .child_write_leases + .lock() + .expect("Workdir tool scope lease mutex poisoned"); + leases.retain(|_, lease| { + lease + .validity + .upgrade() + .is_some_and(|validity| validity.is_active()) + || lease + .cleanup_pending + .upgrade() + .is_some_and(|pending| pending.load(Ordering::Acquire)) + }); + if leases.values().any(|lease| { + lease.rules.iter().any(|rule| { + rule.permission == WorkdirToolScopePermission::Write + && rule_allows_path(rule, path, WorkdirToolScopePermission::Write) + }) + }) { + Err(WorkdirError::Denied(format!( + "logical workdir path `{path}` is leased to child Workdir tools" + ))) + } else { + Ok(()) + } + } + + async fn ensure_source_path_has_no_symlink(&self, path: &FsPath) -> Result<(), WorkdirError> { + let mut current = String::new(); + for component in Path::new(path.as_str()).components() { + let component = component.as_os_str().to_string_lossy(); + if component.is_empty() || component == "." { + continue; + } + if !current.is_empty() { + current.push('/'); + } + current.push_str(&component); + let current = FsPath::new(¤t).map_err(|error| { + WorkdirError::Denied(format!("invalid scoped Workdir path: {error}")) + })?; + match self.source.stat(StatRequest { path: current }).await { + Ok(result) if result.kind == fs_operation::EntryKind::Symlink => { + return Err(WorkdirError::Denied(format!( + "scoped Workdir path `{path}` traverses a symlink" + ))); + } + Ok(_) => {} + Err(WorkdirError::NotFound(_)) => break, + Err(error) => return Err(error), + } + } + Ok(()) + } + + async fn ensure_scope_targets_do_not_traverse_symlinks( + &self, + rules: &[WorkdirToolScopeRule], + ) -> Result<(), WorkdirError> { + for rule in rules { + self.ensure_source_path_has_no_symlink(&rule.target).await?; + } + Ok(()) + } + + async fn resolve_operation_path(&self, path: &FsPath) -> Result { + self.ensure_active()?; + let resolved = self.resolve_path(path)?; + if self.scope.is_some() { + self.ensure_source_path_has_no_symlink(&resolved).await?; + } + Ok(resolved) + } + + fn validate_scope( + &self, + rules: &[WorkdirToolScopeRule], + command: bool, + ) -> Result { + self.ensure_active()?; + if rules.is_empty() { + return Err(WorkdirError::Denied( + "workdir tool scope requires at least one logical scope rule".into(), + )); + } + let writable = rules + .iter() + .any(|rule| rule.permission == WorkdirToolScopePermission::Write); + if !self.capabilities.supports(WorkdirSessionCapability::Read) + || (writable + && (!self.capabilities.supports(WorkdirSessionCapability::Write) + || !self.capabilities.supports(WorkdirSessionCapability::Edit))) + { + return Err(WorkdirError::Denied( + "parent Workdir session cannot scope the requested capabilities".into(), + )); + } + if command { + if !writable { + return Err(WorkdirError::Denied( + "command execution requires a writable scoped path".into(), + )); + } + if !self + .capabilities + .supports(WorkdirSessionCapability::Command) + { + return Err(WorkdirError::Denied( + "parent Workdir session does not support Command".into(), + )); + } + } + for requested in rules { + if let Some(scope) = &self.scope { + if !scope + .iter() + .any(|parent| rule_contains_rule(parent, requested)) + { + return Err(WorkdirError::Denied(format!( + "logical workdir scope `{}` exceeds the parent tool scope", + requested.target + ))); + } + } + } + let mut delegated = vec![WorkdirSessionCapability::Read]; + for capability in [ + WorkdirSessionCapability::Glob, + WorkdirSessionCapability::Grep, + ] { + if self.capabilities.supports(capability) { + delegated.push(capability); + } + } + if writable { + delegated.push(WorkdirSessionCapability::Write); + delegated.push(WorkdirSessionCapability::Edit); + } + if command { + delegated.push(WorkdirSessionCapability::Command); + } + Ok(WorkdirSessionCapabilities::from_capabilities(delegated)) + } + + async fn scope( + self: &Arc, + request: WorkdirToolScope, + ) -> Result { + let capabilities = self.validate_scope(&request.rules, request.command)?; + if !request + .rules + .iter() + .any(|rule| rule_allows_path(rule, &request.cwd, WorkdirToolScopePermission::Read)) + { + return Err(WorkdirError::Denied(format!( + "scoped tool cwd `{}` is outside the readable scope", + request.cwd + ))); + } + self.ensure_scope_targets_do_not_traverse_symlinks(&request.rules) + .await?; + let validity = SessionValidity::child(self.validity.clone()); + let cleanup_pending = Arc::new(AtomicBool::new(true)); + let id = self.next_lease_id.fetch_add(1, Ordering::Relaxed); + if request + .rules + .iter() + .any(|rule| rule.permission == WorkdirToolScopePermission::Write) + { + let mut leases = self + .child_write_leases + .lock() + .expect("Workdir tool scope lease mutex poisoned"); + leases.retain(|_, lease| { + lease + .validity + .upgrade() + .is_some_and(|validity| validity.is_active()) + || lease + .cleanup_pending + .upgrade() + .is_some_and(|pending| pending.load(Ordering::Acquire)) + }); + let requested_write_rules = request + .rules + .iter() + .filter(|rule| rule.permission == WorkdirToolScopePermission::Write); + for requested in requested_write_rules { + if leases.values().any(|lease| { + lease + .rules + .iter() + .any(|active| rules_overlap(active, requested)) + }) { + return Err(WorkdirError::Denied(format!( + "scoped write path `{}` overlaps an active child scope", + requested.target + ))); + } + } + leases.insert( + id, + ActiveWriteLease { + validity: Arc::downgrade(&validity), + cleanup_pending: Arc::downgrade(&cleanup_pending), + rules: request.rules.clone(), + }, + ); + } + let owned_commands = Arc::new(Mutex::new(HashSet::new())); + let pending_command_events = Arc::new(Mutex::new(HashMap::new())); + let starting_tool_calls = Arc::new(Mutex::new(HashSet::new())); + let forwarded_starts = Arc::new(Mutex::new(HashSet::new())); + let forwarded_terminals = Arc::new(Mutex::new(HashSet::new())); + let (command_events, _) = broadcast::channel(64); + let event_forwarder = forward_owned_command_events( + self.source.subscribe_command_events(), + owned_commands.clone(), + pending_command_events.clone(), + starting_tool_calls.clone(), + forwarded_starts.clone(), + forwarded_terminals.clone(), + command_events.clone(), + ) + .map(|handle| Arc::new(Mutex::new(Some(handle)))); + let close_lock = Arc::new(tokio::sync::Mutex::new(())); + let child = Arc::new(ScopedWorkdirSession { + source: self.source.clone(), + cwd: request.cwd, + scope: Some(request.rules), + capabilities, + validity: validity.clone(), + child_write_leases: Mutex::new(HashMap::new()), + next_lease_id: AtomicU64::new(1), + close_lock: close_lock.clone(), + owned_commands, + pending_command_events, + starting_tool_calls, + forwarded_starts, + forwarded_terminals, + command_events, + closes_source: false, + #[cfg(test)] + command_start_gate: Mutex::new(None), + }); + let broker = WorkdirToolBroker { + session: child.clone(), + authority: child, + event_forwarder, + }; + Ok(WorkdirScopeLease { + broker, + capabilities, + validity, + cleanup_pending, + close_lock, + }) + } +} + +#[async_trait] +impl WorkdirSession for ScopedWorkdirSession { + fn workdir(&self) -> &Workdir { + self.source.workdir() + } + + fn capabilities(&self) -> WorkdirSessionCapabilities { + self.capabilities + } + + async fn stat(&self, mut request: StatRequest) -> Result { + let path = self.resolve_operation_path(&request.path).await?; + self.ensure_read(&path, WorkdirSessionCapability::Read)?; + request.path = path; + self.source.stat(request).await + } + + async fn read(&self, mut request: ReadRequest) -> Result { + let path = self.resolve_operation_path(&request.path).await?; + self.ensure_read(&path, WorkdirSessionCapability::Read)?; + request.path = path; + self.source.read(request).await + } + + async fn write(&self, mut request: WriteRequest) -> Result { + let path = self.resolve_operation_path(&request.path).await?; + self.ensure_write(&path, WorkdirSessionCapability::Write)?; + request.path = path; + self.source.write(request).await + } + + async fn edit(&self, mut request: EditRequest) -> Result { + let path = self.resolve_operation_path(&request.path).await?; + self.ensure_write(&path, WorkdirSessionCapability::Edit)?; + request.path = path; + self.source.edit(request).await + } + + async fn list(&self, mut request: ListRequest) -> Result { + let path = self.resolve_operation_path(&request.path).await?; + self.ensure_read(&path, WorkdirSessionCapability::Read)?; + request.path = path; + self.source.list(request).await + } + + async fn glob(&self, mut request: GlobRequest) -> Result { + let path = self.resolve_operation_path(&request.path).await?; + self.ensure_read(&path, WorkdirSessionCapability::Glob)?; + request.path = path; + self.source.glob(request).await + } + + async fn grep(&self, mut request: GrepRequest) -> Result { + let path = self.resolve_operation_path(&request.path).await?; + self.ensure_read(&path, WorkdirSessionCapability::Grep)?; + request.path = path; + self.source.grep(request).await + } + + async fn start_command( + &self, + mut request: CommandRequest, + ) -> Result { + let _admission_guard = self.close_lock.lock().await; + // Command is an explicit capability, not a typed path mutation. We + // intentionally keep an ancestor's Command capability available while + // a child holds a write scope; only typed Write/Edit operations use the + // best-effort overlapping-path guard below. + self.ensure_command()?; + #[cfg(test)] + { + let gate = self + .command_start_gate + .lock() + .expect("command start gate mutex poisoned") + .clone(); + if let Some(gate) = gate { + gate.entered.notify_one(); + gate.release.notified().await; + } + } + if self.scope.is_some() + && self + .owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .len() + >= MAX_SCOPED_COMMANDS + { + return Err(WorkdirError::Unavailable(format!( + "scoped command limit of {MAX_SCOPED_COMMANDS} is reached" + ))); + } + let tool_call_id = request.tool_call_id.clone(); + if self.scope.is_some() { + request.cwd = Some(match request.cwd.as_ref() { + Some(cwd) => self.resolve_path(cwd)?, + None => self.cwd.clone(), + }); + } + if let Some(tool_call_id) = &tool_call_id { + self.starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .insert(tool_call_id.clone()); + } + let handle = match self.source.start_command(request).await { + Ok(handle) => handle, + Err(error) => { + if let Some(tool_call_id) = &tool_call_id { + self.starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .remove(tool_call_id); + } + return Err(error); + } + }; + let mut owned = self + .owned_commands + .lock() + .expect("scoped command set mutex poisoned"); + owned.insert(handle.0.clone()); + if let Some(tool_call_id) = &tool_call_id { + self.starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .remove(tool_call_id); + } + let pending = self + .pending_command_events + .lock() + .expect("pending scoped command event mutex poisoned") + .remove(&handle.0) + .unwrap_or_default(); + drop(owned); + if !pending + .iter() + .any(|event| matches!(event, CommandEvent::Started { .. })) + { + publish_owned_command_event( + &self.command_events, + &self.forwarded_starts, + &self.forwarded_terminals, + CommandEvent::Started { + command_id: handle.0.clone(), + tool_call_id, + observed_at_ms: unix_timestamp_ms(), + }, + ); + } + for event in pending { + publish_owned_command_event( + &self.command_events, + &self.forwarded_starts, + &self.forwarded_terminals, + event, + ); + } + Ok(handle) + } + + async fn command_status(&self, handle: CommandHandle) -> Result { + self.ensure_owned_command(&handle)?; + self.source.command_status(handle).await + } + + async fn command_output( + &self, + request: CommandOutputRequest, + ) -> Result { + self.ensure_owned_command(&request.handle)?; + let command_id = request.handle.0.clone(); + let output = self.source.command_output(request).await?; + if !matches!(output.status, CommandStatus::Running) { + self.publish_terminal_if_missing( + &command_id, + output.status, + output.exit_code, + output.next_cursor.unwrap_or(output.content.len()) as u64, + &output.content, + ); + self.owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .remove(&command_id); + } + Ok(output) + } + + async fn cancel_command(&self, handle: CommandHandle) -> Result<(), WorkdirError> { + self.ensure_owned_command(&handle)?; + self.source.cancel_command(handle).await + } + + fn subscribe_command_events(&self) -> Option> { + if !self + .capabilities + .supports(WorkdirSessionCapability::Command) + { + return None; + } + if self.scope.is_none() { + self.source.subscribe_command_events() + } else { + Some(self.command_events.subscribe()) + } + } + + fn command_snapshot(&self) -> Vec { + if !self + .capabilities + .supports(WorkdirSessionCapability::Command) + { + return Vec::new(); + } + if self.scope.is_none() { + return self.source.command_snapshot(); + } + let owned = self + .owned_commands + .lock() + .expect("scoped command set mutex poisoned"); + self.source + .command_snapshot() + .into_iter() + .filter(|snapshot| owned.contains(&snapshot.command_id)) + .collect() + } + + async fn close(&self) -> Result<(), WorkdirError> { + self.validity.active.store(false, Ordering::Release); + if self.closes_source { + self.source.close().await + } else { + Ok(()) + } + } +} + +/// A fail-closed read-only view over an already scoped scoped tool route. +#[derive(Debug)] +pub struct ReadOnlyWorkdirSession { + inner: WorkdirSessionHandle, +} + +impl ReadOnlyWorkdirSession { + pub fn new(inner: WorkdirSessionHandle) -> Self { + Self { inner } + } +} + +#[async_trait] +impl WorkdirSession for ReadOnlyWorkdirSession { + fn workdir(&self) -> &Workdir { + self.inner.workdir() + } + + fn capabilities(&self) -> WorkdirSessionCapabilities { + WorkdirSessionCapabilities::READ_ONLY + } + + async fn stat(&self, request: StatRequest) -> Result { + self.inner.stat(request).await + } + + async fn read(&self, request: ReadRequest) -> Result { + self.inner.read(request).await + } + + async fn write(&self, _request: WriteRequest) -> Result { + Err(WorkdirError::Denied("read-only workdir session".into())) + } + + async fn edit(&self, _request: EditRequest) -> Result { + Err(WorkdirError::Denied("read-only workdir session".into())) + } + + async fn list(&self, request: ListRequest) -> Result { + self.inner.list(request).await + } + + async fn glob(&self, request: GlobRequest) -> Result { + self.inner.glob(request).await + } + + async fn grep(&self, request: GrepRequest) -> Result { + self.inner.grep(request).await + } + + async fn start_command(&self, _request: CommandRequest) -> Result { + Err(WorkdirError::Denied("read-only workdir session".into())) + } + + async fn command_status(&self, _handle: CommandHandle) -> Result { + Err(WorkdirError::Denied("read-only workdir session".into())) + } + + async fn command_output( + &self, + _request: CommandOutputRequest, + ) -> Result { + Err(WorkdirError::Denied("read-only workdir session".into())) + } + + async fn cancel_command(&self, _handle: CommandHandle) -> Result<(), WorkdirError> { + Err(WorkdirError::Denied("read-only workdir session".into())) + } + + async fn close(&self) -> Result<(), WorkdirError> { + self.inner.close().await + } +} + +fn forward_owned_command_events( + receiver: Option>, + owned_commands: Arc>>, + pending_command_events: Arc>>>, + starting_tool_calls: Arc>>, + forwarded_starts: Arc>>, + forwarded_terminals: Arc>>, + sender: broadcast::Sender, +) -> Option> { + let mut receiver = receiver?; + Some(tokio::spawn(async move { + loop { + let event = match receiver.recv().await { + Ok(event) => event, + Err(broadcast::error::RecvError::Lagged(_)) => continue, + Err(broadcast::error::RecvError::Closed) => break, + }; + let command_id = command_event_id(&event).to_string(); + let mut owned = owned_commands + .lock() + .expect("scoped command set mutex poisoned"); + if !owned.contains(&command_id) { + let claimed = matches!( + &event, + CommandEvent::Started { + tool_call_id: Some(tool_call_id), + .. + } if starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .contains(tool_call_id) + ); + if !claimed { + continue; + } + owned.insert(command_id.clone()); + pending_command_events + .lock() + .expect("pending scoped command event mutex poisoned") + .entry(command_id) + .or_default() + .push(event); + continue; + } + let mut pending = pending_command_events + .lock() + .expect("pending scoped command event mutex poisoned"); + if let Some(events) = pending.get_mut(&command_id) { + if events.len() < 64 { + events.push(event); + } + continue; + } + drop(pending); + drop(owned); + publish_owned_command_event(&sender, &forwarded_starts, &forwarded_terminals, event); + } + })) +} + +fn command_event_id(event: &CommandEvent) -> &str { + match event { + CommandEvent::Started { command_id, .. } + | CommandEvent::Output { command_id, .. } + | CommandEvent::Terminal { command_id, .. } => command_id, + } +} + +fn publish_owned_command_event( + sender: &broadcast::Sender, + forwarded_starts: &Mutex>, + forwarded_terminals: &Mutex>, + event: CommandEvent, +) { + let command_id = command_event_id(&event); + let mut terminals = forwarded_terminals + .lock() + .expect("forwarded terminal command mutex poisoned"); + match &event { + CommandEvent::Terminal { .. } if !terminals.insert(command_id.to_string()) => return, + CommandEvent::Started { .. } if terminals.contains(command_id) => return, + CommandEvent::Started { .. } + if !forwarded_starts + .lock() + .expect("forwarded command start mutex poisoned") + .insert(command_id.to_string()) => + { + return; + } + CommandEvent::Output { .. } if terminals.contains(command_id) => return, + _ => {} + } + drop(terminals); + let _ = sender.send(event); +} + +fn unix_timestamp_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .min(u128::from(u64::MAX)) as u64 +} + +fn rules_overlap(left: &WorkdirToolScopeRule, right: &WorkdirToolScopeRule) -> bool { + left.permission == WorkdirToolScopePermission::Write + && right.permission == WorkdirToolScopePermission::Write + && (rule_allows_path(left, &right.target, WorkdirToolScopePermission::Write) + || rule_allows_path(right, &left.target, WorkdirToolScopePermission::Write)) +} + +fn rule_allows_path( + rule: &WorkdirToolScopeRule, + path: &FsPath, + required: WorkdirToolScopePermission, +) -> bool { + if required == WorkdirToolScopePermission::Write + && rule.permission != WorkdirToolScopePermission::Write + { + return false; + } + path_in_rule(rule, path) +} + +fn path_in_rule(rule: &WorkdirToolScopeRule, path: &FsPath) -> bool { + let target = Path::new(rule.target.as_str()); + let path = Path::new(path.as_str()); + if path == target { + return true; + } + let Ok(suffix) = path.strip_prefix(target) else { + return false; + }; + let depth = suffix.components().count(); + rule.recursive || depth <= 1 +} + +fn rule_contains_rule(parent: &WorkdirToolScopeRule, child: &WorkdirToolScopeRule) -> bool { + if child.permission == WorkdirToolScopePermission::Write + && parent.permission != WorkdirToolScopePermission::Write + { + return false; + } + if !path_in_rule(parent, &child.target) { + return false; + } + if parent.recursive { + return true; + } + !child.recursive && parent.target == child.target +} + +#[cfg(test)] +mod tests { + use std::fs; + + use manifest::{Permission, Scope, ScopeConfig, ScopeRule, SharedScope}; + use tempfile::TempDir; + + use super::*; + use crate::LocalWorkdirSession; + + fn fs_path(path: &str) -> FsPath { + FsPath::new(path).unwrap() + } + + fn session(root: &Path) -> WorkdirToolBroker { + let scope = SharedScope::new( + Scope::from_config(&ScopeConfig { + allow: vec![ScopeRule { + target: root.to_path_buf(), + permission: Permission::Write, + recursive: true, + }], + deny: Vec::new(), + }) + .unwrap(), + ); + WorkdirToolBroker::new(Arc::new(LocalWorkdirSession::materialized_bound( + Workdir::new("delegation-test"), + root.to_path_buf(), + root.to_path_buf(), + scope, + WorkdirSessionCapabilities::ALL, + ))) + } + + fn request(path: &str, permission: WorkdirToolScopePermission) -> WorkdirToolScope { + WorkdirToolScope { + rules: vec![WorkdirToolScopeRule { + target: fs_path(path), + permission, + recursive: true, + }], + cwd: fs_path(path), + command: permission == WorkdirToolScopePermission::Write, + } + } + + fn read(path: &str) -> ReadRequest { + ReadRequest { + path: fs_path(path), + offset: 0, + limit: 20, + max_bytes: 1024, + } + } + + fn write(path: &str, content: &str) -> WriteRequest { + WriteRequest { + path: fs_path(path), + content: content.as_bytes().to_vec(), + expected_hash: None, + } + } + + async fn run_command( + session: &WorkdirSessionHandle, + command: impl Into, + tool_call_id: impl Into, + ) -> CommandOutput { + let handle = session + .start_command(CommandRequest { + command: command.into(), + timeout_secs: 5, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some(tool_call_id.into()), + }) + .await + .unwrap(); + session + .command_output(CommandOutputRequest { + handle, + cursor: 0, + limit: 1024, + wait: true, + }) + .await + .unwrap() + } + + #[tokio::test] + async fn workdir_tool_broker_session_forwards_command_telemetry() { + let root = TempDir::new().unwrap(); + let parent = session(root.path()); + let mut events = parent + .subscribe_command_events() + .expect("delegation wrapper must preserve command observation"); + let handle = parent + .start_command(CommandRequest { + command: "printf ready; sleep 0.2; printf done".into(), + timeout_secs: 5, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some("tool-delegated".into()), + }) + .await + .unwrap(); + + let first_output = loop { + let event = tokio::time::timeout(std::time::Duration::from_secs(2), events.recv()) + .await + .expect("delegated command telemetry should not stall") + .unwrap(); + if let CommandEvent::Output { content, .. } = event { + break content; + } + }; + assert_eq!(first_output, "ready"); + let snapshots = parent.command_snapshot(); + assert_eq!(snapshots.len(), 1); + assert_eq!(snapshots[0].command_id, handle.0); + assert_eq!(snapshots[0].status, CommandStatus::Running); + assert_eq!(snapshots[0].stdout.content, "ready"); + + let output = parent + .command_output(CommandOutputRequest { + handle, + cursor: 0, + limit: 1024, + wait: true, + }) + .await + .unwrap(); + assert_eq!(output.status, CommandStatus::Completed); + assert_eq!(output.content, "readydone"); + assert!(parent.command_snapshot().is_empty()); + } + + #[tokio::test] + async fn write_scope_without_command_grant_has_no_command_capability() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("work")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(WorkdirToolScope { + rules: vec![WorkdirToolScopeRule { + target: fs_path("work"), + permission: WorkdirToolScopePermission::Write, + recursive: true, + }], + cwd: fs_path("work"), + command: false, + }) + .await + .unwrap(); + + assert!(child.capabilities.supports(WorkdirSessionCapability::Write)); + assert!( + !child + .capabilities + .supports(WorkdirSessionCapability::Command) + ); + let error = child + .start_command(CommandRequest { + command: "pwd".into(), + timeout_secs: 5, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: None, + }) + .await + .unwrap_err(); + assert!(matches!(error, WorkdirError::Denied(_))); + } + + #[tokio::test] + async fn scoped_commands_use_child_cwd_and_do_not_leak_between_siblings() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("one")).unwrap(); + fs::create_dir_all(root.path().join("two")).unwrap(); + let parent = session(root.path()); + let first = parent + .scope(request("one", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + let second = parent + .scope(request("two", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + let mut first_events = first.subscribe_command_events().unwrap(); + let mut second_events = second.subscribe_command_events().unwrap(); + + let handle = first + .start_command(CommandRequest { + command: "pwd; sleep 0.2".into(), + timeout_secs: 5, + output_limit: 4096, + cwd: None, + spill_dir: None, + tool_call_id: Some("first-command".into()), + }) + .await + .unwrap(); + assert!(matches!( + first_events.recv().await.unwrap(), + CommandEvent::Started { .. } + )); + assert!(matches!( + tokio::time::timeout(std::time::Duration::from_millis(50), second_events.recv()).await, + Err(_) + )); + assert!(matches!( + second.command_status(handle.clone()).await, + Err(WorkdirError::UnknownCommand(_)) + )); + + let output = first + .command_output(CommandOutputRequest { + handle, + cursor: 0, + limit: 4096, + wait: true, + }) + .await + .unwrap(); + let expected = root.path().join("one").to_string_lossy().into_owned(); + assert!( + output + .content + .lines() + .next() + .is_some_and(|line| line == expected) + ); + } + + #[test] + fn non_recursive_rule_covers_target_and_direct_children_only() { + let rule = WorkdirToolScopeRule { + target: fs_path("docs"), + permission: WorkdirToolScopePermission::Read, + recursive: false, + }; + assert!(path_in_rule(&rule, &fs_path("docs"))); + assert!(path_in_rule(&rule, &fs_path("docs/readme.md"))); + assert!(!path_in_rule(&rule, &fs_path("docs/guides/start.md"))); + } + + #[tokio::test] + async fn read_only_delegation_allows_prefix_and_denies_mutation() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("docs")).unwrap(); + fs::create_dir_all(root.path().join("secret")).unwrap(); + fs::write(root.path().join("docs/readme.md"), "visible").unwrap(); + fs::write(root.path().join("secret/key"), "hidden").unwrap(); + let parent = session(root.path()); + + let child = parent + .scope(request("docs", WorkdirToolScopePermission::Read)) + .await + .unwrap(); + assert_eq!(child.capabilities, WorkdirSessionCapabilities::READ_ONLY); + assert_eq!( + child.read(read("readme.md")).await.unwrap().bytes, + b"visible" + ); + assert!(matches!( + child.write(write("new.md", "no")).await, + Err(WorkdirError::Denied(_)) + )); + assert!( + !child + .capabilities + .supports(WorkdirSessionCapability::Command) + ); + assert!(child.subscribe_command_events().is_none()); + assert!(child.command_snapshot().is_empty()); + assert!(matches!( + child + .start_command(CommandRequest { + command: "printf denied".into(), + timeout_secs: 5, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some("read-only-command".into()), + }) + .await, + Err(WorkdirError::Denied(_)) + )); + } + + #[cfg(unix)] + #[tokio::test] + async fn provider_scope_denies_read_through_symlink_outside_grant() { + use std::os::unix::fs::symlink; + + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("granted")).unwrap(); + fs::create_dir_all(root.path().join("secret")).unwrap(); + fs::write(root.path().join("secret/key"), "hidden").unwrap(); + symlink("../secret/key", root.path().join("granted/link")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(request("granted", WorkdirToolScopePermission::Read)) + .await + .unwrap(); + + let result = child.read(read("link")).await; + assert!( + result.is_err(), + "symlink read escaped provider scope: {result:?}" + ); + } + + #[cfg(unix)] + #[tokio::test] + async fn provider_scope_denies_write_through_symlink_outside_grant() { + use std::os::unix::fs::symlink; + + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("granted")).unwrap(); + fs::create_dir_all(root.path().join("secret")).unwrap(); + symlink("../secret", root.path().join("granted/outside")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(request("granted", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + + let result = child.write(write("outside/new", "forbidden")).await; + assert!( + result.is_err(), + "symlink write escaped provider scope: {result:?}" + ); + assert!(!root.path().join("secret/new").exists()); + } + + #[cfg(unix)] + #[tokio::test] + async fn write_delegation_rejects_symlink_target_before_lease() { + use std::os::unix::fs::symlink; + + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("granted")).unwrap(); + fs::create_dir_all(root.path().join("secret")).unwrap(); + symlink("../secret", root.path().join("granted/outside")).unwrap(); + let parent = session(root.path()); + + assert!(matches!( + parent + .scope(request( + "granted/outside", + WorkdirToolScopePermission::Write + )) + .await, + Err(WorkdirError::Denied(_)) + )); + parent + .write(write("secret/parent", "still-authoritative")) + .await + .unwrap(); + } + + #[tokio::test] + async fn write_lease_keeps_typed_parent_writes_exclusive_without_blocking_commands() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("leased")).unwrap(); + fs::create_dir_all(root.path().join("other")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(request("leased", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + assert!( + child + .capabilities + .supports(WorkdirSessionCapability::Command) + ); + let child_output = + run_command(&child, "printf child-command", "delegated-child-command").await; + assert_eq!(child_output.content, "child-command"); + let parent_output = run_command( + &parent, + "printf parent-write > leased/from-command; printf parent-command", + "parent-command-during-child-write", + ) + .await; + assert_eq!(parent_output.status, CommandStatus::Completed); + assert_eq!(parent_output.content, "parent-command"); + assert_eq!( + fs::read_to_string(root.path().join("leased/from-command")).unwrap(), + "parent-write" + ); + + assert!(matches!( + parent.write(write("leased/file", "parent")).await, + Err(WorkdirError::Denied(_)) + )); + parent.write(write("other/file", "parent")).await.unwrap(); + child.write(write("file", "child")).await.unwrap(); + child.close().await.unwrap(); + assert!(matches!( + child + .start_command(CommandRequest { + command: "printf revoked".into(), + timeout_secs: 5, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some("revoked-child-command".into()), + }) + .await, + Err(WorkdirError::SessionClosed) + )); + parent + .write(write("leased/parent", "parent")) + .await + .unwrap(); + assert!(matches!( + child.read(read("file")).await, + Err(WorkdirError::SessionClosed) + )); + } + + #[tokio::test] + async fn sibling_write_scopes_must_not_overlap() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("shared/one")).unwrap(); + fs::create_dir_all(root.path().join("other")).unwrap(); + let parent = session(root.path()); + let first = parent + .scope(request("shared", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + + assert!(matches!( + parent + .scope(request("shared/one", WorkdirToolScopePermission::Write)) + .await, + Err(WorkdirError::Denied(_)) + )); + let other = parent + .scope(request("other", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + other.close().await.unwrap(); + first.close().await.unwrap(); + } + + #[tokio::test] + async fn fast_command_keeps_started_output_terminal_event_order() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("work")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(request("work", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + let mut events = child.subscribe_command_events().unwrap(); + + let output = run_command(&child.tool_session(), "printf fast-output", "fast-command").await; + assert_eq!(output.content, "fast-output"); + + let mut kinds = Vec::new(); + let mut streamed = String::new(); + while kinds.last().is_none_or(|kind| *kind != "terminal") { + let event = tokio::time::timeout(std::time::Duration::from_secs(1), events.recv()) + .await + .expect("fast command event timeout") + .expect("fast command event channel"); + match event { + CommandEvent::Started { .. } => kinds.push("started"), + CommandEvent::Output { content, .. } => { + kinds.push("output"); + streamed.push_str(&content); + } + CommandEvent::Terminal { .. } => kinds.push("terminal"), + } + } + assert_eq!(kinds.first(), Some(&"started")); + assert_eq!(kinds.last(), Some(&"terminal")); + assert_eq!(kinds.iter().filter(|kind| **kind == "started").count(), 1); + assert!(kinds.contains(&"output")); + assert!(streamed.contains("fast-output")); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), events.recv()) + .await + .is_err(), + "no provider event may follow the terminal event" + ); + assert!(child.command_snapshot().is_empty()); + child.close().await.unwrap(); + } + + #[tokio::test] + async fn scoped_command_ceiling_rejects_the_seventeenth_live_command() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("work")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(request("work", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + for index in 0..MAX_SCOPED_COMMANDS { + child + .start_command(CommandRequest { + command: "sleep 30".into(), + timeout_secs: 60, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some(format!("command-{index}")), + }) + .await + .unwrap(); + } + + let error = child + .start_command(CommandRequest { + command: "sleep 30".into(), + timeout_secs: 60, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some("command-over-limit".into()), + }) + .await + .unwrap_err(); + assert!(matches!(error, WorkdirError::Unavailable(message) if message.contains("limit"))); + child.close().await.unwrap(); + } + + #[tokio::test] + async fn close_serializes_with_inflight_command_admission() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("work")).unwrap(); + let parent = session(root.path()); + let child = Arc::new( + parent + .scope(request("work", WorkdirToolScopePermission::Write)) + .await + .unwrap(), + ); + let gate = Arc::new(TestCommandStartGate { + entered: tokio::sync::Notify::new(), + release: tokio::sync::Notify::new(), + }); + *child.broker.authority.command_start_gate.lock().unwrap() = Some(gate.clone()); + let entered = gate.entered.notified(); + let command_child = child.clone(); + let command = tokio::spawn(async move { + command_child + .start_command(CommandRequest { + command: "sleep 30".into(), + timeout_secs: 60, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some("racing-command".into()), + }) + .await + }); + entered.await; + let close_child = child.clone(); + let mut close = tokio::spawn(async move { close_child.close().await }); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(50), &mut close) + .await + .is_err(), + "close must wait for command admission to commit or fail" + ); + + gate.release.notify_one(); + command.await.unwrap().unwrap(); + close.await.unwrap().unwrap(); + assert!(!child.is_active()); + assert!( + child + .broker + .authority + .owned_commands + .lock() + .unwrap() + .is_empty() + ); + } + + #[tokio::test] + async fn closing_scope_cancels_and_terminalizes_owned_commands() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("work")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(request("work", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + let mut events = child.subscribe_command_events().unwrap(); + let handle = child + .start_command(CommandRequest { + command: "sleep 30; printf leaked > marker".into(), + timeout_secs: 60, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some("owned-command".into()), + }) + .await + .unwrap(); + assert!(matches!( + events.recv().await.unwrap(), + CommandEvent::Started { .. } + )); + + child.close().await.unwrap(); + + assert!(matches!( + parent.command_status(handle).await, + Ok(CommandStatus::Cancelled | CommandStatus::Completed | CommandStatus::Failed) + | Err(WorkdirError::UnknownCommand(_)) + )); + assert!(!root.path().join("work/marker").exists()); + let terminal = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if let CommandEvent::Terminal { .. } = events.recv().await.unwrap() { + break; + } + } + }) + .await; + assert!( + terminal.is_ok(), + "scope close must publish terminal command telemetry" + ); + } + + #[tokio::test] + async fn nested_delegation_is_attenuated_and_parent_revocation_cascades() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("docs/sub")).unwrap(); + fs::create_dir_all(root.path().join("docs/peer")).unwrap(); + fs::write(root.path().join("docs/sub/a"), "a").unwrap(); + fs::write(root.path().join("docs/peer/b"), "b").unwrap(); + let root_session = session(root.path()); + let child = root_session + .scope(request("docs", WorkdirToolScopePermission::Read)) + .await + .unwrap(); + let nested = child + .scope(request("docs/sub", WorkdirToolScopePermission::Read)) + .await + .unwrap(); + + nested.read(read("a")).await.unwrap(); + assert!( + child + .scope(request("other", WorkdirToolScopePermission::Read)) + .await + .is_err() + ); + assert!( + child + .scope(request("docs/sub", WorkdirToolScopePermission::Write)) + .await + .is_err() + ); + + child.close().await.unwrap(); + assert!(matches!( + nested.read(read("a")).await, + Err(WorkdirError::SessionClosed) + )); + } + + #[tokio::test] + async fn nested_write_leases_do_not_block_command_capable_ancestors() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("docs/sub")).unwrap(); + let root_session = session(root.path()); + let child = root_session + .scope(request("docs", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + let nested = child + .scope(request("docs/sub", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + + for (session, label) in [ + (root_session.tool_session(), "root"), + (child.tool_session(), "child"), + (nested.tool_session(), "nested"), + ] { + let output = run_command( + &session, + format!("printf {label}"), + format!("{label}-command-during-nested-write"), + ) + .await; + assert_eq!(output.status, CommandStatus::Completed); + assert_eq!(output.content, label); + } + + assert!(matches!( + root_session.write(write("docs/root", "blocked")).await, + Err(WorkdirError::Denied(_)) + )); + assert!(matches!( + child.write(write("sub/child", "blocked")).await, + Err(WorkdirError::Denied(_)) + )); + nested.write(write("nested", "allowed")).await.unwrap(); + + nested.close().await.unwrap(); + child.close().await.unwrap(); + } + + #[tokio::test] + async fn closing_parent_invalidates_scoped_tools() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("docs")).unwrap(); + fs::write(root.path().join("docs/a"), "a").unwrap(); + let parent = session(root.path()); + let child = parent + .scope(request("docs", WorkdirToolScopePermission::Read)) + .await + .unwrap(); + + parent.close().await.unwrap(); + assert!(matches!( + parent + .start_command(CommandRequest { + command: "printf closed".into(), + timeout_secs: 5, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some("closed-parent-command".into()), + }) + .await, + Err(WorkdirError::SessionClosed) + )); + let child_result = child.read(read("a")).await; + assert!( + matches!(child_result, Err(WorkdirError::SessionClosed)), + "child result after parent close: {child_result:?}" + ); + } +} diff --git a/crates/workdir/src/workspace.rs b/crates/workdir/src/workspace.rs index 3872bc9d..fc696f99 100644 --- a/crates/workdir/src/workspace.rs +++ b/crates/workdir/src/workspace.rs @@ -104,15 +104,5 @@ mod tests { #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct WorkspaceWorkdirSessionOperationRequest { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub expected_session_fence: Option, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub delegations: Vec, pub operation: crate::http::WorkdirSessionOperation, } - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WorkspaceWorkdirSessionFence { - pub value: String, -} diff --git a/crates/worker-runtime/src/execution.rs b/crates/worker-runtime/src/execution.rs index 47660977..c85e9f29 100644 --- a/crates/worker-runtime/src/execution.rs +++ b/crates/worker-runtime/src/execution.rs @@ -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, #[serde(default, skip_serializing_if = "Option::is_none")] - pub input_commit: Option, + pub submission: Option, } /// 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, submission_id: impl Into, + 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::(&json).unwrap(), diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index d1e93598..84eef30b 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -774,8 +774,7 @@ async fn run_workdir_session_operation( .ok_or_else(RuntimeHttpWorkdirError::not_found)?; record.session.clone() }; - let applied = workdir::apply_delegation_chain(source, request.delegations).await?; - let session = applied.scoped_session.as_ref(); + let session = source.as_ref(); let operation = request.operation; let result = match operation { @@ -1240,10 +1239,12 @@ async fn worker_protocol_ws( auth: Option>, Path(worker_id): Path, Query(query): Query, + headers: HeaderMap, ws: WebSocketUpgrade, ) -> Result { 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 +1255,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, 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, worker_ref: WorkerRef, query: RuntimeWorkerEventsWsQuery, + input_source: Option, mut socket: WebSocket, ) { let mut cursor = match query.cursor.as_deref() { @@ -1347,6 +1420,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) @@ -2139,8 +2214,8 @@ mod tests { use manifest::{Scope, SharedScope}; use tower::ServiceExt; use workdir::{ - GrepOutputMode, GrepRequest, LocalWorkdirSession, ReadRequest, StatRequest, Workdir, - WorkdirPath, WorkdirSessionCapabilities, + GrepOutputMode, GrepRequest, LocalWorkdirSession, StatRequest, Workdir, WorkdirPath, + WorkdirSessionCapabilities, }; #[tokio::test] @@ -2219,6 +2294,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!( @@ -2637,16 +2769,6 @@ mod tests { async fn workdir_session_operations_enforce_owner_and_close_terminally() { let temp = tempfile::tempdir().expect("tempdir"); std::fs::write(temp.path().join("hello.txt"), "hello").expect("write fixture"); - #[cfg(unix)] - { - use std::os::unix::fs::symlink; - std::fs::create_dir(temp.path().join("granted")).expect("granted directory"); - std::fs::write(temp.path().join("granted/visible"), "visible") - .expect("visible fixture"); - std::fs::create_dir(temp.path().join("secret")).expect("secret directory"); - std::fs::write(temp.path().join("secret/key"), "hidden").expect("secret fixture"); - symlink("../secret/key", temp.path().join("granted/link")).expect("symlink fixture"); - } let scope = SharedScope::new(Scope::writable(temp.path()).expect("scope")); let session: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::materialized_bound( Workdir::new("wd-1"), @@ -2680,7 +2802,6 @@ mod tests { expires_at: u64::MAX, }; let operation = WorkdirSessionOperationRequest { - delegations: Vec::new(), operation: WorkdirSessionOperation::Stat(StatRequest { path: WorkdirPath::new("hello.txt").expect("logical path"), }), @@ -2697,7 +2818,6 @@ mod tests { assert!(matches!(result, WorkdirSessionOperationResult::Stat(_))); let grep = WorkdirSessionOperationRequest { - delegations: Vec::new(), operation: WorkdirSessionOperation::Grep(GrepRequest { pattern: "hello".into(), path: WorkdirPath::new("hello.txt").unwrap(), @@ -2720,78 +2840,7 @@ mod tests { ) .await .expect("grep direct file through provider operation"); - match result { - WorkdirSessionOperationResult::Grep(result) => { - assert_eq!(result.match_count, 1); - assert_eq!(result.matched_files, 1); - assert!(result.output.starts_with("hello.txt\n")); - assert!(result.output.contains("> 1 │ hello")); - } - other => panic!("unexpected workdir grep result: {other:?}"), - } - - #[cfg(unix)] - { - let delegated_visible = WorkdirSessionOperationRequest { - delegations: vec![workdir::WorkdirDelegationRequest { - rules: vec![workdir::WorkdirDelegationRule { - target: WorkdirPath::new("granted").unwrap(), - permission: workdir::WorkdirDelegationPermission::Read, - recursive: true, - }], - cwd: WorkdirPath::new("granted").unwrap(), - }], - operation: WorkdirSessionOperation::Read(ReadRequest { - path: WorkdirPath::new("visible").unwrap(), - offset: 0, - limit: 20, - max_bytes: 1024, - }), - }; - let visible = run_workdir_session_operation( - State(state.clone()), - Path("session-1".to_string()), - Some(Extension(auth.clone())), - Ok(Json(delegated_visible)), - ) - .await - .expect("non-root delegated cwd should resolve once") - .0; - assert!(matches!( - visible, - WorkdirSessionOperationResult::Read(result) if result.bytes == b"visible" - )); - - let delegated_read = WorkdirSessionOperationRequest { - delegations: vec![workdir::WorkdirDelegationRequest { - rules: vec![workdir::WorkdirDelegationRule { - target: WorkdirPath::new("granted").unwrap(), - permission: workdir::WorkdirDelegationPermission::Read, - recursive: true, - }], - cwd: WorkdirPath::new("granted").unwrap(), - }], - operation: WorkdirSessionOperation::Read(ReadRequest { - path: WorkdirPath::new("link").unwrap(), - offset: 0, - limit: 20, - max_bytes: 1024, - }), - }; - let error = run_workdir_session_operation( - State(state.clone()), - Path("session-1".to_string()), - Some(Extension(auth.clone())), - Ok(Json(delegated_read)), - ) - .await - .expect_err("provider must reject delegated symlink escape"); - assert_ne!(error.status, StatusCode::OK); - assert_eq!( - std::fs::read_to_string(temp.path().join("secret/key")).unwrap(), - "hidden" - ); - } + assert!(matches!(result, WorkdirSessionOperationResult::Grep(_))); let wrong_owner = RuntimeAuthContext { workspace_id: "workspace-b".to_string(), @@ -2870,11 +2919,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 +3245,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( diff --git a/crates/worker-runtime/src/interaction.rs b/crates/worker-runtime/src/interaction.rs index d76c675d..52e0168e 100644 --- a/crates/worker-runtime/src/interaction.rs +++ b/crates/worker-runtime/src/interaction.rs @@ -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, + /// 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, #[serde(default, skip_serializing_if = "Option::is_none")] pub segments: Option>, } @@ -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, } diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index fd5a0b16..9701aae1 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -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 { 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 { 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"); diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 95d0c941..5225dd8b 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -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>>, @@ -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 Drop for WorkerRuntimeExecutionBackend { 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"); diff --git a/crates/worker/Cargo.toml b/crates/worker/Cargo.toml index 87437d17..87046834 100644 --- a/crates/worker/Cargo.toml +++ b/crates/worker/Cargo.toml @@ -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"] } diff --git a/crates/worker/examples/worker_protocol.rs b/crates/worker/examples/worker_protocol.rs index 1e2c618f..cd00370a 100644 --- a/crates/worker/examples/worker_protocol.rs +++ b/crates/worker/examples/worker_protocol.rs @@ -101,7 +101,10 @@ async fn main() -> Result<(), Box> { // 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 diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 70584074..10ff2f67 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -5,7 +5,7 @@ use std::sync::atomic::Ordering; use agen::EngineError; use agen::llm_client::client::LlmClient; use session_store::WorkerMetadataStore; -use session_store::{LogEntry, SessionExtension, Store}; +use session_store::{LogEntry, Store}; use tokio::sync::{broadcast, mpsc, oneshot}; use crate::discovery::WorkerDiscovery; @@ -23,16 +23,12 @@ use crate::shutdown_after_idle::{ }; use crate::spawn::registry::SpawnedWorkerRegistry; use crate::spawn::tool::sub_worker_spawn_tool; -use crate::worker::{ - SystemItemCommitter, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError, - WorkerRunResult, -}; +use crate::worker::{SystemItemCommitter, Worker, WorkerError, WorkerRunResult}; use protocol::{ AlertLevel, AlertSource, CommandEvent as ProtocolCommandEvent, CommandSnapshot as ProtocolCommandSnapshot, CommandStatus as ProtocolCommandStatus, CommandStream as ProtocolCommandStream, CommandStreamSlice as ProtocolCommandStreamSlice, - ErrorCode, Event, Method, RewindTargetId, RunResult, Segment, TurnResult, UploadedFileRef, - WorkerStatus, + ErrorCode, Event, Method, RewindTargetId, RunResult, TurnResult, UploadedFileRef, WorkerStatus, }; use workdir::{ CommandEvent as WorkdirCommandEvent, CommandSnapshot as WorkdirCommandSnapshot, @@ -58,6 +54,7 @@ pub struct WorkerHandle { spawned_registry: Arc, artifact_store: Arc, session_id: session_store::SessionId, + pending_activations: Arc>, } impl WorkerHandle { @@ -131,8 +128,15 @@ impl WorkerHandle { let in_flight = snapshot_from_guard(&in_flight_guard); (entries, entry_rx, in_flight) }; + let mut session = + session_store::public_snapshot::project_current_session_snapshot(&entries); + session.pending_submissions = self + .pending_activations + .lock() + .expect("pending activation state poisoned") + .snapshot(); let event = Event::Snapshot { - session: session_store::public_snapshot::project_current_session_snapshot(&entries), + session, greeting: self.shared_state.greeting.clone(), status: self.shared_state.get_status(), in_flight, @@ -213,21 +217,143 @@ async fn finish_controller_run( /// `Worker::*` entry point — `RunForNotification` carries none because /// `worker.run_for_notification()` drains the NotifyBuffer on its own. enum PendingRun { - Run(Vec), - RunTracked { - input: Vec, - extension: SessionExtension, - }, + Submit(crate::worker::PendingSubmission), /// Self-initiated turn kicked from the notify buffer. The carried /// `InvokeKind` is the trigger that flipped the Worker from IDLE /// (Notify or WorkerEvent) and is recorded by the Invoke marker /// committed at the start of `worker.run_for_notification`. - RunForNotification(protocol::InvokeKind), + RunForNotification { + invoke_kind: protocol::InvokeKind, + notification_request_id: Option, + }, Resume, } +fn resolved_input_source( + pending_submissions: &crate::worker::PendingSubmissionHandle, + source: &protocol::AuthenticatedInputSource, +) -> (String, session_store::LoggedSessionHistoryOrigin) { + if matches!(source, protocol::AuthenticatedInputSource::UntrustedWire) { + return ( + pending_submissions.direct_client_namespace(), + session_store::LoggedSessionHistoryOrigin::LegacyUnknown, + ); + } + ( + source.namespace(), + crate::worker::authenticated_input_provenance(source), + ) +} + +fn durable_parent_notification_target( + pending_submissions: crate::worker::PendingSubmissionHandle, + notify_buffer: NotifyBuffer, +) -> crate::spawn::tool::ParentNotificationTarget { + crate::spawn::tool::ParentNotificationTarget::Durable(Arc::new(move |method| { + let Method::NotifyTracked { + notification_request_id, + message, + auto_run, + source, + } = method + else { + return; + }; + let (source_namespace, provenance) = resolved_input_source(&pending_submissions, &source); + match pending_submissions.accept_notification_from_source( + notification_request_id.clone(), + message, + source_namespace.clone(), + provenance, + auto_run, + ) { + Ok(_) if !auto_run => { + stage_pending_notification( + &pending_submissions, + ¬ify_buffer, + &source_namespace, + ¬ification_request_id, + ); + } + Ok(_) => {} + Err(error) => { + tracing::warn!(%error, "failed to durably accept SubWorker notification"); + } + } + })) +} + +fn stage_pending_notification( + pending_submissions: &crate::worker::PendingSubmissionHandle, + notify_buffer: &NotifyBuffer, + source_namespace: &str, + notification_request_id: &str, +) -> bool { + let Some(notification) = + pending_submissions.prepare_notification(source_namespace, notification_request_id) + else { + return false; + }; + let extension = pending_submissions.notification_activation_extension(); + notify_buffer.push_durable_notify( + notification.message, + notification.auto_run, + notification.provenance, + extension, + ); + true +} + +fn stage_oldest_passive_notification( + pending_submissions: &crate::worker::PendingSubmissionHandle, + notify_buffer: &NotifyBuffer, +) -> bool { + pending_submissions + .next_passive_notification_identity() + .is_some_and(|(source_namespace, request_id)| { + stage_pending_notification( + pending_submissions, + notify_buffer, + &source_namespace, + &request_id, + ) + }) +} + +fn prepare_pending_run( + pending_submissions: &crate::worker::PendingSubmissionHandle, + notify_buffer: &NotifyBuffer, + fence: Option<(u64, &str)>, +) -> Result, crate::worker::PendingSubmissionError> { + let staged_passive_notification = pending_submissions.activating_passive_notification_id(); + Ok(match pending_submissions.prepare_next_activation(fence)? { + Some(crate::worker::PendingActivation::Submission(submission)) => { + if staged_passive_notification.is_some() { + let extension = pending_submissions.notification_activation_extension(); + debug_assert!(notify_buffer.replace_durable_notification_extension(extension)); + } + Some(PendingRun::Submit(submission)) + } + Some(crate::worker::PendingActivation::Notification(notification)) => { + let extension = pending_submissions.notification_activation_extension(); + let notification_request_id = notification.notification_request_id.clone(); + notify_buffer.push_durable_notify( + notification.message, + notification.auto_run, + notification.provenance, + extension, + ); + Some(PendingRun::RunForNotification { + invoke_kind: protocol::InvokeKind::Notify, + notification_request_id: Some(notification_request_id), + }) + } + None => None, + }) +} + impl PendingRun { - /// Whether this turn was kicked off by the parent (via `Method::Run` + /// Whether this turn was kicked off by the parent (via `Method::Submit` /// or `Method::Resume`). Used by [`drive_turn`] to gate upward /// `WorkerEvent::TurnEnded` / `WorkerEvent::Errored` reports so the parent /// only sees completion signals for work it actually delegated. @@ -235,16 +361,12 @@ impl PendingRun { /// notify buffer (Notify / inbound WorkerEvent) and stays silent. fn is_parent_originated(&self) -> bool { match self { - PendingRun::Run(_) | PendingRun::RunTracked { .. } | PendingRun::Resume => true, - PendingRun::RunForNotification(_) => false, + PendingRun::Submit(_) | PendingRun::Resume => true, + PendingRun::RunForNotification { .. } => false, } } } -fn should_auto_run_notification(status: WorkerStatus, auto_run: bool) -> bool { - auto_run && status == WorkerStatus::Idle -} - // --------------------------------------------------------------------------- // WorkerController — actor that owns a Worker // --------------------------------------------------------------------------- @@ -514,6 +636,7 @@ impl WorkerController { runtime_base.to_path_buf(), spawned_registry.clone(), Some(method_tx.downgrade()), + None, ) .await?; if let Some(session) = fs_for_view.as_ref() { @@ -552,6 +675,7 @@ impl WorkerController { let artifact_store: Arc = Arc::new(worker.store().clone()); let session_id = worker.session_id(); + let pending_activations = worker.pending_activation_state(); let handle = WorkerHandle { method_tx, working_event_tx: working_event_tx.clone(), @@ -563,6 +687,7 @@ impl WorkerController { spawned_registry: spawned_registry.clone(), artifact_store, session_id, + pending_activations, }; let socket_server = match transport { @@ -911,6 +1036,7 @@ pub(crate) async fn register_worker_tools( runtime_base: PathBuf, spawned_registry: Arc, parent_method_tx: Option>, + inherited_workdir_tool_broker: Option, ) -> std::io::Result> where C: LlmClient + Clone + 'static, @@ -919,21 +1045,26 @@ where // Worker-immutable snapshots taken before the mutable worker borrow // below so the worker borrow doesn't conflict with reads on `worker`. let feature_config = worker.manifest().feature.clone(); + let mut workdir_tool_broker = inherited_workdir_tool_broker; if feature_config.manage_workdir.enabled && worker.workdir_session().is_none() { let workspace_client = worker.workspace_client_handle(); - worker.bind_workdir_session(Some(workdir::delegation_capable_session( + let broker = workdir::WorkdirToolBroker::new( crate::feature::builtin::manage_workdir::WorkspaceAttachedWorkdirSession::handle( workspace_client, ), - ))); - } - if feature_config.sub_worker.enabled + ); + worker.bind_workdir_session(Some(broker.tool_session())); + workdir_tool_broker = Some(broker); + } else if workdir_tool_broker.is_none() && let Some(existing) = worker.workdir_session().cloned() - && !existing.is_delegation_capable() { - worker.bind_workdir_session(Some(workdir::delegation_capable_session(existing))); + let broker = workdir::WorkdirToolBroker::new(existing); + worker.bind_workdir_session(Some(broker.tool_session())); + workdir_tool_broker = Some(broker); } - let worker_workdir = worker.workdir_session().cloned(); + let worker_workdir = workdir_tool_broker + .as_ref() + .map(workdir::WorkdirToolBroker::tool_session); let local_filesystem = worker.local_working_directory().cloned(); let local_workspace_root = local_filesystem.as_ref().map(|local| local.root.clone()); let task_feature = worker.task_feature(); @@ -942,11 +1073,17 @@ where let spawner_name = worker.manifest().worker.name.clone(); let spawner_manifest = worker.manifest().clone(); let spawner_workspace_context = worker.workspace_context_handle(); - let parent_notifications = parent_method_tx - .map(crate::spawn::tool::ParentNotificationTarget::Controller) - .unwrap_or_else(|| { - crate::spawn::tool::ParentNotificationTarget::Buffer(worker.notify_buffer_handle()) - }); + let pending_submissions = worker.pending_submission_handle(); + let notify_buffer = worker.notify_buffer_handle(); + let durable_parent_notifications = + durable_parent_notification_target(pending_submissions.clone(), notify_buffer.clone()); + let parent_notifications = match parent_method_tx { + Some(sender) => crate::spawn::tool::ParentNotificationTarget::with_controller_fallback( + sender, + durable_parent_notifications, + ), + None => durable_parent_notifications, + }; let prompts = worker.prompts().clone(); let paste_store = worker.store().clone(); let paste_session_id = worker.session_id(); @@ -1094,8 +1231,17 @@ where "manage Workdir tools require Backend Workspace API authority", )); } + let shutdown_registry = spawned_registry.clone(); + let reopen_registry = spawned_registry.clone(); feature_registry.add_module( - crate::feature::builtin::manage_workdir::manage_workdir_feature(workspace_client), + crate::feature::builtin::manage_workdir::ManageWorkdirFeature::with_child_lifecycle( + workspace_client, + Arc::new(move || { + let child_registry = shutdown_registry.clone(); + Box::pin(async move { child_registry.shutdown_internal().await }) + }), + Arc::new(move || reopen_registry.reopen_internal()), + ), ); } if feature_config.workspace_worker_discovery.enabled { @@ -1157,7 +1303,6 @@ where } let host_worker_observation_provider = worker.worker_observation_provider(); - let source_workdir_session = worker.workdir_session().cloned(); { let workspace_client = worker.workspace_client_handle(); let engine = worker.engine_mut(); @@ -1199,7 +1344,7 @@ where runtime_base.clone(), bash_output_dir.clone(), spawner_workspace_root, - source_workdir_session, + workdir_tool_broker, spawned_registry.clone(), spawner_manifest, prompts, @@ -1289,7 +1434,18 @@ async fn controller_loop( discovery_cwd, spawned_registry.clone(), ); - let mut pending: Option = None; + let pending_submissions = worker.pending_submission_handle(); + stage_oldest_passive_notification(&pending_submissions, ¬ify_buffer); + let mut pending = match prepare_pending_run(&pending_submissions, ¬ify_buffer, None) { + Ok(pending) => pending, + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::Internal, + message: error.to_string(), + }); + None + } + }; loop { // Top-of-iteration: if an event handler staged a run, fire it @@ -1306,8 +1462,8 @@ async fn controller_loop( // interrupted/error turn from being carried into the next snapshot. worker.clear_in_flight_events(); let parent_originated = run.is_parent_originated(); - let user_input_run = matches!(&run, PendingRun::Run(_) | PendingRun::RunTracked { .. }); - if !user_input_run { + let user_input_submit = matches!(&run, PendingRun::Submit(_)); + if !user_input_submit { set_controller_status( &shared_state, &runtime_dir, @@ -1316,38 +1472,25 @@ async fn controller_loop( ) .await; } - let (mut new_status, shutdown) = match run { - PendingRun::Run(input) => { + let notification_request_id = match &run { + PendingRun::RunForNotification { + notification_request_id, + .. + } => notification_request_id.clone(), + _ => None, + }; + let passive_notification_request_id = + pending_submissions.activating_passive_notification_id(); + let (mut new_status, shutdown, may_drain_pending) = match run { + PendingRun::Submit(submission) => { let (input_commit_tx, input_commit_rx) = oneshot::channel(); + let committed_submission = submission.clone(); + let extension = pending_submissions.activation_extension(); drive_turn( worker.run_with_input_extensions_and_commit_hook( - input, - Vec::new(), - move || { - let _ = input_commit_tx.send(()); - }, - ), - &mut method_rx, - &working_event_tx, - &cancel_tx, - &pause_tx, - &shared_state, - &runtime_dir, - Some(input_commit_rx), - ¬ify_buffer, - self_parent_socket.as_ref(), - &spawner_name, - &spawned_registry, - parent_originated, - ) - .await - } - PendingRun::RunTracked { input, extension } => { - let (input_commit_tx, input_commit_rx) = oneshot::channel(); - drive_turn( - worker.run_with_input_extensions_and_commit_hook( - input, + submission.input, vec![extension], + submission.provenance, move || { let _ = input_commit_tx.send(()); }, @@ -1358,8 +1501,9 @@ async fn controller_loop( &pause_tx, &shared_state, &runtime_dir, - Some(input_commit_rx), + Some((input_commit_rx, committed_submission)), ¬ify_buffer, + &pending_submissions, self_parent_socket.as_ref(), &spawner_name, &spawned_registry, @@ -1367,9 +1511,9 @@ async fn controller_loop( ) .await } - PendingRun::RunForNotification(kind) => { + PendingRun::RunForNotification { invoke_kind, .. } => { drive_turn( - worker.run_for_notification(kind), + worker.run_for_notification(invoke_kind), &mut method_rx, &working_event_tx, &cancel_tx, @@ -1378,6 +1522,7 @@ async fn controller_loop( &runtime_dir, None, ¬ify_buffer, + &pending_submissions, self_parent_socket.as_ref(), &spawner_name, &spawned_registry, @@ -1396,6 +1541,7 @@ async fn controller_loop( &runtime_dir, None, ¬ify_buffer, + &pending_submissions, self_parent_socket.as_ref(), &spawner_name, &spawned_registry, @@ -1404,10 +1550,35 @@ async fn controller_loop( .await } }; - if !shutdown && new_status == WorkerStatus::Idle && notify_buffer.has_auto_run_pending() + if let Some(notification_request_id) = + notification_request_id.or(passive_notification_request_id) { - pending = Some(PendingRun::RunForNotification(protocol::InvokeKind::Notify)); - new_status = WorkerStatus::Running; + pending_submissions.finish_notification_activation(¬ification_request_id); + stage_oldest_passive_notification(&pending_submissions, ¬ify_buffer); + } + + if !shutdown && may_drain_pending && new_status == WorkerStatus::Idle { + match prepare_pending_run(&pending_submissions, ¬ify_buffer, None) { + Ok(Some(next)) => { + pending = Some(next); + new_status = WorkerStatus::Running; + } + Ok(None) => { + if notify_buffer.has_auto_run_pending() { + pending = Some(PendingRun::RunForNotification { + invoke_kind: protocol::InvokeKind::Notify, + notification_request_id: None, + }); + new_status = WorkerStatus::Running; + } + } + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::Internal, + message: error.to_string(), + }); + } + } } finish_controller_run( &mut worker, @@ -1434,61 +1605,227 @@ async fn controller_loop( }; match method { - Method::Run { input } => { - if shared_state.get_status() == WorkerStatus::Running { - // Defensive: the inner select! inside drive_turn - // already rejects `Run` while a turn is live, so - // this branch is only reachable across a race window - // around status flips. + Method::Submit { + submission_request_id, + input, + } => { + let request_id = submission_request_id.clone(); + match pending_submissions.accept_from_source( + submission_request_id, + input, + pending_submissions.direct_client_namespace(), + session_store::LoggedSessionHistoryOrigin::LegacyUnknown, + true, + ) { + Ok(acceptance) => { + if let Some(activation) = acceptance.activation { + pending = Some(PendingRun::Submit(activation)); + } else { + let _ = working_event_tx.send(Event::SubmissionAccepted { + submission_request_id: acceptance.submission_request_id, + submission_id: acceptance.submission_id, + disposition: acceptance.disposition, + }); + } + } + Err(error) => { + let _ = working_event_tx.send(Event::SubmissionRejected { + submission_request_id: request_id, + message: error.to_string(), + }); + } + } + } + Method::SubmitTracked { + submission_request_id, + input, + source, + } => { + let request_id = submission_request_id.clone(); + let (source_namespace, provenance) = + resolved_input_source(&pending_submissions, &source); + match pending_submissions.accept_from_source( + submission_request_id, + input, + source_namespace, + provenance, + true, + ) { + Ok(acceptance) => { + if let Some(activation) = acceptance.activation { + pending = Some(PendingRun::Submit(activation)); + } else { + let _ = working_event_tx.send(Event::SubmissionAccepted { + submission_request_id: acceptance.submission_request_id, + submission_id: acceptance.submission_id, + disposition: acceptance.disposition, + }); + } + } + Err(error) => { + let _ = working_event_tx.send(Event::SubmissionRejected { + submission_request_id: request_id, + message: error.to_string(), + }); + } + } + } + + Method::Notify { + notification_request_id, + message, + auto_run, + } => { + let request_id = notification_request_id.clone(); + let source_namespace = pending_submissions.direct_client_namespace(); + match pending_submissions.accept_notification_from_source( + notification_request_id, + message, + source_namespace.clone(), + session_store::LoggedSessionHistoryOrigin::LegacyUnknown, + auto_run, + ) { + Ok(_) if auto_run => { + match prepare_pending_run(&pending_submissions, ¬ify_buffer, None) { + Ok(Some(next)) => pending = Some(next), + Ok(None) => {} + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::Internal, + message: error.to_string(), + }); + } + } + } + Ok(_) => { + stage_pending_notification( + &pending_submissions, + ¬ify_buffer, + &source_namespace, + &request_id, + ); + } + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::InvalidRequest, + message: error.to_string(), + }); + } + } + } + + Method::NotifyTracked { + notification_request_id, + message, + auto_run, + source, + } => { + let request_id = notification_request_id.clone(); + let (source_namespace, provenance) = + resolved_input_source(&pending_submissions, &source); + match pending_submissions.accept_notification_from_source( + notification_request_id, + message, + source_namespace.clone(), + provenance, + auto_run, + ) { + Ok(_) if auto_run => { + match prepare_pending_run(&pending_submissions, ¬ify_buffer, None) { + Ok(Some(next)) => pending = Some(next), + Ok(None) => {} + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::Internal, + message: error.to_string(), + }); + } + } + } + Ok(_) => { + stage_pending_notification( + &pending_submissions, + ¬ify_buffer, + &source_namespace, + &request_id, + ); + } + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::InvalidRequest, + message: error.to_string(), + }); + } + } + } + + Method::ListPendingSubmissions => { + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { + pending: pending_submissions.snapshot(), + }); + } + Method::CancelPendingSubmission { + submission_id, + expected_revision, + } => match pending_submissions.cancel(&submission_id, expected_revision) { + Ok(pending_snapshot) => { + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { + pending: pending_snapshot, + }); + } + Err(error) => { let _ = working_event_tx.send(Event::Error { - code: ErrorCode::AlreadyRunning, - message: "Worker is already executing a turn".into(), + code: ErrorCode::InvalidRequest, + message: error.to_string(), + }); + } + }, + Method::ClearPendingSubmissions { expected_revision } => { + match pending_submissions.clear(expected_revision) { + Ok(pending_snapshot) => { + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { + pending: pending_snapshot, + }); + } + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::InvalidRequest, + message: error.to_string(), + }); + } + } + } + Method::ContinuePending { + expected_revision, + expected_head_id, + } => { + if shared_state.get_status() != WorkerStatus::Idle { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::InvalidRequest, + message: "ContinuePending requires an idle Worker; Resume or Cancel a paused run first".into(), }); continue; } - // Stage the run without a speculative user-message echo. - // `Worker::run` validates the input, commits - // `LogEntry::AnnotatedUserInput`, and the session-log sink turns that - // committed entry into the live `Event::UserMessage`. That - // keeps every client ordered against `SegmentStart` replay and - // makes persisted history the single source of visible user - // input. Paused→Run cleanup (orphan tool_result closure + - // interrupt system note) is applied inside `Worker::run` itself - // when the worker's `last_run_interrupted` flag is set. - pending = Some(PendingRun::Run(input)); - } - - Method::RunTracked { - input, - submission_id, - } => { - // Runtime-correlated submissions retain their opaque id in the - // same durable UserInput record used for Flow state. - let extension = SessionExtension::new( - WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, - serde_json::json!({ "submission_id": submission_id }), - ); - pending = Some(PendingRun::RunTracked { input, extension }); - } - - Method::Notify { message, auto_run } => { - // Client-side live echo is delivered as `Event::SystemItem` - // once the interceptor commits the corresponding - // `LogEntry::AnnotatedSystemItem` entry — drained out of the - // notify buffer + broadcast through the sink. No - // separate echo here. - worker.push_notify(message, auto_run); - // RUNNING: the in-flight turn drains the buffer at its next - // pending_history_appends; if an auto-run notification remains - // at turn end, the Controller stages a follow-up notification - // turn. Paused notifications remain queued until Resume/Run. - // IDLE: `auto_run` notifications stage RunForNotification; - // weak progress notices stay queued until an explicit run. - if should_auto_run_notification(shared_state.get_status(), auto_run) { - pending = Some(PendingRun::RunForNotification(protocol::InvokeKind::Notify)); + match prepare_pending_run( + &pending_submissions, + ¬ify_buffer, + Some((expected_revision, &expected_head_id)), + ) { + Ok(Some(next)) => pending = Some(next), + Ok(None) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::InvalidRequest, + message: "pending activation queue is empty".into(), + }); + } + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::Internal, + message: error.to_string(), + }); + } } } - Method::Resume => { if shared_state.get_status() != WorkerStatus::Paused { let _ = working_event_tx.send(Event::Error { @@ -1702,9 +2039,10 @@ async fn controller_loop( // notification is not stranded. Matches the // `Method::Notify` idle path. if shared_state.get_status() == WorkerStatus::Idle { - pending = Some(PendingRun::RunForNotification( - protocol::InvokeKind::WorkerEvent, - )); + pending = Some(PendingRun::RunForNotification { + invoke_kind: protocol::InvokeKind::WorkerEvent, + notification_request_id: None, + }); } } } @@ -1720,7 +2058,16 @@ async fn controller_loop( // Memory/Workdir teardown so they cannot observe a partially closed Worker. worker.stop_feature_runtime("controller shutdown").await; - if let Some(session) = worker.workdir_session() + let child_cleanup_succeeded = match spawned_registry.shutdown_internal().await { + Ok(()) => true, + Err(error) => { + tracing::warn!(%error, "Internal SubWorker cleanup failed before Workdir shutdown"); + false + } + }; + + if child_cleanup_succeeded + && let Some(session) = worker.workdir_session() && let Err(error) = session.close().await { tracing::warn!(%error, "Workdir session close failed"); @@ -1787,12 +2134,12 @@ async fn handle_inbound_worker_event( /// as `Errored` — only the worker-execution `Err` branch below fires. /// /// `parent_originated` further restricts both upward reports to turns -/// the parent actually delegated (`Method::Run` / `Method::Resume`). +/// the parent actually delegated (`Method::Submit` / `Method::Resume`). /// `Method::Notify` / inbound `WorkerEvent` auto-kicks complete silently /// so the parent's history does not get flooded with child-internal /// turn boundaries. #[allow(clippy::too_many_arguments)] -async fn drive_turn( +async fn drive_turn( worker_future: F, method_rx: &mut mpsc::Receiver, working_event_tx: &broadcast::Sender, @@ -1800,15 +2147,17 @@ async fn drive_turn( pause_tx: &mpsc::Sender<()>, shared_state: &Arc, runtime_dir: &RuntimeDir, - mut input_commit_rx: Option>, + mut input_commit: Option<(oneshot::Receiver<()>, crate::worker::PendingSubmission)>, notify_buffer: &NotifyBuffer, + pending_submissions: &crate::worker::PendingSubmissionHandle, parent_socket: Option<&PathBuf>, self_name: &str, spawned_registry: &Arc, parent_originated: bool, -) -> (WorkerStatus, bool) +) -> (WorkerStatus, bool, bool) where F: std::future::Future>, + St: Store + Clone, { tokio::pin!(worker_future); let mut shutdown_requested = false; @@ -1821,13 +2170,25 @@ where // Running snapshot contract deterministic even for immediate clients. biased; committed = async { - input_commit_rx + input_commit .as_mut() + .map(|(receiver, _)| receiver) .expect("input commit receiver guarded by select condition") .await - }, if input_commit_rx.is_some() => { - input_commit_rx = None; + }, if input_commit.is_some() => { + let submission = input_commit.take().map(|(_, submission)| submission); if committed.is_ok() { + if let Some(submission) = submission { + pending_submissions.finish_activation(&submission.submission_id); + let _ = working_event_tx.send(Event::SubmissionAccepted { + submission_request_id: submission.submission_request_id, + submission_id: submission.submission_id, + disposition: protocol::SubmissionDisposition::Started, + }); + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { + pending: pending_submissions.snapshot(), + }); + } set_controller_status( shared_state, runtime_dir, @@ -1835,11 +2196,33 @@ where WorkerStatus::Running, ) .await; + } else if let Some(submission) = submission { + pending_submissions.abort_activation(submission); } } result = &mut worker_future => { + if let Some((mut receiver, submission)) = input_commit.take() { + match receiver.try_recv() { + Ok(()) => { + pending_submissions.finish_activation(&submission.submission_id); + let _ = working_event_tx.send(Event::SubmissionAccepted { + submission_request_id: submission.submission_request_id, + submission_id: submission.submission_id, + disposition: protocol::SubmissionDisposition::Started, + }); + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { + pending: pending_submissions.snapshot(), + }); + } + Err(_) => pending_submissions.abort_activation(submission), + } + } return match result { Ok(r) => { + let may_drain_pending = matches!( + &r, + WorkerRunResult::Finished | WorkerRunResult::LimitReached + ); let (status, run_result) = match r { WorkerRunResult::Finished if pause_requested => { (WorkerStatus::Paused, RunResult::Paused) @@ -1850,7 +2233,7 @@ where WorkerRunResult::RolledBack => (WorkerStatus::Idle, RunResult::RolledBack), WorkerRunResult::Interrupted { .. } if pause_requested => { let _ = working_event_tx.send(Event::RunEnd { result: RunResult::Paused }); - return (WorkerStatus::Paused, shutdown_requested); + return (WorkerStatus::Paused, shutdown_requested, false); } WorkerRunResult::Interrupted { code, message } => { let _ = working_event_tx.send(Event::Error { @@ -1866,7 +2249,7 @@ where }, ); } - return (WorkerStatus::Idle, shutdown_requested); + return (WorkerStatus::Idle, shutdown_requested, false); } }; let _ = working_event_tx.send(Event::RunEnd { result: run_result }); @@ -1878,7 +2261,7 @@ where }, ); } - (status, shutdown_requested) + (status, shutdown_requested, may_drain_pending) } Err(WorkerError::Engine(EngineError::Cancelled)) if pause_requested => { // User-initiated Pause. Report the transition to @@ -1887,7 +2270,7 @@ where // that channel is reserved for worker runtime // failures, not deliberate interruptions. let _ = working_event_tx.send(Event::RunEnd { result: RunResult::Paused }); - (WorkerStatus::Paused, shutdown_requested) + (WorkerStatus::Paused, shutdown_requested, false) } Err(e) => { let code = worker_error_code(&e); @@ -1905,11 +2288,11 @@ where }, ); } - (WorkerStatus::Idle, shutdown_requested) + (WorkerStatus::Idle, shutdown_requested, false) } }; } - method = method_rx.recv() => { + method = method_rx.recv(), if input_commit.is_none() => { match method { Some(Method::Cancel) => { let _ = cancel_tx.try_send(()); @@ -1922,12 +2305,109 @@ where shutdown_requested = true; let _ = cancel_tx.try_send(()); } - Some(Method::Run { .. } | Method::RunTracked { .. } | Method::Resume) => { + Some(Method::Submit { + submission_request_id, + input, + }) => { + let request_id = submission_request_id.clone(); + match pending_submissions.accept_from_source( + submission_request_id, + input, + pending_submissions.direct_client_namespace(), + session_store::LoggedSessionHistoryOrigin::LegacyUnknown, + false, + ) { + Ok(acceptance) => { + let _ = working_event_tx.send(Event::SubmissionAccepted { + submission_request_id: acceptance.submission_request_id, + submission_id: acceptance.submission_id, + disposition: acceptance.disposition, + }); + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { + pending: pending_submissions.snapshot(), + }); + } + Err(error) => { + let _ = working_event_tx.send(Event::SubmissionRejected { + submission_request_id: request_id, + message: error.to_string(), + }); + } + } + } + Some(Method::SubmitTracked { + submission_request_id, + input, + source, + }) => { + let request_id = submission_request_id.clone(); + let (source_namespace, provenance) = + resolved_input_source(pending_submissions, &source); + match pending_submissions.accept_from_source( + submission_request_id, + input, + source_namespace, + provenance, + false, + ) { + Ok(acceptance) => { + let _ = working_event_tx.send(Event::SubmissionAccepted { + submission_request_id: acceptance.submission_request_id, + submission_id: acceptance.submission_id, + disposition: acceptance.disposition, + }); + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { + pending: pending_submissions.snapshot(), + }); + } + Err(error) => { + let _ = working_event_tx.send(Event::SubmissionRejected { + submission_request_id: request_id, + message: error.to_string(), + }); + } + } + } + Some(Method::Resume | Method::ContinuePending { .. }) => { let _ = working_event_tx.send(Event::Error { code: ErrorCode::AlreadyRunning, message: "Worker is already executing a turn".into(), }); } + Some(Method::ListPendingSubmissions) => { + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { + pending: pending_submissions.snapshot(), + }); + } + Some(Method::CancelPendingSubmission { + submission_id, + expected_revision, + }) => { + match pending_submissions.cancel(&submission_id, expected_revision) { + Ok(pending) => { + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { pending }); + } + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::InvalidRequest, + message: error.to_string(), + }); + } + } + } + Some(Method::ClearPendingSubmissions { expected_revision }) => { + match pending_submissions.clear(expected_revision) { + Ok(pending) => { + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { pending }); + } + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::InvalidRequest, + message: error.to_string(), + }); + } + } + } Some(Method::Compact | Method::ListRewindTargets | Method::RewindTo { .. }) => { let _ = working_event_tx.send(Event::Error { code: ErrorCode::AlreadyRunning, @@ -1935,11 +2415,75 @@ where .into(), }); } - Some(Method::Notify { message, auto_run }) => { - // Live echo arrives via `Event::SystemItem` once - // the in-flight turn's next `pending_history_appends` - // drains this entry through the interceptor. - notify_buffer.push_notify(message, auto_run); + Some(Method::Notify { + notification_request_id, + message, + auto_run, + }) => { + let request_id = notification_request_id.clone(); + let source_namespace = pending_submissions.direct_client_namespace(); + match pending_submissions.accept_notification_from_source( + notification_request_id, + message, + source_namespace.clone(), + session_store::LoggedSessionHistoryOrigin::LegacyUnknown, + auto_run, + ) { + Ok(_) if !auto_run => { + stage_pending_notification( + &pending_submissions, + notify_buffer, + &source_namespace, + &request_id, + ); + } + Ok(_) => {} + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::InvalidRequest, + message: error.to_string(), + }); + } + } + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { + pending: pending_submissions.snapshot(), + }); + } + Some(Method::NotifyTracked { + notification_request_id, + message, + auto_run, + source, + }) => { + let request_id = notification_request_id.clone(); + let (source_namespace, provenance) = + resolved_input_source(pending_submissions, &source); + match pending_submissions.accept_notification_from_source( + notification_request_id, + message, + source_namespace.clone(), + provenance, + auto_run, + ) { + Ok(_) if !auto_run => { + stage_pending_notification( + &pending_submissions, + notify_buffer, + &source_namespace, + &request_id, + ); + } + Ok(_) => {} + Err(error) => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::InvalidRequest, + message: error.to_string(), + }); + } + } + let _ = working_event_tx.send(Event::PendingSubmissionsChanged { + pending: pending_submissions.snapshot(), + }); } Some(Method::ListCompletions { .. }) => {} Some(Method::ListWorkers | Method::RestoreWorker { .. } | Method::RegisterPeer { .. }) => { @@ -1968,7 +2512,7 @@ where None => { let _ = cancel_tx.try_send(()); shared_state.set_status(WorkerStatus::Idle); - return (WorkerStatus::Idle, false); + return (WorkerStatus::Idle, false, false); } } } @@ -2117,6 +2661,113 @@ mod tests { use tempfile::TempDir; use tokio::net::UnixListener; + #[test] + fn no_controller_parent_notification_uses_durable_pending_authority() { + let temp = TempDir::new().unwrap(); + let pending = + crate::worker::PendingSubmissionHandle::for_test(&temp.path().join("sessions")); + let target = durable_parent_notification_target(pending.clone(), NotifyBuffer::new()); + let (wire_namespace, wire_provenance) = + resolved_input_source(&pending, &protocol::AuthenticatedInputSource::UntrustedWire); + assert_eq!(wire_namespace, pending.direct_client_namespace()); + assert!(matches!( + wire_provenance, + session_store::LoggedSessionHistoryOrigin::LegacyUnknown + )); + + target.notify("child-session".into(), "completed".into(), true); + + let snapshot = pending.snapshot(); + assert_eq!(snapshot.notification_count, 1); + assert!(snapshot.head_id.is_some()); + let notify_buffer = NotifyBuffer::new(); + assert!(matches!( + prepare_pending_run(&pending, ¬ify_buffer, None).unwrap(), + Some(PendingRun::RunForNotification { + notification_request_id: Some(_), + .. + }) + )); + assert!(notify_buffer.has_auto_run_pending()); + } + + #[test] + fn restored_mixed_activations_preserve_global_fifo_order() { + let temp = TempDir::new().unwrap(); + let pending = + crate::worker::PendingSubmissionHandle::for_test(&temp.path().join("submit-first")); + pending + .accept( + "submit-first".into(), + vec![protocol::Segment::Text { + content: "queued submit".into(), + }], + false, + ) + .unwrap(); + pending + .accept_notification("notify-second".into(), "newer notification".into(), true) + .unwrap(); + let notify_buffer = NotifyBuffer::new(); + + assert!(matches!( + prepare_pending_run(&pending, ¬ify_buffer, None).unwrap(), + Some(PendingRun::Submit(_)) + )); + + let pending = + crate::worker::PendingSubmissionHandle::for_test(&temp.path().join("notify-first")); + pending + .accept_notification("notify-first".into(), "older notification".into(), true) + .unwrap(); + pending + .accept( + "submit-second".into(), + vec![protocol::Segment::Text { + content: "newer submit".into(), + }], + false, + ) + .unwrap(); + let notify_buffer = NotifyBuffer::new(); + + assert!(matches!( + prepare_pending_run(&pending, ¬ify_buffer, None).unwrap(), + Some(PendingRun::RunForNotification { + notification_request_id: Some(request_id), + .. + }) if request_id == "notify-first" + )); + + let pending = + crate::worker::PendingSubmissionHandle::for_test(&temp.path().join("passive-first")); + pending + .accept_notification("passive-first".into(), "passive notification".into(), false) + .unwrap(); + pending + .accept( + "submit-after-passive".into(), + vec![protocol::Segment::Text { + content: "queued after passive".into(), + }], + false, + ) + .unwrap(); + let snapshot = pending.snapshot(); + let head_id = snapshot.head_id.clone().expect("queued Submit is the head"); + let notify_buffer = NotifyBuffer::new(); + assert!(stage_oldest_passive_notification(&pending, ¬ify_buffer)); + assert!(matches!( + prepare_pending_run( + &pending, + ¬ify_buffer, + Some((snapshot.revision + 1, &head_id)), + ) + .unwrap(), + Some(PendingRun::Submit(_)) + )); + } + #[test] fn image_attachment_gate_requires_vision_and_supported_openai_scheme() { let openai = manifest::ModelManifest { @@ -2133,21 +2784,16 @@ mod tests { #[test] fn pending_run_parent_origin_table() { - assert!(PendingRun::Run(Vec::new()).is_parent_originated()); assert!(PendingRun::Resume.is_parent_originated()); assert!( - !PendingRun::RunForNotification(protocol::InvokeKind::Notify).is_parent_originated() + !PendingRun::RunForNotification { + invoke_kind: protocol::InvokeKind::Notify, + notification_request_id: None, + } + .is_parent_originated() ); } - #[test] - fn notification_auto_run_gate_only_allows_idle_auto_run() { - assert!(should_auto_run_notification(WorkerStatus::Idle, true)); - assert!(!should_auto_run_notification(WorkerStatus::Idle, false)); - assert!(!should_auto_run_notification(WorkerStatus::Running, true)); - assert!(!should_auto_run_notification(WorkerStatus::Paused, true)); - } - struct DriveTurnEnv { // Held to keep the channel alive; without this `method_rx.recv()` // would observe channel-closed and confuse the select! arm. @@ -2160,6 +2806,7 @@ mod tests { _pause_rx: mpsc::Receiver<()>, shared_state: Arc, notify_buffer: NotifyBuffer, + pending_submissions: crate::worker::PendingSubmissionHandle, spawned_registry: Arc, parent_socket_path: PathBuf, runtime_dir: Arc, @@ -2193,6 +2840,8 @@ mod tests { }, )); let notify_buffer = NotifyBuffer::new(); + let pending_submissions = + crate::worker::PendingSubmissionHandle::for_test(&temp.path().join("pending-sessions")); let spawned_registry = SpawnedWorkerRegistry::new(runtime_dir.clone()); let parent_socket_path = temp.path().join("parent.sock"); @@ -2206,6 +2855,7 @@ mod tests { _pause_rx: pause_rx, shared_state, notify_buffer, + pending_submissions, spawned_registry, parent_socket_path, runtime_dir, @@ -2224,6 +2874,7 @@ mod tests { writer .write(&Event::Snapshot { session: protocol::SessionSnapshot { + pending_submissions: protocol::PendingSubmissionsSnapshot::default(), entries: Vec::new(), }, greeting: protocol::Greeting { @@ -2258,7 +2909,7 @@ mod tests { let recv = tokio::spawn(recv_worker_event(listener, Duration::from_secs(2))); let worker_future = async { Ok::<_, WorkerError>(WorkerRunResult::Finished) }; - let (status, shutdown) = drive_turn( + let (status, shutdown, _) = drive_turn( worker_future, &mut env.method_rx, &env.working_event_tx, @@ -2268,6 +2919,7 @@ mod tests { &env.runtime_dir, None, &env.notify_buffer, + &env.pending_submissions, Some(&env.parent_socket_path), "child-worker", &env.spawned_registry, @@ -2301,7 +2953,7 @@ mod tests { Ok::<_, WorkerError>(WorkerRunResult::Finished) }; let started_at = std::time::Instant::now(); - let (status, shutdown) = drive_turn( + let (status, shutdown, _) = drive_turn( worker_future, &mut env.method_rx, &env.working_event_tx, @@ -2311,6 +2963,7 @@ mod tests { &env.runtime_dir, None, &env.notify_buffer, + &env.pending_submissions, None, "child-worker", &env.spawned_registry, @@ -2331,7 +2984,7 @@ mod tests { let listener = UnixListener::bind(&env.parent_socket_path).expect("bind listener"); let worker_future = async { Ok::<_, WorkerError>(WorkerRunResult::Finished) }; - let (status, _) = drive_turn( + let (status, _, _) = drive_turn( worker_future, &mut env.method_rx, &env.working_event_tx, @@ -2341,6 +2994,7 @@ mod tests { &env.runtime_dir, None, &env.notify_buffer, + &env.pending_submissions, Some(&env.parent_socket_path), "child-worker", &env.spawned_registry, @@ -2369,7 +3023,7 @@ mod tests { "boom from test".into(), ))) }; - let (status, _) = drive_turn( + let (status, _, _) = drive_turn( worker_future, &mut env.method_rx, &env.working_event_tx, @@ -2379,6 +3033,7 @@ mod tests { &env.runtime_dir, None, &env.notify_buffer, + &env.pending_submissions, Some(&env.parent_socket_path), "child-worker", &env.spawned_registry, @@ -2413,7 +3068,7 @@ mod tests { "boom from notify".into(), ))) }; - let (status, _) = drive_turn( + let (status, _, _) = drive_turn( worker_future, &mut env.method_rx, &env.working_event_tx, @@ -2423,6 +3078,7 @@ mod tests { &env.runtime_dir, None, &env.notify_buffer, + &env.pending_submissions, Some(&env.parent_socket_path), "child-worker", &env.spawned_registry, @@ -2455,7 +3111,7 @@ mod tests { tokio::time::sleep(Duration::from_millis(50)).await; Ok::<_, WorkerError>(WorkerRunResult::Finished) }; - let (status, shutdown) = drive_turn( + let (status, shutdown, _) = drive_turn( worker_future, &mut env.method_rx, &env.working_event_tx, @@ -2465,6 +3121,7 @@ mod tests { &env.runtime_dir, None, &env.notify_buffer, + &env.pending_submissions, Some(&env.parent_socket_path), "parent", &env.spawned_registry, @@ -2494,7 +3151,7 @@ mod tests { tokio::time::sleep(Duration::from_millis(50)).await; Ok::<_, WorkerError>(WorkerRunResult::Finished) }; - let (status, shutdown) = drive_turn( + let (status, shutdown, _) = drive_turn( worker_future, &mut env.method_rx, &env.working_event_tx, @@ -2504,6 +3161,7 @@ mod tests { &env.runtime_dir, None, &env.notify_buffer, + &env.pending_submissions, Some(&env.parent_socket_path), "parent", &env.spawned_registry, @@ -2521,6 +3179,7 @@ mod tests { let mut env = make_env().await; env._method_tx .send(Method::Notify { + notification_request_id: protocol::new_submission_request_id(), message: "continue".into(), auto_run: true, }) @@ -2531,7 +3190,7 @@ mod tests { tokio::time::sleep(Duration::from_millis(50)).await; Ok::<_, WorkerError>(WorkerRunResult::Finished) }; - let (status, shutdown) = drive_turn( + let (status, shutdown, _) = drive_turn( worker_future, &mut env.method_rx, &env.working_event_tx, @@ -2541,6 +3200,7 @@ mod tests { &env.runtime_dir, None, &env.notify_buffer, + &env.pending_submissions, Some(&env.parent_socket_path), "parent", &env.spawned_registry, @@ -2550,8 +3210,8 @@ mod tests { assert_eq!(status, WorkerStatus::Idle); assert!(!shutdown); - assert_eq!(env.notify_buffer.len(), 1); - assert!(env.notify_buffer.has_auto_run_pending()); + assert_eq!(env.notify_buffer.len(), 0); + assert_eq!(env.pending_submissions.snapshot().notification_count, 1); } #[tokio::test] @@ -2567,7 +3227,7 @@ mod tests { tokio::time::sleep(Duration::from_millis(50)).await; Ok::<_, WorkerError>(WorkerRunResult::Finished) }; - let (status, shutdown) = drive_turn( + let (status, shutdown, _) = drive_turn( worker_future, &mut env.method_rx, &env.working_event_tx, @@ -2577,6 +3237,7 @@ mod tests { &env.runtime_dir, None, &env.notify_buffer, + &env.pending_submissions, Some(&env.parent_socket_path), "child-worker", &env.spawned_registry, @@ -2598,4 +3259,21 @@ mod tests { other => panic!("expected compact rejection error, got {other:?}"), } } + + #[test] + fn controller_shutdown_orders_child_cleanup_before_workdir_close() { + let source = include_str!("controller.rs"); + let shutdown_start = source + .rfind("worker.stop_feature_runtime(\"controller shutdown\")") + .expect("controller shutdown block"); + let shutdown = &source[shutdown_start..]; + let children = shutdown + .find("spawned_registry.shutdown_internal().await") + .expect("Internal SubWorker cleanup"); + let workdir = shutdown + .find("session.close().await") + .expect("parent Workdir close"); + assert!(children < workdir); + assert!(shutdown.contains("if child_cleanup_succeeded")); + } } diff --git a/crates/worker/src/discovery.rs b/crates/worker/src/discovery.rs index 32da104d..755180f4 100644 --- a/crates/worker/src/discovery.rs +++ b/crates/worker/src/discovery.rs @@ -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(value: &T) -> Result { @@ -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::().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::().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 { diff --git a/crates/worker/src/feature/builtin/manage_workdir.rs b/crates/worker/src/feature/builtin/manage_workdir.rs index f00093eb..eb950bbc 100644 --- a/crates/worker/src/feature/builtin/manage_workdir.rs +++ b/crates/worker/src/feature/builtin/manage_workdir.rs @@ -5,6 +5,8 @@ //! endpoints, credentials, materializer handles, and operation sessions stay //! behind [`WorkspaceClient`]. +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use agen::tool::{Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput}; @@ -12,7 +14,7 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json::json; use workdir::http::{WorkdirSessionOperation, WorkdirSessionOperationResult}; -use workdir::workspace::{WorkspaceWorkdirSessionFence, WorkspaceWorkdirSessionOperationRequest}; +use workdir::workspace::WorkspaceWorkdirSessionOperationRequest; use workdir::{ CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, @@ -52,16 +54,48 @@ const LIST_DESCRIPTION: &str = "List persistent Workdirs in the current Workspac const CREATE_DESCRIPTION: &str = "Materialize a persistent Workdir on a selected Runtime from a Workspace repository and optional selector. This does not change this Worker's attachment; use WorkdirAttach explicitly after creation."; const ATTACH_DESCRIPTION: &str = "Attach this Worker to one existing Workdir. The Backend enforces one active Workdir per Worker and one active Worker per Workdir, then opens an ephemeral operation session."; const DETACH_DESCRIPTION: &str = "Detach this Worker from its active Workdir and release Workdir occupancy. Any ephemeral operation session is closed."; +pub(crate) type BeforeWorkdirRelease = + Arc Pin> + Send>> + Send + Sync>; +pub(crate) type AfterWorkdirAttach = Arc; + const DELETE_DESCRIPTION: &str = "Request removal of one persistent Workdir by id through durable Backend Workspace authority. The input includes only the Workdir id and a bounded reason. The result reports removed, retained, or attention_required without exposing operation-table or provider internals."; -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct ManageWorkdirFeature { client: Arc, + before_workdir_release: Option, + after_workdir_attach: Option, +} + +impl std::fmt::Debug for ManageWorkdirFeature { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ManageWorkdirFeature") + .field("client_kind", &self.client.kind()) + .field("release_guard", &self.before_workdir_release.is_some()) + .finish() + } } impl ManageWorkdirFeature { pub fn new(client: Arc) -> Self { - Self { client } + Self { + client, + before_workdir_release: None, + after_workdir_attach: None, + } + } + + pub(crate) fn with_child_lifecycle( + client: Arc, + before_workdir_release: BeforeWorkdirRelease, + after_workdir_attach: AfterWorkdirAttach, + ) -> Self { + Self { + client, + before_workdir_release: Some(before_workdir_release), + after_workdir_attach: Some(after_workdir_attach), + } } } @@ -81,7 +115,10 @@ impl FeatureModule for ManageWorkdirFeature { } fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { - let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()); + let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()).with_child_lifecycle( + self.before_workdir_release.clone(), + self.after_workdir_attach.clone(), + ); for (name, definition) in [ ( LIST_TOOL, @@ -142,9 +179,21 @@ impl FeatureModule for ManageWorkdirFeature { } } -#[derive(Clone, Debug)] +#[derive(Clone)] struct WorkspaceHttpWorkdirBackend { client: Arc, + before_workdir_release: Option, + after_workdir_attach: Option, +} + +impl std::fmt::Debug for WorkspaceHttpWorkdirBackend { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkspaceHttpWorkdirBackend") + .field("client_kind", &self.client.kind()) + .field("release_guard", &self.before_workdir_release.is_some()) + .finish() + } } /// Worker-local Workdir handle whose operation authority remains in the Workspace Backend. @@ -156,8 +205,6 @@ struct WorkspaceHttpWorkdirBackend { pub struct WorkspaceAttachedWorkdirSession { client: Arc, workdir: Workdir, - expected_session_fence: Option, - delegations: Vec, } impl WorkspaceAttachedWorkdirSession { @@ -165,8 +212,6 @@ impl WorkspaceAttachedWorkdirSession { Arc::new(Self { client, workdir: Workdir::new("workspace-attachment"), - expected_session_fence: None, - delegations: Vec::new(), }) } @@ -183,16 +228,13 @@ impl WorkspaceAttachedWorkdirSession { "/api/w/{}/workers/self/workdir-session/operations", encode_path_segment(workspace_id) ), - serde_json::to_string(&WorkspaceWorkdirSessionOperationRequest { - expected_session_fence: self.expected_session_fence.clone(), - delegations: self.delegations.clone(), - operation, - }) - .map_err(|error| { - WorkdirError::Transport(format!( - "failed to encode Workspace Workdir operation: {error}" - )) - })?, + serde_json::to_string(&WorkspaceWorkdirSessionOperationRequest { operation }).map_err( + |error| { + WorkdirError::Transport(format!( + "failed to encode Workspace Workdir operation: {error}" + )) + }, + )?, ); let response = self .client @@ -241,59 +283,6 @@ impl WorkdirSession for WorkspaceAttachedWorkdirSession { WorkdirSessionCapabilities::ALL } - fn transports_delegation_context(&self) -> bool { - true - } - - async fn capture_delegation_source( - &self, - request: &workdir::WorkdirDelegationRequest, - ) -> Result { - let expected_session_fence = if let Some(fence) = &self.expected_session_fence { - fence.clone() - } else { - let workspace_id = self.client.workspace_id().ok_or_else(|| { - WorkdirError::Unavailable("Workspace identity is unavailable".to_string()) - })?; - let response = self - .client - .execute(WorkspaceRequest { - method: WorkspaceRequestMethod::Get, - path: format!( - "/api/w/{}/workers/self/workdir-session/fence", - encode_path_segment(workspace_id) - ), - body: None, - }) - .map_err(|error| { - WorkdirError::Unavailable(format!( - "failed to capture Workdir attachment fence: {error}" - )) - })?; - let fence: WorkspaceWorkdirSessionFence = serde_json::from_str(&response.body) - .map_err(|error| { - WorkdirError::Unavailable(format!( - "invalid Workdir attachment fence response: {error}" - )) - })?; - fence.value - }; - let mut delegations = self.delegations.clone(); - delegations.push(request.clone()); - let candidate = Arc::new(Self { - client: self.client.clone(), - workdir: self.workdir.clone(), - expected_session_fence: Some(expected_session_fence), - delegations, - }); - candidate - .stat(StatRequest { - path: workdir::WorkdirPath::new("").expect("empty Workdir path is valid"), - }) - .await?; - Ok(candidate) - } - async fn stat(&self, request: StatRequest) -> Result { match self.operate(WorkdirSessionOperation::Stat(request))? { WorkdirSessionOperationResult::Stat(result) => Ok(result), @@ -387,7 +376,21 @@ impl WorkdirSession for WorkspaceAttachedWorkdirSession { impl WorkspaceHttpWorkdirBackend { fn new(client: Arc) -> Self { - Self { client } + Self { + client, + before_workdir_release: None, + after_workdir_attach: None, + } + } + + fn with_child_lifecycle( + mut self, + before_workdir_release: Option, + after_workdir_attach: Option, + ) -> Self { + self.before_workdir_release = before_workdir_release; + self.after_workdir_attach = after_workdir_attach; + self } fn workspace_id(&self) -> Result<&str, ToolError> { @@ -565,11 +568,26 @@ impl Tool for WorkspaceHttpWorkdirTool { parse_input::(input_json)?, ctx.call_id.to_string(), ), - WorkdirOperation::Attach => self - .backend - .attach(parse_input::(input_json)?), + WorkdirOperation::Attach => { + let result = self + .backend + .attach(parse_input::(input_json)?); + if result.is_ok() + && let Some(after_attach) = &self.backend.after_workdir_attach + { + after_attach(); + } + result + } WorkdirOperation::Detach => { let _input = parse_input::(input_json)?; + if let Some(before_release) = &self.backend.before_workdir_release { + before_release().await.map_err(|error| { + ToolError::ExecutionFailed(format!( + "stop Internal SubWorkers before Workdir detach: {error}" + )) + })?; + } self.backend.detach() } WorkdirOperation::Delete => self @@ -765,6 +783,7 @@ struct WorkdirDeleteInput { #[cfg(test)] mod tests { use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; use super::*; use crate::feature::{FeatureModule, FeatureRegistryBuilder}; @@ -1155,6 +1174,7 @@ mod tests { command: "true".to_string(), timeout_secs: 120, output_limit: 1024, + cwd: None, spill_dir: Some("/worker-local/bash-output".into()), tool_call_id: Some("call-1".to_string()), }) @@ -1178,83 +1198,6 @@ mod tests { ); } - #[tokio::test] - async fn delegated_attached_session_carries_captured_fence_on_operations() { - let client = Arc::new(RecordingWorkspaceClient::new(vec![ - response(json!({"value": "attachment-fence"})), - response(json!({ - "operation": "stat", - "result": {"path": "", "kind": "directory", "size": 0} - })), - response(json!({ - "operation": "stat", - "result": {"path": "visible.txt", "kind": "file", "size": 8} - })), - ])); - let parent = workdir::delegation_capable_session(WorkspaceAttachedWorkdirSession::handle( - client.clone(), - )); - let delegation = parent - .delegate(workdir::WorkdirDelegationRequest { - rules: vec![workdir::WorkdirDelegationRule { - target: workdir::WorkdirPath::new("").unwrap(), - permission: workdir::WorkdirDelegationPermission::Read, - recursive: false, - }], - cwd: workdir::WorkdirPath::new("").unwrap(), - }) - .await - .unwrap(); - delegation - .scoped_session - .stat(StatRequest { - path: workdir::WorkdirPath::new("visible.txt").unwrap(), - }) - .await - .unwrap(); - - let requests = client.requests(); - assert_eq!(requests.len(), 3); - assert_eq!( - requests[0].path, - "/api/w/workspace%2Ftest/workers/self/workdir-session/fence" - ); - let body: serde_json::Value = - serde_json::from_str(requests[2].body.as_deref().unwrap()).unwrap(); - assert_eq!(body["expected_session_fence"], "attachment-fence"); - assert_eq!(body["operation"]["operation"], "stat"); - assert_eq!(body["delegations"][0]["rules"][0]["target"], ""); - } - - #[tokio::test] - async fn attached_provider_rejection_happens_before_delegation_is_returned() { - let client = Arc::new(RecordingWorkspaceClient::new(vec![ - response(json!({"value": "attachment-fence"})), - response(json!({"error": "provider rejected delegated write target"})), - ])); - let parent = workdir::delegation_capable_session(WorkspaceAttachedWorkdirSession::handle( - client.clone(), - )); - let result = parent - .delegate(workdir::WorkdirDelegationRequest { - rules: vec![workdir::WorkdirDelegationRule { - target: workdir::WorkdirPath::new("linked-target").unwrap(), - permission: workdir::WorkdirDelegationPermission::Write, - recursive: true, - }], - cwd: workdir::WorkdirPath::new("linked-target").unwrap(), - }) - .await; - - assert!(result.is_err(), "provider rejection must fail before lease"); - let requests = client.requests(); - assert_eq!(requests.len(), 2); - let validation: serde_json::Value = - serde_json::from_str(requests[1].body.as_deref().unwrap()).unwrap(); - assert_eq!(validation["operation"]["operation"], "stat"); - assert_eq!(validation["delegations"].as_array().unwrap().len(), 1); - } - #[tokio::test] async fn attached_session_preserves_typed_provider_validation_error() { let client = Arc::new(RecordingWorkspaceClient::new(vec![error_response( @@ -1298,73 +1241,52 @@ mod tests { } #[tokio::test] - async fn nested_attached_session_preserves_full_delegation_chain() { + async fn scoped_broker_operations_carry_no_child_context() { let client = Arc::new(RecordingWorkspaceClient::new(vec![ - response(json!({"value": "attachment-fence"})), response(json!({ "operation": "stat", - "result": {"path": "", "kind": "directory", "size": 0} + "result": {"path": "visible.txt", "kind": "file", "size": 8} })), response(json!({ "operation": "stat", - "result": {"path": "nested", "kind": "directory", "size": 0} - })), - response(json!({ - "operation": "stat", - "result": {"path": "nested/file", "kind": "file", "size": 1} + "result": {"path": "visible.txt", "kind": "file", "size": 8} })), ])); - let parent = workdir::delegation_capable_session(WorkspaceAttachedWorkdirSession::handle( + let broker = workdir::WorkdirToolBroker::new(WorkspaceAttachedWorkdirSession::handle( client.clone(), )); - let outer = parent - .delegate(workdir::WorkdirDelegationRequest { - rules: vec![workdir::WorkdirDelegationRule { + let scoped = broker + .scope(workdir::WorkdirToolScope { + rules: vec![workdir::WorkdirToolScopeRule { target: workdir::WorkdirPath::new("").unwrap(), - permission: workdir::WorkdirDelegationPermission::Read, + permission: workdir::WorkdirToolScopePermission::Read, recursive: true, }], cwd: workdir::WorkdirPath::new("").unwrap(), + command: false, }) .await .unwrap(); - let nested = outer - .scoped_session - .delegate(workdir::WorkdirDelegationRequest { - rules: vec![workdir::WorkdirDelegationRule { - target: workdir::WorkdirPath::new("nested").unwrap(), - permission: workdir::WorkdirDelegationPermission::Read, - recursive: true, - }], - cwd: workdir::WorkdirPath::new("nested").unwrap(), - }) - .await - .unwrap(); - nested - .scoped_session + scoped .stat(StatRequest { - path: workdir::WorkdirPath::new("file").unwrap(), + path: workdir::WorkdirPath::new("visible.txt").unwrap(), }) .await .unwrap(); let requests = client.requests(); - assert_eq!(requests.len(), 4); - let outer_validation: serde_json::Value = - serde_json::from_str(requests[1].body.as_deref().unwrap()).unwrap(); - let nested_validation: serde_json::Value = - serde_json::from_str(requests[2].body.as_deref().unwrap()).unwrap(); - assert_eq!(outer_validation["delegations"].as_array().unwrap().len(), 1); - assert_eq!( - nested_validation["delegations"].as_array().unwrap().len(), - 2 - ); - let body: serde_json::Value = - serde_json::from_str(requests[3].body.as_deref().unwrap()).unwrap(); - assert_eq!(body["delegations"].as_array().unwrap().len(), 2); - assert_eq!(body["delegations"][0]["rules"][0]["target"], ""); - assert_eq!(body["delegations"][1]["rules"][0]["target"], "nested"); - assert_eq!(body["operation"]["request"]["path"], "file"); + assert_eq!(requests.len(), 2); + for request in requests { + assert_eq!( + request.path, + "/api/w/workspace%2Ftest/workers/self/workdir-session/operations" + ); + let body: serde_json::Value = + serde_json::from_str(request.body.as_deref().unwrap()).unwrap(); + assert!(body.get("delegations").is_none()); + assert!(body.get("child").is_none()); + assert!(body.get("expected_session_fence").is_none()); + } } #[test] @@ -1416,4 +1338,86 @@ mod tests { assert!(client.requests().is_empty()); assert!(parse_input::(r#"{"path":"/tmp"}"#).is_err()); } + + #[tokio::test] + async fn detach_stops_internal_subworkers_before_backend_release() { + let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({ + "workspace_id": "workspace/test", + "workdir_id": "wd-attached", + "attached": false + }))])); + let cleanup_calls = Arc::new(AtomicUsize::new(0)); + let cleanup_calls_for_guard = cleanup_calls.clone(); + let before_release: BeforeWorkdirRelease = Arc::new(move || { + let cleanup_calls = cleanup_calls_for_guard.clone(); + Box::pin(async move { + cleanup_calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + }); + let tool = WorkspaceHttpWorkdirTool { + backend: WorkspaceHttpWorkdirBackend::new(client.clone()) + .with_child_lifecycle(Some(before_release), None), + operation: WorkdirOperation::Detach, + }; + + tool.execute("{}", ToolExecutionContext::default()) + .await + .unwrap(); + + assert_eq!(cleanup_calls.load(Ordering::SeqCst), 1); + assert_eq!(client.requests().len(), 1); + assert_eq!( + client.requests()[0].path, + "/api/w/workspace%2Ftest/workers/self/workdir-attachment" + ); + } + + #[tokio::test] + async fn detach_does_not_release_backend_when_child_cleanup_fails() { + let client = Arc::new(RecordingWorkspaceClient::new(Vec::new())); + let before_release: BeforeWorkdirRelease = + Arc::new(|| Box::pin(async { Err(std::io::Error::other("child cleanup failed")) })); + let tool = WorkspaceHttpWorkdirTool { + backend: WorkspaceHttpWorkdirBackend::new(client.clone()) + .with_child_lifecycle(Some(before_release), None), + operation: WorkdirOperation::Detach, + }; + + let error = tool + .execute("{}", ToolExecutionContext::default()) + .await + .unwrap_err(); + + assert!(error.to_string().contains("stop Internal SubWorkers")); + assert!(client.requests().is_empty()); + } + + #[tokio::test] + async fn successful_attach_reopens_internal_subworker_admission() { + let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({ + "workspace_id": "workspace/test", + "workdir_id": "wd-attached", + "attached": true + }))])); + let reopen_calls = Arc::new(AtomicUsize::new(0)); + let reopen_calls_for_hook = reopen_calls.clone(); + let after_attach: AfterWorkdirAttach = Arc::new(move || { + reopen_calls_for_hook.fetch_add(1, Ordering::SeqCst); + }); + let tool = WorkspaceHttpWorkdirTool { + backend: WorkspaceHttpWorkdirBackend::new(client) + .with_child_lifecycle(None, Some(after_attach)), + operation: WorkdirOperation::Attach, + }; + + tool.execute( + r#"{"workdir_id":"wd-attached"}"#, + ToolExecutionContext::default(), + ) + .await + .unwrap(); + + assert_eq!(reopen_calls.load(Ordering::SeqCst), 1); + } } diff --git a/crates/worker/src/feature/builtin/worker_observation.rs b/crates/worker/src/feature/builtin/worker_observation.rs index 2dc059f6..bda05e09 100644 --- a/crates/worker/src/feature/builtin/worker_observation.rs +++ b/crates/worker/src/feature/builtin/worker_observation.rs @@ -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, + }, }) } } diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index ec7ef586..52427800 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -709,7 +709,7 @@ pub(crate) fn prepare_internal_worker_from_spec( } Box::pin(prepare_internal_worker_session( - worker, store, visibility, None, None, + worker, store, visibility, None, None, None, )) .await }) @@ -746,13 +746,16 @@ pub(crate) async fn prepare_internal_worker_session( visibility: InternalWorkerVisibility, child_registry: Option>, on_turn_end: Option>, + command_event_broker: Option, ) -> Result { let (event_tx, _event_rx) = broadcast::channel(256); let sink = worker.sink(); spawn_internal_log_event_bridge(sink.clone(), event_tx.clone()); let alerter = Alerter::new(event_tx.clone()); let in_flight = InFlightEvents::new(event_tx.clone()); - if let Some(session) = worker.workdir_session() { + if let Some(broker) = command_event_broker.as_ref() { + wire_workdir_command_events(&broker.tool_session(), &in_flight); + } else if let Some(session) = worker.workdir_session() { wire_workdir_command_events(session, &in_flight); } let actor_in_flight = in_flight.clone(); @@ -887,6 +890,7 @@ pub(crate) async fn spawn_prepared_internal_worker_session( InternalWorkerVisibility::ServicePrivate, None, on_turn_end, + None, ) .await?; handle.send(input).await?; diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 30bf4e12..86f26230 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -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, + Option, + )], + ) -> 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::>(), + ) + } + fn current_turn_index(&self) -> usize { self.next_turn_index .load(Ordering::Relaxed) @@ -327,7 +348,11 @@ impl Interceptor for WorkerInterceptor { projection_digest: projection.catalog_digest.clone(), logical_name: "internal.notify_wrapper".to_string(), }; - let mut system_items: Vec = Vec::with_capacity(drained.len()); + let mut system_items: Vec<( + SystemItem, + Vec, + Option, + )> = Vec::with_capacity(drained.len()); let mut items: Vec = Vec::with_capacity(drained.len()); for entry in &drained { let system_item = match build_system_item_with_provenance( @@ -345,9 +370,9 @@ impl Interceptor 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, diff --git a/crates/worker/src/ipc/notify_buffer.rs b/crates/worker/src/ipc/notify_buffer.rs index 7b97e4ed..2fd2ca03 100644 --- a/crates/worker/src/ipc/notify_buffer.rs +++ b/crates/worker/src/ipc/notify_buffer.rs @@ -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, + history_provenance: Option, + }, + WorkerEvent { + event: WorkerEvent, + }, +} + +impl PendingNotify { + pub(crate) fn extensions(&self) -> Vec { + match self { + PendingNotify::Notify { extensions, .. } => extensions.clone(), + PendingNotify::WorkerEvent { .. } => Vec::new(), + } + } + + pub(crate) fn history_provenance(&self) -> Option { + 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(); diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index d5f7f8d1..c83ef56a 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -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, }; diff --git a/crates/worker/src/segment_log_sink.rs b/crates/worker/src/segment_log_sink.rs index 91c2792b..20233983 100644 --- a/crates/worker/src/segment_log_sink.rs +++ b/crates/worker/src/segment_log_sink.rs @@ -291,6 +291,7 @@ mod tests { prompt_provenance: None, }, ), + extensions: Vec::new(), } } diff --git a/crates/worker/src/spawn/comm_tools.rs b/crates/worker/src/spawn/comm_tools.rs index 27e6845f..0389f288 100644 --- a/crates/worker/src/spawn/comm_tools.rs +++ b/crates/worker/src/spawn/comm_tools.rs @@ -72,6 +72,7 @@ mod tests { fn snapshot(entries: Vec) -> Event { Event::Snapshot { session: protocol::SessionSnapshot { + pending_submissions: protocol::PendingSubmissionsSnapshot::default(), entries: entries .into_iter() .enumerate() diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 62fec8d4..fd946e54 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -12,7 +12,7 @@ use std::collections::{BTreeMap, HashSet}; use std::io; use std::sync::{ Arc, Mutex, - atomic::{AtomicBool, AtomicU64, Ordering}, + atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, }; use std::time::Instant; @@ -23,9 +23,9 @@ use protocol::{Event, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnaps use session_store::{ LoggedItem, WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerStoreError, }; -use tokio::sync::broadcast; +use tokio::sync::{Notify, broadcast}; use tracing::warn; -use workdir::WorkdirDelegation; +use workdir::WorkdirScopeLease; use crate::internal_worker::{InternalWorkerSessionHandle, InternalWorkerVisibility}; use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord}; @@ -68,10 +68,11 @@ pub(crate) struct SubWorkerStopSummary { pub(crate) struct InternalSpawnedWorkerRecord { pub worker_name: String, pub scope_delegated: Vec, - pub workdir_delegation: Arc, + pub workdir_tool_scope: Arc, #[cfg(test)] pub installed_tools: Arc<[String]>, pub session: InternalWorkerSessionHandle, + pub child_registry: Arc, change_tracker: Option, started_at: Instant, stop_lock: Arc>, @@ -86,18 +87,20 @@ impl InternalSpawnedWorkerRecord { pub(crate) fn new( worker_name: String, scope_delegated: Vec, - workdir_delegation: WorkdirDelegation, + workdir_tool_scope: WorkdirScopeLease, #[cfg(test)] installed_tools: Vec, session: InternalWorkerSessionHandle, + child_registry: Arc, change_tracker: Option, ) -> Self { Self { worker_name, scope_delegated, - workdir_delegation: Arc::new(workdir_delegation), + workdir_tool_scope: Arc::new(workdir_tool_scope), #[cfg(test)] installed_tools: installed_tools.into(), session, + child_registry, change_tracker, started_at: Instant::now(), stop_lock: Arc::new(tokio::sync::Mutex::new(())), @@ -235,18 +238,56 @@ pub(crate) struct InternalSpawnReservation { } impl InternalSpawnReservation { - pub(crate) fn commit(mut self, record: InternalSpawnedWorkerRecord) -> io::Result<()> { - if record.worker_name != self.worker_name { - return Err(io::Error::new( + pub(crate) async fn commit(mut self, record: InternalSpawnedWorkerRecord) -> io::Result<()> { + let rejection = if record.worker_name != self.worker_name { + Some(io::Error::new( io::ErrorKind::InvalidInput, "internal SubWorker reservation name does not match record name", - )); + )) + } else { + match self.registry.internal_records.lock() { + Ok(mut records) => { + if self.registry.internal_shutting_down.load(Ordering::Acquire) { + Some(io::Error::new( + io::ErrorKind::Interrupted, + "internal SubWorker registry is shutting down", + )) + } else { + records.push(record.clone()); + None + } + } + Err(_) => Some(io::Error::other( + "internal spawned-worker registry lock poisoned", + )), + } + }; + if let Some(error) = rejection { + let mut cleanup_failures = Vec::new(); + if let Err(cleanup) = record.session.stop().await { + cleanup_failures.push(format!("stop rejected Internal SubWorker: {cleanup}")); + } + if let Err(cleanup) = Box::pin(record.child_registry.shutdown_internal()).await { + cleanup_failures.push(format!( + "stop rejected Internal SubWorker descendants: {cleanup}" + )); + } + if let Err(cleanup) = record.workdir_tool_scope.close().await { + cleanup_failures.push(format!( + "close rejected Internal SubWorker Workdir tools: {cleanup}" + )); + } + if cleanup_failures.is_empty() { + return Err(error); + } + self.registry + .internal_spawn_cleanup_failed + .store(true, Ordering::Release); + return Err(io::Error::other(format!( + "{error}; {}", + cleanup_failures.join("; ") + ))); } - self.registry - .internal_records - .lock() - .map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))? - .push(record.clone()); self.registry.start_protocol_forwarding(record); self.committed = true; Ok(()) @@ -260,6 +301,10 @@ impl Drop for InternalSpawnReservation { names.remove(&self.worker_name); } } + self.registry + .pending_internal_spawns + .fetch_sub(1, Ordering::AcqRel); + self.registry.pending_internal_notify.notify_waiters(); } } @@ -267,6 +312,10 @@ pub struct SpawnedWorkerRegistry { internal_records: std::sync::Mutex>, service_records: std::sync::Mutex>, internal_names: std::sync::Mutex>, + internal_shutting_down: AtomicBool, + pending_internal_spawns: AtomicUsize, + pending_internal_notify: Notify, + internal_spawn_cleanup_failed: AtomicBool, parent_scope: Option, parent_protocol: Mutex, String)>>, } @@ -283,6 +332,10 @@ impl SpawnedWorkerRegistry { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), + internal_shutting_down: AtomicBool::new(false), + pending_internal_spawns: AtomicUsize::new(0), + pending_internal_notify: Notify::new(), + internal_spawn_cleanup_failed: AtomicBool::new(false), parent_scope: None, parent_protocol: Mutex::new(None), }) @@ -294,6 +347,10 @@ impl SpawnedWorkerRegistry { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), + internal_shutting_down: AtomicBool::new(false), + pending_internal_spawns: AtomicUsize::new(0), + pending_internal_notify: Notify::new(), + internal_spawn_cleanup_failed: AtomicBool::new(false), parent_scope: None, parent_protocol: Mutex::new(None), }) @@ -304,6 +361,10 @@ impl SpawnedWorkerRegistry { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), + internal_shutting_down: AtomicBool::new(false), + pending_internal_spawns: AtomicUsize::new(0), + pending_internal_notify: Notify::new(), + internal_spawn_cleanup_failed: AtomicBool::new(false), parent_scope: Some(parent_scope), parent_protocol: Mutex::new(None), }) @@ -383,6 +444,10 @@ impl SpawnedWorkerRegistry { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), + internal_shutting_down: AtomicBool::new(false), + pending_internal_spawns: AtomicUsize::new(0), + pending_internal_notify: Notify::new(), + internal_spawn_cleanup_failed: AtomicBool::new(false), parent_scope, parent_protocol: Mutex::new(None), }), @@ -394,6 +459,16 @@ impl SpawnedWorkerRegistry { self: &Arc, worker_name: String, ) -> io::Result { + let records = self + .internal_records + .lock() + .map_err(|_| io::Error::other("internal Worker registry lock poisoned"))?; + if self.internal_shutting_down.load(Ordering::Acquire) { + return Err(io::Error::new( + io::ErrorKind::Interrupted, + "internal SubWorker registry is shutting down", + )); + } let mut names = self .internal_names .lock() @@ -404,7 +479,9 @@ impl SpawnedWorkerRegistry { format!("spawned worker `{worker_name}` is already registered"), )); } + self.pending_internal_spawns.fetch_add(1, Ordering::AcqRel); drop(names); + drop(records); Ok(InternalSpawnReservation { registry: Arc::clone(self), worker_name, @@ -679,18 +756,11 @@ impl SpawnedWorkerRegistry { .unwrap_or_default() } - pub(crate) fn reclaim_internal_scope(&self, worker_name: &str) -> io::Result { - let record = self.get_internal(worker_name).ok_or_else(|| { - io::Error::new(io::ErrorKind::NotFound, "internal SubWorker not found") - })?; - self.reclaim_record_scope(&record) - } - fn reclaim_record_scope(&self, record: &InternalSpawnedWorkerRecord) -> io::Result { if !record.claim_scope_reclaim() { return Ok(false); } - record.workdir_delegation.release(); + record.workdir_tool_scope.revoke(); let result = if let Some(parent_scope) = &self.parent_scope { parent_scope .update(|current| current.with_removed_deny_rules(delegated_write_rules(record))) @@ -705,6 +775,58 @@ impl SpawnedWorkerRegistry { result } + pub(crate) async fn close_internal_scope(&self, name: &str) -> io::Result { + let Some(record) = self.get_internal(name) else { + return Ok(false); + }; + Box::pin(record.child_registry.shutdown_internal()).await?; + record + .workdir_tool_scope + .close() + .await + .map_err(|error| io::Error::other(error.to_string()))?; + self.reclaim_record_scope(&record) + } + + pub(crate) async fn shutdown_internal(&self) -> io::Result<()> { + let names = { + let records = self + .internal_records + .lock() + .map_err(|_| io::Error::other("internal Worker registry lock poisoned"))?; + self.internal_shutting_down.store(true, Ordering::Release); + records + .iter() + .map(|record| record.worker_name.clone()) + .collect::>() + }; + loop { + let notified = self.pending_internal_notify.notified(); + if self.pending_internal_spawns.load(Ordering::Acquire) == 0 { + break; + } + notified.await; + } + let mut first_error = None; + for name in names { + if let Err(error) = self.remove_internal(&name).await { + first_error.get_or_insert(error); + } + } + if first_error.is_none() && self.internal_spawn_cleanup_failed.load(Ordering::Acquire) { + first_error = Some(io::Error::other( + "an in-flight Internal SubWorker failed cleanup during shutdown", + )); + } + first_error.map_or(Ok(()), Err) + } + + pub(crate) fn reopen_internal(&self) { + self.internal_shutting_down.store(false, Ordering::Release); + self.internal_spawn_cleanup_failed + .store(false, Ordering::Release); + } + /// Stop one direct Internal SubWorker and discard its registry/scope state. /// /// The child actor must acknowledge its stop before the registry is removed. @@ -731,6 +853,12 @@ impl SpawnedWorkerRegistry { .stop() .await .map_err(|error| io::Error::other(error.to_string()))?; + Box::pin(record.child_registry.shutdown_internal()).await?; + record + .workdir_tool_scope + .close() + .await + .map_err(|error| io::Error::other(error.to_string()))?; let summary = record.stop_summary(); self.reclaim_record_scope(&record)?; let removed = @@ -966,7 +1094,7 @@ mod tests { deny: Vec::new(), }) .unwrap(); - let source = workdir::delegation_capable_session(Arc::new( + let source = workdir::WorkdirToolBroker::new(Arc::new( workdir::LocalWorkdirSession::materialized_bound( workdir::Workdir::new("registry-test"), root.clone(), @@ -976,13 +1104,14 @@ mod tests { ), )); let delegation = source - .delegate(workdir::WorkdirDelegationRequest { - rules: vec![workdir::WorkdirDelegationRule { + .scope(workdir::WorkdirToolScope { + rules: vec![workdir::WorkdirToolScopeRule { target: workdir::WorkdirPath::new("").unwrap(), - permission: workdir::WorkdirDelegationPermission::Read, + permission: workdir::WorkdirToolScopePermission::Read, recursive: true, }], cwd: workdir::WorkdirPath::new("").unwrap(), + command: false, }) .await .unwrap(); @@ -993,6 +1122,7 @@ mod tests { delegation, Vec::new(), session, + registry(), None, ), sender, @@ -1230,6 +1360,143 @@ mod tests { } } + #[tokio::test] + async fn parent_shutdown_stops_all_internal_workers_before_returning() { + let registry = registry(); + for name in ["first", "second"] { + let (record, _events) = record(name, InternalWorkerVisibility::ParentClient).await; + record + .session + .force_status(InternalWorkerSessionStatus::Running); + install_record(®istry, record); + } + + registry.shutdown_internal().await.unwrap(); + + assert!(registry.list_internal().is_empty()); + assert!(registry.get_internal("first").is_none()); + assert!(registry.get_internal("second").is_none()); + } + + #[tokio::test] + async fn shutdown_rejects_new_reservations_until_reopened() { + let registry = registry(); + registry.shutdown_internal().await.unwrap(); + assert!(registry.reserve_internal_name("late-child".into()).is_err()); + + registry.reopen_internal(); + let reservation = registry.reserve_internal_name("late-child".into()).unwrap(); + drop(reservation); + } + + #[tokio::test] + async fn concurrent_commit_and_shutdown_leave_no_live_internal_worker() { + let registry = registry(); + let reservation = registry + .reserve_internal_name("racing-child".into()) + .unwrap(); + let (record, _events) = + record("racing-child", InternalWorkerVisibility::ParentClient).await; + let scope = record.workdir_tool_scope.clone(); + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + let commit_barrier = barrier.clone(); + let commit = tokio::spawn(async move { + commit_barrier.wait().await; + reservation.commit(record).await + }); + let shutdown_registry = registry.clone(); + let shutdown = tokio::spawn(async move { + barrier.wait().await; + shutdown_registry.shutdown_internal().await + }); + + let commit = commit.await.unwrap(); + shutdown.await.unwrap().unwrap(); + if let Err(error) = commit { + assert_eq!(error.kind(), io::ErrorKind::Interrupted); + } + + assert!(registry.list_internal().is_empty()); + assert!(!scope.is_active()); + } + + #[tokio::test] + async fn shutdown_fences_a_reservation_that_has_not_committed() { + let registry = registry(); + let reservation = registry + .reserve_internal_name("racing-child".into()) + .unwrap(); + let (record, _events) = + record("racing-child", InternalWorkerVisibility::ParentClient).await; + + let mut shutdown = { + let registry = registry.clone(); + tokio::spawn(async move { registry.shutdown_internal().await }) + }; + while !registry.internal_shutting_down.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + assert!( + tokio::time::timeout(std::time::Duration::from_millis(50), &mut shutdown) + .await + .is_err(), + "shutdown must wait for the pending spawn to roll back" + ); + let error = reservation.commit(record).await.unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::Interrupted); + shutdown.await.unwrap().unwrap(); + } + + #[tokio::test] + async fn rejected_spawn_cleanup_failure_keeps_shutdown_failed_closed() { + let registry = registry(); + let reservation = registry + .reserve_internal_name("cleanup-failure".into()) + .unwrap(); + let (record, _events) = + record("cleanup-failure", InternalWorkerVisibility::ParentClient).await; + record.session.force_stop_failure(); + let shutdown = { + let registry = registry.clone(); + tokio::spawn(async move { registry.shutdown_internal().await }) + }; + while !registry.internal_shutting_down.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + + let error = reservation.commit(record).await.unwrap_err(); + assert!( + error + .to_string() + .contains("stop rejected Internal SubWorker") + ); + let shutdown_error = shutdown.await.unwrap().unwrap_err(); + assert!( + shutdown_error + .to_string() + .contains("failed cleanup during shutdown") + ); + assert!(registry.internal_shutting_down.load(Ordering::Acquire)); + } + + #[tokio::test] + async fn shutdown_recursively_stops_grandchildren_before_parent_scope_release() { + let registry = registry(); + let (child, _child_events) = record("child", InternalWorkerVisibility::ParentClient).await; + let child_registry = child.child_registry.clone(); + let (grandchild, _grandchild_events) = + record("grandchild", InternalWorkerVisibility::ParentClient).await; + let grandchild_scope = grandchild.workdir_tool_scope.clone(); + install_record(&child_registry, grandchild); + install_record(®istry, child); + + registry.shutdown_internal().await.unwrap(); + + assert!(registry.list_internal().is_empty()); + assert!(child_registry.list_internal().is_empty()); + assert!(!grandchild_scope.is_active()); + } + #[tokio::test] async fn running_worker_is_stopped_before_removal() { let registry = registry(); diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index e6b8f872..6ec7c3cc 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -22,8 +22,7 @@ use manifest::{ use serde::Deserialize; use tokio::sync::mpsc; use workdir::{ - WorkdirDelegationPermission, WorkdirDelegationRequest, WorkdirDelegationRule, WorkdirPath, - WorkdirSessionHandle, + WorkdirToolBroker, WorkdirToolScope, WorkdirToolScopePermission, WorkdirToolScopeRule, }; use crate::PromptCatalogSource; @@ -58,12 +57,15 @@ struct SubWorkerSpawnInput { /// a host path and grants no authority. When omitted, the Workdir root is used. #[serde(default)] cwd: Option, - /// 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 /// sufficient. Omit `recursive` for normal workspace/worktree delegation; it defaults to true. scope: Vec, + /// Explicitly grant command execution through the parent-owned Workdir tool broker. + #[serde(default)] + command: bool, /// Binds an actual read-only builtin Reviewer child to the current Merge Request candidate. /// Review capability material is generated by the trusted spawn layer. #[serde(default)] @@ -219,33 +221,50 @@ fn parse_spawn_profile_selector(raw: Option<&str>) -> Result), - Buffer(crate::ipc::notify_buffer::NotifyBuffer), + Controller { + sender: mpsc::WeakSender, + fallback: Arc, + }, + Durable(Arc), } impl ParentNotificationTarget { - fn notify(&self, message: String, auto_run: bool) { + pub(crate) fn with_controller_fallback( + sender: mpsc::WeakSender, + 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), } } } @@ -267,8 +286,8 @@ pub struct SubWorkerSpawnTool { workspace_root: PathBuf, /// Directory the spawned SubWorker's tools should use when the LLM did not /// override it. Defaults to the spawner's cwd. - /// Active provider-backed Workdir session from which child leases are captured. - source_workdir_session: Option, + /// Parent-owned broker for scoped Workdir tool execution. + workdir_tool_broker: Option, /// Parent-owned in-memory registry shared by the five SubWorker tools. registry: Arc, /// Spawner's resolved Manifest. `profile = "inherit"` derives the @@ -295,7 +314,7 @@ impl SubWorkerSpawnTool { runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - source_workdir_session: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompt_loader: PromptCatalogSource, @@ -308,7 +327,7 @@ impl SubWorkerSpawnTool { runtime_base, bash_output_dir, workspace_root, - source_workdir_session, + workdir_tool_broker, registry, spawner_manifest, prompt_loader, @@ -341,6 +360,11 @@ fn validate_reviewer_handoff(input: &SubWorkerSpawnInput) -> Result<(), ToolErro "Merge Request Reviewer SubWorkers must include writable delegated scope".to_string(), )); } + if !input.command { + return Err(ToolError::InvalidArgument( + "Merge Request Reviewer SubWorkers require an explicit command grant".to_string(), + )); + } Ok(()) } @@ -370,7 +394,7 @@ impl Tool for SubWorkerSpawnTool { .reserve_internal_name(input.name.clone()) .map_err(|error| ToolError::InvalidArgument(error.to_string()))?; - let mut workdir_rules = parse_workdir_scope(&input.scope)?; + let workdir_rules = parse_workdir_scope(&input.scope)?; let child_bash_output_dir = self.bash_output_dir.join("sub-workers").join(&input.name); tokio::fs::create_dir_all(&child_bash_output_dir) .await @@ -380,28 +404,15 @@ impl Tool for SubWorkerSpawnTool { child_bash_output_dir.display() )) })?; - let source_workdir_session = - require_active_workdir_session(self.source_workdir_session.as_ref())?; - let transports_delegation_context = source_workdir_session.transports_delegation_context(); - // Provider-transported sessions resolve every delegation rule in the - // receiving Workdir namespace. The Bash spill directory instead belongs - // to this Worker host, so forwarding it would widen the request with a - // foreign absolute path and fail the provider's existing scope check. - if !transports_delegation_context { - workdir_rules.push(WorkdirDelegationRule { - target: WorkdirPath::new_scoped(child_bash_output_dir.to_string_lossy()) - .map_err(|error| ToolError::ExecutionFailed(error.to_string()))?, - permission: WorkdirDelegationPermission::Read, - recursive: true, - }); - } - let delegation_request = workdir_delegation_request(input.cwd.as_deref(), workdir_rules)?; - let workdir_delegation = source_workdir_session - .delegate(delegation_request) + let workdir_tool_broker = require_workdir_tool_broker(self.workdir_tool_broker.as_ref())?; + let tool_scope = workdir_tool_scope(input.cwd.as_deref(), workdir_rules, input.command)?; + let workdir_scope = workdir_tool_broker + .scope(tool_scope) .await .map_err(|error| { - ToolError::InvalidArgument(format!("delegate Workdir session: {error}")) + ToolError::InvalidArgument(format!("scope parent-owned Workdir tools: {error}")) })?; + let child_workdir_tool_broker = workdir_scope.broker(); let spawn_selector = parse_spawn_profile_selector(input.profile.as_deref()).map_err(|msg| { @@ -490,7 +501,6 @@ impl Tool for SubWorkerSpawnTool { ) .await .map_err(|error| ToolError::ExecutionFailed(format!("build Internal Worker: {error}")))?; - child.bind_workdir_session(Some(workdir_delegation.scoped_session.clone())); child .add_scope_rules([ScopeRule { target: child_bash_output_dir.clone(), @@ -510,6 +520,7 @@ impl Tool for SubWorkerSpawnTool { self.runtime_base.clone(), child_registry.clone(), None, + Some(child_workdir_tool_broker.clone()), ) .await .map_err(|error| { @@ -538,20 +549,24 @@ impl Tool for SubWorkerSpawnTool { InternalWorkerSessionStatus::Failed | InternalWorkerSessionStatus::Stopped ) { if let Some(registry) = registry.upgrade() { - if let Err(error) = registry.reclaim_internal_scope(&child_name) { - tracing::warn!( - child_name, - %error, - "failed to reclaim delegated scope after Internal SubWorker failure" - ); - } + let child_name = child_name.clone(); + tokio::spawn(async move { + if let Err(error) = registry.close_internal_scope(&child_name).await { + tracing::warn!( + child_name, + %error, + "failed to close parent-owned Workdir tools after Internal SubWorker failure" + ); + } + }); } } 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); })), + Some(child_workdir_tool_broker.clone()), ) .await; let session = session_result.map_err(|error| { @@ -602,15 +617,19 @@ impl Tool for SubWorkerSpawnTool { ), body.to_string(), ); - let response = self - .workspace_context - .client() - .execute(request) - .map_err(|error| { - ToolError::ExecutionFailed(format!("register review capability: {error}")) - })?; + let response = match self.workspace_context.client().execute(request) { + Ok(response) => response, + Err(error) => { + let _ = session.stop().await; + let _ = workdir_scope.close().await; + return Err(ToolError::ExecutionFailed(format!( + "register review capability: {error}" + ))); + } + }; if !response.is_success() { let _ = session.stop().await; + let _ = workdir_scope.close().await; return Err(ToolError::ExecutionFailed(format!( "register review capability failed with status {}: {}", response.status, response.body @@ -621,14 +640,14 @@ impl Tool for SubWorkerSpawnTool { let record = crate::spawn::registry::InternalSpawnedWorkerRecord::new( input.name.clone(), scope_allow, - workdir_delegation, + workdir_scope, #[cfg(test)] installed_tools, session.clone(), + child_registry, child_change_tracker, ); - if let Err(error) = name_reservation.commit(record) { - let _ = session.stop().await; + if let Err(error) = name_reservation.commit(record).await { return Err(ToolError::ExecutionFailed(format!( "register Internal Worker session: {error}" ))); @@ -674,18 +693,18 @@ fn logical_workdir_path(value: &str, field: &str) -> Result { }) } -fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result, ToolError> { +fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result, ToolError> { if rules.is_empty() { return Err(ToolError::InvalidArgument("scope must not be empty".into())); } rules .iter() .map(|rule| { - Ok(WorkdirDelegationRule { + Ok(WorkdirToolScopeRule { target: logical_workdir_path(&rule.target, "scope.target")?, permission: match rule.permission { - PermissionInput::Read => WorkdirDelegationPermission::Read, - PermissionInput::Write => WorkdirDelegationPermission::Write, + PermissionInput::Read => WorkdirToolScopePermission::Read, + PermissionInput::Write => WorkdirToolScopePermission::Write, }, recursive: rule.recursive, }) @@ -693,22 +712,24 @@ fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result, - rules: Vec, -) -> Result { - Ok(WorkdirDelegationRequest { + rules: Vec, + command: bool, +) -> Result { + Ok(WorkdirToolScope { rules, cwd: logical_workdir_path(cwd.unwrap_or("."), "cwd")?, + command, }) } -fn require_active_workdir_session( - session: Option<&WorkdirSessionHandle>, -) -> Result<&WorkdirSessionHandle, ToolError> { - session.ok_or_else(|| { +fn require_workdir_tool_broker( + broker: Option<&WorkdirToolBroker>, +) -> Result<&WorkdirToolBroker, ToolError> { + broker.ok_or_else(|| { ToolError::InvalidArgument( - "SubWorkerSpawn requires an active Workdir session; attach a Workdir before delegating filesystem access" + "SubWorkerSpawn requires parent-owned Workdir tools; attach a Workdir before granting filesystem access" .to_string(), ) }) @@ -946,7 +967,7 @@ pub(crate) fn sub_worker_spawn_tool( runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - source_workdir_session: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompts: Arc>, @@ -958,7 +979,7 @@ pub(crate) fn sub_worker_spawn_tool( runtime_base, bash_output_dir, workspace_root, - source_workdir_session, + workdir_tool_broker, registry, spawner_manifest, prompts, @@ -972,7 +993,7 @@ fn sub_worker_spawn_tool_impl( runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - source_workdir_session: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompts: Arc>, @@ -1004,7 +1025,7 @@ fn sub_worker_spawn_tool_impl( runtime_base.clone(), bash_output_dir.clone(), workspace_root.clone(), - source_workdir_session.clone(), + workdir_tool_broker.clone(), registry.clone(), spawner_manifest.clone(), prompts.load_full().source(), @@ -1037,12 +1058,12 @@ mod tests { }; #[test] - fn missing_active_workdir_session_fails_deterministically() { - let error = require_active_workdir_session(None).unwrap_err(); + fn missing_parent_workdir_tool_broker_fails_deterministically() { + let error = require_workdir_tool_broker(None).unwrap_err(); assert!(matches!( error, ToolError::InvalidArgument(message) - if message.contains("requires an active Workdir session") + if message.contains("requires parent-owned Workdir tools") )); } @@ -1079,6 +1100,7 @@ mod tests { let valid: SubWorkerSpawnInput = serde_json::from_value(serde_json::json!({ "name":"reviewer","task":"review","profile":"builtin:reviewer", "scope":[{"target":"work","permission":"write"}], + "command":true, "review":{"ticket_id":"T1"} })) .unwrap(); @@ -1134,12 +1156,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] @@ -1173,7 +1224,7 @@ enabled = false let fail_requests = Arc::new(AtomicBool::new(false)); let prompt_loader = PromptCatalogSource::builtins_only(); let (parent_method_tx, mut parent_method_rx) = mpsc::channel(8); - let source_workdir_session = workdir::delegation_capable_session(Arc::new( + let workdir_tool_broker = workdir::WorkdirToolBroker::new(Arc::new( workdir::LocalWorkdirSession::materialized_bound( workdir::Workdir::new("test-workdir"), workspace_root.clone(), @@ -1185,11 +1236,14 @@ 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(), - Some(source_workdir_session), + Some(workdir_tool_broker), registry.clone(), manifest.clone(), prompt_loader, @@ -1212,7 +1266,8 @@ enabled = false "target": ".", "permission": "write", "recursive": true - }] + }], + "command": true }); assert!(spawner_scope.snapshot().is_writable(&workspace_root)); @@ -1247,15 +1302,6 @@ enabled = false let record = registry .get_internal("reviewer-child") .expect("Internal reviewer registry record"); - let child_bash_output_dir = bash_output_dir.join("sub-workers").join("reviewer-child"); - record - .workdir_delegation - .scoped_session - .stat(workdir::StatRequest { - path: WorkdirPath::new_scoped(child_bash_output_dir.to_string_lossy()).unwrap(), - }) - .await - .expect("local child retains read scope for its Bash output directory"); for required in ["Read", "Write", "Edit", "Glob", "Grep", "Bash"] { assert!( record.installed_tools.iter().any(|name| name == required), @@ -1282,10 +1328,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()); @@ -1371,7 +1420,7 @@ enabled = false "Stopped terminal child must release its delegated Workdir session" ); assert!( - !record.workdir_delegation.is_active(), + !record.workdir_tool_scope.is_active(), "stopped child must revoke cloned scoped sessions" ); assert!(registry.get_internal("reviewer-child").is_some()); @@ -1426,7 +1475,7 @@ enabled = false Arc::new(AvailableWorkspaceClient), ); let remote_client = Arc::new(StrictRemoteWorkdirWorkspaceClient::default()); - let source_workdir_session = workdir::delegation_capable_session( + let workdir_tool_broker = workdir::WorkdirToolBroker::new( WorkspaceAttachedWorkdirSession::handle(remote_client.clone()), ); let calls = Arc::new(AtomicUsize::new(0)); @@ -1434,11 +1483,14 @@ 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(), - Some(source_workdir_session), + Some(workdir_tool_broker), registry.clone(), manifest, PromptCatalogSource::builtins_only(), @@ -1478,51 +1530,12 @@ enabled = false record.session.wait_until_idle().await, crate::internal_worker::InternalWorkerSessionStatus::Idle ); + assert!(record.installed_tools.iter().any(|tool| tool == "Write")); + assert!(!record.installed_tools.iter().any(|tool| tool == "Bash")); assert_eq!(calls.load(Ordering::SeqCst), 1); - assert_eq!( - remote_client - .foreign_scope_rejections - .load(Ordering::SeqCst), - 0 - ); - let child_bash_output_dir = bash_output_dir.join("sub-workers").join("remote-child"); - assert!(child_bash_output_dir.is_dir()); - for required in ["Read", "Write", "Edit", "Glob", "Grep", "Bash"] { - assert!( - record.installed_tools.iter().any(|name| name == required), - "remote write-scoped child is missing {required}: {:?}", - record.installed_tools - ); - } - - let remote_requests = remote_client.requests(); - let operate_requests = remote_requests - .iter() - .filter(|request| request.body.is_some()) - .collect::>(); - assert_eq!( - operate_requests.len(), - 1, - "remote requests: {remote_requests:?}" - ); - let operation_body: serde_json::Value = serde_json::from_str( - operate_requests[0] - .body - .as_deref() - .expect("remote operation body"), - ) - .unwrap(); - let rules = operation_body["delegations"][0]["rules"] - .as_array() - .expect("delegation rules"); - assert_eq!(rules.len(), 1, "remote operation body: {operation_body}"); - assert_eq!(rules[0]["target"], ""); assert!( - !operation_body.to_string().contains( - child_bash_output_dir - .to_str() - .expect("UTF-8 test output directory") - ) + remote_client.requests().is_empty(), + "spawning a child must not open or delegate a provider Workdir session" ); } @@ -1534,6 +1547,7 @@ enabled = false .and_then(serde_json::Value::as_object) .expect("schema properties"); assert!(properties.contains_key("cwd"), "schema: {schema}"); + assert!(properties.contains_key("command"), "schema: {schema}"); let required = schema .get("required") .and_then(serde_json::Value::as_array) @@ -1663,7 +1677,6 @@ enabled = false #[derive(Debug, Default)] struct StrictRemoteWorkdirWorkspaceClient { requests: Mutex>, - foreign_scope_rejections: AtomicUsize, } impl StrictRemoteWorkdirWorkspaceClient { @@ -1695,59 +1708,10 @@ enabled = false self.requests .lock() .expect("remote Workdir request lock") - .push(request.clone()); - if request.path.ends_with("/fence") { - return Ok(WorkspaceResponse { - status: 200, - body: serde_json::json!({ "value": "remote-fence-1" }).to_string(), - }); - } - - let body: serde_json::Value = serde_json::from_str( - request - .body - .as_deref() - .ok_or_else(|| WorkspaceClientError::Request("missing request body".into()))?, - ) - .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; - let has_foreign_scope = body - .get("delegations") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .flat_map(|delegation| { - delegation - .get("rules") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - }) - .filter_map(|rule| rule.get("target").and_then(serde_json::Value::as_str)) - .any(|target| Path::new(target).is_absolute()); - if has_foreign_scope { - self.foreign_scope_rejections.fetch_add(1, Ordering::SeqCst); - return Ok(WorkspaceResponse { - status: 403, - body: serde_json::json!({ - "code": "out_of_scope", - "message": "Worker-host path is outside the remote Workdir namespace" - }) - .to_string(), - }); - } - - Ok(WorkspaceResponse { - status: 200, - body: serde_json::json!({ - "operation": "stat", - "result": { - "path": "", - "kind": "directory", - "size": 0 - } - }) - .to_string(), - }) + .push(request); + Err(WorkspaceClientError::Request( + "SubWorker spawn must not call the remote Workdir provider".into(), + )) } } diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 55d041ff..7b051970 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -1,3 +1,4 @@ +use std::collections::VecDeque; #[cfg(test)] use std::path::Path; use std::path::PathBuf; @@ -68,6 +69,243 @@ const LARGE_PASTE_INLINE_MAX_BYTES: usize = 32 * 1024; const WORKER_ORCHESTRATION_INSTRUCTION_ID: &str = "worker.orchestration"; const WORKER_ORCHESTRATION_PROMPT_REF: &str = "common.worker_orchestration"; const FEATURE_HOOK_CHAIN_TIMEOUT: Duration = Duration::from_secs(30); +const SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN: &str = "worker.pending_activations.v1"; +const MAX_PENDING_SUBMISSIONS: usize = 32; +const MAX_PENDING_SUBMISSION_BYTES: u64 = 1024 * 1024; +const MAX_PENDING_ARTIFACT_REFS: usize = 64; +const MAX_ACTIVATION_REQUEST_ID_BYTES: usize = 128; +const MAX_SUBMISSION_RECEIPTS: usize = 128; + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub(crate) struct PendingSubmission { + pub(crate) submission_request_id: String, + source_namespace: String, + pub(crate) submission_id: String, + payload_digest: String, + accepted_at_ms: u64, + activation_sequence: u64, + pub(crate) provenance: WorkerHistoryProvenance, + #[serde(default)] + was_queued: bool, + pub(crate) input: Vec, +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +struct SubmissionReceipt { + submission_request_id: String, + source_namespace: String, + submission_id: String, + payload_digest: String, + disposition: protocol::SubmissionDisposition, +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub(crate) struct PendingNotification { + pub(crate) notification_request_id: String, + source_namespace: String, + pub(crate) message: String, + payload_digest: String, + pub(crate) auto_run: bool, + accepted_at_ms: u64, + activation_sequence: u64, + pub(crate) provenance: WorkerHistoryProvenance, +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +struct NotificationReceipt { + notification_request_id: String, + source_namespace: String, + payload_digest: String, + auto_run: bool, +} + +#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] +pub(crate) struct PendingActivationState { + revision: u64, + next_activation_sequence: u64, + /// A prepared activation remains in checkpoints until the same atomic + /// UserInput record commits the clearing checkpoint. Restore puts it back + /// at the FIFO head. + activating: Option, + activating_notification: Option, + pending: VecDeque, + pending_notifications: VecDeque, + receipts: VecDeque, + notification_receipts: VecDeque, +} + +impl PendingActivationState { + pub(crate) fn snapshot(&self) -> protocol::PendingSubmissionsSnapshot { + let pending_notification = self + .pending_notifications + .iter() + .find(|notification| notification.auto_run); + let head_id = match (self.pending.front(), pending_notification) { + (Some(submission), Some(notification)) + if notification.activation_sequence < submission.activation_sequence => + { + Some(notification_head_id( + ¬ification.source_namespace, + ¬ification.notification_request_id, + )) + } + (Some(submission), _) => Some(submission.submission_id.clone()), + (None, Some(notification)) => Some(notification_head_id( + ¬ification.source_namespace, + ¬ification.notification_request_id, + )), + (None, None) => None, + }; + protocol::PendingSubmissionsSnapshot { + revision: self.revision, + notification_count: u32::try_from(self.pending_notifications.len()).unwrap_or(u32::MAX), + head_id, + submissions: self + .pending + .iter() + .map(|pending| protocol::PendingSubmissionSummary { + submission_id: pending.submission_id.clone(), + accepted_at_ms: pending.accepted_at_ms, + segment_count: u32::try_from(pending.input.len()).unwrap_or(u32::MAX), + byte_len: submission_payload_len(&pending.input), + }) + .collect(), + } + } + + fn live_artifact_pin_owner_ids(&self) -> Vec { + self.activating + .iter() + .chain(self.pending.iter()) + .map(|submission| submission.submission_id.clone()) + .collect() + } + + fn remember_notification_receipt(&mut self, receipt: NotificationReceipt) { + self.notification_receipts.push_back(receipt); + while self.notification_receipts.len() > MAX_SUBMISSION_RECEIPTS { + self.notification_receipts.pop_front(); + } + } + + fn remember_receipt(&mut self, receipt: SubmissionReceipt) { + self.receipts.push_back(receipt); + while self.receipts.len() > MAX_SUBMISSION_RECEIPTS { + self.receipts.pop_front(); + } + } +} + +fn notification_payload_digest(message: &str, auto_run: bool) -> String { + use sha2::Digest as _; + let mut hasher = sha2::Sha256::new(); + hasher.update(if auto_run { + &b"auto\0"[..] + } else { + &b"deferred\0"[..] + }); + hasher.update(message.as_bytes()); + hasher + .finalize() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +pub(crate) fn authenticated_input_provenance( + source: &protocol::AuthenticatedInputSource, +) -> WorkerHistoryProvenance { + match source { + protocol::AuthenticatedInputSource::UntrustedWire => WorkerHistoryProvenance::LegacyUnknown, + protocol::AuthenticatedInputSource::Account { account_id } => { + WorkerHistoryProvenance::HumanInput { + account_id: account_id.clone(), + } + } + protocol::AuthenticatedInputSource::Worker { + runtime_id, + worker_id, + } => WorkerHistoryProvenance::WorkerInput { + actor: session_store::LoggedWorkerSubject { + workspace_id: None, + runtime_id: Some(runtime_id.clone()), + worker_id: worker_id.clone(), + }, + }, + protocol::AuthenticatedInputSource::SubWorker { session_id } => { + WorkerHistoryProvenance::WorkerInput { + actor: session_store::LoggedWorkerSubject { + workspace_id: None, + runtime_id: None, + worker_id: session_id.clone(), + }, + } + } + protocol::AuthenticatedInputSource::Backend { operation_id } => { + WorkerHistoryProvenance::BackendInstruction { + operation_id: Some(operation_id.clone()), + } + } + } +} + +fn notification_head_id(source_namespace: &str, request_id: &str) -> String { + use sha2::Digest as _; + let mut hasher = sha2::Sha256::new(); + hasher.update(source_namespace.as_bytes()); + hasher.update(b"\0"); + hasher.update(request_id.as_bytes()); + format!( + "notification:{}", + hasher + .finalize() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + ) +} + +fn submission_payload_len(input: &[Segment]) -> u64 { + serde_json::to_vec(input) + .map(|bytes| u64::try_from(bytes.len()).unwrap_or(u64::MAX)) + .unwrap_or(u64::MAX) +} + +fn submission_payload_digest(input: &[Segment]) -> String { + use sha2::Digest as _; + sha2::Sha256::digest(serde_json::to_vec(input).unwrap_or_default()) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +fn submission_uploaded_file_refs( + input: &[Segment], +) -> impl Iterator { + input.iter().filter_map(|segment| match segment { + Segment::UploadedFile { file } => Some(file), + _ => None, + }) +} + +fn submission_artifact_ref_count(input: &[Segment]) -> usize { + input + .iter() + .filter(|segment| { + matches!( + segment, + Segment::PasteArtifact { .. } | Segment::UploadedFile { .. } + ) + }) + .count() +} + +fn pending_activation_extension(state: &PendingActivationState) -> SessionExtension { + SessionExtension { + domain: SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.to_owned(), + payload: serde_json::to_value(state).expect("pending activation state must serialize"), + } +} fn hook_run_exit(exit: &EngineRunExit) -> RunCommittedExit { match exit { @@ -970,18 +1208,740 @@ where } } -/// Type-erased commit handle for the interceptor. Lets the -/// interceptor commit `SystemItem`s without being generic over the +#[derive(Debug, Clone)] +pub(crate) enum PendingActivation { + Submission(PendingSubmission), + Notification(PendingNotification), +} + +#[derive(Debug, Clone)] +pub(crate) struct SubmissionAcceptance { + pub(crate) submission_request_id: String, + pub(crate) submission_id: String, + pub(crate) disposition: protocol::SubmissionDisposition, + pub(crate) activation: Option, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum PendingSubmissionError { + #[error("submission_request_id must not be empty")] + EmptyRequestId, + #[error("activation request id exceeds {MAX_ACTIVATION_REQUEST_ID_BYTES} bytes")] + RequestIdLimit, + #[error("submission input must contain at least one typed segment")] + EmptyInput, + #[error("submission request id was already used with a different payload")] + IdempotencyConflict, + #[error("pending submission queue is full (maximum {MAX_PENDING_SUBMISSIONS})")] + CountLimit, + #[error("pending submission bytes exceed {MAX_PENDING_SUBMISSION_BYTES}")] + ByteLimit, + #[error("pending submission artifact references exceed {MAX_PENDING_ARTIFACT_REFS}")] + ArtifactLimit, + #[error("pending queue revision conflict: expected {expected}, current {current}")] + RevisionConflict { expected: u64, current: u64 }, + #[error("pending queue head conflict: expected {expected}, current {current:?}")] + HeadConflict { + expected: String, + current: Option, + }, + #[error("pending submission not found: {0}")] + NotFound(String), + #[error("pending submission state persistence failed: {0}")] + Store(#[from] StoreError), +} + +#[derive(Clone)] +pub(crate) struct PendingSubmissionHandle { + state: Arc>, + writer: LogWriterHandle, +} + +impl PendingSubmissionHandle +where + St: Store + Clone, +{ + fn validate_fence( + state: &PendingActivationState, + expected_revision: u64, + expected_head_id: Option<&str>, + ) -> Result<(), PendingSubmissionError> { + if state.revision != expected_revision { + return Err(PendingSubmissionError::RevisionConflict { + expected: expected_revision, + current: state.revision, + }); + } + if let Some(expected) = expected_head_id { + let current = state.snapshot().head_id; + if current.as_deref() != Some(expected) { + return Err(PendingSubmissionError::HeadConflict { + expected: expected.to_owned(), + current, + }); + } + } + Ok(()) + } + + fn persist_locked(&self, state: &PendingActivationState) -> Result<(), PendingSubmissionError> { + self.writer.append_entry_locked(LogEntry::Extension { + ts: segment_log::now_millis(), + domain: SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.to_owned(), + payload: serde_json::to_value(state).expect("pending activation state must serialize"), + })?; + Ok(()) + } + + fn pin_submission_files( + &self, + pending: &PendingSubmission, + ) -> Result<(), PendingSubmissionError> { + let session_id = self.writer.state.location().session_id; + let mut pinned = Vec::new(); + for reference in submission_uploaded_file_refs(&pending.input) { + if pinned.iter().any(|existing: &protocol::UploadedFileRef| { + existing.artifact_id == reference.artifact_id + }) { + continue; + } + if let Err(pin_error) = + self.writer + .store + .pin_uploaded_file(session_id, reference, &pending.submission_id) + { + let mut rollback_error = None; + for acquired in pinned.iter().rev() { + if let Err(error) = self.writer.store.release_uploaded_file_pin( + session_id, + &acquired.artifact_id, + &pending.submission_id, + ) { + rollback_error.get_or_insert(error); + } + } + return Err(rollback_error.unwrap_or(pin_error).into()); + } + pinned.push(reference.clone()); + } + Ok(()) + } + + fn release_submission_files( + &self, + pending: &PendingSubmission, + ) -> Result<(), PendingSubmissionError> { + let session_id = self.writer.state.location().session_id; + let mut released = Vec::new(); + let mut first_error = None; + for reference in submission_uploaded_file_refs(&pending.input) { + if released + .iter() + .any(|artifact_id: &String| artifact_id == &reference.artifact_id) + { + continue; + } + if let Err(error) = self.writer.store.release_uploaded_file_pin( + session_id, + &reference.artifact_id, + &pending.submission_id, + ) { + first_error.get_or_insert(error); + } + released.push(reference.artifact_id.clone()); + } + match first_error { + Some(error) => Err(error.into()), + None => Ok(()), + } + } + + #[cfg(test)] + pub(crate) fn accept( + &self, + submission_request_id: String, + input: Vec, + activate_now: bool, + ) -> Result { + self.accept_from_source( + submission_request_id, + input, + self.direct_client_namespace(), + WorkerHistoryProvenance::LegacyUnknown, + activate_now, + ) + } + + pub(crate) fn accept_from_source( + &self, + submission_request_id: String, + input: Vec, + source_namespace: String, + provenance: WorkerHistoryProvenance, + activate_now: bool, + ) -> Result { + if submission_request_id.trim().is_empty() { + return Err(PendingSubmissionError::EmptyRequestId); + } + if submission_request_id.len() > MAX_ACTIVATION_REQUEST_ID_BYTES { + return Err(PendingSubmissionError::RequestIdLimit); + } + if input.is_empty() { + return Err(PendingSubmissionError::EmptyInput); + } + let payload_digest = submission_payload_digest(&input); + let _append_guard = self + .writer + .state + .append_lock + .lock() + .expect("segment append lock poisoned"); + let mut current = self + .state + .lock() + .expect("pending activation state poisoned"); + let original = current.clone(); + if let Some(receipt) = current.receipts.iter().find(|receipt| { + receipt.submission_request_id == submission_request_id + && receipt.source_namespace == source_namespace + }) { + if receipt.payload_digest != payload_digest { + return Err(PendingSubmissionError::IdempotencyConflict); + } + return Ok(SubmissionAcceptance { + submission_request_id, + submission_id: receipt.submission_id.clone(), + disposition: receipt.disposition, + activation: None, + }); + } + + let submission_id = uuid::Uuid::now_v7().to_string(); + let pending = PendingSubmission { + submission_request_id: submission_request_id.clone(), + source_namespace: source_namespace.clone(), + submission_id: submission_id.clone(), + payload_digest: payload_digest.clone(), + accepted_at_ms: segment_log::now_millis(), + activation_sequence: current.next_activation_sequence, + provenance, + was_queued: !activate_now, + input, + }; + current.next_activation_sequence = current.next_activation_sequence.saturating_add(1); + let disposition = if activate_now { + protocol::SubmissionDisposition::Started + } else { + protocol::SubmissionDisposition::Queued + }; + current.remember_receipt(SubmissionReceipt { + submission_request_id: submission_request_id.clone(), + source_namespace, + submission_id: submission_id.clone(), + payload_digest, + disposition, + }); + current.revision = current.revision.saturating_add(1); + + if activate_now { + current.activating = Some(pending.clone()); + } else { + let count = current + .pending + .len() + .saturating_add(current.pending_notifications.len()) + .saturating_add(1); + if count > MAX_PENDING_SUBMISSIONS { + *current = original; + return Err(PendingSubmissionError::CountLimit); + } + let bytes = current + .pending + .iter() + .map(|pending| submission_payload_len(&pending.input)) + .sum::() + .saturating_add( + current + .pending_notifications + .iter() + .map(|pending| u64::try_from(pending.message.len()).unwrap_or(u64::MAX)) + .sum::(), + ) + .saturating_add(submission_payload_len(&pending.input)); + if bytes > MAX_PENDING_SUBMISSION_BYTES { + *current = original; + return Err(PendingSubmissionError::ByteLimit); + } + let artifact_refs = current + .pending + .iter() + .map(|pending| submission_artifact_ref_count(&pending.input)) + .sum::() + .saturating_add(submission_artifact_ref_count(&pending.input)); + if artifact_refs > MAX_PENDING_ARTIFACT_REFS { + *current = original; + return Err(PendingSubmissionError::ArtifactLimit); + } + current.pending.push_back(pending.clone()); + } + if !activate_now { + if let Err(error) = self.pin_submission_files(&pending) { + *current = original; + return Err(error); + } + if let Err(error) = self.persist_locked(¤t) { + let _ = self.release_submission_files(&pending); + *current = original; + return Err(error); + } + } + Ok(SubmissionAcceptance { + submission_request_id, + submission_id, + disposition, + activation: activate_now.then_some(pending), + }) + } + + #[cfg(test)] + pub(crate) fn accept_notification( + &self, + notification_request_id: String, + message: String, + auto_run: bool, + ) -> Result { + self.accept_notification_from_source( + notification_request_id, + message, + self.direct_client_namespace(), + WorkerHistoryProvenance::LegacyUnknown, + auto_run, + ) + } + + pub(crate) fn accept_notification_from_source( + &self, + notification_request_id: String, + message: String, + source_namespace: String, + provenance: WorkerHistoryProvenance, + auto_run: bool, + ) -> Result { + if notification_request_id.trim().is_empty() { + return Err(PendingSubmissionError::EmptyRequestId); + } + if notification_request_id.len() > MAX_ACTIVATION_REQUEST_ID_BYTES { + return Err(PendingSubmissionError::RequestIdLimit); + } + let payload_digest = notification_payload_digest(&message, auto_run); + let _append_guard = self + .writer + .state + .append_lock + .lock() + .expect("segment append lock poisoned"); + let mut state = self + .state + .lock() + .expect("pending activation state poisoned"); + if let Some(receipt) = state.notification_receipts.iter().find(|receipt| { + receipt.notification_request_id == notification_request_id + && receipt.source_namespace == source_namespace + }) { + if receipt.payload_digest != payload_digest || receipt.auto_run != auto_run { + return Err(PendingSubmissionError::IdempotencyConflict); + } + return Ok(false); + } + if state + .pending + .len() + .saturating_add(state.pending_notifications.len()) + >= MAX_PENDING_SUBMISSIONS + { + return Err(PendingSubmissionError::CountLimit); + } + let queued_bytes = state + .pending + .iter() + .map(|pending| submission_payload_len(&pending.input)) + .sum::() + .saturating_add( + state + .pending_notifications + .iter() + .map(|pending| u64::try_from(pending.message.len()).unwrap_or(u64::MAX)) + .sum::(), + ) + .saturating_add(u64::try_from(message.len()).unwrap_or(u64::MAX)); + if queued_bytes > MAX_PENDING_SUBMISSION_BYTES { + return Err(PendingSubmissionError::ByteLimit); + } + let original = state.clone(); + let activation_sequence = state.next_activation_sequence; + state.next_activation_sequence = state.next_activation_sequence.saturating_add(1); + state.pending_notifications.push_back(PendingNotification { + notification_request_id: notification_request_id.clone(), + source_namespace: source_namespace.clone(), + message, + payload_digest: payload_digest.clone(), + auto_run, + accepted_at_ms: segment_log::now_millis(), + activation_sequence, + provenance, + }); + state.remember_notification_receipt(NotificationReceipt { + notification_request_id, + source_namespace, + payload_digest, + auto_run, + }); + state.revision = state.revision.saturating_add(1); + if let Err(error) = self.persist_locked(&state) { + *state = original; + return Err(error); + } + Ok(true) + } + + pub(crate) fn activating_passive_notification_id(&self) -> Option { + self.state + .lock() + .expect("pending activation state poisoned") + .activating_notification + .as_ref() + .filter(|notification| !notification.auto_run) + .map(|notification| notification.notification_request_id.clone()) + } + + pub(crate) fn next_passive_notification_identity(&self) -> Option<(String, String)> { + self.state + .lock() + .expect("pending activation state poisoned") + .pending_notifications + .iter() + .find(|notification| !notification.auto_run) + .map(|notification| { + ( + notification.source_namespace.clone(), + notification.notification_request_id.clone(), + ) + }) + } + + pub(crate) fn prepare_notification( + &self, + source_namespace: &str, + notification_request_id: &str, + ) -> Option { + let mut state = self + .state + .lock() + .expect("pending activation state poisoned"); + if state.activating_notification.is_some() { + return None; + } + let index = state + .pending_notifications + .iter() + .position(|notification| { + notification.notification_request_id == notification_request_id + && notification.source_namespace == source_namespace + })?; + let notification = state + .pending_notifications + .remove(index) + .expect("located pending notification must exist"); + state.activating_notification = Some(notification.clone()); + state.revision = state.revision.saturating_add(1); + Some(notification) + } + + pub(crate) fn prepare_next_activation( + &self, + fence: Option<(u64, &str)>, + ) -> Result, PendingSubmissionError> { + let _append_guard = self + .writer + .state + .append_lock + .lock() + .expect("segment append lock poisoned"); + let mut state = self + .state + .lock() + .expect("pending activation state poisoned"); + if let Some((expected_revision, expected_head_id)) = fence { + Self::validate_fence(&state, expected_revision, Some(expected_head_id))?; + } + if state.activating.is_some() + || state + .activating_notification + .as_ref() + .is_some_and(|notification| notification.auto_run) + { + return Ok(None); + } + let has_staged_passive_notification = state.activating_notification.is_some(); + let submission_sequence = state.pending.front().map(|item| item.activation_sequence); + let notification_index = if has_staged_passive_notification { + None + } else { + state + .pending_notifications + .iter() + .position(|item| item.auto_run) + }; + let notification_sequence = notification_index + .and_then(|index| state.pending_notifications.get(index)) + .map(|item| item.activation_sequence); + if notification_sequence.is_some() + && (submission_sequence.is_none() || notification_sequence < submission_sequence) + { + let notification = state + .pending_notifications + .remove(notification_index.expect("notification sequence came from an item")) + .expect("notification sequence came from an existing item"); + state.activating_notification = Some(notification.clone()); + state.revision = state.revision.saturating_add(1); + return Ok(Some(PendingActivation::Notification(notification))); + } + if submission_sequence.is_some() { + let pending = state + .pending + .pop_front() + .expect("submission sequence came from queue head"); + state.activating = Some(pending.clone()); + state.revision = state.revision.saturating_add(1); + return Ok(Some(PendingActivation::Submission(pending))); + } + Ok(None) + } + + pub(crate) fn abort_activation(&self, pending: PendingSubmission) { + let _append_guard = self + .writer + .state + .append_lock + .lock() + .expect("segment append lock poisoned"); + let mut state = self + .state + .lock() + .expect("pending activation state poisoned"); + if state.activating.as_ref().map(|item| &item.submission_id) != Some(&pending.submission_id) + { + return; + } + state.activating = None; + if pending.was_queued { + state.pending.push_front(pending.clone()); + } else { + state + .receipts + .retain(|receipt| receipt.submission_id != pending.submission_id); + } + state.revision = state.revision.saturating_add(1); + if let Err(error) = self.persist_locked(&state) { + tracing::error!(error = %error, "failed to persist aborted pending activation"); + } + } + + pub(crate) fn activation_extension(&self) -> SessionExtension { + let state = self + .state + .lock() + .expect("pending activation state poisoned"); + let mut committed = state.clone(); + if let Some(activating) = &committed.activating + && let Some(receipt) = committed + .receipts + .iter_mut() + .find(|receipt| receipt.submission_id == activating.submission_id) + { + receipt.disposition = protocol::SubmissionDisposition::Started; + } + committed.activating = None; + committed.revision = committed.revision.saturating_add(1); + pending_activation_extension(&committed) + } + + pub(crate) fn finish_activation(&self, submission_id: &str) { + let mut state = self + .state + .lock() + .expect("pending activation state poisoned"); + if state + .activating + .as_ref() + .map(|item| item.submission_id.as_str()) + == Some(submission_id) + { + if let Some(receipt) = state + .receipts + .iter_mut() + .find(|receipt| receipt.submission_id == submission_id) + { + receipt.disposition = protocol::SubmissionDisposition::Started; + } + state.activating = None; + state.revision = state.revision.saturating_add(1); + } + } + + pub(crate) fn notification_activation_extension(&self) -> SessionExtension { + let state = self + .state + .lock() + .expect("pending activation state poisoned"); + let mut committed = state.clone(); + committed.activating = None; + committed.activating_notification = None; + committed.revision = committed.revision.saturating_add(1); + pending_activation_extension(&committed) + } + + pub(crate) fn finish_notification_activation(&self, notification_request_id: &str) { + let mut state = self + .state + .lock() + .expect("pending activation state poisoned"); + if state + .activating_notification + .as_ref() + .map(|item| item.notification_request_id.as_str()) + == Some(notification_request_id) + { + state.activating_notification = None; + state.revision = state.revision.saturating_add(1); + } + } + + pub(crate) fn direct_client_namespace(&self) -> String { + format!("direct:{}", self.writer.state.location().session_id) + } + + pub(crate) fn snapshot(&self) -> protocol::PendingSubmissionsSnapshot { + self.state + .lock() + .expect("pending activation state poisoned") + .snapshot() + } + + fn reconcile_uploaded_file_pins(&self) -> Result { + let live_owner_ids = self + .state + .lock() + .expect("pending activation state poisoned") + .live_artifact_pin_owner_ids(); + Ok(self + .writer + .store + .reconcile_uploaded_file_pins(self.writer.state.session_id(), &live_owner_ids)?) + } + + pub(crate) fn cancel( + &self, + submission_id: &str, + expected_revision: u64, + ) -> Result { + let _append_guard = self + .writer + .state + .append_lock + .lock() + .expect("segment append lock poisoned"); + let mut state = self + .state + .lock() + .expect("pending activation state poisoned"); + Self::validate_fence(&state, expected_revision, None)?; + let original = state.clone(); + let Some(index) = state + .pending + .iter() + .position(|pending| pending.submission_id == submission_id) + else { + return Err(PendingSubmissionError::NotFound(submission_id.to_owned())); + }; + let removed = state + .pending + .remove(index) + .expect("located pending submission must exist"); + state.revision = state.revision.saturating_add(1); + if let Err(error) = self.persist_locked(&state) { + *state = original; + return Err(error); + } + self.release_submission_files(&removed)?; + Ok(state.snapshot()) + } + + pub(crate) fn clear( + &self, + expected_revision: u64, + ) -> Result { + let _append_guard = self + .writer + .state + .append_lock + .lock() + .expect("segment append lock poisoned"); + let mut state = self + .state + .lock() + .expect("pending activation state poisoned"); + Self::validate_fence(&state, expected_revision, None)?; + let original = state.clone(); + let removed = state.pending.drain(..).collect::>(); + state.pending_notifications.clear(); + state.revision = state.revision.saturating_add(1); + if let Err(error) = self.persist_locked(&state) { + *state = original; + return Err(error); + } + for pending in &removed { + self.release_submission_files(pending)?; + } + Ok(state.snapshot()) + } +} + +impl PendingSubmissionHandle { + #[cfg(test)] + pub(crate) fn for_test(root: &std::path::Path) -> Self { + let store = session_store::FsStore::new(root).expect("test session store"); + let session_id = session_store::new_session_id(); + let segment_id = session_store::new_segment_id(); + store + .create_segment(session_id, segment_id, &[]) + .expect("test session segment"); + Self { + state: Arc::new(Mutex::new(PendingActivationState::default())), + writer: LogWriterHandle { + store, + state: SegmentState::new(session_id, segment_id, 0), + sink: SegmentLogSink::new(), + in_flight: None, + }, + } + } +} + +/// Type-erased commit handle for the interceptor. Lets the interceptor commit `SystemItem`s without being generic over the /// concrete `Store` type. pub trait SystemItemCommitter: Send + Sync { fn commit_log_entry(&self, entry: LogEntry) -> Result<(), StoreError>; - fn commit_system_item( + fn commit_system_item_with_extensions( &self, item: SystemItem, + extensions: Vec, + history_provenance: Option, ) -> Result, StoreError> { let metadata = new_history_metadata( - WorkerHistoryProvenance::BackendInstruction { operation_id: None }, + history_provenance + .unwrap_or(WorkerHistoryProvenance::BackendInstruction { operation_id: None }), None, ); let history_item = item.to_history_item(); @@ -991,6 +1951,7 @@ pub trait SystemItemCommitter: Send + Sync { item, metadata: metadata.clone(), }, + extensions, })?; Ok(HistoryEntry::new(history_item, metadata)) } @@ -1027,8 +1988,6 @@ where } } -pub const WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN: &str = "worker.input-submission.v1"; - #[derive(Clone)] struct PreparedFlowProjection { selector: String, @@ -1049,6 +2008,7 @@ pub struct WorkerSession { session_id: SessionId, revision: u64, history: History, + pending_activations: Arc>, } impl WorkerSession { @@ -1058,9 +2018,39 @@ impl WorkerSession { session_id, revision, history: History::from_entries(entries), + pending_activations: Arc::new(Mutex::new(PendingActivationState::default())), } } + fn restore_pending_activations(&mut self, extensions: &[(String, serde_json::Value)]) { + let Some(payload) = extensions.iter().rev().find_map(|(domain, payload)| { + (domain == SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN).then_some(payload) + }) else { + return; + }; + if let Ok(mut state) = serde_json::from_value::(payload.clone()) { + if let Some(activating) = state.activating.take() { + state.pending.push_front(activating); + state.revision = state.revision.saturating_add(1); + } + if let Some(activating) = state.activating_notification.take() { + state.pending_notifications.push_front(activating); + state.revision = state.revision.saturating_add(1); + } + *self + .pending_activations + .lock() + .expect("pending activation state poisoned") = state; + } + } + + pub fn pending_submissions(&self) -> protocol::PendingSubmissionsSnapshot { + self.pending_activations + .lock() + .expect("pending activation state poisoned") + .snapshot() + } + pub fn session_id(&self) -> SessionId { self.session_id } @@ -1314,6 +2304,21 @@ impl Worker { } } + pub(crate) fn pending_activation_state(&self) -> Arc> { + self.session.pending_activations.clone() + } + + pub(crate) fn pending_submission_handle(&self) -> PendingSubmissionHandle { + PendingSubmissionHandle { + state: self.session.pending_activations.clone(), + writer: self.log_writer_handle(), + } + } + + pub fn pending_submissions(&self) -> protocol::PendingSubmissionsSnapshot { + self.session.pending_submissions() + } + /// Attach a type-erased system-item commit handle. The controller /// calls this once during spawn so the interceptor can commit /// `SystemItem`s directly without owning a generic store handle. @@ -1675,6 +2680,7 @@ impl Worker { }, metadata: skill_metadata.clone(), }, + extensions: Vec::new(), })?; let history_entry = HistoryEntry::new(agen::Item::system_message(body), skill_metadata); let mut annotate = history_annotator( @@ -1965,6 +2971,30 @@ impl Worker { .truncate(loc.session_id, loc.segment_id, truncate_entries)?; self.segment_state.set_entries_written(truncate_entries); self.sink.truncate_silent(truncate_entries); + let pending_state = self + .session + .pending_activations + .lock() + .expect("pending activation state poisoned") + .clone(); + if !pending_state.pending.is_empty() + || !pending_state.pending_notifications.is_empty() + || pending_state.activating.is_some() + || pending_state.activating_notification.is_some() + || !pending_state.receipts.is_empty() + || !pending_state.notification_receipts.is_empty() + { + let checkpoint = LogEntry::Extension { + ts: segment_log::now_millis(), + domain: SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.to_owned(), + payload: serde_json::to_value(&pending_state).map_err(|error| { + RewindError::Invalid(format!( + "serialize pending submissions during rewind: {error}" + )) + })?, + }; + self.commit_entry(checkpoint)?; + } let history_entries = restore_history_entries(loc.session_id, loc.segment_id, &retained) .map_err(|error| RewindError::Invalid(error.into()))?; @@ -2534,7 +3564,7 @@ impl Worker { /// Convenience: run with a single `Segment::Text`. /// /// Equivalent to `run(vec![Segment::text(s)])`. The dumb-client - /// counterpart of [`protocol::Method::run_text`]; primarily for + /// counterpart of [`protocol::Method::submit_text`]; primarily for /// tests and tools that have only a string in hand. pub async fn run_text(&mut self, s: impl Into) -> Result where @@ -2785,8 +3815,13 @@ impl Worker { where St: Clone + 'static, { - self.run_with_input_extensions_and_commit_hook(input, input_extensions, || {}) - .await + self.run_with_input_extensions_and_commit_hook( + input, + input_extensions, + WorkerHistoryProvenance::LegacyUnknown, + || {}, + ) + .await } /// Run user input and invoke `on_input_committed` only after the annotated @@ -2797,6 +3832,7 @@ impl Worker { &mut self, input: Vec, mut input_extensions: Vec, + input_provenance: WorkerHistoryProvenance, on_input_committed: F, ) -> Result where @@ -2844,8 +3880,12 @@ impl Worker { trigger: protocol::InvokeKind::UserSend, })?; - let projected_input = - self.projected_input_history(&input, flow_projection.as_ref(), &projected_entry_ids); + let projected_input = self.projected_input_history( + &input, + flow_projection.as_ref(), + &projected_entry_ids, + &input_provenance, + ); // Persist original typed segments together with the exact ordered // model-visible item+origin projection before any entry becomes live. @@ -2858,6 +3898,11 @@ impl Worker { .map(to_logged_history_entry) .collect(), })?; + self.finalize_uploaded_segment_bindings( + &input, + &projected_entry_ids, + flow_projection.is_some(), + ); if let Some(state) = pending_flow_state { *self .flow_runtime_state @@ -3051,6 +4096,7 @@ impl Worker { }, metadata: interrupt_metadata.clone(), }, + extensions: Vec::new(), })?; let interrupt_entry = HistoryEntry::new(agen::Item::system_message(system_note), interrupt_metadata); @@ -3079,6 +4125,36 @@ impl Worker { Ok(()) } + fn finalize_uploaded_segment_bindings( + &self, + input: &[Segment], + projected_entry_ids: &[SessionHistoryEntryId], + one_entry_per_segment: bool, + ) { + for (index, segment) in input.iter().enumerate() { + let Segment::UploadedFile { file } = segment else { + continue; + }; + let entry_index = if one_entry_per_segment { index } else { 0 }; + let source_entry_id = projected_entry_ids + .get(entry_index) + .expect("projected input id exists for every uploaded file") + .0 + .as_str(); + if let Err(error) = self.store.finalize_uploaded_file_binding( + self.session_id(), + &file.artifact_id, + source_entry_id, + ) { + tracing::warn!( + artifact_id = %file.artifact_id, + error = %error, + "deferred uploaded file pin finalization to cleanup reconciliation" + ); + } + } + } + fn materialize_large_pastes( &self, input: &mut [Segment], @@ -3150,6 +4226,7 @@ impl Worker { input: &[Segment], flow_projection: Option<&PreparedFlowProjection>, entry_ids: &[SessionHistoryEntryId], + provenance: &WorkerHistoryProvenance, ) -> Vec> { if let Some(flow) = flow_projection { return input @@ -3170,10 +4247,7 @@ impl Worker { other => history_entry_with_id( Item::user_message(Segment::flatten_to_text(std::slice::from_ref(other))), entry_id.clone(), - // Current public submit transport does not carry a - // trusted account/Worker subject envelope. Fail closed - // instead of promoting role=user to HumanInput. - WorkerHistoryProvenance::LegacyUnknown, + provenance.clone(), ), }) .collect(); @@ -3185,7 +4259,7 @@ impl Worker { .first() .expect("projected Worker input always has one entry id") .clone(), - WorkerHistoryProvenance::LegacyUnknown, + provenance.clone(), )] } @@ -4437,6 +5511,22 @@ impl Worker { { initial_entries.push(checkpoint); } + initial_entries.push(LogEntry::Extension { + ts: segment_log::now_millis(), + domain: SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.to_owned(), + payload: serde_json::to_value( + &*self + .session + .pending_activations + .lock() + .expect("pending activation state poisoned"), + ) + .map_err(|error| { + WorkerError::InvalidState(format!( + "serialize pending submissions during compaction: {error}" + )) + })?, + }); if let Some(flow_state) = self .flow_runtime_state .lock() @@ -5257,6 +6347,12 @@ where history_persistence_wired: false, log_writer: None, }; + worker + .session + .restore_pending_activations(&state.extensions); + worker + .pending_submission_handle() + .reconcile_uploaded_file_pins()?; worker.apply_permissions_from_manifest(); worker.apply_prune_from_manifest(); worker.write_worker_metadata_active(SegmentLocation { @@ -7202,8 +8298,15 @@ mod build_summary_prompt_tests { serde_json::to_value(&state).unwrap(), ); let projected_ids = vec![SessionHistoryEntryId::new(), SessionHistoryEntryId::new()]; - let projected = - worker.projected_input_history(&segments, projection.as_ref(), &projected_ids); + let input_provenance = WorkerHistoryProvenance::HumanInput { + account_id: "account-1".into(), + }; + let projected = worker.projected_input_history( + &segments, + projection.as_ref(), + &projected_ids, + &input_provenance, + ); worker .commit_entry(LogEntry::AnnotatedUserInput { ts: segment_log::now_millis(), @@ -7230,6 +8333,7 @@ mod build_summary_prompt_tests { projected[0].annotation.origin, WorkerHistoryProvenance::FlowInstruction { .. } )); + assert_eq!(projected[1].annotation.origin, input_provenance); assert_eq!(state.instance.definition_revision, 3); assert_eq!(state.instance.current_state.as_str(), "implement"); assert_eq!(workspace_client.requests.lock().unwrap().len(), 1); @@ -7312,7 +8416,12 @@ mod build_summary_prompt_tests { .delete_uploaded_file(worker.session_id(), &file.artifact_id), Err(StoreError::ArtifactAlreadyCommitted) )); - let projected = worker.projected_input_history(&input, None, &[entry_id]); + let projected = worker.projected_input_history( + &input, + None, + &[entry_id], + &WorkerHistoryProvenance::LegacyUnknown, + ); let text = projected[0].item.as_text().unwrap(); assert!(text.contains("notes.md")); assert!(text.contains(&file.artifact_id)); @@ -7394,7 +8503,12 @@ mod build_summary_prompt_tests { if retained.source_entry_id == artifact.source_entry_id )); - let history = worker.projected_input_history(&input, None, &[entry_id]); + let history = worker.projected_input_history( + &input, + None, + &[entry_id], + &WorkerHistoryProvenance::LegacyUnknown, + ); assert!(!history[0].item.as_text().unwrap().contains("終端")); append_test_entry( &worker, @@ -7670,6 +8784,47 @@ mod build_summary_prompt_tests { assert_eq!(worker.history()[0].as_text().unwrap(), "first message"); } + #[tokio::test] + async fn rewind_preserves_notification_only_pending_activation_checkpoint() { + let (_dir, mut worker) = rewind_test_worker().await; + append_user_turn(&worker, 10, "first message"); + append_user_turn(&worker, 20, "second message"); + worker + .pending_submission_handle() + .accept_notification("notification-1".into(), "keep me".into(), true) + .unwrap(); + let (head_entries, targets) = worker.list_rewind_targets().unwrap(); + + worker + .rewind_to(targets.last().unwrap().id.clone(), head_entries) + .await + .unwrap(); + + let location = worker.segment_state.location(); + let entries = worker + .store + .read_all(location.session_id, location.segment_id) + .unwrap(); + let restored: PendingActivationState = entries + .iter() + .rev() + .find_map(|entry| match entry { + LogEntry::Extension { + domain, payload, .. + } if domain == SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN => { + serde_json::from_value(payload.clone()).ok() + } + _ => None, + }) + .unwrap(); + assert_eq!(restored.pending_notifications.len(), 1); + assert_eq!(restored.notification_receipts.len(), 1); + assert_eq!( + restored.pending_notifications[0].notification_request_id, + "notification-1" + ); + } + #[tokio::test] async fn annotated_history_rewind_commits_authoritative_prefix() { let (_dir, mut worker) = rewind_test_worker().await; @@ -8515,6 +9670,501 @@ mod build_summary_prompt_tests { ); } + #[test] + fn submission_retry_identity_is_scoped_to_authenticated_source_and_keeps_provenance() { + let temp = tempfile::tempdir().unwrap(); + let handle = PendingSubmissionHandle::for_test(temp.path()); + let input = vec![Segment::text("same request")]; + let account_a = WorkerHistoryProvenance::HumanInput { + account_id: "account-a".into(), + }; + let account_b = WorkerHistoryProvenance::HumanInput { + account_id: "account-b".into(), + }; + let first = handle + .accept_from_source( + "request-1".into(), + input.clone(), + "account:account-a".into(), + account_a.clone(), + false, + ) + .unwrap(); + let replay = handle + .accept_from_source( + "request-1".into(), + input.clone(), + "account:account-a".into(), + account_a.clone(), + false, + ) + .unwrap(); + let other_source = handle + .accept_from_source( + "request-1".into(), + input, + "account:account-b".into(), + account_b.clone(), + false, + ) + .unwrap(); + + assert_eq!(replay.submission_id, first.submission_id); + assert_ne!(other_source.submission_id, first.submission_id); + let state = handle.state.lock().unwrap(); + assert_eq!(state.pending[0].provenance, account_a); + assert_eq!(state.pending[1].provenance, account_b); + } + + #[test] + fn restore_reconciliation_clears_interrupted_acceptance_pin_for_retry() { + let temp = tempfile::tempdir().unwrap(); + let handle = PendingSubmissionHandle::for_test(temp.path()); + let session_id = handle.writer.state.session_id(); + let segment_id = handle.writer.state.location().segment_id; + let limits = session_store::UploadedFileLimits { + max_file_bytes: 1024, + max_session_bytes: 2048, + }; + let file = handle + .writer + .store + .write_uploaded_file(session_id, "retry.txt", "text/plain", b"retry", limits) + .unwrap(); + handle + .writer + .store + .pin_uploaded_file(session_id, &file, "interrupted-before-checkpoint") + .unwrap(); + drop(handle); + let handle = PendingSubmissionHandle { + state: Arc::new(Mutex::new(PendingActivationState::default())), + writer: LogWriterHandle { + store: session_store::FsStore::new(temp.path()).unwrap(), + state: SegmentState::new(session_id, segment_id, 0), + sink: SegmentLogSink::new(), + in_flight: None, + }, + }; + + assert_eq!(handle.reconcile_uploaded_file_pins().unwrap(), 1); + let accepted = handle + .accept( + "request-after-restore".into(), + vec![Segment::UploadedFile { file: file.clone() }], + false, + ) + .unwrap(); + assert!(!accepted.submission_id.is_empty()); + assert_eq!(handle.reconcile_uploaded_file_pins().unwrap(), 0); + assert!(matches!( + handle + .writer + .store + .delete_uploaded_file(session_id, &file.artifact_id), + Err(StoreError::ArtifactAlreadyCommitted) + )); + } + + #[test] + fn rejected_submission_rolls_back_uploaded_file_pins_acquired_before_conflict() { + let temp = tempfile::tempdir().unwrap(); + let handle = PendingSubmissionHandle::for_test(temp.path()); + let session_id = handle.writer.state.session_id(); + let limits = session_store::UploadedFileLimits { + max_file_bytes: 1024, + max_session_bytes: 2048, + }; + let first = handle + .writer + .store + .write_uploaded_file(session_id, "first.txt", "text/plain", b"first", limits) + .unwrap(); + let second = handle + .writer + .store + .write_uploaded_file(session_id, "second.txt", "text/plain", b"second", limits) + .unwrap(); + handle + .writer + .store + .pin_uploaded_file(session_id, &second, "other-submission") + .unwrap(); + + assert!( + handle + .accept( + "request-partial-pin".into(), + vec![ + Segment::UploadedFile { + file: first.clone(), + }, + Segment::UploadedFile { + file: second.clone(), + }, + ], + false, + ) + .is_err() + ); + assert!(handle.snapshot().submissions.is_empty()); + assert!( + handle + .writer + .store + .delete_uploaded_file(session_id, &first.artifact_id) + .unwrap() + ); + assert!(matches!( + handle + .writer + .store + .delete_uploaded_file(session_id, &second.artifact_id), + Err(StoreError::ArtifactAlreadyCommitted) + )); + handle + .writer + .store + .release_uploaded_file_pin(session_id, &second.artifact_id, "other-submission") + .unwrap(); + assert!( + handle + .writer + .store + .delete_uploaded_file(session_id, &second.artifact_id) + .unwrap() + ); + } + + #[test] + fn queued_submission_pins_uploaded_file_until_cancelled() { + let temp = tempfile::tempdir().unwrap(); + let handle = PendingSubmissionHandle::for_test(temp.path()); + let session_id = handle.writer.state.session_id(); + let reference = handle + .writer + .store + .write_uploaded_file( + session_id, + "queued.txt", + "text/plain", + b"queued artifact", + session_store::UploadedFileLimits { + max_file_bytes: 1024, + max_session_bytes: 2048, + }, + ) + .unwrap(); + let accepted = handle + .accept( + "artifact-request".into(), + vec![Segment::UploadedFile { + file: reference.clone(), + }], + false, + ) + .unwrap(); + + assert_eq!( + handle + .writer + .store + .delete_uncommitted_uploaded_files(session_id) + .unwrap(), + 0 + ); + assert!( + handle + .writer + .store + .read_uploaded_file_by_id(session_id, &reference.artifact_id) + .is_ok() + ); + + handle + .cancel(&accepted.submission_id, handle.snapshot().revision) + .unwrap(); + assert_eq!( + handle + .writer + .store + .delete_uncommitted_uploaded_files(session_id) + .unwrap(), + 1 + ); + } + + #[test] + fn pending_submission_queue_is_durable_idempotent_and_bounded() { + let temp = tempfile::tempdir().unwrap(); + let handle = PendingSubmissionHandle::for_test(temp.path()); + let input = vec![Segment::text("queued")]; + let accepted = handle + .accept("request-1".into(), input.clone(), false) + .unwrap(); + assert_eq!( + accepted.disposition, + protocol::SubmissionDisposition::Queued + ); + assert_eq!(handle.snapshot().submissions.len(), 1); + + let replay = handle + .accept("request-1".into(), input.clone(), false) + .unwrap(); + assert_eq!(replay.submission_id, accepted.submission_id); + assert!(replay.activation.is_none()); + assert_eq!(handle.snapshot().submissions.len(), 1); + assert!(matches!( + handle.accept("request-1".into(), vec![Segment::text("different")], false), + Err(PendingSubmissionError::IdempotencyConflict) + )); + + let entries = handle + .writer + .store + .read_all( + handle.writer.state.session_id(), + handle.writer.state.segment_id(), + ) + .unwrap(); + let payload = entries + .iter() + .rev() + .find_map(|entry| match entry { + LogEntry::Extension { + domain, payload, .. + } if domain == SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN => { + Some(payload.clone()) + } + _ => None, + }) + .unwrap(); + let restored: PendingActivationState = serde_json::from_value(payload).unwrap(); + assert_eq!(restored.pending.len(), 1); + assert_eq!(restored.pending[0].submission_id, accepted.submission_id); + + let fence = handle.snapshot(); + assert!(matches!( + handle.cancel(&accepted.submission_id, fence.revision.saturating_sub(1)), + Err(PendingSubmissionError::RevisionConflict { .. }) + )); + assert!(matches!( + handle.prepare_next_activation(Some((fence.revision, "wrong-head"))), + Err(PendingSubmissionError::HeadConflict { .. }) + )); + assert_eq!(handle.snapshot(), fence); + + let snapshot = handle + .cancel(&accepted.submission_id, handle.snapshot().revision) + .unwrap(); + assert!(snapshot.submissions.is_empty()); + assert!(matches!( + handle.cancel(&accepted.submission_id, handle.snapshot().revision), + Err(PendingSubmissionError::NotFound(_)) + )); + + for index in 0..MAX_PENDING_SUBMISSIONS { + handle + .accept( + format!("limit-{index}"), + vec![Segment::text(format!("value-{index}"))], + false, + ) + .unwrap(); + } + assert!(matches!( + handle.accept("over-limit".into(), vec![Segment::text("too much")], false), + Err(PendingSubmissionError::CountLimit) + )); + assert_eq!(handle.snapshot().submissions.len(), MAX_PENDING_SUBMISSIONS); + let cleared = handle.clear(handle.snapshot().revision).unwrap(); + assert!(cleared.submissions.is_empty()); + assert_eq!(cleared.notification_count, 0); + } + + #[test] + fn durable_notification_commits_authenticated_history_provenance() { + let temp = tempfile::tempdir().unwrap(); + let handle = PendingSubmissionHandle::for_test(temp.path()); + let provenance = WorkerHistoryProvenance::HumanInput { + account_id: "account-1".into(), + }; + let committed = handle + .writer + .commit_system_item_with_extensions( + SystemItem::Notification { + message: "notice".into(), + body: "notice".into(), + prompt_provenance: None, + }, + Vec::new(), + Some(provenance.clone()), + ) + .unwrap(); + assert_eq!(committed.annotation.origin, provenance); + } + + #[test] + fn notification_retry_identity_is_scoped_to_authenticated_source() { + let temp = tempfile::tempdir().unwrap(); + let handle = PendingSubmissionHandle::for_test(temp.path()); + let account_a = WorkerHistoryProvenance::HumanInput { + account_id: "account-a".into(), + }; + let account_b = WorkerHistoryProvenance::HumanInput { + account_id: "account-b".into(), + }; + assert!( + handle + .accept_notification_from_source( + "request-1".into(), + "notice".into(), + "account:account-a".into(), + account_a.clone(), + false, + ) + .unwrap() + ); + assert!( + !handle + .accept_notification_from_source( + "request-1".into(), + "notice".into(), + "account:account-a".into(), + account_a.clone(), + false, + ) + .unwrap() + ); + assert!( + handle + .accept_notification_from_source( + "request-1".into(), + "notice".into(), + "account:account-b".into(), + account_b.clone(), + false, + ) + .unwrap() + ); + let state = handle.state.lock().unwrap(); + assert_eq!(state.pending_notifications[0].provenance, account_a); + assert_eq!(state.pending_notifications[1].provenance, account_b); + } + + #[test] + fn notification_and_submit_share_activation_order_and_notification_dedupes() { + let temp = tempfile::tempdir().unwrap(); + let handle = PendingSubmissionHandle::for_test(temp.path()); + assert!( + handle + .accept_notification("notification-1".into(), "notice".into(), true) + .unwrap() + ); + assert!( + !handle + .accept_notification("notification-1".into(), "notice".into(), true) + .unwrap() + ); + assert!(matches!( + handle.accept_notification("notification-1".into(), "different".into(), true), + Err(PendingSubmissionError::IdempotencyConflict) + )); + assert!(matches!( + handle.accept_notification("notification-1".into(), "notice".into(), false), + Err(PendingSubmissionError::IdempotencyConflict) + )); + handle + .accept("request-1".into(), vec![Segment::text("submit")], false) + .unwrap(); + + let first = handle.prepare_next_activation(None).unwrap().unwrap(); + assert!(matches!( + first, + PendingActivation::Notification(PendingNotification { ref message, .. }) + if message == "notice" + )); + let committed = handle.notification_activation_extension(); + let committed_state: PendingActivationState = + serde_json::from_value(committed.payload).unwrap(); + assert!(committed_state.pending_notifications.is_empty()); + assert!(committed_state.activating_notification.is_none()); + handle.finish_notification_activation("notification-1"); + let second = handle.prepare_next_activation(None).unwrap().unwrap(); + assert!(matches!(second, PendingActivation::Submission(_))); + } + + #[test] + fn restoring_an_in_flight_activation_requeues_it_at_the_fifo_head() { + let mut session = WorkerSession::new(session_store::new_session_id(), Vec::new()); + let state = PendingActivationState { + revision: 4, + next_activation_sequence: 3, + activating: Some(PendingSubmission { + submission_request_id: "request-1".into(), + source_namespace: "direct:test".into(), + submission_id: "submission-1".into(), + payload_digest: submission_payload_digest(&[Segment::text("first")]), + accepted_at_ms: 1, + activation_sequence: 0, + provenance: WorkerHistoryProvenance::LegacyUnknown, + was_queued: false, + input: vec![Segment::text("first")], + }), + activating_notification: Some(PendingNotification { + notification_request_id: "notification-1".into(), + source_namespace: "account:account-1".into(), + message: "deferred notice".into(), + payload_digest: notification_payload_digest("deferred notice", false), + auto_run: false, + accepted_at_ms: 3, + activation_sequence: 2, + provenance: WorkerHistoryProvenance::HumanInput { + account_id: "account-1".into(), + }, + }), + pending: VecDeque::from([PendingSubmission { + submission_request_id: "request-2".into(), + source_namespace: "direct:test".into(), + submission_id: "submission-2".into(), + payload_digest: submission_payload_digest(&[Segment::text("second")]), + accepted_at_ms: 2, + activation_sequence: 1, + provenance: WorkerHistoryProvenance::LegacyUnknown, + was_queued: true, + input: vec![Segment::text("second")], + }]), + pending_notifications: VecDeque::new(), + receipts: VecDeque::new(), + notification_receipts: VecDeque::from([NotificationReceipt { + notification_request_id: "notification-1".into(), + source_namespace: "account:account-1".into(), + payload_digest: notification_payload_digest("deferred notice", false), + auto_run: false, + }]), + }; + session.restore_pending_activations(&[( + SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.into(), + serde_json::to_value(state).unwrap(), + )]); + let state = session + .pending_activations + .lock() + .expect("pending activation state poisoned"); + assert!(state.activating.is_none()); + assert_eq!(state.pending.len(), 2); + assert_eq!(state.pending[0].submission_id, "submission-1"); + assert_eq!(state.pending[1].submission_id, "submission-2"); + assert!(state.activating_notification.is_none()); + assert_eq!(state.pending_notifications.len(), 1); + assert!(!state.pending_notifications[0].auto_run); + assert!(matches!( + state.pending_notifications[0].provenance, + WorkerHistoryProvenance::HumanInput { ref account_id } if account_id == "account-1" + )); + assert_eq!(state.notification_receipts.len(), 1); + } + fn minimal_manifest() -> WorkerManifest { let toml_str = r#" [worker] diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index 47fbd351..8265f9e2 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -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 { diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 143c1da9..0db13a0c 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -332,6 +332,7 @@ async fn shutdown_closes_bound_workdir_session() { command: "sleep 30".to_owned(), timeout_secs: 60, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: None, }) @@ -376,6 +377,7 @@ async fn controller_projects_workdir_command_events_and_snapshot_state() { command: "printf ready; sleep 0.3; printf done".to_owned(), timeout_secs: 5, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: Some("tool-command-1".into()), }) @@ -484,6 +486,7 @@ async fn controller_refreshes_command_snapshot_after_high_output_provider_lag() .to_owned(), timeout_secs: 10, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: Some("tool-high-output".into()), }) @@ -560,6 +563,7 @@ async fn controller_startup_failure_closes_bound_workdir_session() { command: "printf unreachable".to_owned(), timeout_secs: 5, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: None, }) @@ -617,7 +621,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 +682,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 +774,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 +848,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 +903,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 +962,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 +1015,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 +1063,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 +1118,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 +1186,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 +1230,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 +1250,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 +1274,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 +1315,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 +1326,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 +1474,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 +1547,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 +1601,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 +1655,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 +1709,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 +1750,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 +1808,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 +1828,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 +2014,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 +2056,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 +2150,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 +2463,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 +2558,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 +2595,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 +2632,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 +2769,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 +2843,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 +2935,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 +2974,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 +3011,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 +3031,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 +3081,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!( diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 4a672748..4a90e192 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -538,7 +538,7 @@ fn initial_worker_input(segments: &[Segment]) -> Option { 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( diff --git a/crates/workspace-server/src/runtime_subscription_tests.rs b/crates/workspace-server/src/runtime_subscription_tests.rs index a42c0c02..6c14c56a 100644 --- a/crates/workspace-server/src/runtime_subscription_tests.rs +++ b/crates/workspace-server/src/runtime_subscription_tests.rs @@ -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( diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index f34f0ae1..41492447 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -48,8 +48,7 @@ use workdir::http::{ }; use workdir::workspace::{ MaterializerKind, WorkingDirectoryCleanupTarget, WorkingDirectoryOccupancy, - WorkingDirectoryStatusKind, WorkingDirectorySummary, WorkspaceWorkdirSessionFence, - WorkspaceWorkdirSessionOperationRequest, + WorkingDirectoryStatusKind, WorkingDirectorySummary, WorkspaceWorkdirSessionOperationRequest, }; use workdir::{CommandHandle, WorkdirSessionHandle}; use worker::feature::builtin::{WorkerObservationSubject, WorkerObservationSubjectRef}; @@ -355,7 +354,6 @@ static EMBEDDED_RUNTIME_REQUEST_IDENTITY: std::sync::LazyLock< struct WorkdirCommandSession { source: WorkdirSessionHandle, provider_handle: CommandHandle, - delegations: Vec, } enum RegisteredWorkdirSession { @@ -398,7 +396,6 @@ impl WorkdirSessionRegistry { worker: RuntimeWorkerRef, source: WorkdirSessionHandle, provider_handle: CommandHandle, - delegations: Vec, ) -> CommandHandle { let external_handle = loop { let candidate = CommandHandle(Uuid::now_v7().to_string()); @@ -414,7 +411,6 @@ impl WorkdirSessionRegistry { WorkdirCommandSession { source, provider_handle, - delegations, }, ); external_handle @@ -2653,10 +2649,6 @@ fn build_inner_router(api: WorkspaceApi) -> Router { post(scoped_attach_current_worker_workdir) .delete(scoped_detach_current_worker_workdir), ) - .route( - "/api/w/{workspace_id}/workers/self/workdir-session/fence", - get(scoped_current_worker_workdir_session_fence), - ) .route( "/api/w/{workspace_id}/workers/self/workdir-session/operations", post(scoped_execute_current_worker_workdir_operation), @@ -7412,46 +7404,11 @@ async fn scoped_detach_current_worker_workdir( })) } -async fn scoped_current_worker_workdir_session_fence( - State(api): State, - AxumPath(path): AxumPath, - headers: HeaderMap, -) -> ApiResult> { - validate_workspace_scope(&api, &path.workspace_id)?; - let worker = current_worker_identity(&api, &path.workspace_id, &headers)?; - let session_lock = current_worker_session_lock(&api, &worker); - let _session_guard = session_lock.lock().await; - let link = current_worker_active_attachment(&api, &worker)?; - Ok(Json(WorkspaceWorkdirSessionFence { - value: current_worker_workdir_session_fence(&link), - })) -} - -fn current_worker_workdir_session_fence(link: &WorkerWorkdirLinkRecord) -> String { - format!("v1:{}\0{}", link.workdir_id, link.linked_at) -} - -fn validate_current_worker_workdir_session_fence( - link: &WorkerWorkdirLinkRecord, - expected: Option<&str>, -) -> Result<()> { - if expected.is_some_and(|expected| expected != current_worker_workdir_session_fence(link)) { - Err(Error::WorkdirAttachmentConflict( - "delegated Workdir session attachment changed".to_string(), - )) - } else { - Ok(()) - } -} - fn validated_current_worker_attachment( api: &WorkspaceApi, worker: &RuntimeWorkerRef, - expected_session_fence: Option<&str>, ) -> ApiResult { - let link = current_worker_active_attachment(api, worker)?; - validate_current_worker_workdir_session_fence(&link, expected_session_fence)?; - Ok(link) + current_worker_active_attachment(api, worker) } #[derive(Debug)] @@ -7506,23 +7463,13 @@ async fn scoped_execute_current_worker_workdir_operation( ) -> std::result::Result, WorkdirOperationApiError> { validate_workspace_scope(&api, &path.workspace_id)?; let worker = current_worker_identity(&api, &path.workspace_id, &headers)?; - let expected_session_fence = request.expected_session_fence; - let delegations = request.delegations; let result = match request.operation { WorkdirSessionOperation::CommandStart(command) => { let session_lock = current_worker_session_lock(&api, &worker); let _session_guard = session_lock.lock().await; - let link = validated_current_worker_attachment( - &api, - &worker, - expected_session_fence.as_deref(), - )?; + let link = validated_current_worker_attachment(&api, &worker)?; let source = open_current_worker_workdir_session_locked(&api, &worker, &link).await?; - let applied = - apply_current_worker_delegations(&worker, source.clone(), delegations.clone()) - .await?; - let provider_handle = applied - .scoped_session + let provider_handle = source .start_command(command) .await .map_err(|error| current_worker_workdir_operation_error(&worker, error))?; @@ -7541,58 +7488,32 @@ async fn scoped_execute_current_worker_workdir_operation( .workdir_sessions .lock() .expect("Workdir session registry lock poisoned") - .register_command( - worker.clone(), - registered_source, - provider_handle, - delegations, - ); + .register_command(worker.clone(), registered_source, provider_handle); WorkdirSessionOperationResult::CommandStart(external_handle) } WorkdirSessionOperation::CommandStatus(external_handle) => { - let (session, provider_handle) = current_worker_command_session( - &api, - &worker, - &external_handle, - &delegations, - expected_session_fence.as_deref(), - ) - .await?; + let (session, provider_handle) = + current_worker_command_session(&api, &worker, &external_handle)?; session - .scoped_session .command_status(provider_handle) .await .map(WorkdirSessionOperationResult::CommandStatus) .map_err(|error| current_worker_workdir_operation_error(&worker, error))? } WorkdirSessionOperation::CommandOutput(mut output) => { - let (session, provider_handle) = current_worker_command_session( - &api, - &worker, - &output.handle, - &delegations, - expected_session_fence.as_deref(), - ) - .await?; + let (session, provider_handle) = + current_worker_command_session(&api, &worker, &output.handle)?; output.handle = provider_handle; session - .scoped_session .command_output(output) .await .map(WorkdirSessionOperationResult::CommandOutput) .map_err(|error| current_worker_workdir_operation_error(&worker, error))? } WorkdirSessionOperation::CommandCancel(external_handle) => { - let (session, provider_handle) = current_worker_command_session( - &api, - &worker, - &external_handle, - &delegations, - expected_session_fence.as_deref(), - ) - .await?; + let (session, provider_handle) = + current_worker_command_session(&api, &worker, &external_handle)?; session - .scoped_session .cancel_command(provider_handle) .await .map(|()| WorkdirSessionOperationResult::CommandCancel) @@ -7607,14 +7528,9 @@ async fn scoped_execute_current_worker_workdir_operation( | WorkdirSessionOperation::Grep(_)) => { let session_lock = current_worker_session_lock(&api, &worker); let _session_guard = session_lock.lock().await; - let link = validated_current_worker_attachment( - &api, - &worker, - expected_session_fence.as_deref(), - )?; + let link = validated_current_worker_attachment(&api, &worker)?; let source = open_current_worker_workdir_session_locked(&api, &worker, &link).await?; - let applied = apply_current_worker_delegations(&worker, source, delegations).await?; - execute_workdir_session_operation(&applied.scoped_session, operation) + execute_workdir_session_operation(&source, operation) .await .map_err(|error| current_worker_workdir_operation_error(&worker, error))? } @@ -7622,29 +7538,12 @@ async fn scoped_execute_current_worker_workdir_operation( Ok(Json(result)) } -async fn apply_current_worker_delegations( - worker: &RuntimeWorkerRef, - source: WorkdirSessionHandle, - delegations: Vec, -) -> Result { - workdir::apply_delegation_chain(source, delegations) - .await - .map_err(|error| Error::RuntimeOperationFailed { - runtime_id: worker.runtime_id.clone(), - code: "workdir_session_delegation_failed".to_string(), - message: error.to_string(), - }) -} - -async fn current_worker_command_session( +fn current_worker_command_session( api: &WorkspaceApi, worker: &RuntimeWorkerRef, external_handle: &CommandHandle, - delegations: &[workdir::WorkdirDelegationRequest], - expected_session_fence: Option<&str>, -) -> std::result::Result<(workdir::AppliedWorkdirDelegation, CommandHandle), WorkdirOperationApiError> -{ - let _link = validated_current_worker_attachment(api, worker, expected_session_fence)?; +) -> std::result::Result<(WorkdirSessionHandle, CommandHandle), WorkdirOperationApiError> { + let _link = validated_current_worker_attachment(api, worker)?; let command = api .workdir_sessions .lock() @@ -7656,15 +7555,7 @@ async fn current_worker_command_session( workdir::WorkdirError::UnknownCommand(external_handle.0.clone()), )) })?; - if command.delegations != delegations { - return Err(Error::WorkdirAttachmentConflict( - "command lifecycle delegation differs from CommandStart".to_string(), - ) - .into()); - } - let session = - apply_current_worker_delegations(worker, command.source, command.delegations).await?; - Ok((session, command.provider_handle)) + Ok((command.source, command.provider_handle)) } fn current_worker_workdir_operation_error( @@ -8606,13 +8497,15 @@ async fn scoped_list_runtimes( async fn scoped_workspace_protocol_ws( State(api): State, + Extension(actor): Extension, AxumPath(workspace_id): AxumPath, ws: axum::extract::ws::WebSocketUpgrade, ) -> std::result::Result { 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()) } @@ -9318,7 +9211,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(|_| { @@ -11741,6 +11634,7 @@ async fn scoped_cancel_runtime_worker( async fn scoped_worker_protocol_ws( ws: WebSocketUpgrade, State(api): State, + Extension(actor): Extension, AxumPath(path): AxumPath, ) -> Response { if let Err(err) = validate_workspace_scope(&api, &path.workspace_id) { @@ -11748,6 +11642,7 @@ async fn scoped_worker_protocol_ws( } worker_protocol_ws( State(api), + Extension(actor), AxumPath((path.worker.runtime_id, path.worker.worker_id)), ws, ) @@ -14140,8 +14035,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 { + 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, + Extension(actor): Extension, AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>, ws: WebSocketUpgrade, ) -> impl IntoResponse { @@ -14165,7 +14097,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 { @@ -14176,6 +14109,7 @@ pub(crate) struct WorkspaceWorkerProtocolConnection { pub(crate) async fn connect_workspace_worker_protocol( api: &WorkspaceApi, worker: &RuntimeWorkerRef, + input_source: Option<&protocol::AuthenticatedInputSource>, ) -> Result { let source = match api.observation_proxy.source(worker) { Ok(source) => source, @@ -14192,15 +14126,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 { let mut request = config .endpoint @@ -14215,6 +14173,7 @@ async fn connect_remote_worker_protocol( })?, ); } + insert_authenticated_input_source_header(request.headers_mut(), input_source)?; let (socket, _) = connect_async(request) .await @@ -14292,13 +14251,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; } } } @@ -14306,6 +14269,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, @@ -14333,6 +14297,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, @@ -14354,14 +14328,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; @@ -14417,6 +14410,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(), @@ -14436,24 +14430,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; } @@ -16593,6 +16595,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"); @@ -16788,6 +16832,7 @@ mod tests { command: "printf ready; sleep 30".to_string(), timeout_secs: 60, output_limit: 4096, + cwd: None, spill_dir: None, tool_call_id: Some("tool-call-command-session".to_string()), }) @@ -16797,12 +16842,8 @@ mod tests { let mut registry = WorkdirSessionRegistry::default(); registry.insert_attachment(worker.clone(), source.clone()); let registered_source = registry.remove_attachment(&worker).unwrap(); - let external_handle = registry.register_command( - worker.clone(), - registered_source, - provider_handle.clone(), - Vec::new(), - ); + let external_handle = + registry.register_command(worker.clone(), registered_source, provider_handle.clone()); assert_ne!(external_handle, provider_handle); let refreshed: WorkdirSessionHandle = Arc::new(workdir::LocalWorkdirSession::new( @@ -18843,7 +18884,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)); @@ -18851,11 +18892,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( @@ -23747,30 +23790,6 @@ mod tests { assert_eq!(response.status(), StatusCode::BAD_REQUEST); } - #[test] - fn delegated_workdir_session_fence_rejects_reattached_link() { - let first = WorkerWorkdirLinkRecord { - workspace_id: "workspace-a".to_string(), - worker: workdir::workspace::RuntimeWorkerRef::new("runtime-a", "worker-a"), - workdir_id: "workdir-a".to_string(), - role: "primary".to_string(), - linked_at: "2026-01-01T00:00:00Z".to_string(), - unlinked_at: None, - }; - let expected = current_worker_workdir_session_fence(&first); - assert!(validate_current_worker_workdir_session_fence(&first, None).is_ok()); - assert!(validate_current_worker_workdir_session_fence(&first, Some(&expected)).is_ok()); - - let reattached = WorkerWorkdirLinkRecord { - linked_at: "2026-01-01T00:00:01Z".to_string(), - ..first - }; - assert!(matches!( - validate_current_worker_workdir_session_fence(&reattached, Some(&expected)), - Err(Error::WorkdirAttachmentConflict(_)) - )); - } - #[tokio::test] async fn backend_workdir_session_proxy_executes_typed_operations() { use manifest::Scope; @@ -27638,6 +27657,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) { @@ -27656,11 +27685,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, @@ -27672,7 +27698,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; }); @@ -27720,7 +27747,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; }); diff --git a/crates/workspace-server/src/workspace_subscription.rs b/crates/workspace-server/src/workspace_subscription.rs index fe26f7f9..018a9622 100644 --- a/crates/workspace-server/src/workspace_subscription.rs +++ b/crates/workspace-server/src/workspace_subscription.rs @@ -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>, } -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::(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; } } diff --git a/docs/design/flow-state-graph.md b/docs/design/flow-state-graph.md index 3fc49249..bd0cea00 100644 --- a/docs/design/flow-state-graph.md +++ b/docs/design/flow-state-graph.md @@ -55,11 +55,12 @@ Workspace Server schema migration v26 removes the legacy `flow_instances`, `flow ## Worker boundary -Flow invocation uses the normal Submit/Run segment vector rather than a Worker-create field: +Flow invocation uses the normal Submit segment vector rather than a Worker-create field: ```json { - "method": "run", + "method": "submit", + "submission_request_id": "018f4f15-5c41-7d3a-8a72-2e755bc71681", "input": [ { "kind": "flow", "selector": "builtin:coder-review" }, { "kind": "text", "content": "Ticket 00001... implementation" } @@ -69,7 +70,7 @@ Flow invocation uses the normal Submit/Run segment vector rather than a Worker-c Runtime accepts exactly one Flow segment only when the resolved Profile enables `feature.flow` and a Workspace client is available. The Worker asks Workspace authority only for an immutable source snapshot, creates the instance locally, replaces the Flow segment with the entered state's instructions, and commits that runtime state atomically with the remaining Submit segments before LLM execution. A Worker with an active Flow rejects the duplicate input without changing its local state or events. -The generic model-facing `WorkerSpawn` accepts `initial_submit: Vec` and routes them unchanged through the shared Workspace spawn request into Runtime `CreateWorkerRequest.initial_input`. It does not have a parallel `initial_text` or a role-specific `SpawnCoder` wrapper. Backend derives the flat content projection from the canonical segment vector, validates Flow shape before spawn, and includes the segment vector in lifecycle idempotency fingerprints. Runtime does not commit Worker creation or report spawn success merely because the initial Run method entered the Worker's in-memory channel: Runtime assigns the Submit an opaque id, the Worker commits that id as an extension on the same `UserInput` entry as any initial `FlowRuntimeState`, and the execution backend must return a matching typed input-commit acknowledgement. Restoring the same Worker never replays spawn initial segments. +The generic model-facing `WorkerSpawn` accepts `initial_submit: Vec` and routes them unchanged through the shared Workspace spawn request into Runtime `CreateWorkerRequest.initial_input`. It does not have a parallel `initial_text` or a role-specific `SpawnCoder` wrapper. Backend derives the flat content projection from the canonical segment vector, validates Flow shape before spawn, and includes the segment vector in lifecycle idempotency fingerprints. Runtime does not commit Worker creation or report spawn success merely because the initial Submit request entered the Worker's in-memory channel: Runtime assigns the Submit an opaque id, the Worker commits that id as an extension on the same `UserInput` entry as any initial `FlowRuntimeState`, and the execution backend must return a matching typed input-commit acknowledgement. Restoring the same Worker never replays spawn initial segments. When an Orchestrator supplies `ticket_id` to generic `WorkerSpawn`, the Worker tool derives the assignment operation id from the durable tool-call id rather than accepting lifecycle authority from model input. The shared Workspace worker-create route projects that request into a Coder Ticket-role intent and atomically applies the existing queued-Ticket assignment operation only after Runtime has returned the input-commit acknowledgement. A spawn or pre-commit input failure therefore leaves the Ticket queued and unassigned. diff --git a/docs/development/work-items.md b/docs/development/work-items.md index 61c174ff..34bec74f 100644 --- a/docs/development/work-items.md +++ b/docs/development/work-items.md @@ -286,7 +286,7 @@ User triggers a Ticket action in yoi panel -> client Ticket role launcher reads .yoi/workspace.toml [ticket] settings -> launcher selects the role Profile -> launcher spawns the role Worker - -> launcher sends Method::Run with Text segments + -> launcher sends Method::Submit with Text segments -> launcher waits for run-acceptance evidence -> Dashboard reports success/failure ``` diff --git a/docs/report/test-validity-20260612/pod.md b/docs/report/test-validity-20260612/pod.md index d5c5cfc8..087ea4d2 100644 --- a/docs/report/test-validity-20260612/pod.md +++ b/docs/report/test-validity-20260612/pod.md @@ -34,7 +34,7 @@ - いくつかの async/socket tests はまだ fixed sleeps と process-wide environment mutation (`YOI_RUNTIME_DIR`, `YOI_HOME`, `XDG_RUNTIME_DIR`) に依存している。一部ファイルでは env-changing tests を guards で serialize しているが、fixed timing と global env は高負荷 CI や parallel execution 下で flakiness risk のまま。 - Real provider wire behavior は意図的に `pod` の外側だが、crate と streaming edge cases の interaction はまだ大部分が mock されている。tests は重要な `Worker` outcomes をカバーしているが、malformed/partial provider streams を real provider adapter 経由では exercise していない。 - Prompt tests は数が多く有用だが、一部は behavior-coupled というより prose-coupled である。critical safety wording に対して、その文字列を意図的に stable contract として扱う場合だけこれは許容できる。そうでなければ maintenance noise になる。 -- startup profile resolution、socket server、`Method::Run`、session persistence、shutdown、restore を跨ぐ full lifecycle integration は slice ごとにしかカバーされておらず、1 つの scenario としてはカバーされていない。 +- startup profile resolution、socket server、`Method::Submit`、session persistence、shutdown、restore を跨ぐ full lifecycle integration は slice ごとにしかカバーされておらず、1 つの scenario としてはカバーされていない。 ## 追加を提案するもの diff --git a/resources/flows/coder-review.dcdl b/resources/flows/coder-review.dcdl index b42dbd76..4de10077 100644 --- a/resources/flows/coder-review.dcdl +++ b/resources/flows/coder-review.dcdl @@ -15,7 +15,7 @@ }; review = { - instructions = "Use the current Ticket Merge Request as review authority. Call `ShowMergeRequest` and confirm its source selector resolves to exact committed implementation HEAD, then spawn one actual direct-child SubWorker with profile builtin:reviewer, write scope for Workdir inspection and command validation, and only the Ticket id in the structured review handoff. The trusted spawn layer records `ReviewRequested` with the exact source ref and injects review capability; do not place commit/ref identity, capability material, or a prewritten verdict in model input. The child must commit `ReviewMergeRequest`; prose output and Worker observation are not approval authority. After the structured result for the exact current source ref exists, request a Flow transition."; + instructions = "Use the current Ticket Merge Request as review authority. Call `ShowMergeRequest` and confirm its source selector resolves to exact committed implementation HEAD, then spawn one actual direct-child SubWorker with profile builtin:reviewer, write scope plus an explicit command grant for Workdir inspection and command validation, and only the Ticket id in the structured review handoff. The trusted spawn layer records `ReviewRequested` with the exact source ref and injects review capability; do not place commit/ref identity, capability material, or a prewritten verdict in model input. The child must commit `ReviewMergeRequest`; prose output and Worker observation are not approval authority. After the structured result for the exact current source ref exists, request a Flow transition."; transitions = { approved = { target = "complete"; diff --git a/resources/prompts/internal/sub_worker_spawn_tool_description.md b/resources/prompts/internal/sub_worker_spawn_tool_description.md index a652f551..d6d50df1 100644 --- a/resources/prompts/internal/sub_worker_spawn_tool_description.md +++ b/resources/prompts/internal/sub_worker_spawn_tool_description.md @@ -1,8 +1,8 @@ Spawn a parent-owned Internal SubWorker session to split context for a delegated task. The parent Worker's write scope is reduced by the scope passed here; the Internal SubWorker starts running `task` immediately without creating a Runtime Worker record, OS process, PID, or Unix socket. It remains available for follow-up turns until explicitly stopped or its parent exits. -Optional `cwd`: when provided, the spawned SubWorker's tool default working directory only. It must be an absolute existing directory covered by the child's delegated readable scope, and it does not change workspace/Profile/memory/Ticket roots or grant authority. `name` must be unique among this Worker's direct children. +Optional `cwd`: when provided, the spawned SubWorker's tool default working directory only. It must be a Workdir-relative existing directory covered by the child's readable scope, and it does not change workspace/Profile/memory/Ticket roots or grant authority. `name` must be unique among this Worker's direct children. -Profile selection: `profile` may be omitted or set to `default` to use the effective child default profile, set to `inherit` to derive reusable child configuration from this Worker, or set to one of the registry selectors below. Raw/path profile selectors are not accepted by SubWorkerSpawn. `scope` is always the only delegated filesystem capability; profile scope is replaced by the explicit SubWorkerSpawn scope. +Profile selection: `profile` may be omitted or set to `default` to use the effective child default profile, set to `inherit` to derive reusable child configuration from this Worker, or set to one of the registry selectors below. Raw/path profile selectors are not accepted by SubWorkerSpawn. `scope` is the child's only filesystem capability and replaces profile scope. `command` is a separate explicit grant, defaults to false, and is accepted only with a writable scope; writable scope alone does not grant command execution. Default profile: {{ default_profile }} Special selector: inherit — derive reusable model/worker/tool policy from the spawner while replacing worker.name and scope. diff --git a/web/workspace/deno.json b/web/workspace/deno.json index eff2525d..4fda14c1 100644 --- a/web/workspace/deno.json +++ b/web/workspace/deno.json @@ -6,7 +6,7 @@ "dev": "deno run -A npm:vite@7.2.7 dev", "dev:backend": "cd ../.. && cargo run -p yoi-workspace-server --bin yoi-server -- serve --listen 127.0.0.1:8787", "check": "deno run -A npm:@sveltejs/kit@2.49.4 sync && deno run -A npm:svelte-check@4.3.4 --tsconfig ./tsconfig.json", - "test": "deno test --allow-read=src,test,tests --allow-env=LOG,VSCODE_TEXTMATE_DEBUG,NODE_ENV tests/workspace-model.test.ts tests/workspace-catalog.test.ts tests/profile-api.test.ts tests/skill-api.test.ts src/lib/workspace/auth/model.test.ts tests/auth-api.test.ts src/lib/workspace/api/http.test.ts src/lib/workspace/api/workers.test.ts src/lib/workspace/header/breadcrumb-model.test.ts src/lib/workspace/console/chat-submit.test.ts test/composer-history.test.ts tests/composer-paste.test.ts src/lib/workspace/console/composer-command.test.ts src/lib/workspace/console/composer-draft.test.ts src/lib/workspace/console/composer-completion.test.ts src/lib/workspace/console/markdown.test.ts test/console/ansi.test.ts src/lib/workspace/console/model.test.ts src/lib/workspace/companion/api.test.ts tests/workdir-api.test.ts src/lib/workspace/console/tasks.test.ts test/ticket-detail-route-reuse.test.ts test/repositories/ui.test.ts src/lib/workspace/console/worker-console.ui.test.ts src/lib/workspace/settings/model.test.ts tests/runtime-connection.test.ts tests/runtime-management.test.ts tests/runtime-management-source.test.ts src/lib/workspace/sidebar/override-stack.test.ts src/lib/workspace/sidebar/workers.test.ts src/lib/workspace/sidebar/workspace-switcher.test.ts src/lib/workspace/sidebar/worker-subscription.test.ts src/lib/workspace/sidebar/worker-launch.test.ts test/sidebar/worker-actions.test.ts src/lib/workspace/tickets/merge-request-resources.test.ts src/lib/workspace/tickets/ticket-panel.test.ts test/merge-request-status.test.ts test/config-source/decodal-grammar.test.ts test/config-source/editor-state.test.ts test/config-source/fixed-schema-wrapper.test.ts test/config-source/toolchain.test.ts test/config-source/wasm-parity.test.ts test/repository-access/api.test.ts test/repository-access/loader.test.ts test/repository-access/ui.test.ts", + "test": "deno test --allow-read=src,test,tests --allow-env=LOG,VSCODE_TEXTMATE_DEBUG,NODE_ENV tests/workspace-model.test.ts tests/workspace-catalog.test.ts tests/profile-api.test.ts tests/skill-api.test.ts src/lib/workspace/auth/model.test.ts tests/auth-api.test.ts src/lib/workspace/api/http.test.ts src/lib/workspace/api/workers.test.ts src/lib/workspace/header/breadcrumb-model.test.ts src/lib/workspace/console/chat-submit.test.ts src/lib/workspace/console/composer-delivery.test.ts test/composer-history.test.ts tests/composer-paste.test.ts src/lib/workspace/console/composer-command.test.ts src/lib/workspace/console/composer-draft.test.ts src/lib/workspace/console/composer-completion.test.ts src/lib/workspace/console/markdown.test.ts test/console/ansi.test.ts src/lib/workspace/console/model.test.ts src/lib/workspace/companion/api.test.ts tests/workdir-api.test.ts src/lib/workspace/console/tasks.test.ts test/ticket-detail-route-reuse.test.ts test/repositories/ui.test.ts src/lib/workspace/console/worker-console.ui.test.ts src/lib/workspace/settings/model.test.ts tests/runtime-connection.test.ts tests/runtime-management.test.ts tests/runtime-management-source.test.ts src/lib/workspace/sidebar/override-stack.test.ts src/lib/workspace/sidebar/workers.test.ts src/lib/workspace/sidebar/workspace-switcher.test.ts src/lib/workspace/sidebar/worker-subscription.test.ts src/lib/workspace/sidebar/worker-launch.test.ts test/sidebar/worker-actions.test.ts src/lib/workspace/tickets/merge-request-resources.test.ts src/lib/workspace/tickets/ticket-panel.test.ts test/merge-request-status.test.ts test/config-source/decodal-grammar.test.ts test/config-source/editor-state.test.ts test/config-source/fixed-schema-wrapper.test.ts test/config-source/toolchain.test.ts test/config-source/wasm-parity.test.ts test/repository-access/api.test.ts test/repository-access/loader.test.ts test/repository-access/ui.test.ts", "build": "deno run -A npm:vite@7.2.7 build", "preview": "deno run -A npm:vite@7.2.7 preview" }, diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index 8cd81901..6055d3e7 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -103,7 +103,13 @@ entry_id: string, */ timestamp: number, provenance: SessionEntryProvenance, derived_from?: Array, } & ({ "kind": "user_input", segments: Array, } | { "kind": "message", role: SessionMessageRole, content: Array, } | { "kind": "tool_call", call_id: string, name: string, arguments: string, } | { "kind": "tool_result", call_id: string, summary: string, content?: string | null, is_error: boolean, attachments?: Array, } | { "kind": "system_item", item_kind: string, content: string, data?: unknown, } | { "kind": "run_error", message: string, }); -export type SessionSnapshot = { entries: Array, }; +export type PendingSubmissionSummary = { submission_id: string, accepted_at_ms: number, segment_count: number, byte_len: number, }; + +export type PendingSubmissionsSnapshot = { revision: number, notification_count: number, head_id: string | null, submissions: Array, }; + +export type SubmissionDisposition = "started" | "queued"; + +export type SessionSnapshot = { pending_submissions: PendingSubmissionsSnapshot, entries: Array, }; export type InternalWorkerKind = "sub_worker" | { "service": { kind: string, } }; @@ -225,9 +231,9 @@ export type SubscriptionFramePayload = { "frame": "request", "message": Subscrip export type SubscriptionFrame = { protocol_version: number, } & ({ "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent } | { "frame": "worker_protocol", "message": SubscriptionWorkerProtocolMethod }); -export type Method = { "method": "run", "params": { input: Array, } } | { "method": "notify", "params": { message: string, auto_run?: boolean, } } | { "method": "worker_event", "params": WorkerEvent } | { "method": "resume" } | { "method": "cancel" } | { "method": "pause" } | { "method": "compact" } | { "method": "list_rewind_targets" } | { "method": "rewind_to", "params": { target: RewindTargetId, expected_head_entries: number, } } | { "method": "shutdown" } | { "method": "list_completions", "params": { kind: CompletionKind, prefix: string, } } | { "method": "list_workers" } | { "method": "restore_worker", "params": { name: string, } } | { "method": "register_peer", "params": { name: string, } }; +export type Method = { "method": "submit", "params": { submission_request_id: string, input: Array, } } | { "method": "notify", "params": { notification_request_id: string, message: string, auto_run?: boolean, } } | { "method": "worker_event", "params": WorkerEvent } | { "method": "list_pending_submissions" } | { "method": "cancel_pending_submission", "params": { submission_id: string, expected_revision: number, } } | { "method": "clear_pending_submissions", "params": { expected_revision: number, } } | { "method": "continue_pending", "params": { expected_revision: number, expected_head_id: string, } } | { "method": "resume" } | { "method": "cancel" } | { "method": "pause" } | { "method": "compact" } | { "method": "list_rewind_targets" } | { "method": "rewind_to", "params": { target: RewindTargetId, expected_head_entries: number, } } | { "method": "shutdown" } | { "method": "list_completions", "params": { kind: CompletionKind, prefix: string, } } | { "method": "list_workers" } | { "method": "restore_worker", "params": { name: string, } } | { "method": "register_peer", "params": { name: string, } }; -export type Event = { "event": "user_message", "data": { segments: Array, } } | { "event": "system_item", "data": { item: unknown, } } | { "event": "invoke_start", "data": { kind: InvokeKind, } } | { "event": "turn_start", "data": { turn: number, } } | { "event": "turn_end", "data": { turn: number, result: TurnResult, } } | { "event": "llm_call_start", "data": { llm_call: number, } } | { "event": "llm_call_end", "data": { llm_call: number, } } | { "event": "llm_retry", "data": { llm_call: number, +export type Event = { "event": "submission_accepted", "data": { submission_request_id: string, submission_id: string, disposition: SubmissionDisposition, } } | { "event": "submission_rejected", "data": { submission_request_id: string, message: string, } } | { "event": "pending_submissions_changed", "data": { pending: PendingSubmissionsSnapshot, } } | { "event": "user_message", "data": { segments: Array, } } | { "event": "system_item", "data": { item: unknown, } } | { "event": "invoke_start", "data": { kind: InvokeKind, } } | { "event": "turn_start", "data": { turn: number, } } | { "event": "turn_end", "data": { turn: number, result: TurnResult, } } | { "event": "llm_call_start", "data": { llm_call: number, } } | { "event": "llm_call_end", "data": { llm_call: number, } } | { "event": "llm_retry", "data": { llm_call: number, /** * The attempt that just failed. 1 origin. */ diff --git a/web/workspace/src/lib/workspace/console/composer-delivery.test.ts b/web/workspace/src/lib/workspace/console/composer-delivery.test.ts new file mode 100644 index 00000000..4d052215 --- /dev/null +++ b/web/workspace/src/lib/workspace/console/composer-delivery.test.ts @@ -0,0 +1,136 @@ +declare const Deno: { + test(name: string, fn: () => void): void; +}; + +import { + canDeliverComposerDraft, + sendComposerDelivery, +} from "./composer-delivery.ts"; + +function assertEquals(actual: unknown, expected: unknown): void { + if (actual !== expected) { + throw new Error(`Expected ${String(expected)}, got ${String(actual)}`); + } +} + +const base = { + protocolOpen: true, + sending: false, + hasText: true, + hasAttachments: false, +}; + +Deno.test("running Composer enables Queue Submit and Notify but not immediate Submit", () => { + assertEquals( + canDeliverComposerDraft({ + ...base, + delivery: "queue", + workerState: "running", + }), + true, + ); + assertEquals( + canDeliverComposerDraft({ + ...base, + delivery: "notify", + workerState: "running", + }), + true, + ); + assertEquals( + canDeliverComposerDraft({ + ...base, + delivery: "submit", + workerState: "running", + }), + false, + ); +}); + +Deno.test("running Queue Submit and Notify dispatch their protocol methods", () => { + const sent: string[] = []; + assertEquals( + sendComposerDelivery( + { ...base, delivery: "queue", workerState: "running" }, + "submit", + (method) => sent.push(method), + ), + true, + ); + assertEquals( + sendComposerDelivery( + { ...base, delivery: "notify", workerState: "running" }, + "notify", + (method) => sent.push(method), + ), + true, + ); + assertEquals(sent.join(","), "submit,notify"); +}); + +Deno.test("idle Composer enables only immediate Submit", () => { + assertEquals( + canDeliverComposerDraft({ + ...base, + delivery: "submit", + workerState: "idle", + }), + true, + ); + assertEquals( + canDeliverComposerDraft({ + ...base, + delivery: "queue", + workerState: "idle", + }), + false, + ); + assertEquals( + canDeliverComposerDraft({ + ...base, + delivery: "notify", + workerState: "idle", + }), + false, + ); +}); + +Deno.test("running delivery remains fenced by protocol, send state, and payload kind", () => { + assertEquals( + canDeliverComposerDraft({ + ...base, + delivery: "queue", + workerState: "running", + protocolOpen: false, + }), + false, + ); + assertEquals( + canDeliverComposerDraft({ + ...base, + delivery: "notify", + workerState: "running", + sending: true, + }), + false, + ); + assertEquals( + canDeliverComposerDraft({ + ...base, + delivery: "notify", + workerState: "running", + hasAttachments: true, + }), + false, + ); + assertEquals( + canDeliverComposerDraft({ + ...base, + delivery: "queue", + workerState: "running", + hasText: false, + hasAttachments: true, + }), + true, + ); +}); diff --git a/web/workspace/src/lib/workspace/console/composer-delivery.ts b/web/workspace/src/lib/workspace/console/composer-delivery.ts new file mode 100644 index 00000000..adf17883 --- /dev/null +++ b/web/workspace/src/lib/workspace/console/composer-delivery.ts @@ -0,0 +1,39 @@ +export type ComposerDelivery = "submit" | "queue" | "notify"; + +export type ComposerDeliveryState = { + delivery: ComposerDelivery; + workerState: string; + protocolOpen: boolean; + sending: boolean; + hasText: boolean; + hasAttachments: boolean; +}; + +/** + * Resolve whether the current Composer draft can use one delivery action. + * Immediate Submit is idle-only; Queue and Notify are running-only. + */ +export function canDeliverComposerDraft(state: ComposerDeliveryState): boolean { + if (!state.protocolOpen || state.sending) return false; + + const hasInput = state.hasText || state.hasAttachments; + switch (state.delivery) { + case "submit": + return state.workerState === "idle" && hasInput; + case "queue": + return state.workerState === "running" && hasInput; + case "notify": + return state.workerState === "running" && state.hasText && + !state.hasAttachments; + } +} + +export function sendComposerDelivery( + state: ComposerDeliveryState, + method: T, + send: (method: T) => void, +): boolean { + if (!canDeliverComposerDraft(state)) return false; + send(method); + return true; +} diff --git a/web/workspace/src/lib/workspace/console/model.test.ts b/web/workspace/src/lib/workspace/console/model.test.ts index e2dfdceb..60b98633 100644 --- a/web/workspace/src/lib/workspace/console/model.test.ts +++ b/web/workspace/src/lib/workspace/console/model.test.ts @@ -2150,6 +2150,12 @@ Deno.test("snapshot restores TaskStore state from system history", () => { const event = snapshotEvent("/repo"); if (event.event !== "snapshot") throw new Error("snapshot fixture expected"); event.data.session = { + pending_submissions: { + revision: 0, + notification_count: 0, + head_id: null, + submissions: [], + }, entries: [{ entry_id: "task-reminder-1", timestamp: 1, diff --git a/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts b/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts index e7effcb7..abd53396 100644 --- a/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts +++ b/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts @@ -1064,3 +1064,46 @@ Deno.test("Web Console switches main and direct SubWorker views from the Tasks r "Worker view selection should expose only direct SubWorker session identities with main fallback", ); }); + +Deno.test("Web Console uses Notify while running and exposes durable pending controls", async () => { + const consolePage = await Deno.readTextFile( + new URL( + "./../../../routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte", + import.meta.url, + ), + ); + + for ( + const token of [ + 'method: "submit"', + 'method: "notify"', + "notification_request_id: crypto.randomUUID()", + "submission_request_id: crypto.randomUUID()", + 'payload.event === "pending_submissions_changed"', + 'method: "cancel_pending_submission"', + 'method: "clear_pending_submissions"', + 'method: "continue_pending"', + "handleQueueSubmit", + "handleNotifySubmit", + 'submitDraft(composerInputElement?.snapshot() ?? draft, "queue")', + "disabled={!canQueueDraft}", + "disabled={!canNotifyDraft}", + ">Queue Submit", + ">Notify", + ] + ) { + assert( + consolePage.includes(token), + `missing durable pending control token: ${token}`, + ); + } + + const userCase = consolePage.slice( + consolePage.indexOf('case "user":'), + consolePage.indexOf('case "compact":'), + ); + assert( + !userCase.includes("workerRunning"), + "ordinary text must remain Submit instead of being implicitly converted to Notify", + ); +}); diff --git a/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte b/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte index 03efcf8d..c7852975 100644 --- a/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte +++ b/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte @@ -5,6 +5,11 @@ import ConsoleTimeline from "$lib/workspace/console/ConsoleTimeline.svelte"; import ComposerInput from "$lib/workspace/console/ComposerInput.svelte"; import type { ComposerDraftSnapshot } from "$lib/workspace/console/composer-draft"; + import { + canDeliverComposerDraft, + sendComposerDelivery, + type ComposerDelivery, + } from "$lib/workspace/console/composer-delivery"; import { buildComposerSegmentsRequest, type WorkerConsoleInputRequest, @@ -31,7 +36,13 @@ type ConsoleViewMode, type ConsoleViewScroll, } from "$lib/workspace/console/model"; - import type { Event as ProtocolEvent, Method as ProtocolMethod, RewindTarget, Segment } from "$lib/generated/protocol"; + import type { + Event as ProtocolEvent, + Method as ProtocolMethod, + PendingSubmissionsSnapshot, + RewindTarget, + Segment, + } from "$lib/generated/protocol"; import { MAX_FILES_PER_SUBMISSION, uploadAttachment, @@ -152,6 +163,13 @@ "connecting", ); let protocolSubscription: WorkspaceMultiplexerSubscription | null = null; + let pendingSubmissions = $state({ + revision: 0, + notification_count: 0, + head_id: null, + submissions: [], + }); + let pendingSubmissionItems = $derived(pendingSubmissions.submissions ?? []); let pendingCompletionRequest: { resolve: (entries: ComposerCompletionEntry[]) => void; reject: (error: Error) => void; @@ -234,13 +252,42 @@ const workerState = $derived(liveWorkerState ?? worker?.state ?? "loading"); const workerRunning = $derived(workerState === "running"); const workerPaused = $derived(workerState === "paused"); - const inputReady = $derived(workerState === "idle"); const composerEditable = $derived(protocolState === "open" && !sending); - const canSubmitDraft = $derived(inputReady && composerEditable); - const canSend = $derived(canSubmitDraft && draft.content.trim().length > 0); + const draftHasText = $derived(draft.content.trim().length > 0); + const draftHasAttachments = $derived(attachments.length > 0); + const canSubmitDraft = $derived( + canDeliverComposerDraft({ + delivery: "submit", + workerState, + protocolOpen: protocolState === "open", + sending, + hasText: draftHasText, + hasAttachments: draftHasAttachments, + }), + ); + const canQueueDraft = $derived( + canDeliverComposerDraft({ + delivery: "queue", + workerState, + protocolOpen: protocolState === "open", + sending, + hasText: draftHasText, + hasAttachments: draftHasAttachments, + }), + ); + const canNotifyDraft = $derived( + canDeliverComposerDraft({ + delivery: "notify", + workerState, + protocolOpen: protocolState === "open", + sending, + hasText: draftHasText, + hasAttachments: draftHasAttachments, + }), + ); const canStopFromComposer = $derived(workerRunning && composerEditable); const composerSubmitDisabled = $derived( - workerRunning ? !canStopFromComposer : !canSend, + workerRunning ? !canStopFromComposer : !canSubmitDraft, ); async function getJson(path: string): Promise { @@ -334,6 +381,13 @@ function handleIncomingProtocolEvent(payload: ProtocolEvent) { handleProtocolCommandEvent(payload); + if (payload.event === "snapshot") { + pendingSubmissions = payload.data.session.pending_submissions; + } else if (payload.event === "segment_rotated") { + pendingSubmissions = payload.data.session.pending_submissions; + } else if (payload.event === "pending_submissions_changed") { + pendingSubmissions = payload.data.pending; + } if (payload.event === "error") { queueObservationDiagnostic({ code: payload.data.code, @@ -556,8 +610,9 @@ switch (request.kind) { case "user": return { - method: "run", + method: "submit", params: { + submission_request_id: crypto.randomUUID(), input: request.segments ?? [ { kind: "text", content: request.content }, ], @@ -566,7 +621,11 @@ case "notify": return { method: "notify", - params: { message: request.content, auto_run: true }, + params: { + notification_request_id: crypto.randomUUID(), + message: request.content, + auto_run: true, + }, }; case "compact": return { method: "compact" }; @@ -638,6 +697,14 @@ void submitDraft(composerInputElement?.snapshot() ?? draft); } + function handleQueueSubmit() { + void submitDraft(composerInputElement?.snapshot() ?? draft, "queue"); + } + + function handleNotifySubmit() { + void submitDraft(composerInputElement?.snapshot() ?? draft, "notify"); + } + function attachmentPath(): string { return `/api/w/${encodeURIComponent(workspaceId)}/runtimes/${encodeURIComponent(runtimeId)}/workers/${encodeURIComponent(workerId)}`; } @@ -737,7 +804,15 @@ if (event.dataTransfer?.files) addAttachmentFiles(event.dataTransfer.files); } - async function submitDraft(value: ComposerDraftSnapshot) { + async function submitDraft( + value: ComposerDraftSnapshot, + delivery: ComposerDelivery = "submit", + ) { + if (delivery === "notify" && attachments.length > 0) { + composerNotice = null; + sendError = "Notify accepts text only; remove attachments or queue a Submit."; + return; + } const incompleteAttachment = attachments.find((attachment) => attachment.state !== "uploaded" || !attachment.reference ); @@ -767,19 +842,38 @@ composerInputElement?.clear(); return; } - if (sending || !inputReady) { + const deliveryState = { + delivery, + workerState, + protocolOpen: protocolState === "open", + sending, + hasText: value.content.trim().length > 0, + hasAttachments: attachments.length > 0, + }; + if (!canDeliverComposerDraft(deliveryState)) { return; } + let request: WorkerConsoleInputRequest = command.request; + if (delivery === "notify") { + if (request.kind !== "user") { + composerNotice = null; + sendError = "Notify accepts ordinary text, not a Composer command."; + return; + } + request = { kind: "notify", content: request.content }; + } sending = true; sendError = null; try { - const method = composerRequestToProtocolMethod(command.request); - sendProtocolMethod(method); + const method = composerRequestToProtocolMethod(request); + if (!sendComposerDelivery(deliveryState, method, sendProtocolMethod)) { + return; + } composerInputElement?.recordHistory(value); composerInputElement?.clear(); attachments = []; - if (method.method === "run" || method.method === "notify") { + if (method.method === "submit" || method.method === "notify") { liveWorkerState = "running"; } composerNotice = "Sent through Worker protocol."; @@ -1722,6 +1816,62 @@ {/if} + {#if pendingSubmissionItems.length > 0 || pendingSubmissions.notification_count > 0} +
+ + Pending activations ({pendingSubmissionItems.length} submissions · {pendingSubmissions.notification_count} notifications) + +
    + {#each pendingSubmissionItems as submission (submission.submission_id)} +
  1. + {submission.submission_id} + {submission.segment_count} segments · {submission.byte_len} bytes + +
    2. + {/each} +
+ + +
+ {/if} + {#if workerRunning}
+ {#if workerRunning} + + + {/if} {#if composerNotice} {composerNotice} {/if} @@ -2035,6 +2197,31 @@ display: none; } + .pending-submissions { + margin: 0 var(--space-3); + color: var(--muted); + font-size: 0.75rem; + } + + .pending-submissions ol { + display: grid; + gap: var(--space-1); + margin: var(--space-2) 0; + padding-left: var(--space-5); + } + + .pending-submissions li { + display: flex; + gap: var(--space-2); + align-items: center; + } + + .pending-submissions code { + max-width: 16rem; + overflow: hidden; + text-overflow: ellipsis; + } + .console-log { display: grid; align-content: start;