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)
// ---------------------------------------------------------------------------
/// 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)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(tag = "method", content = "params", rename_all = "snake_case")]
@@ -45,13 +77,13 @@ pub enum Method {
submission_request_id: String,
input: Vec<Segment>,
},
/// Runtime-internal Submit with the same request identity contract. This
/// variant is not serializable on the public Client → Worker protocol.
#[serde(skip)]
/// Authenticated transport form of Submit. Trusted adapters replace
/// public Submit before forwarding it to the Worker.
#[cfg_attr(feature = "typescript", ts(skip))]
SubmitTracked {
submission_request_id: String,
input: Vec<Segment>,
source: AuthenticatedInputSource,
},
/// Human-readable text injected into the target Worker's LLM context
/// 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")]
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.
WorkerEvent(WorkerEvent),
/// Return the authoritative FIFO summary without exposing queued payloads.
@@ -1497,15 +1538,26 @@ mod tests {
}
#[test]
fn runtime_tracked_submit_is_not_public_protocol_json() {
fn authenticated_submit_round_trips_trusted_source() {
let method = Method::SubmitTracked {
input: vec![Segment::text("private")],
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!(
serde_json::from_str::<Method>(
r#"{"method":"submit_tracked","input":[],"submission_id":"forged"}"#,
r#"{"method":"submit_tracked","input":[],"submission_request_id":"forged"}"#,
)
.is_err()
);
+116 -1
View File
@@ -22,7 +22,8 @@ use crate::store::{Store, StoreError};
use crate::uploaded_file::{
bind_uploaded_file, clear_uploaded_file_binding, copy_committed_uploaded_files,
delete_uncommitted_uploaded_files, delete_uploaded_file, list_uploaded_file_refs,
read_uploaded_file, read_uploaded_file_by_id, write_uploaded_file,
pin_uploaded_file, read_uploaded_file, read_uploaded_file_by_id, release_uploaded_file_pin,
write_uploaded_file,
};
use crate::{
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(
&self,
session_id: SessionId,
@@ -865,6 +892,94 @@ mod tests {
assert!(store.read_uploaded_file(owner, &reference).is_err());
}
#[test]
fn pending_upload_pin_survives_cleanup_until_release_or_history_binding() {
let tmp = tempfile::TempDir::new().unwrap();
let store = FsStore::new(tmp.path()).unwrap();
let session_id = new_session_id();
let limits = UploadedFileLimits {
max_file_bytes: 64,
max_session_bytes: 128,
};
let pending = store
.write_uploaded_file(session_id, "pending.txt", "text/plain", b"pending", limits)
.unwrap();
store
.pin_uploaded_file(session_id, &pending, "submission-1")
.unwrap();
assert!(matches!(
store.pin_uploaded_file(session_id, &pending, "submission-other"),
Err(StoreError::ArtifactAlreadyCommitted)
));
drop(store);
let store = FsStore::new(tmp.path()).unwrap();
assert_eq!(
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
0
);
assert_eq!(
store
.read_uploaded_file_by_id(session_id, &pending.artifact_id)
.unwrap()
.1,
b"pending"
);
let fork_session_id = new_session_id();
assert_eq!(
store
.copy_committed_uploaded_files(session_id, fork_session_id)
.unwrap(),
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]
fn uploaded_file_validation_and_shared_quota_fail_closed() {
let tmp = tempfile::TempDir::new().unwrap();
+20
View File
@@ -226,6 +226,26 @@ pub trait Store: Send + Sync {
Err(StoreError::PasteArtifactUnsupported)
}
/// Retain an uploaded file while a durable pending operation owns it.
fn pin_uploaded_file(
&self,
_session_id: SessionId,
_reference: &UploadedFileRef,
_owner_id: &str,
) -> Result<(), StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
/// Release a pending-operation pin without changing committed ownership.
fn release_uploaded_file_pin(
&self,
_session_id: SessionId,
_artifact_id: &str,
_owner_id: &str,
) -> Result<(), StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
/// Delete an uncommitted uploaded file owned by `session_id`.
fn delete_uploaded_file(
&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;
const MAX_FILE_NAME_CHARS: usize = 255;
const MAX_MEDIA_TYPE_BYTES: usize = 127;
fn validate_pending_owner_id(owner_id: &str) -> Result<()> {
if owner_id.is_empty() || owner_id.len() > 256 {
return Err(StoreError::ArtifactIntegrityMismatch);
}
Ok(())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct UploadedFileLimits {
@@ -59,6 +65,8 @@ struct StoredUploadedFile {
#[serde(default, skip_serializing_if = "Option::is_none")]
source_entry_id: Option<String>,
#[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>,
content_base64: String,
}
@@ -291,6 +299,7 @@ pub(crate) fn write_uploaded_file(
byte_len,
sha256: sha256.clone(),
source_entry_id: None,
pending_owner_id: None,
upload_context: context.cloned(),
content_base64: BASE64.encode(content),
};
@@ -376,6 +385,72 @@ pub(crate) fn clear_uploaded_file_binding(
Ok(())
}
pub(crate) fn pin_uploaded_file(
dir: &Path,
reference: &UploadedFileRef,
owner_id: &str,
) -> Result<()> {
validate_pending_owner_id(owner_id)?;
if reference.source_entry_id.is_some() {
return Err(StoreError::ArtifactAlreadyCommitted);
}
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let path = record_path(dir, &reference.artifact_id)?;
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.file_name != reference.file_name
|| stored.media_type != reference.media_type
|| stored.created_at_ms != reference.created_at_ms
|| stored.byte_len != reference.byte_len
|| stored.sha256 != reference.sha256
{
return Err(StoreError::ArtifactIntegrityMismatch);
}
if stored.source_entry_id.is_some() {
return Err(StoreError::ArtifactAlreadyCommitted);
}
if let Some(existing_owner) = stored.pending_owner_id.as_deref() {
return if existing_owner == owner_id {
Ok(())
} else {
Err(StoreError::ArtifactAlreadyCommitted)
};
}
stored.pending_owner_id = Some(owner_id.to_owned());
let temp = dir.join(format!(".{}.file.pin.tmp", reference.artifact_id));
fs::write(&temp, serde_json::to_vec(&stored)?)?;
fs::rename(temp, path)?;
Ok(())
}
pub(crate) fn release_uploaded_file_pin(
dir: &Path,
artifact_id: &str,
owner_id: &str,
) -> Result<()> {
validate_pending_owner_id(owner_id)?;
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let path = record_path(dir, artifact_id)?;
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.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(
dir: &Path,
reference: &UploadedFileRef,
@@ -405,6 +480,7 @@ pub(crate) fn bind_uploaded_file(
return Err(StoreError::ArtifactAlreadyCommitted);
}
stored.source_entry_id = Some(source_entry_id.to_owned());
stored.pending_owner_id = None;
let temp = dir.join(format!(".{}.file.bind.tmp", reference.artifact_id));
fs::write(&temp, serde_json::to_vec(&stored)?)?;
fs::rename(&temp, path)?;
@@ -455,7 +531,7 @@ pub(crate) fn copy_committed_uploaded_files(source_dir: &Path, target_dir: &Path
}
let bytes = fs::read(&path)?;
let stored: StoredUploadedFile = serde_json::from_slice(&bytes)?;
if stored.source_entry_id.is_none() {
if stored.source_entry_id.is_none() && stored.pending_owner_id.is_none() {
continue;
}
let target = target_dir.join(name);
@@ -499,7 +575,7 @@ pub(crate) fn delete_uncommitted_uploaded_files(dir: &Path) -> Result<u64> {
continue;
}
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.source_entry_id.is_none() {
if stored.source_entry_id.is_none() && stored.pending_owner_id.is_none() {
fs::remove_file(path)?;
removed = removed
.checked_add(1)
@@ -523,7 +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) => return Err(error.into()),
};
if stored.source_entry_id.is_some() {
if stored.source_entry_id.is_some() || stored.pending_owner_id.is_some() {
return Err(StoreError::ArtifactAlreadyCommitted);
}
match fs::remove_file(path) {
@@ -1522,6 +1522,7 @@ fn method_starts_turn(method: &Method) -> bool {
Method::Submit { .. }
| Method::SubmitTracked { .. }
| Method::Notify { auto_run: true, .. }
| Method::NotifyTracked { auto_run: true, .. }
| Method::Resume
| Method::Compact
)
@@ -1549,6 +1550,7 @@ fn accepted_run_state_for_method(method: &Method) -> WorkerExecutionRunState {
Method::Submit { .. }
| Method::SubmitTracked { .. }
| Method::Notify { auto_run: true, .. }
| Method::NotifyTracked { auto_run: true, .. }
| Method::Resume
| Method::Compact => WorkerExecutionRunState::Busy,
Method::Shutdown => WorkerExecutionRunState::Stopped,
@@ -1946,6 +1948,9 @@ where
input: input.segments.unwrap_or_else(|| {
vec![Segment::text(input.content.trim().to_string())]
}),
source: protocol::AuthenticatedInputSource::Backend {
operation_id: submission_id.clone(),
},
},
Some(submission_id),
)
@@ -1966,6 +1971,7 @@ where
Method::Submit { .. }
| Method::SubmitTracked { .. }
| Method::Notify { .. }
| Method::NotifyTracked { .. }
| Method::Compact => WorkerExecutionRunState::Busy,
_ => WorkerExecutionRunState::Idle,
};
+259 -43
View File
@@ -229,6 +229,43 @@ enum PendingRun {
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>(
pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
notify_buffer: &NotifyBuffer,
@@ -241,7 +278,12 @@ fn prepare_pending_run<St: Store + Clone>(
Some(crate::worker::PendingActivation::Notification(notification)) => {
let extension = pending_submissions.notification_activation_extension();
let notification_request_id = notification.notification_request_id.clone();
notify_buffer.push_durable_notify(notification.message, extension);
notify_buffer.push_durable_notify(
notification.message,
notification.auto_run,
notification.provenance,
extension,
);
Some(PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::Notify,
notification_request_id: Some(notification_request_id),
@@ -1314,6 +1356,7 @@ async fn controller_loop<C, St>(
);
let mut pending: Option<PendingRun> = None;
let pending_submissions = worker.pending_submission_handle();
stage_oldest_passive_notification(&pending_submissions, &notify_buffer);
loop {
// 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(),
_ => None,
};
let passive_notification_request_id =
pending_submissions.activating_passive_notification_id();
let (mut new_status, shutdown, may_drain_pending) = match run {
PendingRun::Submit(submission) => {
let (input_commit_tx, input_commit_rx) = oneshot::channel();
@@ -1356,6 +1401,7 @@ async fn controller_loop<C, St>(
worker.run_with_input_extensions_and_commit_hook(
submission.input,
vec![extension],
submission.provenance,
move || {
let _ = input_commit_tx.send(());
},
@@ -1415,8 +1461,11 @@ async fn controller_loop<C, St>(
.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);
stage_oldest_passive_notification(&pending_submissions, &notify_buffer);
}
if !shutdown && may_drain_pending && new_status == WorkerStatus::Idle {
@@ -1470,13 +1519,47 @@ async fn controller_loop<C, St>(
Method::Submit {
submission_request_id,
input,
}
| Method::SubmitTracked {
submission_request_id,
input,
} => {
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) => {
if let Some(activation) = acceptance.activation {
pending = Some(PendingRun::Submit(activation));
@@ -1502,31 +1585,85 @@ async fn controller_loop<C, St>(
message,
auto_run,
} => {
if auto_run {
match pending_submissions.accept_notification(notification_request_id, message)
{
Ok(true) => {
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(),
});
}
let request_id = notification_request_id.clone();
let source_namespace = pending_submissions.direct_client_namespace();
match pending_submissions.accept_notification_from_source(
notification_request_id,
message,
source_namespace.clone(),
session_store::LoggedSessionHistoryOrigin::LegacyUnknown,
auto_run,
) {
Ok(_) if auto_run => {
match prepare_pending_run(&pending_submissions, &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(false) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
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(),
});
}
}
}
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(),
});
}
}
}
} else {
worker.push_notify(message, false);
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 {
submission_request_id,
input,
}
| Method::SubmitTracked {
submission_request_id,
input,
}) => {
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) => {
let _ = working_event_tx.send(Event::SubmissionAccepted {
submission_request_id: acceptance.submission_request_id,
@@ -2147,23 +2317,69 @@ where
message,
auto_run,
}) => {
if auto_run {
if let Err(error) = pending_submissions.accept_notification(
notification_request_id,
message,
) {
let request_id = notification_request_id.clone();
let source_namespace = pending_submissions.direct_client_namespace();
match pending_submissions.accept_notification_from_source(
notification_request_id,
message,
source_namespace.clone(),
session_store::LoggedSessionHistoryOrigin::LegacyUnknown,
auto_run,
) {
Ok(_) if !auto_run => {
stage_pending_notification(
&pending_submissions,
notify_buffer,
&source_namespace,
&request_id,
);
}
Ok(_) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
} else {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_submissions.snapshot(),
}
}
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_submissions.snapshot(),
});
}
Some(Method::NotifyTracked {
notification_request_id,
message,
auto_run,
source,
}) => {
let request_id = notification_request_id.clone();
let source_namespace = 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(),
});
}
} else {
notify_buffer.push_notify(message, false);
}
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_submissions.snapshot(),
});
}
Some(Method::ListCompletions { .. }) => {}
Some(Method::ListWorkers | Method::RestoreWorker { .. } | Method::RegisterPeer { .. }) => {
+18 -8
View File
@@ -178,14 +178,21 @@ impl WorkerInterceptor {
/// matches worker-history order.
fn commit_system_items_with_extensions(
&self,
items: &[(SystemItem, Vec<session_store::SessionExtension>)],
items: &[(
SystemItem,
Vec<session_store::SessionExtension>,
Option<session_store::LoggedSessionHistoryOrigin>,
)],
) -> Result<(), session_store::StoreError> {
let Some(writer) = self.log_writer.as_ref() else {
return Ok(());
};
for (item, extensions) in items {
let entry =
writer.commit_system_item_with_extensions(item.clone(), extensions.clone())?;
for (item, extensions, history_provenance) in items {
let entry = writer.commit_system_item_with_extensions(
item.clone(),
extensions.clone(),
history_provenance.clone(),
)?;
self.pending_committed_history
.lock()
.expect("pending committed history poisoned")
@@ -199,7 +206,7 @@ impl WorkerInterceptor {
&items
.iter()
.cloned()
.map(|item| (item, Vec::new()))
.map(|item| (item, Vec::new(), None))
.collect::<Vec<_>>(),
)
}
@@ -341,8 +348,11 @@ impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
projection_digest: projection.catalog_digest.clone(),
logical_name: "internal.notify_wrapper".to_string(),
};
let mut system_items: Vec<(SystemItem, Vec<session_store::SessionExtension>)> =
Vec::with_capacity(drained.len());
let mut system_items: Vec<(
SystemItem,
Vec<session_store::SessionExtension>,
Option<session_store::LoggedSessionHistoryOrigin>,
)> = Vec::with_capacity(drained.len());
let mut items: Vec<Item> = Vec::with_capacity(drained.len());
for entry in &drained {
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());
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) {
self.pending_notifies.requeue_front(drained);
+22 -3
View File
@@ -25,7 +25,7 @@ use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use protocol::WorkerEvent;
use session_store::{SessionExtension, SystemItem};
use session_store::{LoggedSessionHistoryOrigin, SessionExtension, SystemItem};
use tracing::warn;
use crate::prompt::catalog::{CatalogError, PromptCatalog};
@@ -45,6 +45,7 @@ pub enum PendingNotify {
message: String,
auto_run: bool,
extensions: Vec<SessionExtension>,
history_provenance: Option<LoggedSessionHistoryOrigin>,
},
WorkerEvent {
event: WorkerEvent,
@@ -58,6 +59,15 @@ impl PendingNotify {
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.
@@ -81,14 +91,22 @@ impl NotifyBuffer {
message,
auto_run,
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 {
message,
auto_run: true,
auto_run,
extensions: vec![extension],
history_provenance: Some(history_provenance),
});
}
@@ -230,6 +248,7 @@ mod tests {
message: "hello".into(),
auto_run: false,
extensions: Vec::new(),
history_provenance: None,
};
let catalog = PromptCatalog::builtins_only().unwrap();
let item = build_system_item(&entry, &catalog).unwrap();
+557 -48
View File
@@ -79,11 +79,12 @@ const MAX_SUBMISSION_RECEIPTS: usize = 128;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub(crate) struct PendingSubmission {
pub(crate) submission_request_id: String,
source_namespace: String,
pub(crate) submission_id: String,
payload_digest: String,
accepted_at_ms: u64,
activation_sequence: u64,
provenance: WorkerHistoryProvenance,
pub(crate) provenance: WorkerHistoryProvenance,
#[serde(default)]
was_queued: bool,
pub(crate) input: Vec<Segment>,
@@ -92,6 +93,7 @@ pub(crate) struct PendingSubmission {
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
struct SubmissionReceipt {
submission_request_id: String,
source_namespace: String,
submission_id: String,
payload_digest: String,
disposition: protocol::SubmissionDisposition,
@@ -100,17 +102,21 @@ struct SubmissionReceipt {
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub(crate) struct PendingNotification {
pub(crate) notification_request_id: String,
source_namespace: String,
pub(crate) message: String,
payload_digest: String,
pub(crate) auto_run: bool,
accepted_at_ms: u64,
activation_sequence: u64,
provenance: WorkerHistoryProvenance,
pub(crate) provenance: WorkerHistoryProvenance,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
struct NotificationReceipt {
notification_request_id: String,
source_namespace: String,
payload_digest: String,
auto_run: bool,
}
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
@@ -130,14 +136,24 @@ pub(crate) struct PendingActivationState {
impl PendingActivationState {
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))
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()),
(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,
};
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 {
serde_json::to_vec(input)
.map(|bytes| u64::try_from(bytes.len()).unwrap_or(u64::MAX))
@@ -186,6 +261,15 @@ fn submission_payload_digest(input: &[Segment]) -> String {
.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 {
input
.iter()
@@ -1191,11 +1275,57 @@ where
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(
&self,
submission_request_id: String,
input: Vec<Segment>,
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> {
if submission_request_id.trim().is_empty() {
return Err(PendingSubmissionError::EmptyRequestId);
@@ -1218,11 +1348,10 @@ where
.lock()
.expect("pending activation state poisoned");
let original = current.clone();
if let Some(receipt) = current
.receipts
.iter()
.find(|receipt| receipt.submission_request_id == submission_request_id)
{
if let Some(receipt) = current.receipts.iter().find(|receipt| {
receipt.submission_request_id == submission_request_id
&& receipt.source_namespace == source_namespace
}) {
if receipt.payload_digest != payload_digest {
return Err(PendingSubmissionError::IdempotencyConflict);
}
@@ -1237,11 +1366,12 @@ where
let submission_id = uuid::Uuid::now_v7().to_string();
let pending = PendingSubmission {
submission_request_id: submission_request_id.clone(),
source_namespace: source_namespace.clone(),
submission_id: submission_id.clone(),
payload_digest: payload_digest.clone(),
accepted_at_ms: segment_log::now_millis(),
activation_sequence: current.next_activation_sequence,
provenance: WorkerHistoryProvenance::LegacyUnknown,
provenance,
was_queued: !activate_now,
input,
};
@@ -1253,6 +1383,7 @@ where
};
current.remember_receipt(SubmissionReceipt {
submission_request_id: submission_request_id.clone(),
source_namespace,
submission_id: submission_id.clone(),
payload_digest,
disposition,
@@ -1299,7 +1430,14 @@ where
return Err(PendingSubmissionError::ArtifactLimit);
}
current.pending.push_back(pending.clone());
}
if !activate_now {
if let Err(error) = self.pin_submission_files(&pending) {
*current = original;
return Err(error);
}
if let Err(error) = self.persist_locked(&current) {
let _ = self.release_submission_files(&pending);
*current = original;
return Err(error);
}
@@ -1312,10 +1450,29 @@ where
})
}
#[cfg(test)]
pub(crate) fn accept_notification(
&self,
notification_request_id: 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> {
if notification_request_id.trim().is_empty() {
return Err(PendingSubmissionError::EmptyRequestId);
@@ -1323,7 +1480,7 @@ where
if notification_request_id.len() > MAX_ACTIVATION_REQUEST_ID_BYTES {
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
.writer
.state
@@ -1334,12 +1491,11 @@ where
.state
.lock()
.expect("pending activation state poisoned");
if let Some(receipt) = state
.notification_receipts
.iter()
.find(|receipt| receipt.notification_request_id == notification_request_id)
{
if receipt.payload_digest != payload_digest {
if let Some(receipt) = state.notification_receipts.iter().find(|receipt| {
receipt.notification_request_id == notification_request_id
&& receipt.source_namespace == source_namespace
}) {
if receipt.payload_digest != payload_digest || receipt.auto_run != auto_run {
return Err(PendingSubmissionError::IdempotencyConflict);
}
return Ok(false);
@@ -1373,17 +1529,19 @@ where
state.next_activation_sequence = state.next_activation_sequence.saturating_add(1);
state.pending_notifications.push_back(PendingNotification {
notification_request_id: notification_request_id.clone(),
source_namespace: source_namespace.clone(),
message,
payload_digest: payload_digest.clone(),
auto_run,
accepted_at_ms: segment_log::now_millis(),
activation_sequence,
provenance: WorkerHistoryProvenance::BackendInstruction {
operation_id: Some(notification_request_id.clone()),
},
provenance,
});
state.remember_notification_receipt(NotificationReceipt {
notification_request_id,
source_namespace,
payload_digest,
auto_run,
});
state.revision = state.revision.saturating_add(1);
if let Err(error) = self.persist_locked(&state) {
@@ -1393,6 +1551,59 @@ where
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(
&self,
fence: Option<(u64, &str)>,
@@ -1414,17 +1625,20 @@ where
return Ok(None);
}
let submission_sequence = state.pending.front().map(|item| item.activation_sequence);
let notification_sequence = state
let notification_index = state
.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);
if notification_sequence.is_some()
&& (submission_sequence.is_none() || notification_sequence < submission_sequence)
{
let notification = state
.pending_notifications
.pop_front()
.expect("notification sequence came from queue head");
.remove(notification_index.expect("notification sequence came from an item"))
.expect("notification sequence came from an existing item");
state.activating_notification = Some(notification.clone());
state.revision = state.revision.saturating_add(1);
return Ok(Some(PendingActivation::Notification(notification)));
@@ -1442,6 +1656,12 @@ where
}
pub(crate) fn abort_activation(&self, pending: PendingSubmission) {
let _append_guard = self
.writer
.state
.append_lock
.lock()
.expect("segment append lock poisoned");
let mut state = self
.state
.lock()
@@ -1452,13 +1672,16 @@ where
}
state.activating = None;
if pending.was_queued {
state.pending.push_front(pending);
state.pending.push_front(pending.clone());
} else {
state
.receipts
.retain(|receipt| receipt.submission_id != pending.submission_id);
}
state.revision = state.revision.saturating_add(1);
if let Err(error) = self.persist_locked(&state) {
tracing::error!(error = %error, "failed to persist aborted pending activation");
}
}
pub(crate) fn activation_extension(&self) -> SessionExtension {
@@ -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 {
self.state
.lock()
@@ -1561,12 +1788,16 @@ where
else {
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);
if let Err(error) = self.persist_locked(&state) {
*state = original;
return Err(error);
}
self.release_submission_files(&removed)?;
Ok(state.snapshot())
}
@@ -1586,13 +1817,16 @@ where
.expect("pending activation state poisoned");
Self::validate_fence(&state, expected_revision, None)?;
let original = state.clone();
state.pending.clear();
let removed = state.pending.drain(..).collect::<Vec<_>>();
state.pending_notifications.clear();
state.revision = state.revision.saturating_add(1);
if let Err(error) = self.persist_locked(&state) {
*state = original;
return Err(error);
}
for pending in &removed {
self.release_submission_files(pending)?;
}
Ok(state.snapshot())
}
}
@@ -1627,9 +1861,11 @@ pub trait SystemItemCommitter: Send + Sync {
&self,
item: SystemItem,
extensions: Vec<SessionExtension>,
history_provenance: Option<WorkerHistoryProvenance>,
) -> Result<HistoryEntry<SessionHistoryMetadata>, StoreError> {
let metadata = new_history_metadata(
WorkerHistoryProvenance::BackendInstruction { operation_id: None },
history_provenance
.unwrap_or(WorkerHistoryProvenance::BackendInstruction { operation_id: None }),
None,
);
let history_item = item.to_history_item();
@@ -2666,9 +2902,11 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.expect("pending activation state poisoned")
.clone();
if !pending_state.pending.is_empty()
|| !pending_state.pending_notifications.is_empty()
|| pending_state.activating.is_some()
|| pending_state.activating_notification.is_some()
|| !pending_state.receipts.is_empty()
|| !pending_state.notification_receipts.is_empty()
{
let checkpoint = LogEntry::Extension {
ts: segment_log::now_millis(),
@@ -3501,8 +3739,13 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
where
St: Clone + 'static,
{
self.run_with_input_extensions_and_commit_hook(input, input_extensions, || {})
.await
self.run_with_input_extensions_and_commit_hook(
input,
input_extensions,
WorkerHistoryProvenance::LegacyUnknown,
|| {},
)
.await
}
/// Run user input and invoke `on_input_committed` only after the annotated
@@ -3513,6 +3756,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
&mut self,
input: Vec<Segment>,
mut input_extensions: Vec<SessionExtension>,
input_provenance: WorkerHistoryProvenance,
on_input_committed: F,
) -> Result<WorkerRunResult, WorkerError>
where
@@ -3560,8 +3804,12 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
trigger: protocol::InvokeKind::UserSend,
})?;
let projected_input =
self.projected_input_history(&input, flow_projection.as_ref(), &projected_entry_ids);
let projected_input = self.projected_input_history(
&input,
flow_projection.as_ref(),
&projected_entry_ids,
&input_provenance,
);
// Persist original typed segments together with the exact ordered
// model-visible item+origin projection before any entry becomes live.
@@ -3867,6 +4115,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
input: &[Segment],
flow_projection: Option<&PreparedFlowProjection>,
entry_ids: &[SessionHistoryEntryId],
provenance: &WorkerHistoryProvenance,
) -> Vec<HistoryEntry<SessionHistoryMetadata>> {
if let Some(flow) = flow_projection {
return input
@@ -3887,10 +4136,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
other => history_entry_with_id(
Item::user_message(Segment::flatten_to_text(std::slice::from_ref(other))),
entry_id.clone(),
// Current public submit transport does not carry a
// trusted account/Worker subject envelope. Fail closed
// instead of promoting role=user to HumanInput.
WorkerHistoryProvenance::LegacyUnknown,
provenance.clone(),
),
})
.collect();
@@ -3902,7 +4148,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.first()
.expect("projected Worker input always has one entry id")
.clone(),
WorkerHistoryProvenance::LegacyUnknown,
provenance.clone(),
)]
}
@@ -7938,8 +8184,15 @@ mod build_summary_prompt_tests {
serde_json::to_value(&state).unwrap(),
);
let projected_ids = vec![SessionHistoryEntryId::new(), SessionHistoryEntryId::new()];
let projected =
worker.projected_input_history(&segments, projection.as_ref(), &projected_ids);
let input_provenance = WorkerHistoryProvenance::HumanInput {
account_id: "account-1".into(),
};
let projected = worker.projected_input_history(
&segments,
projection.as_ref(),
&projected_ids,
&input_provenance,
);
worker
.commit_entry(LogEntry::AnnotatedUserInput {
ts: segment_log::now_millis(),
@@ -7966,6 +8219,7 @@ mod build_summary_prompt_tests {
projected[0].annotation.origin,
WorkerHistoryProvenance::FlowInstruction { .. }
));
assert_eq!(projected[1].annotation.origin, input_provenance);
assert_eq!(state.instance.definition_revision, 3);
assert_eq!(state.instance.current_state.as_str(), "implement");
assert_eq!(workspace_client.requests.lock().unwrap().len(), 1);
@@ -8048,7 +8302,12 @@ mod build_summary_prompt_tests {
.delete_uploaded_file(worker.session_id(), &file.artifact_id),
Err(StoreError::ArtifactAlreadyCommitted)
));
let projected = worker.projected_input_history(&input, None, &[entry_id]);
let projected = worker.projected_input_history(
&input,
None,
&[entry_id],
&WorkerHistoryProvenance::LegacyUnknown,
);
let text = projected[0].item.as_text().unwrap();
assert!(text.contains("notes.md"));
assert!(text.contains(&file.artifact_id));
@@ -8130,7 +8389,12 @@ mod build_summary_prompt_tests {
if retained.source_entry_id == artifact.source_entry_id
));
let history = worker.projected_input_history(&input, None, &[entry_id]);
let history = worker.projected_input_history(
&input,
None,
&[entry_id],
&WorkerHistoryProvenance::LegacyUnknown,
);
assert!(!history[0].item.as_text().unwrap().contains("終端"));
append_test_entry(
&worker,
@@ -8406,6 +8670,47 @@ mod build_summary_prompt_tests {
assert_eq!(worker.history()[0].as_text().unwrap(), "first message");
}
#[tokio::test]
async fn rewind_preserves_notification_only_pending_activation_checkpoint() {
let (_dir, mut worker) = rewind_test_worker().await;
append_user_turn(&worker, 10, "first message");
append_user_turn(&worker, 20, "second message");
worker
.pending_submission_handle()
.accept_notification("notification-1".into(), "keep me".into(), true)
.unwrap();
let (head_entries, targets) = worker.list_rewind_targets().unwrap();
worker
.rewind_to(targets.last().unwrap().id.clone(), head_entries)
.await
.unwrap();
let location = worker.segment_state.location();
let entries = worker
.store
.read_all(location.session_id, location.segment_id)
.unwrap();
let restored: PendingActivationState = entries
.iter()
.rev()
.find_map(|entry| match entry {
LogEntry::Extension {
domain, payload, ..
} if domain == SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN => {
serde_json::from_value(payload.clone()).ok()
}
_ => None,
})
.unwrap();
assert_eq!(restored.pending_notifications.len(), 1);
assert_eq!(restored.notification_receipts.len(), 1);
assert_eq!(
restored.pending_notifications[0].notification_request_id,
"notification-1"
);
}
#[tokio::test]
async fn annotated_history_rewind_commits_authoritative_prefix() {
let (_dir, mut worker) = rewind_test_worker().await;
@@ -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]
fn pending_submission_queue_is_durable_idempotent_and_bounded() {
let temp = tempfile::tempdir().unwrap();
@@ -9339,22 +9748,96 @@ mod build_summary_prompt_tests {
assert_eq!(cleared.notification_count, 0);
}
#[test]
fn durable_notification_commits_authenticated_history_provenance() {
let temp = tempfile::tempdir().unwrap();
let handle = PendingSubmissionHandle::for_test(temp.path());
let provenance = WorkerHistoryProvenance::HumanInput {
account_id: "account-1".into(),
};
let committed = handle
.writer
.commit_system_item_with_extensions(
SystemItem::Notification {
message: "notice".into(),
body: "notice".into(),
prompt_provenance: None,
},
Vec::new(),
Some(provenance.clone()),
)
.unwrap();
assert_eq!(committed.annotation.origin, provenance);
}
#[test]
fn notification_retry_identity_is_scoped_to_authenticated_source() {
let temp = tempfile::tempdir().unwrap();
let handle = PendingSubmissionHandle::for_test(temp.path());
let account_a = WorkerHistoryProvenance::HumanInput {
account_id: "account-a".into(),
};
let account_b = WorkerHistoryProvenance::HumanInput {
account_id: "account-b".into(),
};
assert!(
handle
.accept_notification_from_source(
"request-1".into(),
"notice".into(),
"account:account-a".into(),
account_a.clone(),
false,
)
.unwrap()
);
assert!(
!handle
.accept_notification_from_source(
"request-1".into(),
"notice".into(),
"account:account-a".into(),
account_a.clone(),
false,
)
.unwrap()
);
assert!(
handle
.accept_notification_from_source(
"request-1".into(),
"notice".into(),
"account:account-b".into(),
account_b.clone(),
false,
)
.unwrap()
);
let state = handle.state.lock().unwrap();
assert_eq!(state.pending_notifications[0].provenance, account_a);
assert_eq!(state.pending_notifications[1].provenance, account_b);
}
#[test]
fn notification_and_submit_share_activation_order_and_notification_dedupes() {
let temp = tempfile::tempdir().unwrap();
let handle = PendingSubmissionHandle::for_test(temp.path());
assert!(
handle
.accept_notification("notification-1".into(), "notice".into())
.accept_notification("notification-1".into(), "notice".into(), true)
.unwrap()
);
assert!(
!handle
.accept_notification("notification-1".into(), "notice".into())
.accept_notification("notification-1".into(), "notice".into(), true)
.unwrap()
);
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)
));
handle
@@ -9382,9 +9865,10 @@ mod build_summary_prompt_tests {
let mut session = WorkerSession::new(session_store::new_session_id(), Vec::new());
let state = PendingActivationState {
revision: 4,
next_activation_sequence: 2,
next_activation_sequence: 3,
activating: Some(PendingSubmission {
submission_request_id: "request-1".into(),
source_namespace: "direct:test".into(),
submission_id: "submission-1".into(),
payload_digest: submission_payload_digest(&[Segment::text("first")]),
accepted_at_ms: 1,
@@ -9393,9 +9877,21 @@ mod build_summary_prompt_tests {
was_queued: false,
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 {
submission_request_id: "request-2".into(),
source_namespace: "direct:test".into(),
submission_id: "submission-2".into(),
payload_digest: submission_payload_digest(&[Segment::text("second")]),
accepted_at_ms: 2,
@@ -9406,7 +9902,12 @@ mod build_summary_prompt_tests {
}]),
pending_notifications: 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_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.into(),
@@ -9420,6 +9921,14 @@ mod build_summary_prompt_tests {
assert_eq!(state.pending.len(), 2);
assert_eq!(state.pending[0].submission_id, "submission-1");
assert_eq!(state.pending[1].submission_id, "submission-2");
assert!(state.activating_notification.is_none());
assert_eq!(state.pending_notifications.len(), 1);
assert!(!state.pending_notifications[0].auto_run);
assert!(matches!(
state.pending_notifications[0].provenance,
WorkerHistoryProvenance::HumanInput { ref account_id } if account_id == "account-1"
));
assert_eq!(state.notification_receipts.len(), 1);
}
fn minimal_manifest() -> WorkerManifest {
+71 -8
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 worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
let notification_request_id = protocol::new_submission_request_id();
handle
.send(Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message: "progress snapshot".into(),
auto_run: false,
})
.await
.unwrap();
for _ in 0..2 {
handle
.send(Method::Notify {
notification_request_id: notification_request_id.clone(),
message: "progress snapshot".into(),
auto_run: false,
})
.await
.unwrap();
}
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Idle);
@@ -2049,6 +2052,66 @@ async fn notify_while_running_does_not_emit_already_running_error() {
wait_for_status(&handle, WorkerStatus::Idle).await;
}
#[tokio::test]
async fn weak_notify_while_running_is_deduped_and_survives_until_next_submit() {
let client = MockClient::sequential(vec![
MockResponse::Hang(Vec::new()),
MockResponse::Complete(simple_text_events()),
]);
let client_for_assert = client.clone();
let worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"first",
))
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Running).await;
let notification_request_id = protocol::new_submission_request_id();
for _ in 0..2 {
handle
.send(Method::Notify {
notification_request_id: notification_request_id.clone(),
message: "durable weak notice".into(),
auto_run: false,
})
.await
.unwrap();
}
handle.send(Method::Cancel).await.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let mut rx = handle.subscribe();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"second",
))
.await
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if matches!(rx.recv().await, Ok(Event::TurnEnd { .. })) {
break;
}
}
})
.await
.expect("second submit completes");
let requests = client_for_assert.captured_requests();
let notice_count = requests[1]
.items
.iter()
.filter_map(|item| item.as_text())
.filter(|text| text.contains("durable weak notice"))
.count();
assert_eq!(notice_count, 1);
}
#[tokio::test]
async fn status_json_reflects_worker_name() {
let client = MockClient::new(simple_text_events());
+6 -4
View File
@@ -4987,7 +4987,7 @@ mod tests {
"missing test context",
);
};
let submission_id = input.submission_id.clone();
let submission_request_id = input.submission_request_id.clone();
let content = input.content;
std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(10));
@@ -5004,11 +5004,13 @@ mod tests {
status: protocol::WorkerStatus::Idle,
});
});
if let Some(submission_id) = submission_id {
worker_runtime::execution::WorkerExecutionResult::accepted_input_committed(
if let Some(submission_request_id) = submission_request_id {
worker_runtime::execution::WorkerExecutionResult::accepted_submission(
worker_runtime::execution::WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
submission_id,
submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
)
} else {
worker_runtime::execution::WorkerExecutionResult::accepted(
@@ -32,11 +32,13 @@ impl WorkerExecutionBackend for TestExecutionBackend {
_handle: &WorkerExecutionHandle,
input: worker_runtime::interaction::WorkerInput,
) -> WorkerExecutionResult {
if let Some(submission_id) = input.submission_id {
WorkerExecutionResult::accepted_input_committed(
if let Some(submission_request_id) = input.submission_request_id {
WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
submission_id,
submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
)
} else {
WorkerExecutionResult::accepted(
+139 -31
View File
@@ -8537,13 +8537,15 @@ async fn scoped_list_runtimes(
async fn scoped_workspace_protocol_ws(
State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath(workspace_id): AxumPath<String>,
ws: axum::extract::ws::WebSocketUpgrade,
) -> std::result::Result<Response, Response> {
validate_workspace_scope(&api, &workspace_id).map_err(|error| error.into_response())?;
let input_source = authenticated_browser_input_source(&actor);
Ok(ws
.on_upgrade(move |socket| {
crate::workspace_subscription::serve_workspace_subscription(api, socket)
crate::workspace_subscription::serve_workspace_subscription(api, socket, input_source)
})
.into_response())
}
@@ -11458,6 +11460,7 @@ async fn scoped_cancel_runtime_worker(
async fn scoped_worker_protocol_ws(
ws: WebSocketUpgrade,
State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
) -> Response {
if let Err(err) = validate_workspace_scope(&api, &path.workspace_id) {
@@ -11465,6 +11468,7 @@ async fn scoped_worker_protocol_ws(
}
worker_protocol_ws(
State(api),
Extension(actor),
AxumPath((path.worker.runtime_id, path.worker.worker_id)),
ws,
)
@@ -13877,8 +13881,45 @@ async fn cancel_runtime_worker(
Ok(Json(result))
}
fn authenticated_browser_input_source(actor: &RequestActor) -> protocol::AuthenticatedInputSource {
protocol::AuthenticatedInputSource::Account {
account_id: actor.account_id.clone(),
}
}
pub(crate) fn authorize_browser_worker_method(
method: protocol::Method,
source: &protocol::AuthenticatedInputSource,
) -> std::result::Result<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(
State(api): State<WorkspaceApi>,
Extension(actor): Extension<RequestActor>,
AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>,
ws: WebSocketUpgrade,
) -> impl IntoResponse {
@@ -13902,7 +13943,8 @@ async fn worker_protocol_ws(
.into_response();
}
};
ws.on_upgrade(move |socket| worker_protocol_ws_session(source, socket))
let input_source = authenticated_browser_input_source(&actor);
ws.on_upgrade(move |socket| worker_protocol_ws_session(source, socket, input_source))
}
pub(crate) struct WorkspaceWorkerProtocolConnection {
@@ -14029,13 +14071,17 @@ async fn connect_embedded_worker_protocol(
Ok(WorkspaceWorkerProtocolConnection { methods, events })
}
async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: WebSocket) {
async fn worker_protocol_ws_session(
source: RuntimeObservationSource,
socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
match source {
RuntimeObservationSource::RemoteWs(config) => {
remote_worker_protocol_ws_session(config, socket).await;
remote_worker_protocol_ws_session(config, socket, input_source).await;
}
RuntimeObservationSource::Embedded(source) => {
embedded_worker_protocol_ws_session(source, socket).await;
embedded_worker_protocol_ws_session(source, socket, input_source).await;
}
}
}
@@ -14043,6 +14089,7 @@ async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: We
async fn remote_worker_protocol_ws_session(
config: RuntimeObservationSourceConfig,
socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
let mut request = match config.endpoint.clone().into_client_request() {
Ok(request) => request,
@@ -14091,14 +14138,33 @@ async fn remote_worker_protocol_ws_session(
inbound = client_stream.next() => {
match inbound {
Some(Ok(WsMessage::Text(text))) => {
if upstream_sink.send(TungsteniteMessage::Text(text.to_string().into())).await.is_err() {
let method = match protocol::stream::decode_method(text.as_ref()) {
Ok(method) => match authorize_browser_worker_method(method, &input_source) {
Ok(method) => method,
Err(message) => {
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event(message)) {
let _ = client_sink.send(WsMessage::Text(event.into())).await;
}
break;
}
},
Err(error) => {
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event(error.to_string())) {
let _ = client_sink.send(WsMessage::Text(event.into())).await;
}
break;
}
};
let Ok(method) = protocol::stream::encode_method(&method) else { break };
if upstream_sink.send(TungsteniteMessage::Text(method.into())).await.is_err() {
break;
}
}
Some(Ok(WsMessage::Binary(binary))) => {
if upstream_sink.send(TungsteniteMessage::Binary(binary.to_vec().into())).await.is_err() {
break;
Some(Ok(WsMessage::Binary(_))) => {
if let Ok(event) = protocol::stream::encode_event(&protocol_error_event("binary Worker methods are not accepted")) {
let _ = client_sink.send(WsMessage::Text(event.into())).await;
}
break;
}
Some(Ok(WsMessage::Close(_))) | None => {
let _ = upstream_sink.send(TungsteniteMessage::Close(None)).await;
@@ -14154,6 +14220,7 @@ async fn remote_worker_protocol_ws_session(
async fn embedded_worker_protocol_ws_session(
source: crate::observation::EmbeddedRuntimeObservationSource,
mut socket: WebSocket,
input_source: protocol::AuthenticatedInputSource,
) {
let mut upstream = match RuntimeObservationClient::connect(&RuntimeObservationSource::Embedded(
source.clone(),
@@ -14173,24 +14240,32 @@ async fn embedded_worker_protocol_ws_session(
inbound = socket.next() => {
match inbound {
Some(Ok(WsMessage::Text(text))) => match decode_method(&text) {
Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) {
Ok(events) => {
for event in events {
Ok(method) => match authorize_browser_worker_method(method, &input_source) {
Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) {
Ok(events) => {
for event in events {
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
}
Err(error) => {
let event = protocol_error_event(error.to_string());
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
}
Err(error) => {
let event = protocol_error_event(error.to_string());
if !send_protocol_event(&mut socket, &event).await {
return;
}
},
Err(message) => {
let event = protocol_error_event(message);
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
},
Err(error) => {
let event =
protocol_error_event(format!("malformed protocol method frame: {error}"));
let event = protocol_error_event(format!(
"malformed protocol method frame: {error}"
));
if !send_protocol_event(&mut socket, &event).await {
return;
}
@@ -16580,6 +16655,29 @@ mod tests {
&tail[..end]
}
#[test]
fn browser_worker_methods_receive_server_owned_account_source() {
let source = protocol::AuthenticatedInputSource::Account {
account_id: "account-1".into(),
};
let method = authorize_browser_worker_method(
protocol::Method::Submit {
submission_request_id: "request-1".into(),
input: vec![protocol::Segment::text("hello")],
},
&source,
)
.unwrap();
assert!(matches!(
method,
protocol::Method::SubmitTracked {
source: protocol::AuthenticatedInputSource::Account { ref account_id },
..
} if account_id == "account-1"
));
assert!(authorize_browser_worker_method(method, &source).is_err());
}
#[test]
fn merge_request_http_paths_observe_refs_through_runtime_provider_authority() {
let source = include_str!("server.rs");
@@ -18825,7 +18923,7 @@ mod tests {
.get(handle.worker_ref())
.cloned()
.expect("execution context");
let submission_id = input.submission_id.clone();
let submission_request_id = input.submission_request_id.clone();
let content = input.content.clone();
std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(25));
@@ -18833,11 +18931,13 @@ mod tests {
text: format!("server companion echoed: {content}"),
});
});
if let Some(submission_id) = submission_id {
worker_runtime::execution::WorkerExecutionResult::accepted_input_committed(
if let Some(submission_request_id) = submission_request_id {
worker_runtime::execution::WorkerExecutionResult::accepted_submission(
worker_runtime::execution::WorkerExecutionOperation::Input,
worker_runtime::execution::WorkerExecutionRunState::Idle,
submission_id,
submission_request_id,
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
)
} else {
worker_runtime::execution::WorkerExecutionResult::accepted(
@@ -27223,6 +27323,16 @@ mod tests {
(runtime, worker_ref, endpoint)
}
fn test_browser_request_actor() -> RequestActor {
RequestActor {
user_id: "test-user".into(),
account_id: format!("account-{TEST_WORKSPACE_ID}"),
handle: "test".into(),
display_name: "Test".into(),
auth_method: ActorAuthMethod::BrowserSession,
}
}
async fn spawn_workspace_proxy(
source: RuntimeObservationSourceConfig,
) -> (String, tempfile::TempDir) {
@@ -27241,11 +27351,8 @@ mod tests {
.unwrap();
let app_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let app_addr = app_listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(app_listener, build_inner_router(api))
.await
.unwrap()
});
let app = build_inner_router(api).layer(Extension(test_browser_request_actor()));
tokio::spawn(async move { axum::serve(app_listener, app).await.unwrap() });
(
format!("ws://{app_addr}/api/runtimes/{runtime_id}/workers/{worker_id}/protocol/ws"),
dir,
@@ -27257,7 +27364,8 @@ mod tests {
let dir = tempfile::tempdir().unwrap();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = build_inner_router(test_api(dir.path()).await);
let app = build_inner_router(test_api(dir.path()).await)
.layer(Extension(test_browser_request_actor()));
let server = tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
});
@@ -27305,7 +27413,7 @@ mod tests {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = build_inner_router(api);
let app = build_inner_router(api).layer(Extension(test_browser_request_actor()));
let server = tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
});
@@ -11,7 +11,9 @@ use tokio::sync::mpsc;
use worker_runtime::identity::RuntimeWorkerRef;
use crate::runtime_subscription::{BrokerSubscriptionEvent, RuntimeSubscriptionBroker};
use crate::server::{WorkspaceApi, connect_workspace_worker_protocol};
use crate::server::{
WorkspaceApi, authorize_browser_worker_method, connect_workspace_worker_protocol,
};
use crate::store::WorkspaceResourceKind;
const OUTBOUND_CAPACITY: usize = 256;
@@ -21,7 +23,11 @@ struct ActiveSubscription {
methods: Option<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 (mut socket_sender, mut socket_receiver) = socket.split();
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 {
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;
}
}