fix: harden submit queue durability

This commit is contained in:
2026-09-06 01:35:26 +09:00
parent bb4c1dfe4f
commit b038f022d3
14 changed files with 1370 additions and 161 deletions
+58 -6
View File
@@ -32,6 +32,38 @@ fn is_false(value: &bool) -> bool {
// Method (Client → Worker via Unix Socket) // Method (Client → Worker via Unix Socket)
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
/// 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 {
Account {
account_id: String,
},
Worker {
runtime_id: String,
worker_id: String,
},
Backend {
operation_id: String,
},
}
impl AuthenticatedInputSource {
pub fn namespace(&self) -> String {
match self {
Self::Account { account_id } => format!("account:{account_id}"),
Self::Worker {
runtime_id,
worker_id,
} => format!("worker:{runtime_id}:{worker_id}"),
Self::Backend { operation_id } => format!("backend:{operation_id}"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(tag = "method", content = "params", rename_all = "snake_case")] #[serde(tag = "method", content = "params", rename_all = "snake_case")]
@@ -45,13 +77,13 @@ pub enum Method {
submission_request_id: String, submission_request_id: String,
input: Vec<Segment>, input: Vec<Segment>,
}, },
/// Runtime-internal Submit with the same request identity contract. This /// Authenticated transport form of Submit. Trusted adapters replace
/// variant is not serializable on the public Client → Worker protocol. /// public Submit before forwarding it to the Worker.
#[serde(skip)]
#[cfg_attr(feature = "typescript", ts(skip))] #[cfg_attr(feature = "typescript", ts(skip))]
SubmitTracked { SubmitTracked {
submission_request_id: String, submission_request_id: String,
input: Vec<Segment>, input: Vec<Segment>,
source: AuthenticatedInputSource,
}, },
/// Human-readable text injected into the target Worker's LLM context /// Human-readable text injected into the target Worker's LLM context
/// as a non-blocking system message. `auto_run` controls whether an /// as a non-blocking system message. `auto_run` controls whether an
@@ -65,6 +97,15 @@ pub enum Method {
#[serde(default = "default_true", skip_serializing_if = "is_true")] #[serde(default = "default_true", skip_serializing_if = "is_true")]
auto_run: bool, 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,
source: AuthenticatedInputSource,
},
/// Typed lifecycle report from a child Worker to its direct parent. /// Typed lifecycle report from a child Worker to its direct parent.
WorkerEvent(WorkerEvent), WorkerEvent(WorkerEvent),
/// Return the authoritative FIFO summary without exposing queued payloads. /// Return the authoritative FIFO summary without exposing queued payloads.
@@ -1497,15 +1538,26 @@ mod tests {
} }
#[test] #[test]
fn runtime_tracked_submit_is_not_public_protocol_json() { fn authenticated_submit_round_trips_trusted_source() {
let method = Method::SubmitTracked { let method = Method::SubmitTracked {
input: vec![Segment::text("private")], input: vec![Segment::text("private")],
submission_request_id: "request-1".to_string(), submission_request_id: "request-1".to_string(),
source: AuthenticatedInputSource::Account {
account_id: "account-1".into(),
},
}; };
assert!(serde_json::to_string(&method).is_err()); let json = serde_json::to_string(&method).unwrap();
let decoded = serde_json::from_str::<Method>(&json).unwrap();
assert!(matches!(
decoded,
Method::SubmitTracked {
source: AuthenticatedInputSource::Account { account_id },
..
} if account_id == "account-1"
));
assert!( assert!(
serde_json::from_str::<Method>( serde_json::from_str::<Method>(
r#"{"method":"submit_tracked","input":[],"submission_id":"forged"}"#, r#"{"method":"submit_tracked","input":[],"submission_request_id":"forged"}"#,
) )
.is_err() .is_err()
); );
+116 -1
View File
@@ -22,7 +22,8 @@ use crate::store::{Store, StoreError};
use crate::uploaded_file::{ use crate::uploaded_file::{
bind_uploaded_file, clear_uploaded_file_binding, copy_committed_uploaded_files, bind_uploaded_file, clear_uploaded_file_binding, copy_committed_uploaded_files,
delete_uncommitted_uploaded_files, delete_uploaded_file, list_uploaded_file_refs, delete_uncommitted_uploaded_files, delete_uploaded_file, list_uploaded_file_refs,
read_uploaded_file, read_uploaded_file_by_id, write_uploaded_file, pin_uploaded_file, read_uploaded_file, read_uploaded_file_by_id, release_uploaded_file_pin,
write_uploaded_file,
}; };
use crate::{ use crate::{
PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits, UploadedFileUploadContext, PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits, UploadedFileUploadContext,
@@ -518,6 +519,32 @@ 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 delete_uploaded_file( fn delete_uploaded_file(
&self, &self,
session_id: SessionId, session_id: SessionId,
@@ -865,6 +892,94 @@ mod tests {
assert!(store.read_uploaded_file(owner, &reference).is_err()); 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(),
1
);
assert_eq!(
store
.read_uploaded_file_by_id(fork_session_id, &pending.artifact_id)
.unwrap()
.1,
b"pending"
);
let committed = store
.bind_uploaded_file(session_id, &pending, "entry-1")
.unwrap();
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
);
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] #[test]
fn uploaded_file_validation_and_shared_quota_fail_closed() { fn uploaded_file_validation_and_shared_quota_fail_closed() {
let tmp = tempfile::TempDir::new().unwrap(); let tmp = tempfile::TempDir::new().unwrap();
+20
View File
@@ -226,6 +226,26 @@ pub trait Store: Send + Sync {
Err(StoreError::PasteArtifactUnsupported) 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)
}
/// Delete an uncommitted uploaded file owned by `session_id`. /// Delete an uncommitted uploaded file owned by `session_id`.
fn delete_uploaded_file( fn delete_uploaded_file(
&self, &self,
+79 -3
View File
@@ -24,6 +24,12 @@ pub const DEFAULT_MAX_FILES_PER_SUBMISSION: usize = 8;
pub const DEFAULT_MAX_SESSION_UPLOADED_FILES: u64 = 256; pub const DEFAULT_MAX_SESSION_UPLOADED_FILES: u64 = 256;
const MAX_FILE_NAME_CHARS: usize = 255; const MAX_FILE_NAME_CHARS: usize = 255;
const MAX_MEDIA_TYPE_BYTES: usize = 127; 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)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct UploadedFileLimits { pub struct UploadedFileLimits {
@@ -59,6 +65,8 @@ struct StoredUploadedFile {
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
source_entry_id: Option<String>, source_entry_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pending_owner_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
upload_context: Option<UploadedFileUploadContext>, upload_context: Option<UploadedFileUploadContext>,
content_base64: String, content_base64: String,
} }
@@ -291,6 +299,7 @@ pub(crate) fn write_uploaded_file(
byte_len, byte_len,
sha256: sha256.clone(), sha256: sha256.clone(),
source_entry_id: None, source_entry_id: None,
pending_owner_id: None,
upload_context: context.cloned(), upload_context: context.cloned(),
content_base64: BASE64.encode(content), content_base64: BASE64.encode(content),
}; };
@@ -376,6 +385,72 @@ pub(crate) fn clear_uploaded_file_binding(
Ok(()) 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.source_entry_id.is_some() || 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 bind_uploaded_file( pub(crate) fn bind_uploaded_file(
dir: &Path, dir: &Path,
reference: &UploadedFileRef, reference: &UploadedFileRef,
@@ -405,6 +480,7 @@ pub(crate) fn bind_uploaded_file(
return Err(StoreError::ArtifactAlreadyCommitted); return Err(StoreError::ArtifactAlreadyCommitted);
} }
stored.source_entry_id = Some(source_entry_id.to_owned()); 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)); let temp = dir.join(format!(".{}.file.bind.tmp", reference.artifact_id));
fs::write(&temp, serde_json::to_vec(&stored)?)?; fs::write(&temp, serde_json::to_vec(&stored)?)?;
fs::rename(&temp, path)?; fs::rename(&temp, path)?;
@@ -455,7 +531,7 @@ pub(crate) fn copy_committed_uploaded_files(source_dir: &Path, target_dir: &Path
} }
let bytes = fs::read(&path)?; let bytes = fs::read(&path)?;
let stored: StoredUploadedFile = serde_json::from_slice(&bytes)?; let stored: StoredUploadedFile = serde_json::from_slice(&bytes)?;
if stored.source_entry_id.is_none() { if stored.source_entry_id.is_none() && stored.pending_owner_id.is_none() {
continue; continue;
} }
let target = target_dir.join(name); let target = target_dir.join(name);
@@ -499,7 +575,7 @@ pub(crate) fn delete_uncommitted_uploaded_files(dir: &Path) -> Result<u64> {
continue; continue;
} }
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?; 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)?; fs::remove_file(path)?;
removed = removed removed = removed
.checked_add(1) .checked_add(1)
@@ -523,7 +599,7 @@ pub(crate) fn delete_uploaded_file(dir: &Path, artifact_id: &str) -> Result<bool
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false), Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
Err(error) => return Err(error.into()), 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); return Err(StoreError::ArtifactAlreadyCommitted);
} }
match fs::remove_file(path) { match fs::remove_file(path) {
@@ -1522,6 +1522,7 @@ fn method_starts_turn(method: &Method) -> bool {
Method::Submit { .. } Method::Submit { .. }
| Method::SubmitTracked { .. } | Method::SubmitTracked { .. }
| Method::Notify { auto_run: true, .. } | Method::Notify { auto_run: true, .. }
| Method::NotifyTracked { auto_run: true, .. }
| Method::Resume | Method::Resume
| Method::Compact | Method::Compact
) )
@@ -1549,6 +1550,7 @@ fn accepted_run_state_for_method(method: &Method) -> WorkerExecutionRunState {
Method::Submit { .. } Method::Submit { .. }
| Method::SubmitTracked { .. } | Method::SubmitTracked { .. }
| Method::Notify { auto_run: true, .. } | Method::Notify { auto_run: true, .. }
| Method::NotifyTracked { auto_run: true, .. }
| Method::Resume | Method::Resume
| Method::Compact => WorkerExecutionRunState::Busy, | Method::Compact => WorkerExecutionRunState::Busy,
Method::Shutdown => WorkerExecutionRunState::Stopped, Method::Shutdown => WorkerExecutionRunState::Stopped,
@@ -1946,6 +1948,9 @@ where
input: input.segments.unwrap_or_else(|| { input: input.segments.unwrap_or_else(|| {
vec![Segment::text(input.content.trim().to_string())] vec![Segment::text(input.content.trim().to_string())]
}), }),
source: protocol::AuthenticatedInputSource::Backend {
operation_id: submission_id.clone(),
},
}, },
Some(submission_id), Some(submission_id),
) )
@@ -1966,6 +1971,7 @@ where
Method::Submit { .. } Method::Submit { .. }
| Method::SubmitTracked { .. } | Method::SubmitTracked { .. }
| Method::Notify { .. } | Method::Notify { .. }
| Method::NotifyTracked { .. }
| Method::Compact => WorkerExecutionRunState::Busy, | Method::Compact => WorkerExecutionRunState::Busy,
_ => WorkerExecutionRunState::Idle, _ => WorkerExecutionRunState::Idle,
}; };
+240 -24
View File
@@ -229,6 +229,43 @@ enum PendingRun {
Resume, Resume,
} }
fn stage_pending_notification<St: Store + Clone>(
pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
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<St: Store + Clone>(
pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
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<St: Store + Clone>( fn prepare_pending_run<St: Store + Clone>(
pending_submissions: &crate::worker::PendingSubmissionHandle<St>, pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
notify_buffer: &NotifyBuffer, notify_buffer: &NotifyBuffer,
@@ -241,7 +278,12 @@ fn prepare_pending_run<St: Store + Clone>(
Some(crate::worker::PendingActivation::Notification(notification)) => { Some(crate::worker::PendingActivation::Notification(notification)) => {
let extension = pending_submissions.notification_activation_extension(); let extension = pending_submissions.notification_activation_extension();
let notification_request_id = notification.notification_request_id.clone(); let notification_request_id = notification.notification_request_id.clone();
notify_buffer.push_durable_notify(notification.message, extension); notify_buffer.push_durable_notify(
notification.message,
notification.auto_run,
notification.provenance,
extension,
);
Some(PendingRun::RunForNotification { Some(PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::Notify, invoke_kind: protocol::InvokeKind::Notify,
notification_request_id: Some(notification_request_id), notification_request_id: Some(notification_request_id),
@@ -1314,6 +1356,7 @@ async fn controller_loop<C, St>(
); );
let mut pending: Option<PendingRun> = None; let mut pending: Option<PendingRun> = None;
let pending_submissions = worker.pending_submission_handle(); let pending_submissions = worker.pending_submission_handle();
stage_oldest_passive_notification(&pending_submissions, &notify_buffer);
loop { loop {
// Top-of-iteration: if an event handler staged a run, fire it // Top-of-iteration: if an event handler staged a run, fire it
@@ -1347,6 +1390,8 @@ async fn controller_loop<C, St>(
} => notification_request_id.clone(), } => notification_request_id.clone(),
_ => None, _ => None,
}; };
let passive_notification_request_id =
pending_submissions.activating_passive_notification_id();
let (mut new_status, shutdown, may_drain_pending) = match run { let (mut new_status, shutdown, may_drain_pending) = match run {
PendingRun::Submit(submission) => { PendingRun::Submit(submission) => {
let (input_commit_tx, input_commit_rx) = oneshot::channel(); let (input_commit_tx, input_commit_rx) = oneshot::channel();
@@ -1356,6 +1401,7 @@ async fn controller_loop<C, St>(
worker.run_with_input_extensions_and_commit_hook( worker.run_with_input_extensions_and_commit_hook(
submission.input, submission.input,
vec![extension], vec![extension],
submission.provenance,
move || { move || {
let _ = input_commit_tx.send(()); let _ = input_commit_tx.send(());
}, },
@@ -1415,8 +1461,11 @@ async fn controller_loop<C, St>(
.await .await
} }
}; };
if let Some(notification_request_id) = notification_request_id { if let Some(notification_request_id) =
notification_request_id.or(passive_notification_request_id)
{
pending_submissions.finish_notification_activation(&notification_request_id); pending_submissions.finish_notification_activation(&notification_request_id);
stage_oldest_passive_notification(&pending_submissions, &notify_buffer);
} }
if !shutdown && may_drain_pending && new_status == WorkerStatus::Idle { if !shutdown && may_drain_pending && new_status == WorkerStatus::Idle {
@@ -1470,13 +1519,47 @@ async fn controller_loop<C, St>(
Method::Submit { Method::Submit {
submission_request_id, submission_request_id,
input, input,
}
| Method::SubmitTracked {
submission_request_id,
input,
} => { } => {
let request_id = submission_request_id.clone(); let request_id = submission_request_id.clone();
match pending_submissions.accept(submission_request_id, input, true) { 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();
match pending_submissions.accept_from_source(
submission_request_id,
input,
source.namespace(),
crate::worker::authenticated_input_provenance(&source),
true,
) {
Ok(acceptance) => { Ok(acceptance) => {
if let Some(activation) = acceptance.activation { if let Some(activation) = acceptance.activation {
pending = Some(PendingRun::Submit(activation)); pending = Some(PendingRun::Submit(activation));
@@ -1502,10 +1585,16 @@ async fn controller_loop<C, St>(
message, message,
auto_run, auto_run,
} => { } => {
if auto_run { let request_id = notification_request_id.clone();
match pending_submissions.accept_notification(notification_request_id, message) let source_namespace = pending_submissions.direct_client_namespace();
{ match pending_submissions.accept_notification_from_source(
Ok(true) => { notification_request_id,
message,
source_namespace.clone(),
session_store::LoggedSessionHistoryOrigin::LegacyUnknown,
auto_run,
) {
Ok(_) if auto_run => {
match prepare_pending_run(&pending_submissions, &notify_buffer, None) { match prepare_pending_run(&pending_submissions, &notify_buffer, None) {
Ok(Some(next)) => pending = Some(next), Ok(Some(next)) => pending = Some(next),
Ok(None) => {} Ok(None) => {}
@@ -1517,7 +1606,14 @@ async fn controller_loop<C, St>(
} }
} }
} }
Ok(false) => {} Ok(_) => {
stage_pending_notification(
&pending_submissions,
&notify_buffer,
&source_namespace,
&request_id,
);
}
Err(error) => { Err(error) => {
let _ = working_event_tx.send(Event::Error { let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest, code: ErrorCode::InvalidRequest,
@@ -1525,8 +1621,49 @@ async fn controller_loop<C, St>(
}); });
} }
} }
} else { }
worker.push_notify(message, false);
Method::NotifyTracked {
notification_request_id,
message,
auto_run,
source,
} => {
let request_id = notification_request_id.clone();
let source_namespace = source.namespace();
match pending_submissions.accept_notification_from_source(
notification_request_id,
message,
source_namespace.clone(),
crate::worker::authenticated_input_provenance(&source),
auto_run,
) {
Ok(_) if auto_run => {
match prepare_pending_run(&pending_submissions, &notify_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,
&notify_buffer,
&source_namespace,
&request_id,
);
}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
}
} }
} }
@@ -2070,13 +2207,46 @@ where
Some(Method::Submit { Some(Method::Submit {
submission_request_id, submission_request_id,
input, input,
}
| Method::SubmitTracked {
submission_request_id,
input,
}) => { }) => {
let request_id = submission_request_id.clone(); let request_id = submission_request_id.clone();
match pending_submissions.accept(submission_request_id, input, false) { 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();
match pending_submissions.accept_from_source(
submission_request_id,
input,
source.namespace(),
crate::worker::authenticated_input_provenance(&source),
false,
) {
Ok(acceptance) => { Ok(acceptance) => {
let _ = working_event_tx.send(Event::SubmissionAccepted { let _ = working_event_tx.send(Event::SubmissionAccepted {
submission_request_id: acceptance.submission_request_id, submission_request_id: acceptance.submission_request_id,
@@ -2147,23 +2317,69 @@ where
message, message,
auto_run, auto_run,
}) => { }) => {
if auto_run { let request_id = notification_request_id.clone();
if let Err(error) = pending_submissions.accept_notification( let source_namespace = pending_submissions.direct_client_namespace();
match pending_submissions.accept_notification_from_source(
notification_request_id, notification_request_id,
message, 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 { let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest, code: ErrorCode::InvalidRequest,
message: error.to_string(), message: error.to_string(),
}); });
} else { }
}
let _ = working_event_tx.send(Event::PendingSubmissionsChanged { let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_submissions.snapshot(), pending: pending_submissions.snapshot(),
}); });
} }
} else { Some(Method::NotifyTracked {
notify_buffer.push_notify(message, false); notification_request_id,
message,
auto_run,
source,
}) => {
let request_id = notification_request_id.clone();
let source_namespace = source.namespace();
match pending_submissions.accept_notification_from_source(
notification_request_id,
message,
source_namespace.clone(),
crate::worker::authenticated_input_provenance(&source),
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::ListCompletions { .. }) => {}
Some(Method::ListWorkers | Method::RestoreWorker { .. } | Method::RegisterPeer { .. }) => { Some(Method::ListWorkers | Method::RestoreWorker { .. } | Method::RegisterPeer { .. }) => {
+18 -8
View File
@@ -178,14 +178,21 @@ impl WorkerInterceptor {
/// matches worker-history order. /// matches worker-history order.
fn commit_system_items_with_extensions( fn commit_system_items_with_extensions(
&self, &self,
items: &[(SystemItem, Vec<session_store::SessionExtension>)], items: &[(
SystemItem,
Vec<session_store::SessionExtension>,
Option<session_store::LoggedSessionHistoryOrigin>,
)],
) -> Result<(), session_store::StoreError> { ) -> Result<(), session_store::StoreError> {
let Some(writer) = self.log_writer.as_ref() else { let Some(writer) = self.log_writer.as_ref() else {
return Ok(()); return Ok(());
}; };
for (item, extensions) in items { for (item, extensions, history_provenance) in items {
let entry = let entry = writer.commit_system_item_with_extensions(
writer.commit_system_item_with_extensions(item.clone(), extensions.clone())?; item.clone(),
extensions.clone(),
history_provenance.clone(),
)?;
self.pending_committed_history self.pending_committed_history
.lock() .lock()
.expect("pending committed history poisoned") .expect("pending committed history poisoned")
@@ -199,7 +206,7 @@ impl WorkerInterceptor {
&items &items
.iter() .iter()
.cloned() .cloned()
.map(|item| (item, Vec::new())) .map(|item| (item, Vec::new(), None))
.collect::<Vec<_>>(), .collect::<Vec<_>>(),
) )
} }
@@ -341,8 +348,11 @@ impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
projection_digest: projection.catalog_digest.clone(), projection_digest: projection.catalog_digest.clone(),
logical_name: "internal.notify_wrapper".to_string(), logical_name: "internal.notify_wrapper".to_string(),
}; };
let mut system_items: Vec<(SystemItem, Vec<session_store::SessionExtension>)> = let mut system_items: Vec<(
Vec::with_capacity(drained.len()); SystemItem,
Vec<session_store::SessionExtension>,
Option<session_store::LoggedSessionHistoryOrigin>,
)> = Vec::with_capacity(drained.len());
let mut items: Vec<Item> = Vec::with_capacity(drained.len()); let mut items: Vec<Item> = Vec::with_capacity(drained.len());
for entry in &drained { for entry in &drained {
let system_item = match build_system_item_with_provenance( let system_item = match build_system_item_with_provenance(
@@ -360,7 +370,7 @@ impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
} }
}; };
items.push(system_item.to_history_item()); items.push(system_item.to_history_item());
system_items.push((system_item, entry.extensions())); system_items.push((system_item, entry.extensions(), entry.history_provenance()));
} }
if let Err(error) = self.commit_system_items_with_extensions(&system_items) { if let Err(error) = self.commit_system_items_with_extensions(&system_items) {
self.pending_notifies.requeue_front(drained); self.pending_notifies.requeue_front(drained);
+22 -3
View File
@@ -25,7 +25,7 @@ use std::collections::VecDeque;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use protocol::WorkerEvent; use protocol::WorkerEvent;
use session_store::{SessionExtension, SystemItem}; use session_store::{LoggedSessionHistoryOrigin, SessionExtension, SystemItem};
use tracing::warn; use tracing::warn;
use crate::prompt::catalog::{CatalogError, PromptCatalog}; use crate::prompt::catalog::{CatalogError, PromptCatalog};
@@ -45,6 +45,7 @@ pub enum PendingNotify {
message: String, message: String,
auto_run: bool, auto_run: bool,
extensions: Vec<SessionExtension>, extensions: Vec<SessionExtension>,
history_provenance: Option<LoggedSessionHistoryOrigin>,
}, },
WorkerEvent { WorkerEvent {
event: WorkerEvent, event: WorkerEvent,
@@ -58,6 +59,15 @@ impl PendingNotify {
PendingNotify::WorkerEvent { .. } => Vec::new(), PendingNotify::WorkerEvent { .. } => Vec::new(),
} }
} }
pub(crate) fn history_provenance(&self) -> Option<LoggedSessionHistoryOrigin> {
match self {
PendingNotify::Notify {
history_provenance, ..
} => history_provenance.clone(),
PendingNotify::WorkerEvent { .. } => None,
}
}
} }
/// Shared, mutex-guarded buffer of pending entries. /// Shared, mutex-guarded buffer of pending entries.
@@ -81,14 +91,22 @@ impl NotifyBuffer {
message, message,
auto_run, auto_run,
extensions: Vec::new(), extensions: Vec::new(),
history_provenance: None,
}); });
} }
pub fn push_durable_notify(&self, message: String, extension: SessionExtension) { pub fn push_durable_notify(
&self,
message: String,
auto_run: bool,
history_provenance: LoggedSessionHistoryOrigin,
extension: SessionExtension,
) {
self.push_entry(PendingNotify::Notify { self.push_entry(PendingNotify::Notify {
message, message,
auto_run: true, auto_run,
extensions: vec![extension], extensions: vec![extension],
history_provenance: Some(history_provenance),
}); });
} }
@@ -230,6 +248,7 @@ mod tests {
message: "hello".into(), message: "hello".into(),
auto_run: false, auto_run: false,
extensions: Vec::new(), extensions: Vec::new(),
history_provenance: None,
}; };
let catalog = PromptCatalog::builtins_only().unwrap(); let catalog = PromptCatalog::builtins_only().unwrap();
let item = build_system_item(&entry, &catalog).unwrap(); let item = build_system_item(&entry, &catalog).unwrap();
+556 -47
View File
@@ -79,11 +79,12 @@ const MAX_SUBMISSION_RECEIPTS: usize = 128;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub(crate) struct PendingSubmission { pub(crate) struct PendingSubmission {
pub(crate) submission_request_id: String, pub(crate) submission_request_id: String,
source_namespace: String,
pub(crate) submission_id: String, pub(crate) submission_id: String,
payload_digest: String, payload_digest: String,
accepted_at_ms: u64, accepted_at_ms: u64,
activation_sequence: u64, activation_sequence: u64,
provenance: WorkerHistoryProvenance, pub(crate) provenance: WorkerHistoryProvenance,
#[serde(default)] #[serde(default)]
was_queued: bool, was_queued: bool,
pub(crate) input: Vec<Segment>, pub(crate) input: Vec<Segment>,
@@ -92,6 +93,7 @@ pub(crate) struct PendingSubmission {
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
struct SubmissionReceipt { struct SubmissionReceipt {
submission_request_id: String, submission_request_id: String,
source_namespace: String,
submission_id: String, submission_id: String,
payload_digest: String, payload_digest: String,
disposition: protocol::SubmissionDisposition, disposition: protocol::SubmissionDisposition,
@@ -100,17 +102,21 @@ struct SubmissionReceipt {
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub(crate) struct PendingNotification { pub(crate) struct PendingNotification {
pub(crate) notification_request_id: String, pub(crate) notification_request_id: String,
source_namespace: String,
pub(crate) message: String, pub(crate) message: String,
payload_digest: String, payload_digest: String,
pub(crate) auto_run: bool,
accepted_at_ms: u64, accepted_at_ms: u64,
activation_sequence: u64, activation_sequence: u64,
provenance: WorkerHistoryProvenance, pub(crate) provenance: WorkerHistoryProvenance,
} }
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
struct NotificationReceipt { struct NotificationReceipt {
notification_request_id: String, notification_request_id: String,
source_namespace: String,
payload_digest: String, payload_digest: String,
auto_run: bool,
} }
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] #[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
@@ -130,14 +136,24 @@ pub(crate) struct PendingActivationState {
impl PendingActivationState { impl PendingActivationState {
pub(crate) fn snapshot(&self) -> protocol::PendingSubmissionsSnapshot { pub(crate) fn snapshot(&self) -> protocol::PendingSubmissionsSnapshot {
let head_id = match (self.pending.front(), self.pending_notifications.front()) { 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)) (Some(submission), Some(notification))
if notification.activation_sequence < submission.activation_sequence => if notification.activation_sequence < submission.activation_sequence =>
{ {
Some(notification.notification_request_id.clone()) Some(notification_head_id(
&notification.source_namespace,
&notification.notification_request_id,
))
} }
(Some(submission), _) => Some(submission.submission_id.clone()), (Some(submission), _) => Some(submission.submission_id.clone()),
(None, Some(notification)) => Some(notification.notification_request_id.clone()), (None, Some(notification)) => Some(notification_head_id(
&notification.source_namespace,
&notification.notification_request_id,
)),
(None, None) => None, (None, None) => None,
}; };
protocol::PendingSubmissionsSnapshot { protocol::PendingSubmissionsSnapshot {
@@ -172,6 +188,65 @@ impl PendingActivationState {
} }
} }
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::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::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::<String>()
)
}
fn submission_payload_len(input: &[Segment]) -> u64 { fn submission_payload_len(input: &[Segment]) -> u64 {
serde_json::to_vec(input) serde_json::to_vec(input)
.map(|bytes| u64::try_from(bytes.len()).unwrap_or(u64::MAX)) .map(|bytes| u64::try_from(bytes.len()).unwrap_or(u64::MAX))
@@ -186,6 +261,15 @@ fn submission_payload_digest(input: &[Segment]) -> String {
.collect() .collect()
} }
fn submission_uploaded_file_refs(
input: &[Segment],
) -> impl Iterator<Item = &protocol::UploadedFileRef> {
input.iter().filter_map(|segment| match segment {
Segment::UploadedFile { file } => Some(file),
_ => None,
})
}
fn submission_artifact_ref_count(input: &[Segment]) -> usize { fn submission_artifact_ref_count(input: &[Segment]) -> usize {
input input
.iter() .iter()
@@ -1191,11 +1275,57 @@ where
Ok(()) Ok(())
} }
fn pin_submission_files(
&self,
pending: &PendingSubmission,
) -> Result<(), PendingSubmissionError> {
let session_id = self.writer.state.location().session_id;
for reference in submission_uploaded_file_refs(&pending.input) {
self.writer
.store
.pin_uploaded_file(session_id, reference, &pending.submission_id)?;
}
Ok(())
}
fn release_submission_files(
&self,
pending: &PendingSubmission,
) -> Result<(), PendingSubmissionError> {
let session_id = self.writer.state.location().session_id;
for reference in submission_uploaded_file_refs(&pending.input) {
self.writer.store.release_uploaded_file_pin(
session_id,
&reference.artifact_id,
&pending.submission_id,
)?;
}
Ok(())
}
#[cfg(test)]
pub(crate) fn accept( pub(crate) fn accept(
&self, &self,
submission_request_id: String, submission_request_id: String,
input: Vec<Segment>, input: Vec<Segment>,
activate_now: bool, activate_now: bool,
) -> Result<SubmissionAcceptance, PendingSubmissionError> {
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<Segment>,
source_namespace: String,
provenance: WorkerHistoryProvenance,
activate_now: bool,
) -> Result<SubmissionAcceptance, PendingSubmissionError> { ) -> Result<SubmissionAcceptance, PendingSubmissionError> {
if submission_request_id.trim().is_empty() { if submission_request_id.trim().is_empty() {
return Err(PendingSubmissionError::EmptyRequestId); return Err(PendingSubmissionError::EmptyRequestId);
@@ -1218,11 +1348,10 @@ where
.lock() .lock()
.expect("pending activation state poisoned"); .expect("pending activation state poisoned");
let original = current.clone(); let original = current.clone();
if let Some(receipt) = current if let Some(receipt) = current.receipts.iter().find(|receipt| {
.receipts receipt.submission_request_id == submission_request_id
.iter() && receipt.source_namespace == source_namespace
.find(|receipt| receipt.submission_request_id == submission_request_id) }) {
{
if receipt.payload_digest != payload_digest { if receipt.payload_digest != payload_digest {
return Err(PendingSubmissionError::IdempotencyConflict); return Err(PendingSubmissionError::IdempotencyConflict);
} }
@@ -1237,11 +1366,12 @@ where
let submission_id = uuid::Uuid::now_v7().to_string(); let submission_id = uuid::Uuid::now_v7().to_string();
let pending = PendingSubmission { let pending = PendingSubmission {
submission_request_id: submission_request_id.clone(), submission_request_id: submission_request_id.clone(),
source_namespace: source_namespace.clone(),
submission_id: submission_id.clone(), submission_id: submission_id.clone(),
payload_digest: payload_digest.clone(), payload_digest: payload_digest.clone(),
accepted_at_ms: segment_log::now_millis(), accepted_at_ms: segment_log::now_millis(),
activation_sequence: current.next_activation_sequence, activation_sequence: current.next_activation_sequence,
provenance: WorkerHistoryProvenance::LegacyUnknown, provenance,
was_queued: !activate_now, was_queued: !activate_now,
input, input,
}; };
@@ -1253,6 +1383,7 @@ where
}; };
current.remember_receipt(SubmissionReceipt { current.remember_receipt(SubmissionReceipt {
submission_request_id: submission_request_id.clone(), submission_request_id: submission_request_id.clone(),
source_namespace,
submission_id: submission_id.clone(), submission_id: submission_id.clone(),
payload_digest, payload_digest,
disposition, disposition,
@@ -1299,7 +1430,14 @@ where
return Err(PendingSubmissionError::ArtifactLimit); return Err(PendingSubmissionError::ArtifactLimit);
} }
current.pending.push_back(pending.clone()); 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(&current) { if let Err(error) = self.persist_locked(&current) {
let _ = self.release_submission_files(&pending);
*current = original; *current = original;
return Err(error); return Err(error);
} }
@@ -1312,10 +1450,29 @@ where
}) })
} }
#[cfg(test)]
pub(crate) fn accept_notification( pub(crate) fn accept_notification(
&self, &self,
notification_request_id: String, notification_request_id: String,
message: String, message: String,
auto_run: bool,
) -> Result<bool, PendingSubmissionError> {
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<bool, PendingSubmissionError> { ) -> Result<bool, PendingSubmissionError> {
if notification_request_id.trim().is_empty() { if notification_request_id.trim().is_empty() {
return Err(PendingSubmissionError::EmptyRequestId); return Err(PendingSubmissionError::EmptyRequestId);
@@ -1323,7 +1480,7 @@ where
if notification_request_id.len() > MAX_ACTIVATION_REQUEST_ID_BYTES { if notification_request_id.len() > MAX_ACTIVATION_REQUEST_ID_BYTES {
return Err(PendingSubmissionError::RequestIdLimit); return Err(PendingSubmissionError::RequestIdLimit);
} }
let payload_digest = submission_payload_digest(&[Segment::text(message.clone())]); let payload_digest = notification_payload_digest(&message, auto_run);
let _append_guard = self let _append_guard = self
.writer .writer
.state .state
@@ -1334,12 +1491,11 @@ where
.state .state
.lock() .lock()
.expect("pending activation state poisoned"); .expect("pending activation state poisoned");
if let Some(receipt) = state if let Some(receipt) = state.notification_receipts.iter().find(|receipt| {
.notification_receipts receipt.notification_request_id == notification_request_id
.iter() && receipt.source_namespace == source_namespace
.find(|receipt| receipt.notification_request_id == notification_request_id) }) {
{ if receipt.payload_digest != payload_digest || receipt.auto_run != auto_run {
if receipt.payload_digest != payload_digest {
return Err(PendingSubmissionError::IdempotencyConflict); return Err(PendingSubmissionError::IdempotencyConflict);
} }
return Ok(false); return Ok(false);
@@ -1373,17 +1529,19 @@ where
state.next_activation_sequence = state.next_activation_sequence.saturating_add(1); state.next_activation_sequence = state.next_activation_sequence.saturating_add(1);
state.pending_notifications.push_back(PendingNotification { state.pending_notifications.push_back(PendingNotification {
notification_request_id: notification_request_id.clone(), notification_request_id: notification_request_id.clone(),
source_namespace: source_namespace.clone(),
message, message,
payload_digest: payload_digest.clone(), payload_digest: payload_digest.clone(),
auto_run,
accepted_at_ms: segment_log::now_millis(), accepted_at_ms: segment_log::now_millis(),
activation_sequence, activation_sequence,
provenance: WorkerHistoryProvenance::BackendInstruction { provenance,
operation_id: Some(notification_request_id.clone()),
},
}); });
state.remember_notification_receipt(NotificationReceipt { state.remember_notification_receipt(NotificationReceipt {
notification_request_id, notification_request_id,
source_namespace,
payload_digest, payload_digest,
auto_run,
}); });
state.revision = state.revision.saturating_add(1); state.revision = state.revision.saturating_add(1);
if let Err(error) = self.persist_locked(&state) { if let Err(error) = self.persist_locked(&state) {
@@ -1393,6 +1551,59 @@ where
Ok(true) Ok(true)
} }
pub(crate) fn activating_passive_notification_id(&self) -> Option<String> {
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<PendingNotification> {
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( pub(crate) fn prepare_next_activation(
&self, &self,
fence: Option<(u64, &str)>, fence: Option<(u64, &str)>,
@@ -1414,17 +1625,20 @@ where
return Ok(None); return Ok(None);
} }
let submission_sequence = state.pending.front().map(|item| item.activation_sequence); let submission_sequence = state.pending.front().map(|item| item.activation_sequence);
let notification_sequence = state let notification_index = state
.pending_notifications .pending_notifications
.front() .iter()
.position(|item| item.auto_run);
let notification_sequence = notification_index
.and_then(|index| state.pending_notifications.get(index))
.map(|item| item.activation_sequence); .map(|item| item.activation_sequence);
if notification_sequence.is_some() if notification_sequence.is_some()
&& (submission_sequence.is_none() || notification_sequence < submission_sequence) && (submission_sequence.is_none() || notification_sequence < submission_sequence)
{ {
let notification = state let notification = state
.pending_notifications .pending_notifications
.pop_front() .remove(notification_index.expect("notification sequence came from an item"))
.expect("notification sequence came from queue head"); .expect("notification sequence came from an existing item");
state.activating_notification = Some(notification.clone()); state.activating_notification = Some(notification.clone());
state.revision = state.revision.saturating_add(1); state.revision = state.revision.saturating_add(1);
return Ok(Some(PendingActivation::Notification(notification))); return Ok(Some(PendingActivation::Notification(notification)));
@@ -1442,6 +1656,12 @@ where
} }
pub(crate) fn abort_activation(&self, pending: PendingSubmission) { 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 let mut state = self
.state .state
.lock() .lock()
@@ -1452,13 +1672,16 @@ where
} }
state.activating = None; state.activating = None;
if pending.was_queued { if pending.was_queued {
state.pending.push_front(pending); state.pending.push_front(pending.clone());
} else { } else {
state state
.receipts .receipts
.retain(|receipt| receipt.submission_id != pending.submission_id); .retain(|receipt| receipt.submission_id != pending.submission_id);
} }
state.revision = state.revision.saturating_add(1); 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 { pub(crate) fn activation_extension(&self) -> SessionExtension {
@@ -1530,6 +1753,10 @@ where
} }
} }
pub(crate) fn direct_client_namespace(&self) -> String {
format!("direct:{}", self.writer.state.location().session_id)
}
pub(crate) fn snapshot(&self) -> protocol::PendingSubmissionsSnapshot { pub(crate) fn snapshot(&self) -> protocol::PendingSubmissionsSnapshot {
self.state self.state
.lock() .lock()
@@ -1561,12 +1788,16 @@ where
else { else {
return Err(PendingSubmissionError::NotFound(submission_id.to_owned())); return Err(PendingSubmissionError::NotFound(submission_id.to_owned()));
}; };
state.pending.remove(index); let removed = state
.pending
.remove(index)
.expect("located pending submission must exist");
state.revision = state.revision.saturating_add(1); state.revision = state.revision.saturating_add(1);
if let Err(error) = self.persist_locked(&state) { if let Err(error) = self.persist_locked(&state) {
*state = original; *state = original;
return Err(error); return Err(error);
} }
self.release_submission_files(&removed)?;
Ok(state.snapshot()) Ok(state.snapshot())
} }
@@ -1586,13 +1817,16 @@ where
.expect("pending activation state poisoned"); .expect("pending activation state poisoned");
Self::validate_fence(&state, expected_revision, None)?; Self::validate_fence(&state, expected_revision, None)?;
let original = state.clone(); let original = state.clone();
state.pending.clear(); let removed = state.pending.drain(..).collect::<Vec<_>>();
state.pending_notifications.clear(); state.pending_notifications.clear();
state.revision = state.revision.saturating_add(1); state.revision = state.revision.saturating_add(1);
if let Err(error) = self.persist_locked(&state) { if let Err(error) = self.persist_locked(&state) {
*state = original; *state = original;
return Err(error); return Err(error);
} }
for pending in &removed {
self.release_submission_files(pending)?;
}
Ok(state.snapshot()) Ok(state.snapshot())
} }
} }
@@ -1627,9 +1861,11 @@ pub trait SystemItemCommitter: Send + Sync {
&self, &self,
item: SystemItem, item: SystemItem,
extensions: Vec<SessionExtension>, extensions: Vec<SessionExtension>,
history_provenance: Option<WorkerHistoryProvenance>,
) -> Result<HistoryEntry<SessionHistoryMetadata>, StoreError> { ) -> Result<HistoryEntry<SessionHistoryMetadata>, StoreError> {
let metadata = new_history_metadata( let metadata = new_history_metadata(
WorkerHistoryProvenance::BackendInstruction { operation_id: None }, history_provenance
.unwrap_or(WorkerHistoryProvenance::BackendInstruction { operation_id: None }),
None, None,
); );
let history_item = item.to_history_item(); let history_item = item.to_history_item();
@@ -2666,9 +2902,11 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.expect("pending activation state poisoned") .expect("pending activation state poisoned")
.clone(); .clone();
if !pending_state.pending.is_empty() if !pending_state.pending.is_empty()
|| !pending_state.pending_notifications.is_empty()
|| pending_state.activating.is_some() || pending_state.activating.is_some()
|| pending_state.activating_notification.is_some() || pending_state.activating_notification.is_some()
|| !pending_state.receipts.is_empty() || !pending_state.receipts.is_empty()
|| !pending_state.notification_receipts.is_empty()
{ {
let checkpoint = LogEntry::Extension { let checkpoint = LogEntry::Extension {
ts: segment_log::now_millis(), ts: segment_log::now_millis(),
@@ -3501,7 +3739,12 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
where where
St: Clone + 'static, St: Clone + 'static,
{ {
self.run_with_input_extensions_and_commit_hook(input, input_extensions, || {}) self.run_with_input_extensions_and_commit_hook(
input,
input_extensions,
WorkerHistoryProvenance::LegacyUnknown,
|| {},
)
.await .await
} }
@@ -3513,6 +3756,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
&mut self, &mut self,
input: Vec<Segment>, input: Vec<Segment>,
mut input_extensions: Vec<SessionExtension>, mut input_extensions: Vec<SessionExtension>,
input_provenance: WorkerHistoryProvenance,
on_input_committed: F, on_input_committed: F,
) -> Result<WorkerRunResult, WorkerError> ) -> Result<WorkerRunResult, WorkerError>
where where
@@ -3560,8 +3804,12 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
trigger: protocol::InvokeKind::UserSend, trigger: protocol::InvokeKind::UserSend,
})?; })?;
let projected_input = let projected_input = self.projected_input_history(
self.projected_input_history(&input, flow_projection.as_ref(), &projected_entry_ids); &input,
flow_projection.as_ref(),
&projected_entry_ids,
&input_provenance,
);
// Persist original typed segments together with the exact ordered // Persist original typed segments together with the exact ordered
// model-visible item+origin projection before any entry becomes live. // model-visible item+origin projection before any entry becomes live.
@@ -3867,6 +4115,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
input: &[Segment], input: &[Segment],
flow_projection: Option<&PreparedFlowProjection>, flow_projection: Option<&PreparedFlowProjection>,
entry_ids: &[SessionHistoryEntryId], entry_ids: &[SessionHistoryEntryId],
provenance: &WorkerHistoryProvenance,
) -> Vec<HistoryEntry<SessionHistoryMetadata>> { ) -> Vec<HistoryEntry<SessionHistoryMetadata>> {
if let Some(flow) = flow_projection { if let Some(flow) = flow_projection {
return input return input
@@ -3887,10 +4136,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
other => history_entry_with_id( other => history_entry_with_id(
Item::user_message(Segment::flatten_to_text(std::slice::from_ref(other))), Item::user_message(Segment::flatten_to_text(std::slice::from_ref(other))),
entry_id.clone(), entry_id.clone(),
// Current public submit transport does not carry a provenance.clone(),
// trusted account/Worker subject envelope. Fail closed
// instead of promoting role=user to HumanInput.
WorkerHistoryProvenance::LegacyUnknown,
), ),
}) })
.collect(); .collect();
@@ -3902,7 +4148,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.first() .first()
.expect("projected Worker input always has one entry id") .expect("projected Worker input always has one entry id")
.clone(), .clone(),
WorkerHistoryProvenance::LegacyUnknown, provenance.clone(),
)] )]
} }
@@ -7938,8 +8184,15 @@ mod build_summary_prompt_tests {
serde_json::to_value(&state).unwrap(), serde_json::to_value(&state).unwrap(),
); );
let projected_ids = vec![SessionHistoryEntryId::new(), SessionHistoryEntryId::new()]; let projected_ids = vec![SessionHistoryEntryId::new(), SessionHistoryEntryId::new()];
let projected = let input_provenance = WorkerHistoryProvenance::HumanInput {
worker.projected_input_history(&segments, projection.as_ref(), &projected_ids); account_id: "account-1".into(),
};
let projected = worker.projected_input_history(
&segments,
projection.as_ref(),
&projected_ids,
&input_provenance,
);
worker worker
.commit_entry(LogEntry::AnnotatedUserInput { .commit_entry(LogEntry::AnnotatedUserInput {
ts: segment_log::now_millis(), ts: segment_log::now_millis(),
@@ -7966,6 +8219,7 @@ mod build_summary_prompt_tests {
projected[0].annotation.origin, projected[0].annotation.origin,
WorkerHistoryProvenance::FlowInstruction { .. } WorkerHistoryProvenance::FlowInstruction { .. }
)); ));
assert_eq!(projected[1].annotation.origin, input_provenance);
assert_eq!(state.instance.definition_revision, 3); assert_eq!(state.instance.definition_revision, 3);
assert_eq!(state.instance.current_state.as_str(), "implement"); assert_eq!(state.instance.current_state.as_str(), "implement");
assert_eq!(workspace_client.requests.lock().unwrap().len(), 1); assert_eq!(workspace_client.requests.lock().unwrap().len(), 1);
@@ -8048,7 +8302,12 @@ mod build_summary_prompt_tests {
.delete_uploaded_file(worker.session_id(), &file.artifact_id), .delete_uploaded_file(worker.session_id(), &file.artifact_id),
Err(StoreError::ArtifactAlreadyCommitted) 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(); let text = projected[0].item.as_text().unwrap();
assert!(text.contains("notes.md")); assert!(text.contains("notes.md"));
assert!(text.contains(&file.artifact_id)); assert!(text.contains(&file.artifact_id));
@@ -8130,7 +8389,12 @@ mod build_summary_prompt_tests {
if retained.source_entry_id == artifact.source_entry_id 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("終端")); assert!(!history[0].item.as_text().unwrap().contains("終端"));
append_test_entry( append_test_entry(
&worker, &worker,
@@ -8406,6 +8670,47 @@ mod build_summary_prompt_tests {
assert_eq!(worker.history()[0].as_text().unwrap(), "first message"); 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] #[tokio::test]
async fn annotated_history_rewind_commits_authoritative_prefix() { async fn annotated_history_rewind_commits_authoritative_prefix() {
let (_dir, mut worker) = rewind_test_worker().await; let (_dir, mut worker) = rewind_test_worker().await;
@@ -9251,6 +9556,110 @@ 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 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] #[test]
fn pending_submission_queue_is_durable_idempotent_and_bounded() { fn pending_submission_queue_is_durable_idempotent_and_bounded() {
let temp = tempfile::tempdir().unwrap(); let temp = tempfile::tempdir().unwrap();
@@ -9339,22 +9748,96 @@ mod build_summary_prompt_tests {
assert_eq!(cleared.notification_count, 0); 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] #[test]
fn notification_and_submit_share_activation_order_and_notification_dedupes() { fn notification_and_submit_share_activation_order_and_notification_dedupes() {
let temp = tempfile::tempdir().unwrap(); let temp = tempfile::tempdir().unwrap();
let handle = PendingSubmissionHandle::for_test(temp.path()); let handle = PendingSubmissionHandle::for_test(temp.path());
assert!( assert!(
handle handle
.accept_notification("notification-1".into(), "notice".into()) .accept_notification("notification-1".into(), "notice".into(), true)
.unwrap() .unwrap()
); );
assert!( assert!(
!handle !handle
.accept_notification("notification-1".into(), "notice".into()) .accept_notification("notification-1".into(), "notice".into(), true)
.unwrap() .unwrap()
); );
assert!(matches!( assert!(matches!(
handle.accept_notification("notification-1".into(), "different".into()), 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) Err(PendingSubmissionError::IdempotencyConflict)
)); ));
handle handle
@@ -9382,9 +9865,10 @@ mod build_summary_prompt_tests {
let mut session = WorkerSession::new(session_store::new_session_id(), Vec::new()); let mut session = WorkerSession::new(session_store::new_session_id(), Vec::new());
let state = PendingActivationState { let state = PendingActivationState {
revision: 4, revision: 4,
next_activation_sequence: 2, next_activation_sequence: 3,
activating: Some(PendingSubmission { activating: Some(PendingSubmission {
submission_request_id: "request-1".into(), submission_request_id: "request-1".into(),
source_namespace: "direct:test".into(),
submission_id: "submission-1".into(), submission_id: "submission-1".into(),
payload_digest: submission_payload_digest(&[Segment::text("first")]), payload_digest: submission_payload_digest(&[Segment::text("first")]),
accepted_at_ms: 1, accepted_at_ms: 1,
@@ -9393,9 +9877,21 @@ mod build_summary_prompt_tests {
was_queued: false, was_queued: false,
input: vec![Segment::text("first")], input: vec![Segment::text("first")],
}), }),
activating_notification: None, 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 { pending: VecDeque::from([PendingSubmission {
submission_request_id: "request-2".into(), submission_request_id: "request-2".into(),
source_namespace: "direct:test".into(),
submission_id: "submission-2".into(), submission_id: "submission-2".into(),
payload_digest: submission_payload_digest(&[Segment::text("second")]), payload_digest: submission_payload_digest(&[Segment::text("second")]),
accepted_at_ms: 2, accepted_at_ms: 2,
@@ -9406,7 +9902,12 @@ mod build_summary_prompt_tests {
}]), }]),
pending_notifications: VecDeque::new(), pending_notifications: VecDeque::new(),
receipts: VecDeque::new(), receipts: VecDeque::new(),
notification_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.restore_pending_activations(&[(
SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.into(), SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.into(),
@@ -9420,6 +9921,14 @@ mod build_summary_prompt_tests {
assert_eq!(state.pending.len(), 2); assert_eq!(state.pending.len(), 2);
assert_eq!(state.pending[0].submission_id, "submission-1"); assert_eq!(state.pending[0].submission_id, "submission-1");
assert_eq!(state.pending[1].submission_id, "submission-2"); 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 { fn minimal_manifest() -> WorkerManifest {
+64 -1
View File
@@ -1804,15 +1804,18 @@ async fn notify_while_idle_with_auto_run_false_waits_for_explicit_run() {
let client_for_assert = client.clone(); let client_for_assert = client.clone();
let worker = make_worker(client).await; let worker = make_worker(client).await;
let handle = spawn_controller(worker).await; let handle = spawn_controller(worker).await;
let notification_request_id = protocol::new_submission_request_id();
for _ in 0..2 {
handle handle
.send(Method::Notify { .send(Method::Notify {
notification_request_id: protocol::new_submission_request_id(), notification_request_id: notification_request_id.clone(),
message: "progress snapshot".into(), message: "progress snapshot".into(),
auto_run: false, auto_run: false,
}) })
.await .await
.unwrap(); .unwrap();
}
tokio::time::sleep(std::time::Duration::from_millis(100)).await; tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Idle); assert_eq!(handle.shared_state.get_status(), WorkerStatus::Idle);
@@ -2049,6 +2052,66 @@ async fn notify_while_running_does_not_emit_already_running_error() {
wait_for_status(&handle, WorkerStatus::Idle).await; 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] #[tokio::test]
async fn status_json_reflects_worker_name() { async fn status_json_reflects_worker_name() {
let client = MockClient::new(simple_text_events()); let client = MockClient::new(simple_text_events());
+6 -4
View File
@@ -4987,7 +4987,7 @@ mod tests {
"missing test context", "missing test context",
); );
}; };
let submission_id = input.submission_id.clone(); let submission_request_id = input.submission_request_id.clone();
let content = input.content; let content = input.content;
std::thread::spawn(move || { std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(10)); std::thread::sleep(std::time::Duration::from_millis(10));
@@ -5004,11 +5004,13 @@ mod tests {
status: protocol::WorkerStatus::Idle, status: protocol::WorkerStatus::Idle,
}); });
}); });
if let Some(submission_id) = submission_id { if let Some(submission_request_id) = submission_request_id {
worker_runtime::execution::WorkerExecutionResult::accepted_input_committed( worker_runtime::execution::WorkerExecutionResult::accepted_submission(
worker_runtime::execution::WorkerExecutionOperation::Input, worker_runtime::execution::WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy, WorkerExecutionRunState::Busy,
submission_id, submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
) )
} else { } else {
worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionResult::accepted(
@@ -32,11 +32,13 @@ impl WorkerExecutionBackend for TestExecutionBackend {
_handle: &WorkerExecutionHandle, _handle: &WorkerExecutionHandle,
input: worker_runtime::interaction::WorkerInput, input: worker_runtime::interaction::WorkerInput,
) -> WorkerExecutionResult { ) -> WorkerExecutionResult {
if let Some(submission_id) = input.submission_id { if let Some(submission_request_id) = input.submission_request_id {
WorkerExecutionResult::accepted_input_committed( WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input, WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy, WorkerExecutionRunState::Busy,
submission_id, submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
) )
} else { } else {
WorkerExecutionResult::accepted( WorkerExecutionResult::accepted(
+130 -22
View File
@@ -8537,13 +8537,15 @@ async fn scoped_list_runtimes(
async fn scoped_workspace_protocol_ws( async fn scoped_workspace_protocol_ws(
State(api): State<WorkspaceApi>, State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath(workspace_id): AxumPath<String>, AxumPath(workspace_id): AxumPath<String>,
ws: axum::extract::ws::WebSocketUpgrade, ws: axum::extract::ws::WebSocketUpgrade,
) -> std::result::Result<Response, Response> { ) -> std::result::Result<Response, Response> {
validate_workspace_scope(&api, &workspace_id).map_err(|error| error.into_response())?; validate_workspace_scope(&api, &workspace_id).map_err(|error| error.into_response())?;
let input_source = authenticated_browser_input_source(&actor);
Ok(ws Ok(ws
.on_upgrade(move |socket| { .on_upgrade(move |socket| {
crate::workspace_subscription::serve_workspace_subscription(api, socket) crate::workspace_subscription::serve_workspace_subscription(api, socket, input_source)
}) })
.into_response()) .into_response())
} }
@@ -11458,6 +11460,7 @@ async fn scoped_cancel_runtime_worker(
async fn scoped_worker_protocol_ws( async fn scoped_worker_protocol_ws(
ws: WebSocketUpgrade, ws: WebSocketUpgrade,
State(api): State<WorkspaceApi>, State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>, AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
) -> Response { ) -> Response {
if let Err(err) = validate_workspace_scope(&api, &path.workspace_id) { if let Err(err) = validate_workspace_scope(&api, &path.workspace_id) {
@@ -11465,6 +11468,7 @@ async fn scoped_worker_protocol_ws(
} }
worker_protocol_ws( worker_protocol_ws(
State(api), State(api),
Extension(actor),
AxumPath((path.worker.runtime_id, path.worker.worker_id)), AxumPath((path.worker.runtime_id, path.worker.worker_id)),
ws, ws,
) )
@@ -13877,8 +13881,45 @@ async fn cancel_runtime_worker(
Ok(Json(result)) Ok(Json(result))
} }
fn authenticated_browser_input_source(actor: &RequestActor) -> protocol::AuthenticatedInputSource {
protocol::AuthenticatedInputSource::Account {
account_id: actor.account_id.clone(),
}
}
pub(crate) fn authorize_browser_worker_method(
method: protocol::Method,
source: &protocol::AuthenticatedInputSource,
) -> std::result::Result<protocol::Method, &'static str> {
match method {
protocol::Method::Submit {
submission_request_id,
input,
} => Ok(protocol::Method::SubmitTracked {
submission_request_id,
input,
source: source.clone(),
}),
protocol::Method::Notify {
notification_request_id,
message,
auto_run,
} => Ok(protocol::Method::NotifyTracked {
notification_request_id,
message,
auto_run,
source: source.clone(),
}),
protocol::Method::SubmitTracked { .. } | protocol::Method::NotifyTracked { .. } => {
Err("authenticated Worker input source is server-owned")
}
other => Ok(other),
}
}
async fn worker_protocol_ws( async fn worker_protocol_ws(
State(api): State<WorkspaceApi>, State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>, AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>,
ws: WebSocketUpgrade, ws: WebSocketUpgrade,
) -> impl IntoResponse { ) -> impl IntoResponse {
@@ -13902,7 +13943,8 @@ async fn worker_protocol_ws(
.into_response(); .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 { pub(crate) struct WorkspaceWorkerProtocolConnection {
@@ -14029,13 +14071,17 @@ async fn connect_embedded_worker_protocol(
Ok(WorkspaceWorkerProtocolConnection { methods, events }) 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 { match source {
RuntimeObservationSource::RemoteWs(config) => { RuntimeObservationSource::RemoteWs(config) => {
remote_worker_protocol_ws_session(config, socket).await; remote_worker_protocol_ws_session(config, socket, input_source).await;
} }
RuntimeObservationSource::Embedded(source) => { RuntimeObservationSource::Embedded(source) => {
embedded_worker_protocol_ws_session(source, socket).await; embedded_worker_protocol_ws_session(source, socket, input_source).await;
} }
} }
} }
@@ -14043,6 +14089,7 @@ async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: We
async fn remote_worker_protocol_ws_session( async fn remote_worker_protocol_ws_session(
config: RuntimeObservationSourceConfig, config: RuntimeObservationSourceConfig,
socket: WebSocket, socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) { ) {
let mut request = match config.endpoint.clone().into_client_request() { let mut request = match config.endpoint.clone().into_client_request() {
Ok(request) => request, Ok(request) => request,
@@ -14091,14 +14138,33 @@ async fn remote_worker_protocol_ws_session(
inbound = client_stream.next() => { inbound = client_stream.next() => {
match inbound { match inbound {
Some(Ok(WsMessage::Text(text))) => { 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; break;
} }
} }
Some(Ok(WsMessage::Binary(binary))) => { Some(Ok(WsMessage::Binary(_))) => {
if upstream_sink.send(TungsteniteMessage::Binary(binary.to_vec().into())).await.is_err() { if let Ok(event) = protocol::stream::encode_event(&protocol_error_event("binary Worker methods are not accepted")) {
break; let _ = client_sink.send(WsMessage::Text(event.into())).await;
} }
break;
} }
Some(Ok(WsMessage::Close(_))) | None => { Some(Ok(WsMessage::Close(_))) | None => {
let _ = upstream_sink.send(TungsteniteMessage::Close(None)).await; let _ = upstream_sink.send(TungsteniteMessage::Close(None)).await;
@@ -14154,6 +14220,7 @@ async fn remote_worker_protocol_ws_session(
async fn embedded_worker_protocol_ws_session( async fn embedded_worker_protocol_ws_session(
source: crate::observation::EmbeddedRuntimeObservationSource, source: crate::observation::EmbeddedRuntimeObservationSource,
mut socket: WebSocket, mut socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) { ) {
let mut upstream = match RuntimeObservationClient::connect(&RuntimeObservationSource::Embedded( let mut upstream = match RuntimeObservationClient::connect(&RuntimeObservationSource::Embedded(
source.clone(), source.clone(),
@@ -14173,6 +14240,7 @@ async fn embedded_worker_protocol_ws_session(
inbound = socket.next() => { inbound = socket.next() => {
match inbound { match inbound {
Some(Ok(WsMessage::Text(text))) => match decode_method(&text) { Some(Ok(WsMessage::Text(text))) => match decode_method(&text) {
Ok(method) => match authorize_browser_worker_method(method, &input_source) {
Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) { Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) {
Ok(events) => { Ok(events) => {
for event in events { for event in events {
@@ -14188,9 +14256,16 @@ async fn embedded_worker_protocol_ws_session(
} }
} }
}, },
Err(message) => {
let event = protocol_error_event(message);
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
},
Err(error) => { Err(error) => {
let event = let event = protocol_error_event(format!(
protocol_error_event(format!("malformed protocol method frame: {error}")); "malformed protocol method frame: {error}"
));
if !send_protocol_event(&mut socket, &event).await { if !send_protocol_event(&mut socket, &event).await {
return; return;
} }
@@ -16580,6 +16655,29 @@ mod tests {
&tail[..end] &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] #[test]
fn merge_request_http_paths_observe_refs_through_runtime_provider_authority() { fn merge_request_http_paths_observe_refs_through_runtime_provider_authority() {
let source = include_str!("server.rs"); let source = include_str!("server.rs");
@@ -18825,7 +18923,7 @@ mod tests {
.get(handle.worker_ref()) .get(handle.worker_ref())
.cloned() .cloned()
.expect("execution context"); .expect("execution context");
let submission_id = input.submission_id.clone(); let submission_request_id = input.submission_request_id.clone();
let content = input.content.clone(); let content = input.content.clone();
std::thread::spawn(move || { std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(25)); std::thread::sleep(std::time::Duration::from_millis(25));
@@ -18833,11 +18931,13 @@ mod tests {
text: format!("server companion echoed: {content}"), text: format!("server companion echoed: {content}"),
}); });
}); });
if let Some(submission_id) = submission_id { if let Some(submission_request_id) = submission_request_id {
worker_runtime::execution::WorkerExecutionResult::accepted_input_committed( worker_runtime::execution::WorkerExecutionResult::accepted_submission(
worker_runtime::execution::WorkerExecutionOperation::Input, worker_runtime::execution::WorkerExecutionOperation::Input,
worker_runtime::execution::WorkerExecutionRunState::Idle, worker_runtime::execution::WorkerExecutionRunState::Idle,
submission_id, submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
) )
} else { } else {
worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionResult::accepted(
@@ -27223,6 +27323,16 @@ mod tests {
(runtime, worker_ref, endpoint) (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( async fn spawn_workspace_proxy(
source: RuntimeObservationSourceConfig, source: RuntimeObservationSourceConfig,
) -> (String, tempfile::TempDir) { ) -> (String, tempfile::TempDir) {
@@ -27241,11 +27351,8 @@ mod tests {
.unwrap(); .unwrap();
let app_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let app_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let app_addr = app_listener.local_addr().unwrap(); let app_addr = app_listener.local_addr().unwrap();
tokio::spawn(async move { let app = build_inner_router(api).layer(Extension(test_browser_request_actor()));
axum::serve(app_listener, build_inner_router(api)) tokio::spawn(async move { axum::serve(app_listener, app).await.unwrap() });
.await
.unwrap()
});
( (
format!("ws://{app_addr}/api/runtimes/{runtime_id}/workers/{worker_id}/protocol/ws"), format!("ws://{app_addr}/api/runtimes/{runtime_id}/workers/{worker_id}/protocol/ws"),
dir, dir,
@@ -27257,7 +27364,8 @@ mod tests {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().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 server = tokio::spawn(async move {
let _ = axum::serve(listener, app).await; let _ = axum::serve(listener, app).await;
}); });
@@ -27305,7 +27413,7 @@ mod tests {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().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 server = tokio::spawn(async move {
let _ = axum::serve(listener, app).await; let _ = axum::serve(listener, app).await;
}); });
@@ -11,7 +11,9 @@ use tokio::sync::mpsc;
use worker_runtime::identity::RuntimeWorkerRef; use worker_runtime::identity::RuntimeWorkerRef;
use crate::runtime_subscription::{BrokerSubscriptionEvent, RuntimeSubscriptionBroker}; 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; use crate::store::WorkspaceResourceKind;
const OUTBOUND_CAPACITY: usize = 256; const OUTBOUND_CAPACITY: usize = 256;
@@ -21,7 +23,11 @@ struct ActiveSubscription {
methods: Option<mpsc::Sender<protocol::Method>>, methods: Option<mpsc::Sender<protocol::Method>>,
} }
pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebSocket) { pub(crate) async fn serve_workspace_subscription(
api: WorkspaceApi,
socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
let broker = api.runtime_subscription_broker().clone(); let broker = api.runtime_subscription_broker().clone();
let (mut socket_sender, mut socket_receiver) = socket.split(); let (mut socket_sender, mut socket_receiver) = socket.split();
let (control_outbound, mut control_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY); let (control_outbound, mut control_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
@@ -153,7 +159,12 @@ pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebS
else { else {
break; 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; break;
} }
} }