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

fix: secure submit activation handoff

Browse Source
This commit is contained in:
Hare
2026-09-06 02:31:10 +09:00
parent b038f022d3
commit dea5bd581d
10 changed files with 464 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/protocol/src/lib.rs
+19 -3
View File
@@ -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<Segment>,
#[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::<Method>(
crates/session-store/src/fs_store.rs
+49 -15
View File
@@ -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)
crates/session-store/src/store.rs
+10
View File
@@ -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,
crates/session-store/src/uploaded_file.rs
+34 -3
View File
@@ -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<bool> {
let path = record_path(dir, artifact_id)?;
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(path)?)?;
Ok(stored.pending_owner_id.is_some())
}
pub(crate) fn read_uploaded_file(dir: &Path, reference: &UploadedFileRef) -> Result<Vec<u8>> {
let (stored_reference, content) = read_uploaded_file_by_id(dir, &reference.artifact_id)?;
if stored_reference.file_name != reference.file_name
@@ -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);
crates/worker-runtime/src/http_server.rs
+59
View File
@@ -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!(
crates/worker-runtime/src/worker_backend.rs
+14 -6
View File
@@ -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
crates/worker/src/controller.rs
+129 -14
View File
@@ -229,6 +229,60 @@ enum PendingRun {
Resume,
}
fn resolved_input_source<St: Store + Clone>(
pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
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<St: Store + Clone + Send + Sync + 'static>(
pending_submissions: crate::worker::PendingSubmissionHandle<St>,
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,
&notify_buffer,
&source_namespace,
&notification_request_id,
);
}
Ok(_) => {}
Err(error) => {
tracing::warn!(%error, "failed to durably accept SubWorker notification");
}
}
}))
}
fn stage_pending_notification<St: Store + Clone>(
pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
notify_buffer: &NotifyBuffer,
@@ -266,6 +320,25 @@ fn stage_oldest_passive_notification<St: Store + Clone>(
})
}
fn prepare_restored_auto_notification<St: Store + Clone>(
pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
notify_buffer: &NotifyBuffer,
) -> Option<PendingRun> {
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<St: Store + Clone>(
pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
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<C, St>(
discovery_cwd,
spawned_registry.clone(),
);
let mut pending: Option<PendingRun> = None;
let pending_submissions = worker.pending_submission_handle();
stage_oldest_passive_notification(&pending_submissions, &notify_buffer);
let mut pending = prepare_restored_auto_notification(&pending_submissions, &notify_buffer);
loop {
// Top-of-iteration: if an event handler staged a run, fire it
@@ -1553,11 +1632,13 @@ async fn controller_loop<C, St>(
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<C, St>(
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, &notify_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 {
crates/worker/src/discovery.rs
+8 -4
View File
@@ -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::<Method>().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::<Method>().await.unwrap().unwrap();
if let Method::Notify {
if let Method::NotifyTracked {
message, auto_run, ..
} = method
{
crates/worker/src/spawn/tool.rs
+76 -26
View File
@@ -219,37 +219,50 @@ fn parse_spawn_profile_selector(raw: Option<&str>) -> Result<SpawnProfileSelecto
#[derive(Clone)]
pub(crate) enum ParentNotificationTarget {
Controller(mpsc::WeakSender<Method>),
Buffer(crate::ipc::notify_buffer::NotifyBuffer),
Controller {
sender: mpsc::WeakSender<Method>,
fallback: Arc<dyn Fn(Method) + Send + Sync>,
},
Durable(Arc<dyn Fn(Method) + Send + Sync>),
}
impl ParentNotificationTarget {
fn notify(&self, message: String, auto_run: bool) {
pub(crate) fn with_controller_fallback(
sender: mpsc::WeakSender<Method>,
fallback: ParentNotificationTarget,
) -> Self {
let ParentNotificationTarget::Durable(fallback) = fallback else {
unreachable!("controller fallback must use durable pending authority");
};
Self::Controller { sender, fallback }
}
pub(crate) fn notify(&self, child_session_id: String, message: String, auto_run: bool) {
let method = Method::NotifyTracked {
notification_request_id: protocol::new_submission_request_id(),
message,
auto_run,
source: protocol::AuthenticatedInputSource::SubWorker {
session_id: child_session_id,
},
};
match self {
Self::Controller(parent_method_tx) => {
let Some(parent_method_tx) = parent_method_tx.upgrade() else {
tracing::warn!(
"parent Worker controller closed before Internal SubWorker completion notification"
);
Self::Controller { sender, fallback } => {
let Some(parent_method_tx) = sender.upgrade() else {
fallback(method);
return;
};
let fallback = fallback.clone();
tokio::spawn(async move {
if let Err(error) = parent_method_tx
.send(Method::Notify {
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(),
crates/worker/src/worker.rs
+66
View File
@@ -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<PendingNotification> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
.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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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],