Make notifications realtime and coalesced

This commit is contained in:
2026-09-18 05:04:47 +09:00
parent ae7f0f0446
commit 9cde75b70f
16 changed files with 326 additions and 536 deletions
+8 -31
View File
@@ -16,14 +16,6 @@ pub fn new_submission_request_id() -> String {
uuid::Uuid::now_v7().to_string()
}
fn default_true() -> bool {
true
}
fn is_true(value: &bool) -> bool {
*value
}
fn is_false(value: &bool) -> bool {
!*value
}
@@ -231,25 +223,20 @@ pub enum Method {
#[serde(skip_deserializing, default)]
source: AuthenticatedInputSource,
},
/// Human-readable text injected into the target Worker's LLM context
/// as a non-blocking system message. `auto_run` controls whether an
/// idle target is kicked into `RunForNotification`; weak notifications
/// (`auto_run: false`) are only queued for the next turn/resume/run.
/// No side effects beyond LLM context; use `WorkerEvent` for typed
/// lifecycle reports.
/// Human-readable text injected into the target Worker's LLM context as a
/// non-blocking system message. Accepted notifications are delivered at the
/// next between-turn boundary, or start a coalesced notification run while
/// the Worker is idle. No side effects beyond LLM context; use
/// `WorkerEvent` for typed lifecycle reports.
Notify {
notification_request_id: String,
message: String,
#[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,
#[serde(skip_deserializing, default)]
source: AuthenticatedInputSource,
},
@@ -2040,26 +2027,16 @@ mod tests {
}
#[test]
fn method_notify_json_roundtrip_defaults_to_auto_run() {
fn method_notify_json_roundtrip_has_no_delivery_mode() {
let json = r#"{"method":"notify","params":{"notification_request_id":"notification-1","message":"turn done"}}"#;
let method: Method = serde_json::from_str(json).unwrap();
assert!(matches!(
method,
Method::Notify { ref message, auto_run: true, .. } if message == "turn done"
Method::Notify { ref message, .. } if message == "turn done"
));
let serialized = serde_json::to_string(&method).unwrap();
assert_eq!(serialized, json);
}
#[test]
fn method_notify_weak_json_roundtrip_serializes_auto_run_false() {
let json = r#"{"method":"notify","params":{"notification_request_id":"notification-1","message":"progress","auto_run":false}}"#;
let method: Method = serde_json::from_str(json).unwrap();
assert!(matches!(
method,
Method::Notify { ref message, auto_run: false, .. } if message == "progress"
));
assert_eq!(serde_json::to_string(&method).unwrap(), json);
assert!(!serialized.contains("auto_run"));
}
#[test]
-1
View File
@@ -349,7 +349,6 @@ async fn standalone_restore_preserves_history_tasks_notifications_and_cwd_scope(
.send(&Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message: "persisted notification".to_string(),
auto_run: true,
})
.await?;
wait_for_run_end(&mut protocol_client).await?;
-1
View File
@@ -782,7 +782,6 @@ impl App {
Some(Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message,
auto_run: true,
})
}
-1
View File
@@ -2009,7 +2009,6 @@ mod tests {
method,
Some(Method::Notify {
ref message,
auto_run: true,
..
}) if message == "progress"
));
-3
View File
@@ -1565,7 +1565,6 @@ fn authorize_runtime_protocol_method(
protocol::Method::NotifyTracked {
notification_request_id,
message,
auto_run,
..
} => protocol::Method::NotifyTracked {
source: transport_source.cloned().unwrap_or_else(|| {
@@ -1575,7 +1574,6 @@ fn authorize_runtime_protocol_method(
}),
notification_request_id,
message,
auto_run,
},
other => other,
}
@@ -2987,7 +2985,6 @@ mod tests {
let wire = serde_json::to_string(&protocol::Method::NotifyTracked {
notification_request_id: "notification-1".into(),
message: "hello".into(),
auto_run: true,
source: protocol::AuthenticatedInputSource::Account {
account_id: "forged".into(),
},
@@ -2374,7 +2374,6 @@ where
Method::NotifyTracked {
notification_request_id: notification_request_id.clone(),
message: input.content,
auto_run: true,
source: protocol::AuthenticatedInputSource::Backend {
operation_id: notification_request_id,
},
+111 -207
View File
@@ -2,6 +2,7 @@ use std::collections::VecDeque;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;
use agen::EngineError;
use agen::llm_client::client::LlmClient;
@@ -307,7 +308,6 @@ enum PendingRun {
/// committed at the start of `worker.run_for_notification`.
RunForNotification {
invoke_kind: protocol::InvokeKind,
notification_request_id: Option<String>,
},
Resume,
}
@@ -336,7 +336,6 @@ fn durable_parent_notification_target<St: Store + Clone + Send + Sync + 'static>
let Method::NotifyTracked {
notification_request_id,
message,
auto_run,
source,
} = method
else {
@@ -344,21 +343,14 @@ fn durable_parent_notification_target<St: Store + Clone + Send + Sync + 'static>
};
let (source_namespace, provenance) = resolved_input_source(&pending_submissions, &source);
match pending_submissions.accept_notification_from_source(
notification_request_id.clone(),
notification_request_id,
message,
source_namespace.clone(),
source_namespace,
provenance,
auto_run,
) {
Ok(_) if !auto_run => {
stage_pending_notification(
&pending_submissions,
&notify_buffer,
&source_namespace,
&notification_request_id,
);
Ok(_) => {
stage_pending_notifications(&pending_submissions, &notify_buffer);
}
Ok(_) => {}
Err(error) => {
tracing::warn!(%error, "failed to durably accept SubWorker notification");
}
@@ -366,69 +358,60 @@ fn durable_parent_notification_target<St: Store + Clone + Send + Sync + 'static>
}))
}
fn stage_pending_notification<St: Store + Clone>(
fn stage_pending_notifications<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
) -> usize {
let notifications = pending_submissions.prepare_notification_batch();
if notifications.is_empty() {
return 0;
}
let count = notifications.len();
for (notification, extension) in notifications {
notify_buffer.push_durable_notify(
notification.message,
notification.provenance,
extension,
);
}
count
}
fn stage_oldest_passive_notification<St: Store + Clone>(
fn notification_coalesce_remaining<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,
)
})
delay: Duration,
) -> Option<Duration> {
let Some(accepted_at_ms) = pending_submissions.oldest_pending_notification_accepted_at_ms()
else {
return notify_buffer
.has_notification_pending()
.then_some(Duration::ZERO);
};
let now_ms = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
let elapsed = Duration::from_millis(now_ms.saturating_sub(accepted_at_ms));
Some(delay.saturating_sub(elapsed))
}
fn prepare_pending_run<St: Store + Clone>(
pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
notify_buffer: &NotifyBuffer,
fence: Option<(u64, &str)>,
allow_notification_run: bool,
) -> Result<Option<PendingRun>, crate::worker::PendingSubmissionError> {
let staged_passive_notification = pending_submissions.activating_passive_notification_id();
Ok(match pending_submissions.prepare_next_activation(fence)? {
Some(crate::worker::PendingActivation::Submission(submission)) => {
if staged_passive_notification.is_some() {
let extension = pending_submissions.notification_activation_extension();
debug_assert!(notify_buffer.replace_durable_notification_extension(extension));
}
stage_pending_notifications(pending_submissions, notify_buffer);
Some(PendingRun::Submit(submission))
}
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,
notification.auto_run,
notification.provenance,
extension,
);
None if allow_notification_run
&& (stage_pending_notifications(pending_submissions, notify_buffer) > 0
|| notify_buffer.has_notification_pending()) => {
Some(PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::Notify,
notification_request_id: Some(notification_request_id),
})
}
None => None,
@@ -702,7 +685,6 @@ impl WorkerController {
worker.push_notify(
"Restored Worker state contained unreachable delegated child Workers; their delegated write scopes were reclaimed before resume."
.to_string(),
false,
);
}
@@ -1556,9 +1538,9 @@ async fn controller_loop<C, St>(
discovery_cwd,
spawned_registry.clone(),
);
let notification_coalesce_delay = worker.notification_coalesce_delay();
let pending_submissions = worker.pending_submission_handle();
stage_oldest_passive_notification(&pending_submissions, &notify_buffer);
let mut pending = match prepare_pending_run(&pending_submissions, &notify_buffer, None) {
let mut pending = match prepare_pending_run(&pending_submissions, &notify_buffer, None, false) {
Ok(pending) => pending,
Err(error) => {
let _ = working_event_tx.send(Event::Error {
@@ -1575,6 +1557,10 @@ async fn controller_loop<C, St>(
// here so the status flip → drive_turn → finish sequence lives
// in one place, regardless of which Method caused it.
if let Some(run) = pending.take() {
// A submit/resume/run always provides an immediate model boundary,
// so any coalescing window ends and the whole FIFO notification
// batch is injected into that run.
stage_pending_notifications(&pending_submissions, &notify_buffer);
// Cancellation is meaningful only for an accepted running turn. Clear
// idle/stale signals before the status flip; any Cancel/Pause received
// after this point is delivered to the turn and must not be discarded by
@@ -1595,15 +1581,6 @@ async fn controller_loop<C, St>(
)
.await;
}
let notification_request_id = match &run {
PendingRun::RunForNotification {
notification_request_id,
..
} => 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();
@@ -1673,28 +1650,17 @@ async fn controller_loop<C, St>(
.await
}
};
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 notify_buffer.is_empty() {
pending_submissions.finish_notification_batch();
}
if !shutdown && may_drain_pending && new_status == WorkerStatus::Idle {
match prepare_pending_run(&pending_submissions, &notify_buffer, None) {
match prepare_pending_run(&pending_submissions, &notify_buffer, None, false) {
Ok(Some(next)) => {
pending = Some(next);
new_status = WorkerStatus::Running;
}
Ok(None) => {
if notify_buffer.has_auto_run_pending() {
pending = Some(PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::Notify,
notification_request_id: None,
});
new_status = WorkerStatus::Running;
}
}
Ok(None) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
@@ -1724,10 +1690,45 @@ async fn controller_loop<C, St>(
let method = if let Some(method) = deferred_methods.pop_front() {
method
} else if let Some(remaining) =
notification_coalesce_remaining(
&pending_submissions,
&notify_buffer,
notification_coalesce_delay,
)
{
tokio::select! {
method = method_rx.recv() => match method {
Some(method) => method,
None => break,
},
_ = notify_buffer.notified() => continue,
_ = tokio::time::sleep(remaining) => {
match prepare_pending_run(
&pending_submissions,
&notify_buffer,
None,
true,
) {
Ok(Some(next)) => pending = Some(next),
Ok(None) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: error.to_string(),
});
}
}
continue;
}
}
} else {
match method_rx.recv().await {
Some(method) => method,
None => break,
tokio::select! {
method = method_rx.recv() => match method {
Some(method) => method,
None => break,
},
_ = notify_buffer.notified() => continue,
}
};
@@ -1799,37 +1800,15 @@ async fn controller_loop<C, St>(
Method::Notify {
notification_request_id,
message,
auto_run,
} => {
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(_) => {
stage_pending_notification(
&pending_submissions,
&notify_buffer,
&source_namespace,
&request_id,
);
}
Ok(_) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
@@ -1842,10 +1821,8 @@ async fn controller_loop<C, St>(
Method::NotifyTracked {
notification_request_id,
message,
auto_run,
source,
} => {
let request_id = notification_request_id.clone();
let (source_namespace, provenance) =
resolved_input_source(&pending_submissions, &source);
match pending_submissions.accept_notification_from_source(
@@ -1853,28 +1830,8 @@ async fn controller_loop<C, St>(
message,
source_namespace.clone(),
provenance,
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,
);
}
Ok(_) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
@@ -1935,6 +1892,7 @@ async fn controller_loop<C, St>(
&pending_submissions,
&notify_buffer,
Some((expected_revision, &expected_head_id)),
true,
) {
Ok(Some(next)) => pending = Some(next),
Ok(None) => {
@@ -2361,7 +2319,6 @@ async fn controller_loop<C, St>(
if shared_state.catalog_status() == WorkerStatus::Idle {
pending = Some(PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::WorkerEvent,
notification_request_id: None,
});
}
}
@@ -2893,26 +2850,17 @@ where
Some(Method::Notify {
notification_request_id,
message,
auto_run,
}) => {
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(_) => {
stage_pending_notifications(&pending_submissions, notify_buffer);
}
Ok(_) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
@@ -2927,10 +2875,8 @@ where
Some(Method::NotifyTracked {
notification_request_id,
message,
auto_run,
source,
}) => {
let request_id = notification_request_id.clone();
let (source_namespace, provenance) =
resolved_input_source(pending_submissions, &source);
match pending_submissions.accept_notification_from_source(
@@ -2938,17 +2884,10 @@ where
message,
source_namespace.clone(),
provenance,
auto_run,
) {
Ok(_) if !auto_run => {
stage_pending_notification(
&pending_submissions,
notify_buffer,
&source_namespace,
&request_id,
);
Ok(_) => {
stage_pending_notifications(&pending_submissions, notify_buffer);
}
Ok(_) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
@@ -3141,7 +3080,8 @@ mod tests {
let temp = TempDir::new().unwrap();
let pending =
crate::worker::PendingSubmissionHandle::for_test(&temp.path().join("sessions"));
let target = durable_parent_notification_target(pending.clone(), NotifyBuffer::new());
let notify_buffer = NotifyBuffer::new();
let target = durable_parent_notification_target(pending.clone(), notify_buffer.clone());
let (wire_namespace, wire_provenance) =
resolved_input_source(&pending, &protocol::AuthenticatedInputSource::UntrustedWire);
assert_eq!(wire_namespace, pending.direct_client_namespace());
@@ -3150,20 +3090,15 @@ mod tests {
session_store::LoggedSessionHistoryOrigin::LegacyUnknown
));
target.notify("child-session".into(), "completed".into(), true);
target.notify("child-session".into(), "completed".into());
let snapshot = pending.snapshot();
assert_eq!(snapshot.notification_count, 1);
assert!(snapshot.head_id.is_some());
let notify_buffer = NotifyBuffer::new();
assert_eq!(snapshot.notification_count, 0);
assert!(matches!(
prepare_pending_run(&pending, &notify_buffer, None).unwrap(),
Some(PendingRun::RunForNotification {
notification_request_id: Some(_),
..
})
prepare_pending_run(&pending, &notify_buffer, None, true).unwrap(),
Some(PendingRun::RunForNotification { .. })
));
assert!(notify_buffer.has_auto_run_pending());
assert!(notify_buffer.has_notification_pending());
}
#[test]
@@ -3181,19 +3116,19 @@ mod tests {
)
.unwrap();
pending
.accept_notification("notify-second".into(), "newer notification".into(), true)
.accept_notification("notify-second".into(), "newer notification".into())
.unwrap();
let notify_buffer = NotifyBuffer::new();
assert!(matches!(
prepare_pending_run(&pending, &notify_buffer, None).unwrap(),
prepare_pending_run(&pending, &notify_buffer, None, true).unwrap(),
Some(PendingRun::Submit(_))
));
let pending =
crate::worker::PendingSubmissionHandle::for_test(&temp.path().join("notify-first"));
pending
.accept_notification("notify-first".into(), "older notification".into(), true)
.accept_notification("notify-first".into(), "older notification".into())
.unwrap();
pending
.accept(
@@ -3207,40 +3142,11 @@ mod tests {
let notify_buffer = NotifyBuffer::new();
assert!(matches!(
prepare_pending_run(&pending, &notify_buffer, None).unwrap(),
Some(PendingRun::RunForNotification {
notification_request_id: Some(request_id),
..
}) if request_id == "notify-first"
));
let pending =
crate::worker::PendingSubmissionHandle::for_test(&temp.path().join("passive-first"));
pending
.accept_notification("passive-first".into(), "passive notification".into(), false)
.unwrap();
pending
.accept(
"submit-after-passive".into(),
vec![protocol::Segment::Text {
content: "queued after passive".into(),
}],
false,
)
.unwrap();
let snapshot = pending.snapshot();
let head_id = snapshot.head_id.clone().expect("queued Submit is the head");
let notify_buffer = NotifyBuffer::new();
assert!(stage_oldest_passive_notification(&pending, &notify_buffer));
assert!(matches!(
prepare_pending_run(
&pending,
&notify_buffer,
Some((snapshot.revision + 1, &head_id)),
)
.unwrap(),
prepare_pending_run(&pending, &notify_buffer, None, true).unwrap(),
Some(PendingRun::Submit(_))
));
assert_eq!(notify_buffer.len(), 1);
}
#[test]
@@ -3263,7 +3169,6 @@ mod tests {
assert!(
!PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::Notify,
notification_request_id: None,
}
.is_parent_originated()
);
@@ -3658,13 +3563,12 @@ mod tests {
}
#[tokio::test]
async fn running_auto_run_notify_remains_staged_for_followup_turn() {
async fn running_notify_remains_staged_for_the_next_model_boundary() {
let mut env = make_env().await;
env._method_tx
.send(Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message: "continue".into(),
auto_run: true,
})
.await
.expect("send notify");
@@ -3693,8 +3597,8 @@ mod tests {
assert_eq!(status, WorkerStatus::Idle);
assert!(!shutdown);
assert_eq!(env.notify_buffer.len(), 0);
assert_eq!(env.pending_submissions.snapshot().notification_count, 1);
assert_eq!(env.notify_buffer.len(), 1);
assert_eq!(env.pending_submissions.snapshot().notification_count, 0);
}
#[tokio::test]
+6 -15
View File
@@ -389,7 +389,7 @@ where
reason: detail.live.error,
};
}
match send_notify(&detail.live.socket_path, message, false).await {
match send_notify(&detail.live.socket_path, message).await {
Ok(()) => WeakNotifyDelivery::Delivered,
Err(error) => WeakNotifyDelivery::SendFailed {
error: error.to_string(),
@@ -1009,17 +1009,16 @@ where
}
async fn send_peer_notify(socket_path: &Path, message: String) -> io::Result<()> {
send_notify(socket_path, message, true).await
send_notify(socket_path, message).await
}
async fn send_notify(socket_path: &Path, message: String, auto_run: bool) -> io::Result<()> {
async fn send_notify(socket_path: &Path, message: String) -> io::Result<()> {
let notification_request_id = protocol::new_submission_request_id();
connect_and_send(
socket_path,
&Method::NotifyTracked {
notification_request_id: notification_request_id.clone(),
message,
auto_run,
source: protocol::AuthenticatedInputSource::Backend {
operation_id: notification_request_id,
},
@@ -1552,11 +1551,7 @@ mod tests {
.await
.unwrap();
let method = reader.next::<Method>().await.unwrap().unwrap();
if let Method::NotifyTracked {
message, auto_run, ..
} = method
{
assert!(auto_run);
if let Method::NotifyTracked { message, .. } = method {
tx.send(message).await.unwrap();
} else {
panic!("expected Notify, got {method:?}");
@@ -1575,7 +1570,7 @@ mod tests {
}
#[tokio::test(flavor = "current_thread")]
async fn weak_notify_to_live_peer_uses_notify_without_auto_run_and_noops_when_missing() {
async fn weak_notify_to_live_peer_uses_realtime_notify_and_noops_when_missing() {
let root = TempDir::new().unwrap();
let store_dir = root.path().join("store");
let runtime_base = root.path().join("runtime");
@@ -1674,11 +1669,7 @@ mod tests {
.await
.unwrap();
let method = reader.next::<Method>().await.unwrap().unwrap();
if let Method::NotifyTracked {
message, auto_run, ..
} = method
{
assert!(!auto_run);
if let Method::NotifyTracked { message, .. } = method {
tx.send(message).await.unwrap();
} else {
panic!("expected Notify, got {method:?}");
+5 -5
View File
@@ -1543,7 +1543,7 @@ mod tests {
PromptCatalog::from_projection(projection).unwrap(),
));
buffer.push_notify("updated".to_string(), false);
buffer.push_notify("updated".to_string());
let appends = interceptor
.pending_history_appends(PendingHistoryAppendsContext {
invocation: Default::default(),
@@ -1598,7 +1598,7 @@ mod tests {
prompts.store(Arc::new(
PromptCatalog::from_projection(projection).unwrap(),
));
buffer.push_notify("must persist".to_string(), false);
buffer.push_notify("must persist".to_string());
let error = interceptor
.pending_history_appends(PendingHistoryAppendsContext {
@@ -1621,8 +1621,8 @@ mod tests {
async fn pending_history_appends_drains_buffer_into_items() {
let registry = Arc::new(HookRegistryBuilder::new().build());
let buffer = NotifyBuffer::new();
buffer.push_notify("first".into(), false);
buffer.push_notify("second".into(), false);
buffer.push_notify("first".into());
buffer.push_notify("second".into());
let interceptor = WorkerInterceptor::new(
registry,
@@ -1671,7 +1671,7 @@ mod tests {
// anything itself.
let registry = Arc::new(HookRegistryBuilder::new().build());
let buffer = NotifyBuffer::new();
buffer.push_notify("msg".into(), false);
buffer.push_notify("msg".into());
let interceptor = WorkerInterceptor::new(
registry,
+27 -60
View File
@@ -26,13 +26,9 @@ use std::sync::{Arc, Mutex};
use protocol::WorkerEvent;
use session_store::{LoggedSessionHistoryOrigin, SessionExtension, SystemItem};
use tracing::warn;
use crate::prompt::catalog::{CatalogError, PromptCatalog};
/// Maximum queued pending entries. Oldest entries are dropped beyond this.
const CAPACITY: usize = 128;
/// One pending entry awaiting drain into the next LLM request.
///
/// The buffer keeps the raw input shape so the drain step can decide
@@ -43,7 +39,6 @@ const CAPACITY: usize = 128;
pub enum PendingNotify {
Notify {
message: String,
auto_run: bool,
extensions: Vec<SessionExtension>,
history_provenance: Option<LoggedSessionHistoryOrigin>,
},
@@ -76,6 +71,7 @@ impl PendingNotify {
#[derive(Clone, Default)]
pub struct NotifyBuffer {
inner: Arc<Mutex<VecDeque<PendingNotify>>>,
wake: Arc<tokio::sync::Notify>,
}
impl NotifyBuffer {
@@ -86,10 +82,9 @@ impl NotifyBuffer {
/// Push a notify entry onto the queue. If the queue is full, the
/// oldest entry is dropped and a `tracing::warn` is emitted — the
/// caller should never hit this in normal operation.
pub fn push_notify(&self, message: String, auto_run: bool) {
pub fn push_notify(&self, message: String) {
self.push_entry(PendingNotify::Notify {
message,
auto_run,
extensions: Vec::new(),
history_provenance: None,
});
@@ -98,53 +93,32 @@ impl NotifyBuffer {
pub fn push_durable_notify(
&self,
message: String,
auto_run: bool,
history_provenance: LoggedSessionHistoryOrigin,
extension: SessionExtension,
) {
self.push_entry(PendingNotify::Notify {
message,
auto_run,
extensions: vec![extension],
history_provenance: Some(history_provenance),
});
}
pub(crate) fn replace_durable_notification_extension(
&self,
extension: SessionExtension,
) -> bool {
let mut queue = self.inner.lock().expect("notify buffer poisoned");
let Some(extensions) = queue.iter_mut().rev().find_map(|pending| match pending {
PendingNotify::Notify {
auto_run: false,
extensions,
..
} if !extensions.is_empty() => Some(extensions),
_ => None,
}) else {
return false;
};
*extensions = vec![extension];
true
}
/// Push a typed worker-event entry onto the queue.
pub fn push_worker_event(&self, event: WorkerEvent) {
self.push_entry(PendingNotify::WorkerEvent { event });
}
fn push_entry(&self, entry: PendingNotify) {
let mut q = self.inner.lock().expect("notify buffer poisoned");
if q.len() >= CAPACITY {
let dropped = q.pop_front();
warn!(
capacity = CAPACITY,
dropped = ?dropped,
"notify buffer overflow; dropped oldest"
);
}
q.push_back(entry);
self.inner
.lock()
.expect("notify buffer poisoned")
.push_back(entry);
self.wake.notify_one();
}
/// Wait until a producer adds work outside the controller method channel.
pub async fn notified(&self) {
self.wake.notified().await;
}
/// Remove and return all pending entries in FIFO order.
@@ -160,23 +134,17 @@ impl NotifyBuffer {
for entry in entries.into_iter().rev() {
q.push_front(entry);
}
while q.len() > CAPACITY {
let dropped = q.pop_front();
warn!(
capacity = CAPACITY,
dropped = ?dropped,
"notify buffer overflow while restoring failed drain; dropped oldest"
);
}
drop(q);
self.wake.notify_one();
}
/// Whether an undrained `Method::Notify { auto_run: true }` remains.
pub fn has_auto_run_pending(&self) -> bool {
/// Whether an undrained notification remains.
pub fn has_notification_pending(&self) -> bool {
self.inner
.lock()
.expect("notify buffer poisoned")
.iter()
.any(|entry| matches!(entry, PendingNotify::Notify { auto_run: true, .. }))
.any(|entry| matches!(entry, PendingNotify::Notify { .. }))
}
/// Number of pending entries. Primarily for tests.
@@ -233,12 +201,12 @@ mod tests {
#[test]
fn push_then_drain_preserves_order() {
let buf = NotifyBuffer::new();
buf.push_notify("one".into(), false);
assert!(!buf.has_auto_run_pending());
buf.push_notify("two".into(), true);
assert!(buf.has_auto_run_pending());
buf.push_notify("one".into());
assert!(buf.has_notification_pending());
buf.push_notify("two".into());
assert!(buf.has_notification_pending());
let drained = buf.drain();
assert!(!buf.has_auto_run_pending());
assert!(!buf.has_notification_pending());
assert_eq!(drained.len(), 2);
match &drained[0] {
PendingNotify::Notify { message, .. } => assert_eq!(message, "one"),
@@ -248,15 +216,15 @@ mod tests {
}
#[test]
fn capacity_drops_oldest() {
fn notification_burst_is_not_truncated() {
let buf = NotifyBuffer::new();
for i in 0..(CAPACITY + 5) {
buf.push_notify(format!("msg{i}"), false);
for i in 0..300 {
buf.push_notify(format!("msg{i}"));
}
let drained = buf.drain();
assert_eq!(drained.len(), CAPACITY);
assert_eq!(drained.len(), 300);
match &drained[0] {
PendingNotify::Notify { message, .. } => assert_eq!(message, "msg5"),
PendingNotify::Notify { message, .. } => assert_eq!(message, "msg0"),
other => panic!("unexpected: {other:?}"),
}
}
@@ -265,7 +233,6 @@ mod tests {
fn build_system_item_for_notify_carries_wrapper_body() {
let entry = PendingNotify::Notify {
message: "hello".into(),
auto_run: false,
extensions: Vec::new(),
history_provenance: None,
};
+4 -7
View File
@@ -262,11 +262,10 @@ impl ParentNotificationTarget {
Self::Controller { sender, fallback }
}
pub(crate) fn notify(&self, child_session_id: String, message: String, auto_run: bool) {
pub(crate) fn notify(&self, child_session_id: String, message: String) {
let method = Method::NotifyTracked {
notification_request_id: protocol::new_submission_request_id(),
message,
auto_run,
source: protocol::AuthenticatedInputSource::SubWorker {
session_id: child_session_id,
},
@@ -588,7 +587,7 @@ impl Tool for SubWorkerSpawnTool {
let message = format!(
"SubWorker `{child_name}` turn ended with status {status:?}. Inspect its committed session with worker-observation tools before making completion decisions."
);
parent_notifications.notify(child_name.clone(), message, true);
parent_notifications.notify(child_name.clone(), message);
})),
Some(child_workdir_tool_broker.clone()),
)
@@ -1198,7 +1197,7 @@ enabled = false
drop(parent_method_tx);
assert!(parent_method_rx.recv().await.is_none());
target.notify("child-session".into(), "late completion".to_string(), true);
target.notify("child-session".into(), "late completion".to_string());
assert!(*captured.lock().unwrap());
}
@@ -1210,13 +1209,12 @@ enabled = false
*captured_for_target.lock().unwrap() = Some(method);
}));
target.notify("child-session".into(), "completed".into(), true);
target.notify("child-session".into(), "completed".into());
assert!(matches!(
captured.lock().unwrap().take(),
Some(Method::NotifyTracked {
message,
auto_run: true,
source: protocol::AuthenticatedInputSource::SubWorker { session_id },
..
}) if session_id == "child-session" && message == "completed"
@@ -1360,7 +1358,6 @@ enabled = false
completion,
Method::NotifyTracked {
message,
auto_run: true,
source: protocol::AuthenticatedInputSource::SubWorker { session_id },
..
} if session_id == "reviewer-child"
+140 -161
View File
@@ -111,8 +111,7 @@ pub(crate) struct PendingNotification {
source_namespace: String,
pub(crate) message: String,
payload_digest: String,
pub(crate) auto_run: bool,
accepted_at_ms: u64,
pub(crate) accepted_at_ms: u64,
activation_sequence: u64,
pub(crate) provenance: WorkerHistoryProvenance,
}
@@ -122,7 +121,6 @@ struct NotificationReceipt {
notification_request_id: String,
source_namespace: String,
payload_digest: String,
auto_run: bool,
}
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
@@ -133,19 +131,41 @@ pub(crate) struct PendingActivationState {
/// UserInput record commits the clearing checkpoint. Restore puts it back
/// at the FIFO head.
activating: Option<PendingSubmission>,
activating_notification: Option<PendingNotification>,
#[serde(
default,
rename = "activating_notification",
deserialize_with = "deserialize_activating_notifications"
)]
activating_notifications: VecDeque<PendingNotification>,
pending: VecDeque<PendingSubmission>,
pending_notifications: VecDeque<PendingNotification>,
receipts: VecDeque<SubmissionReceipt>,
notification_receipts: VecDeque<NotificationReceipt>,
}
fn deserialize_activating_notifications<'de, D>(
deserializer: D,
) -> Result<VecDeque<PendingNotification>, D::Error>
where
D: serde::Deserializer<'de>,
{
#[derive(serde::Deserialize)]
#[serde(untagged)]
enum OneOrMany {
One(PendingNotification),
Many(VecDeque<PendingNotification>),
}
Ok(match <Option<OneOrMany> as serde::Deserialize>::deserialize(deserializer)? {
Some(OneOrMany::One(notification)) => VecDeque::from([notification]),
Some(OneOrMany::Many(notifications)) => notifications,
None => VecDeque::new(),
})
}
impl PendingActivationState {
pub(crate) fn snapshot(&self) -> protocol::PendingSubmissionsSnapshot {
let pending_notification = self
.pending_notifications
.iter()
.find(|notification| notification.auto_run);
let pending_notification = self.pending_notifications.front();
let head_id = match (self.pending.front(), pending_notification) {
(Some(submission), Some(notification))
if notification.activation_sequence < submission.activation_sequence =>
@@ -202,14 +222,14 @@ impl PendingActivationState {
}
}
fn notification_payload_digest(message: &str, auto_run: bool) -> String {
pub const DEFAULT_NOTIFICATION_COALESCE_DELAY: Duration = Duration::from_secs(2);
fn notification_payload_digest(message: &str) -> String {
use sha2::Digest as _;
let mut hasher = sha2::Sha256::new();
hasher.update(if auto_run {
&b"auto\0"[..]
} else {
&b"deferred\0"[..]
});
// Keep the original real-time notification digest domain so receipts
// accepted before the delivery-mode field was removed still dedupe.
hasher.update(b"auto\0");
hasher.update(message.as_bytes());
hasher
.finalize()
@@ -1233,7 +1253,6 @@ where
#[derive(Debug, Clone)]
pub(crate) enum PendingActivation {
Submission(PendingSubmission),
Notification(PendingNotification),
}
#[derive(Debug, Clone)]
@@ -1530,14 +1549,12 @@ where
&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,
)
}
@@ -1547,7 +1564,6 @@ where
message: String,
source_namespace: String,
provenance: WorkerHistoryProvenance,
auto_run: bool,
) -> Result<bool, PendingSubmissionError> {
if notification_request_id.trim().is_empty() {
return Err(PendingSubmissionError::EmptyRequestId);
@@ -1555,7 +1571,7 @@ where
if notification_request_id.len() > MAX_ACTIVATION_REQUEST_ID_BYTES {
return Err(PendingSubmissionError::RequestIdLimit);
}
let payload_digest = notification_payload_digest(&message, auto_run);
let payload_digest = notification_payload_digest(&message);
let _append_guard = self
.writer
.state
@@ -1570,7 +1586,7 @@ where
receipt.notification_request_id == notification_request_id
&& receipt.source_namespace == source_namespace
}) {
if receipt.payload_digest != payload_digest || receipt.auto_run != auto_run {
if receipt.payload_digest != payload_digest {
return Err(PendingSubmissionError::IdempotencyConflict);
}
return Ok(false);
@@ -1607,7 +1623,6 @@ where
source_namespace: source_namespace.clone(),
message,
payload_digest: payload_digest.clone(),
auto_run,
accepted_at_ms: segment_log::now_millis(),
activation_sequence,
provenance,
@@ -1616,7 +1631,6 @@ where
notification_request_id,
source_namespace,
payload_digest,
auto_run,
});
state.revision = state.revision.saturating_add(1);
if let Err(error) = self.persist_locked(&state) {
@@ -1626,57 +1640,53 @@ where
Ok(true)
}
pub(crate) fn activating_passive_notification_id(&self) -> Option<String> {
self.state
pub(crate) fn oldest_pending_notification_accepted_at_ms(&self) -> Option<u64> {
let state = self
.state
.lock()
.expect("pending activation state poisoned")
.activating_notification
.as_ref()
.filter(|notification| !notification.auto_run)
.map(|notification| notification.notification_request_id.clone())
.expect("pending activation state poisoned");
state
.activating_notifications
.front()
.into_iter()
.chain(state.pending_notifications.front())
.map(|notification| notification.accepted_at_ms)
.min()
}
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(
pub(crate) fn prepare_notification_batch(
&self,
source_namespace: &str,
notification_request_id: &str,
) -> Option<PendingNotification> {
) -> Vec<(PendingNotification, SessionExtension)> {
let mut state = self
.state
.lock()
.expect("pending activation state poisoned");
if state.activating_notification.is_some() {
return None;
let notifications = state.pending_notifications.drain(..).collect::<Vec<_>>();
if notifications.is_empty() {
return Vec::new();
}
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
.activating_notifications
.extend(notifications.iter().cloned());
state.revision = state.revision.saturating_add(1);
Some(notification)
notifications
.into_iter()
.map(|notification| {
let index = state
.activating_notifications
.iter()
.position(|active| {
active.notification_request_id == notification.notification_request_id
&& active.source_namespace == notification.source_namespace
})
.expect("newly staged notification must be active");
let mut committed = state.clone();
committed.activating_notifications.drain(..=index);
committed.revision = committed.revision.saturating_add(1);
(notification, pending_activation_extension(&committed))
})
.collect()
}
pub(crate) fn prepare_next_activation(
@@ -1696,48 +1706,15 @@ where
if let Some((expected_revision, expected_head_id)) = fence {
Self::validate_fence(&state, expected_revision, Some(expected_head_id))?;
}
if state.activating.is_some()
|| state
.activating_notification
.as_ref()
.is_some_and(|notification| notification.auto_run)
{
if state.activating.is_some() {
return Ok(None);
}
let has_staged_passive_notification = state.activating_notification.is_some();
let submission_sequence = state.pending.front().map(|item| item.activation_sequence);
let notification_index = if has_staged_passive_notification {
None
} else {
state
.pending_notifications
.iter()
.position(|item| item.auto_run)
let Some(pending) = state.pending.pop_front() else {
return Ok(None);
};
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
.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)));
}
if submission_sequence.is_some() {
let pending = state
.pending
.pop_front()
.expect("submission sequence came from queue head");
state.activating = Some(pending.clone());
state.revision = state.revision.saturating_add(1);
return Ok(Some(PendingActivation::Submission(pending)));
}
Ok(None)
state.activating = Some(pending.clone());
state.revision = state.revision.saturating_add(1);
Ok(Some(PendingActivation::Submission(pending)))
}
pub(crate) fn abort_activation(&self, pending: PendingSubmission) {
@@ -1811,30 +1788,13 @@ where
}
}
pub(crate) fn notification_activation_extension(&self) -> SessionExtension {
let state = self
.state
.lock()
.expect("pending activation state poisoned");
let mut committed = state.clone();
committed.activating = None;
committed.activating_notification = None;
committed.revision = committed.revision.saturating_add(1);
pending_activation_extension(&committed)
}
pub(crate) fn finish_notification_activation(&self, notification_request_id: &str) {
pub(crate) fn finish_notification_batch(&self) {
let mut state = self
.state
.lock()
.expect("pending activation state poisoned");
if state
.activating_notification
.as_ref()
.map(|item| item.notification_request_id.as_str())
== Some(notification_request_id)
{
state.activating_notification = None;
if !state.activating_notifications.is_empty() {
state.activating_notifications.clear();
state.revision = state.revision.saturating_add(1);
}
}
@@ -2055,8 +2015,10 @@ impl WorkerSession {
state.pending.push_front(activating);
state.revision = state.revision.saturating_add(1);
}
if let Some(activating) = state.activating_notification.take() {
state.pending_notifications.push_front(activating);
if !state.activating_notifications.is_empty() {
let mut restored = std::mem::take(&mut state.activating_notifications);
restored.append(&mut state.pending_notifications);
state.pending_notifications = restored;
state.revision = state.revision.saturating_add(1);
}
*self
@@ -2213,6 +2175,9 @@ pub struct Worker<C: LlmClient, St: Store> {
/// injection into the next LLM request. Shared with the
/// WorkerInterceptor installed in `ensure_interceptor_installed`.
pending_notifies: NotifyBuffer,
/// Maximum time an idle Worker waits from the first accepted notification
/// so a burst can be delivered as one FIFO batch.
notification_coalesce_delay: Duration,
/// Submit-scoped stash for resolver-produced system messages
/// (currently `@<path>` file content). `Worker::run` fills this
/// before handing off to the worker; `WorkerInterceptor::on_prompt_submit`
@@ -2479,6 +2444,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
in_flight: None,
ai_activity_counter: Arc::new(AtomicUsize::new(0)),
pending_notifies: NotifyBuffer::new(),
notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY,
pending_attachments: Arc::new(Mutex::new(Vec::<SystemItem>::new())),
pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())),
scope_allocation: None,
@@ -3004,7 +2970,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
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.activating_notifications.is_empty()
|| !pending_state.receipts.is_empty()
|| !pending_state.notification_receipts.is_empty()
{
@@ -3381,8 +3347,21 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// `Item::system_message` just before the next LLM request, via
/// `WorkerInterceptor::pending_history_appends`. See [`NotifyBuffer`]
/// for overflow behaviour and the lane-of-record rationale.
pub fn push_notify(&self, message: String, auto_run: bool) {
self.pending_notifies.push_notify(message, auto_run);
/// Configure the fixed idle notification coalescing window.
///
/// The deadline is derived from the oldest durable notification receipt,
/// so restart restores the remaining window rather than starting it over.
pub fn with_notification_coalesce_delay(mut self, delay: Duration) -> Self {
self.notification_coalesce_delay = delay;
self
}
pub(crate) fn notification_coalesce_delay(&self) -> Duration {
self.notification_coalesce_delay
}
pub fn push_notify(&self, message: String) {
self.pending_notifies.push_notify(message);
}
/// Push an agent-visible typed `WorkerEvent` entry onto the pending buffer.
@@ -5964,6 +5943,7 @@ where
in_flight: None,
ai_activity_counter: Arc::new(AtomicUsize::new(0)),
pending_notifies: NotifyBuffer::new(),
notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY,
pending_attachments: Arc::new(Mutex::new(Vec::<SystemItem>::new())),
pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())),
scope_allocation: Some(scope_allocation),
@@ -6049,6 +6029,7 @@ where
in_flight: None,
ai_activity_counter: Arc::new(AtomicUsize::new(0)),
pending_notifies: NotifyBuffer::new(),
notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY,
pending_attachments: Arc::new(Mutex::new(Vec::<SystemItem>::new())),
pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())),
scope_allocation: None,
@@ -6169,6 +6150,7 @@ where
in_flight: None,
ai_activity_counter: Arc::new(AtomicUsize::new(0)),
pending_notifies: NotifyBuffer::new(),
notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY,
pending_attachments: Arc::new(Mutex::new(Vec::<SystemItem>::new())),
pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())),
scope_allocation: Some(scope_allocation),
@@ -6546,6 +6528,7 @@ where
in_flight: None,
ai_activity_counter: Arc::new(AtomicUsize::new(0)),
pending_notifies: NotifyBuffer::new(),
notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY,
pending_attachments: Arc::new(Mutex::new(Vec::<SystemItem>::new())),
pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())),
scope_allocation: Some(scope_allocation),
@@ -6626,7 +6609,6 @@ where
self.push_notify(
"Restored Worker state contained missing or unreachable delegated child Workers; their delegated write scopes were reclaimed before resume."
.to_string(),
false,
);
Ok(())
}
@@ -8981,7 +8963,7 @@ mod build_summary_prompt_tests {
append_user_turn(&worker, 20, "second message");
worker
.pending_submission_handle()
.accept_notification("notification-1".into(), "keep me".into(), true)
.accept_notification("notification-1".into(), "keep me".into())
.unwrap();
let (head_entries, targets) = worker.list_rewind_targets().unwrap();
@@ -10211,7 +10193,6 @@ mod build_summary_prompt_tests {
"notice".into(),
"account:account-a".into(),
account_a.clone(),
false,
)
.unwrap()
);
@@ -10222,7 +10203,6 @@ mod build_summary_prompt_tests {
"notice".into(),
"account:account-a".into(),
account_a.clone(),
false,
)
.unwrap()
);
@@ -10233,7 +10213,6 @@ mod build_summary_prompt_tests {
"notice".into(),
"account:account-b".into(),
account_b.clone(),
false,
)
.unwrap()
);
@@ -10243,45 +10222,41 @@ mod build_summary_prompt_tests {
}
#[test]
fn notification_and_submit_share_activation_order_and_notification_dedupes() {
fn notification_batch_preserves_fifo_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(), true)
.accept_notification("notification-1".into(), "notice".into())
.unwrap()
);
assert!(
!handle
.accept_notification("notification-1".into(), "notice".into(), true)
.accept_notification("notification-1".into(), "notice".into())
.unwrap()
);
assert!(matches!(
handle.accept_notification("notification-1".into(), "different".into(), true),
Err(PendingSubmissionError::IdempotencyConflict)
));
assert!(matches!(
handle.accept_notification("notification-1".into(), "notice".into(), false),
handle.accept_notification("notification-1".into(), "different".into()),
Err(PendingSubmissionError::IdempotencyConflict)
));
assert!(handle
.accept_notification("notification-2".into(), "second notice".into())
.unwrap());
handle
.accept("request-1".into(), vec![Segment::text("submit")], false)
.unwrap();
let first = handle.prepare_next_activation(None).unwrap().unwrap();
assert!(matches!(
first,
PendingActivation::Notification(PendingNotification { ref message, .. })
if message == "notice"
));
let committed = handle.notification_activation_extension();
let notifications = handle.prepare_notification_batch();
assert_eq!(notifications.len(), 2);
assert_eq!(notifications[0].0.message, "notice");
assert_eq!(notifications[1].0.message, "second notice");
let committed_state: PendingActivationState =
serde_json::from_value(committed.payload).unwrap();
serde_json::from_value(notifications[1].1.payload.clone()).unwrap();
assert!(committed_state.pending_notifications.is_empty());
assert!(committed_state.activating_notification.is_none());
handle.finish_notification_activation("notification-1");
let second = handle.prepare_next_activation(None).unwrap().unwrap();
assert!(matches!(second, PendingActivation::Submission(_)));
assert!(committed_state.activating_notifications.is_empty());
handle.finish_notification_batch();
let submission = handle.prepare_next_activation(None).unwrap().unwrap();
assert!(matches!(submission, PendingActivation::Submission(_)));
}
#[test]
@@ -10301,18 +10276,17 @@ mod build_summary_prompt_tests {
was_queued: false,
input: vec![Segment::text("first")],
}),
activating_notification: Some(PendingNotification {
activating_notifications: VecDeque::from([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,
payload_digest: notification_payload_digest("deferred notice"),
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(),
@@ -10329,13 +10303,19 @@ mod build_summary_prompt_tests {
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,
payload_digest: notification_payload_digest("deferred notice"),
}]),
};
let mut payload = serde_json::to_value(state).unwrap();
let legacy_activating = payload["activating_notification"]
.as_array_mut()
.expect("new state serializes an activating notification array")
.pop()
.expect("one activating notification");
payload["activating_notification"] = legacy_activating;
session.restore_pending_activations(&[(
SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.into(),
serde_json::to_value(state).unwrap(),
payload,
)]);
let state = session
.pending_activations
@@ -10345,9 +10325,8 @@ 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!(state.activating_notifications.is_empty());
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"
+24 -39
View File
@@ -1783,7 +1783,9 @@ async fn run_with_unresolved_segment_emits_alert_and_placeholder() {
async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
let client = MockClient::new(simple_text_events());
let client_for_assert = client.clone();
let worker = make_worker(client).await;
let worker = make_worker(client)
.await
.with_notification_coalesce_delay(std::time::Duration::from_millis(20));
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
@@ -1791,7 +1793,6 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
.send(Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message: "turn finished".into(),
auto_run: true,
})
.await
.unwrap();
@@ -1883,10 +1884,12 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
}
#[tokio::test]
async fn notify_while_idle_with_auto_run_false_waits_for_explicit_run() {
async fn repeated_notify_while_idle_coalesces_and_auto_starts_one_turn() {
let client = MockClient::new(simple_text_events());
let client_for_assert = client.clone();
let worker = make_worker(client).await;
let worker = make_worker(client)
.await
.with_notification_coalesce_delay(std::time::Duration::from_millis(50));
let handle = spawn_controller(worker).await;
let notification_request_id = protocol::new_submission_request_id();
@@ -1895,26 +1898,18 @@ async fn notify_while_idle_with_auto_run_false_waits_for_explicit_run() {
.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.catalog_status(), WorkerStatus::Idle);
assert!(
client_for_assert.captured_requests().is_empty(),
"weak Notify must not stage RunForNotification while idle"
);
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"continue",
))
.send(Method::Notify {
notification_request_id: "notify-coalesced-2".into(),
message: "second coalesced notification".into(),
})
.await
.unwrap();
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
if !client_for_assert.captured_requests().is_empty() {
@@ -1922,30 +1917,22 @@ async fn notify_while_idle_with_auto_run_false_waits_for_explicit_run() {
}
assert!(
tokio::time::Instant::now() < deadline,
"explicit run did not reach the mock LLM"
"coalescing deadline did not start a notification turn"
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
wait_for_status(&handle, WorkerStatus::Idle).await;
let requests = client_for_assert.captured_requests();
assert_eq!(
requests.len(),
1,
"explicit run should drain the queued notification"
);
let notify_in_request = requests[0].items.iter().any(|i| {
i.as_text()
.is_some_and(|t| t.contains("[Notification]") && t.contains("progress snapshot"))
});
assert!(
notify_in_request,
"queued weak notification must be history-backed on the next explicit run; got items: {:?}",
requests[0]
.items
.iter()
.filter_map(|i| i.as_text())
.collect::<Vec<_>>()
);
assert_eq!(requests.len(), 1, "coalesced notifications need one turn");
let notifications = requests[0]
.items
.iter()
.filter_map(|item| item.as_text())
.filter(|text| text.contains("[Notification]"))
.collect::<Vec<_>>();
assert_eq!(notifications.len(), 2, "duplicate receipt must not duplicate content");
assert!(notifications[0].contains("progress snapshot"));
assert!(notifications[1].contains("second coalesced notification"));
}
#[tokio::test]
@@ -2105,7 +2092,6 @@ async fn notify_while_running_does_not_emit_already_running_error() {
.send(Method::Notify {
notification_request_id: protocol::new_submission_request_id(),
message: "ping".into(),
auto_run: true,
})
.await
.unwrap();
@@ -2137,7 +2123,7 @@ async fn notify_while_running_does_not_emit_already_running_error() {
}
#[tokio::test]
async fn weak_notify_while_running_is_deduped_and_survives_until_next_submit() {
async fn notify_while_running_is_deduped_and_survives_until_next_model_boundary() {
let client = MockClient::sequential(vec![
MockResponse::Hang(Vec::new()),
MockResponse::Complete(simple_text_events()),
@@ -2160,7 +2146,6 @@ async fn weak_notify_while_running_is_deduped_and_survives_until_next_submit() {
.send(Method::Notify {
notification_request_id: notification_request_id.clone(),
message: "durable weak notice".into(),
auto_run: false,
})
.await
.unwrap();
-2
View File
@@ -16008,11 +16008,9 @@ pub(crate) fn authorize_browser_worker_method(
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 { .. } => {
+1 -1
View File
@@ -286,7 +286,7 @@ export type SubscriptionFramePayload = { "frame": "request", "message": Subscrip
export type SubscriptionFrame = { protocol_version: number, } & ({ "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent } | { "frame": "worker_protocol", "message": SubscriptionWorkerProtocolMethod });
export type Method = { "method": "submit", "params": { submission_request_id: string, input: Array<Segment>, } } | { "method": "notify", "params": { notification_request_id: string, message: string, auto_run?: boolean, } } | { "method": "worker_event", "params": WorkerEvent } | { "method": "list_pending_submissions" } | { "method": "cancel_pending_submission", "params": { submission_id: string, expected_revision: number, } } | { "method": "clear_pending_submissions", "params": { expected_revision: number, } } | { "method": "continue_pending", "params": { expected_revision: number, expected_head_id: string, } } | { "method": "resume", "params": { command: WorkerCommandEnvelope, } } | { "method": "cancel", "params": { command: WorkerCommandEnvelope, } } | { "method": "pause", "params": { command: WorkerCommandEnvelope, } } | { "method": "compact", "params": { command: WorkerCommandEnvelope, } } | { "method": "list_rewind_targets" } | { "method": "rewind_to", "params": { target: RewindTargetId, expected_head_entries: number, } } | { "method": "shutdown", "params": { command: WorkerCommandEnvelope, } } | { "method": "list_completions", "params": { kind: CompletionKind, prefix: string, } } | { "method": "list_workers" } | { "method": "restore_worker", "params": { name: string, } } | { "method": "register_peer", "params": { name: string, } };
export type Method = { "method": "submit", "params": { submission_request_id: string, input: Array<Segment>, } } | { "method": "notify", "params": { notification_request_id: string, message: string, } } | { "method": "worker_event", "params": WorkerEvent } | { "method": "list_pending_submissions" } | { "method": "cancel_pending_submission", "params": { submission_id: string, expected_revision: number, } } | { "method": "clear_pending_submissions", "params": { expected_revision: number, } } | { "method": "continue_pending", "params": { expected_revision: number, expected_head_id: string, } } | { "method": "resume", "params": { command: WorkerCommandEnvelope, } } | { "method": "cancel", "params": { command: WorkerCommandEnvelope, } } | { "method": "pause", "params": { command: WorkerCommandEnvelope, } } | { "method": "compact", "params": { command: WorkerCommandEnvelope, } } | { "method": "list_rewind_targets" } | { "method": "rewind_to", "params": { target: RewindTargetId, expected_head_entries: number, } } | { "method": "shutdown", "params": { command: WorkerCommandEnvelope, } } | { "method": "list_completions", "params": { kind: CompletionKind, prefix: string, } } | { "method": "list_workers" } | { "method": "restore_worker", "params": { name: string, } } | { "method": "register_peer", "params": { name: string, } };
export type Event = { "event": "submission_accepted", "data": { submission_request_id: string, submission_id: string, disposition: SubmissionDisposition, } } | { "event": "submission_rejected", "data": { submission_request_id: string, message: string, } } | { "event": "pending_submissions_changed", "data": { pending: PendingSubmissionsSnapshot, } } | { "event": "user_message", "data": { segments: Array<Segment>, } } | { "event": "system_item", "data": { item: unknown, } } | { "event": "invoke_start", "data": { kind: InvokeKind, } } | { "event": "turn_start", "data": { turn: number, } } | { "event": "turn_end", "data": { turn: number, result: TurnResult, } } | { "event": "llm_call_start", "data": { llm_call: number, } } | { "event": "llm_call_end", "data": { llm_call: number, } } | { "event": "llm_retry", "data": { llm_call: number,
/**
@@ -647,7 +647,6 @@
params: {
notification_request_id: crypto.randomUUID(),
message: request.content,
auto_run: true,
},
};
case "compact": {