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/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],