From dea5bd581d1be6d5f80674af436ffe46c2852985 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 02:31:10 +0900 Subject: [PATCH] fix: secure submit activation handoff --- crates/protocol/src/lib.rs | 22 ++- crates/session-store/src/fs_store.rs | 64 +++++++-- crates/session-store/src/store.rs | 10 ++ crates/session-store/src/uploaded_file.rs | 37 ++++- crates/worker-runtime/src/http_server.rs | 59 ++++++++ crates/worker-runtime/src/worker_backend.rs | 20 ++- crates/worker/src/controller.rs | 143 ++++++++++++++++++-- crates/worker/src/discovery.rs | 12 +- crates/worker/src/spawn/tool.rs | 102 ++++++++++---- crates/worker/src/worker.rs | 66 +++++++++ 10 files changed, 464 insertions(+), 71 deletions(-) diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index e6d790db..de5e8549 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -39,6 +39,9 @@ fn is_false(value: &bool) -> bool { #[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, }, @@ -46,19 +49,30 @@ pub enum AuthenticatedInputSource { 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}"), } } @@ -83,6 +97,7 @@ pub enum Method { SubmitTracked { submission_request_id: String, input: Vec, + #[serde(skip_deserializing, default)] source: AuthenticatedInputSource, }, /// Human-readable text injected into the target Worker's LLM context @@ -104,6 +119,7 @@ pub enum Method { 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. @@ -1538,7 +1554,7 @@ mod tests { } #[test] - fn authenticated_submit_round_trips_trusted_source() { + fn authenticated_submit_replaces_wire_source_with_transport_identity() { let method = Method::SubmitTracked { input: vec![Segment::text("private")], submission_request_id: "request-1".to_string(), @@ -1551,9 +1567,9 @@ mod tests { assert!(matches!( decoded, Method::SubmitTracked { - source: AuthenticatedInputSource::Account { account_id }, + source: AuthenticatedInputSource::UntrustedWire, .. - } if account_id == "account-1" + } )); assert!( serde_json::from_str::( diff --git a/crates/session-store/src/fs_store.rs b/crates/session-store/src/fs_store.rs index abdf14e4..0ebb1498 100644 --- a/crates/session-store/src/fs_store.rs +++ b/crates/session-store/src/fs_store.rs @@ -21,9 +21,9 @@ 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, - pin_uploaded_file, read_uploaded_file, read_uploaded_file_by_id, release_uploaded_file_pin, - 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, + release_uploaded_file_pin, uploaded_file_has_pending_owner, write_uploaded_file, }; use crate::{ PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits, UploadedFileUploadContext, @@ -545,6 +545,23 @@ impl Store for FsStore { 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 delete_uploaded_file( &self, session_id: SessionId, @@ -568,13 +585,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) @@ -930,19 +952,26 @@ mod tests { store .copy_committed_uploaded_files(session_id, fork_session_id) .unwrap(), - 1 + 0 ); - assert_eq!( + assert!( store .read_uploaded_file_by_id(fork_session_id, &pending.artifact_id) - .unwrap() - .1, - b"pending" + .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, @@ -959,6 +988,11 @@ mod tests { 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) diff --git a/crates/session-store/src/store.rs b/crates/session-store/src/store.rs index 4c98fa33..b14839fa 100644 --- a/crates/session-store/src/store.rs +++ b/crates/session-store/src/store.rs @@ -246,6 +246,16 @@ pub trait Store: Send + Sync { 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) + } + /// 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 a7c544ea..c8474b0f 100644 --- a/crates/session-store/src/uploaded_file.rs +++ b/crates/session-store/src/uploaded_file.rs @@ -347,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 @@ -441,7 +447,7 @@ pub(crate) fn release_uploaded_file_pin( 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.is_some() || stored.pending_owner_id.as_deref() != Some(owner_id) { + if stored.pending_owner_id.as_deref() != Some(owner_id) { return Err(StoreError::ArtifactIntegrityMismatch); } stored.pending_owner_id = None; @@ -451,6 +457,32 @@ pub(crate) fn release_uploaded_file_pin( 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, @@ -480,7 +512,6 @@ pub(crate) fn bind_uploaded_file( return Err(StoreError::ArtifactAlreadyCommitted); } stored.source_entry_id = Some(source_entry_id.to_owned()); - stored.pending_owner_id = None; let temp = dir.join(format!(".{}.file.bind.tmp", reference.artifact_id)); fs::write(&temp, serde_json::to_vec(&stored)?)?; fs::rename(&temp, path)?; @@ -531,7 +562,7 @@ pub(crate) fn copy_committed_uploaded_files(source_dir: &Path, target_dir: &Path } let bytes = fs::read(&path)?; let stored: StoredUploadedFile = serde_json::from_slice(&bytes)?; - if stored.source_entry_id.is_none() && stored.pending_owner_id.is_none() { + if stored.source_entry_id.is_none() { continue; } let target = target_dir.join(name); diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index 575df269..5c849c14 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -1203,6 +1203,37 @@ async fn worker_protocol_ws( .into_response()) } +#[cfg(feature = "ws-server")] +fn authorize_runtime_protocol_method(method: protocol::Method) -> protocol::Method { + match method { + protocol::Method::SubmitTracked { + submission_request_id, + input, + .. + } => protocol::Method::SubmitTracked { + source: 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: 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, @@ -1291,6 +1322,7 @@ 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); let result = match scope.as_ref() { Some(scope) => { runtime.send_protocol_method_scoped(scope, &worker_ref, method) @@ -2084,6 +2116,33 @@ mod tests { WorkdirPath, WorkdirSessionCapabilities, }; + #[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), + protocol::Method::SubmitTracked { + source: protocol::AuthenticatedInputSource::Backend { operation_id }, + .. + } if operation_id == "request-1" + )); + } + #[test] fn attachment_routes_require_worker_input_permission() { assert_eq!( diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 408d8155..5225dd8b 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -1898,15 +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 { - notification_request_id: input - .submission_request_id - .unwrap_or_else(protocol::new_submission_request_id), + 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, ); @@ -2066,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 diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 14fd26f6..e39c337d 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -229,6 +229,60 @@ enum PendingRun { 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, @@ -266,6 +320,25 @@ fn stage_oldest_passive_notification( }) } +fn prepare_restored_auto_notification( + pending_submissions: &crate::worker::PendingSubmissionHandle, + notify_buffer: &NotifyBuffer, +) -> Option { + let notification = pending_submissions.prepare_oldest_auto_notification()?; + let extension = pending_submissions.notification_activation_extension(); + let notification_request_id = notification.notification_request_id.clone(); + notify_buffer.push_durable_notify( + notification.message, + true, + notification.provenance, + extension, + ); + Some(PendingRun::RunForNotification { + invoke_kind: protocol::InvokeKind::Notify, + notification_request_id: Some(notification_request_id), + }) +} + fn prepare_pending_run( pending_submissions: &crate::worker::PendingSubmissionHandle, notify_buffer: &NotifyBuffer, @@ -1007,11 +1080,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(); @@ -1354,9 +1433,9 @@ 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 = prepare_restored_auto_notification(&pending_submissions, ¬ify_buffer); loop { // Top-of-iteration: if an event handler staged a run, fire it @@ -1553,11 +1632,13 @@ async fn controller_loop( 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(), - crate::worker::authenticated_input_provenance(&source), + source_namespace, + provenance, true, ) { Ok(acceptance) => { @@ -1630,12 +1711,13 @@ async fn controller_loop( source, } => { let request_id = notification_request_id.clone(); - let source_namespace = source.namespace(); + let (source_namespace, provenance) = + resolved_input_source(&pending_submissions, &source); match pending_submissions.accept_notification_from_source( notification_request_id, message, source_namespace.clone(), - crate::worker::authenticated_input_provenance(&source), + provenance, auto_run, ) { Ok(_) if auto_run => { @@ -2240,11 +2322,13 @@ where 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(), - crate::worker::authenticated_input_provenance(&source), + source_namespace, + provenance, false, ) { Ok(acceptance) => { @@ -2353,12 +2437,13 @@ where source, }) => { let request_id = notification_request_id.clone(); - let source_namespace = source.namespace(); + let (source_namespace, provenance) = + resolved_input_source(pending_submissions, &source); match pending_submissions.accept_notification_from_source( notification_request_id, message, source_namespace.clone(), - crate::worker::authenticated_input_provenance(&source), + provenance, auto_run, ) { Ok(_) if !auto_run => { @@ -2557,6 +2642,36 @@ 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_restored_auto_notification(&pending, ¬ify_buffer), + Some(PendingRun::RunForNotification { + notification_request_id: Some(_), + .. + }) + )); + assert!(notify_buffer.has_auto_run_pending()); + } + #[test] fn image_attachment_gate_requires_vision_and_supported_openai_scheme() { let openai = manifest::ModelManifest { diff --git a/crates/worker/src/discovery.rs b/crates/worker/src/discovery.rs index effc1f6f..755180f4 100644 --- a/crates/worker/src/discovery.rs +++ b/crates/worker/src/discovery.rs @@ -1012,12 +1012,16 @@ 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<()> { + let notification_request_id = protocol::new_submission_request_id(); connect_and_send( socket_path, - &Method::Notify { - notification_request_id: protocol::new_submission_request_id(), + &Method::NotifyTracked { + notification_request_id: notification_request_id.clone(), message, auto_run, + source: protocol::AuthenticatedInputSource::Backend { + operation_id: notification_request_id, + }, }, ) .await @@ -1546,7 +1550,7 @@ mod tests { .await .unwrap(); let method = reader.next::().await.unwrap().unwrap(); - if let Method::Notify { + if let Method::NotifyTracked { message, auto_run, .. } = method { @@ -1668,7 +1672,7 @@ mod tests { .await .unwrap(); let method = reader.next::().await.unwrap().unwrap(); - if let Method::Notify { + if let Method::NotifyTracked { message, auto_run, .. } = method { diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index b071ee29..6bb1b47b 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -219,37 +219,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 { - notification_request_id: protocol::new_submission_request_id(), - 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), } } } @@ -554,7 +567,7 @@ impl Tool for SubWorkerSpawnTool { let message = format!( "SubWorker `{child_name}` turn ended with status {status:?}. Inspect its committed session with worker-observation tools before making completion decisions." ); - parent_notifications.notify(message, true); + parent_notifications.notify(child_name.clone(), message, true); })), ) .await; @@ -1138,12 +1151,41 @@ enabled = false #[tokio::test] async fn parent_controller_notification_target_does_not_keep_channel_open() { let (parent_method_tx, mut parent_method_rx) = mpsc::channel(1); - let target = ParentNotificationTarget::Controller(parent_method_tx.downgrade()); + let captured = Arc::new(std::sync::Mutex::new(false)); + let captured_for_fallback = captured.clone(); + let target = ParentNotificationTarget::with_controller_fallback( + parent_method_tx.downgrade(), + ParentNotificationTarget::Durable(Arc::new(move |_| { + *captured_for_fallback.lock().unwrap() = true; + })), + ); drop(parent_method_tx); assert!(parent_method_rx.recv().await.is_none()); - target.notify("late completion".to_string(), true); + target.notify("child-session".into(), "late completion".to_string(), true); + assert!(*captured.lock().unwrap()); + } + + #[test] + fn durable_parent_notification_target_preserves_child_source() { + let captured = Arc::new(std::sync::Mutex::new(None)); + let captured_for_target = captured.clone(); + let target = ParentNotificationTarget::Durable(Arc::new(move |method| { + *captured_for_target.lock().unwrap() = Some(method); + })); + + target.notify("child-session".into(), "completed".into(), true); + + assert!(matches!( + captured.lock().unwrap().take(), + Some(Method::NotifyTracked { + message, + auto_run: true, + source: protocol::AuthenticatedInputSource::SubWorker { session_id }, + .. + }) if session_id == "child-session" && message == "completed" + )); } #[tokio::test] @@ -1189,7 +1231,10 @@ enabled = false let tool = SubWorkerSpawnTool::new( "parent".into(), workspace_context, - ParentNotificationTarget::Controller(parent_method_tx.downgrade()), + ParentNotificationTarget::with_controller_fallback( + parent_method_tx.downgrade(), + ParentNotificationTarget::Durable(Arc::new(|_| {})), + ), runtime.path().to_path_buf(), bash_output_dir.clone(), workspace_root.clone(), @@ -1286,11 +1331,13 @@ enabled = false .expect("parent method channel remains open"); assert!(matches!( completion, - Method::Notify { + Method::NotifyTracked { message, auto_run: true, + source: protocol::AuthenticatedInputSource::SubWorker { session_id }, .. - } if message.contains("SubWorker `reviewer-child` turn ended with status Idle") + } if session_id == "reviewer-child" + && message.contains("SubWorker `reviewer-child` turn ended with status Idle") )); assert!(!runtime.path().join("reviewer-child/sock").exists()); @@ -1439,7 +1486,10 @@ enabled = false let tool = SubWorkerSpawnTool::new( "parent".into(), workspace_context, - ParentNotificationTarget::Controller(parent_method_tx.downgrade()), + ParentNotificationTarget::with_controller_fallback( + parent_method_tx.downgrade(), + ParentNotificationTarget::Durable(Arc::new(|_| {})), + ), runtime.path().to_path_buf(), bash_output_dir.clone(), workspace_root.clone(), diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index d24658d0..2aecd80d 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -208,6 +208,7 @@ 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(), @@ -223,6 +224,15 @@ pub(crate) fn authenticated_input_provenance( 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()), @@ -1604,6 +1614,27 @@ where Some(notification) } + pub(crate) fn prepare_oldest_auto_notification(&self) -> Option { + let mut state = self + .state + .lock() + .expect("pending activation state poisoned"); + if state.activating.is_some() || state.activating_notification.is_some() { + return None; + } + let index = state + .pending_notifications + .iter() + .position(|notification| notification.auto_run)?; + let notification = state + .pending_notifications + .remove(index) + .expect("located auto-run 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)>, @@ -3822,6 +3853,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 @@ -4044,6 +4080,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],