fix: preserve console alert order and improve alert colors

This commit is contained in:
2026-10-16 05:28:17 +09:00
parent b67aa2de7d
commit a2cec5afd1
25 changed files with 2098 additions and 90 deletions
+4
View File
@@ -1532,6 +1532,8 @@ mod tests {
let mut writer = JsonLineWriter::new(writer_half);
writer
.write(&Event::Alert(Alert {
id: None,
position: None,
level: AlertLevel::Warn,
source: AlertSource::Worker,
message: "connect-time alert".into(),
@@ -1763,6 +1765,8 @@ mod tests {
let mut writer = JsonLineWriter::new(stream);
writer
.write(&Event::Alert(Alert {
id: None,
position: None,
level: AlertLevel::Warn,
source: AlertSource::Worker,
message: "replayed alert".into(),
+142 -14
View File
@@ -1,11 +1,11 @@
use std::{
collections::HashSet,
sync::{Arc, Mutex, MutexGuard},
sync::{Arc, Mutex, MutexGuard, Weak},
};
use protocol::{
CommandEvent, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice, Event,
InFlightBlock, InFlightSnapshot, InFlightToolCallState,
AlertPosition, CommandEvent, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice,
Event, InFlightBlock, InFlightSnapshot, InFlightToolCallState,
};
use session_store::{LoggedContentPart, LoggedItem};
use tokio::sync::broadcast;
@@ -38,6 +38,7 @@ enum TrackedBlock {
block_id: InFlightBlockId,
text: String,
finished: bool,
alert_positions: Vec<Weak<Mutex<AlertPosition>>>,
},
Thinking {
block_id: InFlightBlockId,
@@ -78,6 +79,7 @@ impl InFlightEvents {
block_id,
text: String::new(),
finished: false,
alert_positions: Vec::new(),
});
block_id
}
@@ -213,10 +215,11 @@ impl InFlightEvents {
pub(crate) fn clear_for_committed_item_then<R>(
&self,
item: &LoggedItem,
entry_id: Option<&str>,
f: impl FnOnce() -> R,
) -> R {
let mut inner = self.lock();
inner.clear_for_committed_item(item);
inner.clear_for_committed_item(item, entry_id);
f()
}
@@ -378,7 +381,7 @@ impl InFlightInner {
}
}
fn clear_for_committed_item(&mut self, item: &LoggedItem) {
fn clear_for_committed_item(&mut self, item: &LoggedItem, entry_id: Option<&str>) {
match item {
LoggedItem::Message { role, content }
if matches!(role, session_store::LoggedRole::Assistant) =>
@@ -391,7 +394,7 @@ impl InFlightInner {
})
.collect::<String>();
if !text.is_empty() {
self.remove_first_text_matching(&text);
self.remove_first_text_matching(&text, entry_id);
}
}
LoggedItem::Reasoning {
@@ -420,13 +423,55 @@ impl InFlightInner {
}
}
pub(crate) fn capture_alert_position(&mut self) -> Option<Arc<Mutex<AlertPosition>>> {
// Prune evicted alerts even when their streamed block has not committed yet.
for block in &mut self.blocks {
if let TrackedBlock::Text {
alert_positions, ..
} = block
{
alert_positions.retain(|position| position.strong_count() > 0);
}
}
for block in self.blocks.iter_mut().rev() {
match block {
TrackedBlock::Text {
block_id,
text,
alert_positions,
..
} if !text.is_empty() => {
let position = Arc::new(Mutex::new(AlertPosition::AfterInFlight {
block_id: block_id.0,
}));
alert_positions.push(Arc::downgrade(&position));
return Some(position);
}
TrackedBlock::ToolCall { id, .. } => {
return Some(Arc::new(Mutex::new(AlertPosition::AfterToolCall {
call_id: id.clone(),
})));
}
// Hidden reasoning and empty text are not conversation rows.
_ => {}
}
}
None
}
fn snapshot(&self) -> InFlightSnapshot {
let (block_ids, blocks) = self
.blocks
.iter()
.filter_map(|block| {
block
.to_snapshot_block()
.map(|snapshot| (block.block_id().0, snapshot))
})
.unzip();
InFlightSnapshot {
blocks: self
.blocks
.iter()
.filter_map(TrackedBlock::to_snapshot_block)
.collect(),
block_ids,
blocks,
commands: self.commands.clone(),
compaction: self.compaction.clone(),
}
@@ -444,11 +489,25 @@ impl InFlightInner {
}
}
fn remove_first_text_matching(&mut self, committed: &str) -> bool {
fn remove_first_text_matching(&mut self, committed: &str, entry_id: Option<&str>) -> bool {
if let Some(index) = self.blocks.iter().position(|block| match block {
TrackedBlock::Text { text, .. } => text == committed,
_ => false,
}) {
if let (
TrackedBlock::Text {
alert_positions, ..
},
Some(entry_id),
) = (&self.blocks[index], entry_id)
{
for position in alert_positions.iter().filter_map(Weak::upgrade) {
*position.lock().expect("alert position mutex poisoned") =
AlertPosition::AfterEntry {
entry_id: entry_id.to_owned(),
};
}
}
self.blocks.remove(index);
true
} else {
@@ -559,6 +618,72 @@ mod tests {
}
}
#[test]
fn alert_watchers_are_pruned_and_removed_with_their_text_block() {
let (tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(tx);
let text = in_flight.start_text_block();
in_flight.text_delta(text, "same".into());
let first = in_flight.snapshot_guard().capture_alert_position().unwrap();
drop(first);
let second = in_flight.snapshot_guard().capture_alert_position().unwrap();
{
let guard = in_flight.snapshot_guard();
let TrackedBlock::Text {
alert_positions, ..
} = &guard.blocks[0]
else {
panic!("expected text")
};
assert_eq!(alert_positions.len(), 1);
}
let later = in_flight.start_text_block();
in_flight.text_delta(later, "same".into());
let third = in_flight.snapshot_guard().capture_alert_position().unwrap();
let item = LoggedItem::Message {
role: session_store::LoggedRole::Assistant,
content: vec![LoggedContentPart::Text {
text: "same".into(),
}],
};
in_flight.clear_for_committed_item_then(&item, Some("first-entry"), || ());
assert_eq!(
*second.lock().unwrap(),
AlertPosition::AfterEntry {
entry_id: "first-entry".into()
}
);
assert_eq!(
*third.lock().unwrap(),
AlertPosition::AfterInFlight { block_id: later.0 }
);
in_flight.clear();
assert!(in_flight.snapshot_guard().blocks.is_empty());
assert_eq!(
*third.lock().unwrap(),
AlertPosition::AfterInFlight { block_id: later.0 }
);
}
#[test]
fn snapshot_ids_are_parallel_to_emitted_blocks_not_filtered_placeholders() {
let (tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(tx);
in_flight.start_text_block();
let hidden = in_flight.thinking_start();
in_flight.thinking_done(hidden, String::new());
let visible = in_flight.start_text_block();
in_flight.text_delta(visible, "shown".into());
let tool = in_flight.tool_call_start("call-1".into(), "Read".into());
let snapshot = snapshot_from_guard(&in_flight.snapshot_guard());
assert_eq!(snapshot.blocks.len(), 2);
assert_eq!(snapshot.block_ids, vec![visible.0, tool.0]);
assert_eq!(
snapshot.block_ids,
snapshot_from_guard(&in_flight.snapshot_guard()).block_ids
);
}
#[test]
fn snapshot_boundary_does_not_duplicate_or_gap_delta_sent_after_subscribe() {
let (working_event_tx, _) = broadcast::channel(16);
@@ -621,7 +746,7 @@ mod tests {
// This mirrors Worker::append_entry ordering: clear in-flight first,
// then publish the finalized AssistantItem. Identity-aware clients
// reconcile this durable commit with the preceding streamed block.
in_flight_for_commit.clear_for_committed_item_then(&assistant_item, || {
in_flight_for_commit.clear_for_committed_item_then(&assistant_item, None, || {
sink_for_commit.publish(assistant_entry);
});
committed_tx.send(()).unwrap();
@@ -670,7 +795,7 @@ mod tests {
entry: crate::session_history::test_logged_history_entry(assistant_item.clone()),
};
in_flight.clear_for_committed_item_then(&assistant_item, || {
in_flight.clear_for_committed_item_then(&assistant_item, None, || {
sink.publish(assistant_entry);
});
@@ -698,6 +823,7 @@ mod tests {
text: "done".into(),
}],
},
None,
|| (),
);
@@ -723,6 +849,7 @@ mod tests {
encrypted_content: Some("opaque".into()),
signature: None,
},
None,
|| (),
);
@@ -746,6 +873,7 @@ mod tests {
encrypted_content: Some("opaque".into()),
signature: None,
},
None,
|| (),
);
+380 -8
View File
@@ -6,7 +6,7 @@
//!
//! Each alert is broadcast on the shared `Event` channel and
//! also appended to an in-memory buffer so that clients connecting
//! after the fact still see everything emitted during the session.
//! after the fact can replay the bounded recent prefix from the same live Worker.
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
@@ -14,7 +14,9 @@ use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::broadcast;
use protocol::{Alert, AlertLevel, AlertSource, Event};
use protocol::{Alert, AlertLevel, AlertPosition, AlertSource, Event};
use crate::{in_flight::InFlightEvents, segment_log_sink::SegmentLogSink};
/// Upper bound on buffered alerts. When exceeded, the oldest
/// entries are discarded so a long-running session cannot leak
@@ -27,9 +29,53 @@ pub struct Alerter {
inner: Arc<Inner>,
}
/// A queued identified Alert is no longer retained by this live Worker.
/// Transports must disconnect/resnapshot instead of forwarding its stale anchor.
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("alert is no longer retained; reconnect for a fresh snapshot")]
pub struct AlertReplayRequired;
struct Inner {
working_event_tx: broadcast::Sender<Event>,
buffer: Mutex<VecDeque<Alert>>,
buffer: Mutex<VecDeque<BufferedAlert>>,
position_source: Mutex<Option<PositionSource>>,
}
#[derive(Clone)]
struct PositionSource {
sink: SegmentLogSink,
in_flight: Option<InFlightEvents>,
}
#[derive(Clone)]
struct BufferedAlert {
alert: Alert,
position: Option<Arc<Mutex<AlertPosition>>>,
}
impl BufferedAlert {
fn snapshot(&self) -> Alert {
let mut alert = self.alert.clone();
alert.position = self.position.as_ref().map(|position| {
position
.lock()
.expect("alert position mutex poisoned")
.clone()
});
alert
}
}
impl PositionSource {
fn capture(&self) -> Arc<Mutex<AlertPosition>> {
// Same order as commit/snapshot: in-flight -> sink. Never hold the
// alerter buffer lock while acquiring either of these locks.
let mut in_flight = self.in_flight.as_ref().map(InFlightEvents::snapshot_guard);
in_flight
.as_mut()
.and_then(|guard| guard.capture_alert_position())
.unwrap_or_else(|| Arc::new(Mutex::new(self.sink.alert_position())))
}
}
impl Alerter {
@@ -38,10 +84,23 @@ impl Alerter {
inner: Arc::new(Inner {
working_event_tx,
buffer: Mutex::new(VecDeque::with_capacity(MAX_BUFFERED_ALERTS)),
position_source: Mutex::new(None),
}),
}
}
pub(crate) fn attach_position_source(
&self,
sink: SegmentLogSink,
in_flight: Option<InFlightEvents>,
) {
*self
.inner
.position_source
.lock()
.expect("alert source mutex poisoned") = Some(PositionSource { sink, in_flight });
}
/// Record and broadcast an alert.
///
/// The broadcast may have no subscribers (e.g. during Worker
@@ -55,35 +114,88 @@ impl Alerter {
/// exactly once: older ones from the snapshot, newer ones from
/// the freshly-subscribed receiver.
pub fn alert(&self, level: AlertLevel, source: AlertSource, message: String) {
self.alert_at(level, source, message, now_ms());
}
fn alert_at(&self, level: AlertLevel, source: AlertSource, message: String, timestamp_ms: i64) {
let source_handle = self
.inner
.position_source
.lock()
.expect("alert source mutex poisoned")
.clone();
let position = source_handle.map(|source| source.capture());
let alert = Alert {
id: Some(uuid::Uuid::now_v7().to_string()),
position: None,
level,
source,
message,
timestamp_ms: now_ms(),
timestamp_ms,
};
if let Ok(mut buf) = self.inner.buffer.lock() {
if buf.len() >= MAX_BUFFERED_ALERTS {
buf.pop_front();
}
buf.push_back(alert.clone());
let _ = self.inner.working_event_tx.send(Event::Alert(alert));
let buffered = BufferedAlert { alert, position };
let event = Event::Alert(buffered.snapshot());
buf.push_back(buffered);
let _ = self.inner.working_event_tx.send(event);
}
}
/// Refresh a queued Alert immediately before transport forwarding.
///
/// The independent log lane may have already delivered the text commit, so
/// the broadcast's copied in-flight anchor is not authoritative anymore.
/// Missing identified alerts require reconnect/resnapshot; legacy alerts
/// without IDs and non-Alert events pass through unchanged.
pub fn freshen_event(&self, event: Event) -> Result<Event, AlertReplayRequired> {
let Event::Alert(ref alert) = event else {
return Ok(event);
};
let Some(id) = alert.id.as_ref() else {
return Ok(event);
};
let buffer = self
.inner
.buffer
.lock()
.expect("alerter buffer mutex poisoned");
buffer
.iter()
.find(|buffered| buffered.alert.id.as_ref() == Some(id))
.map(|buffered| Event::Alert(buffered.snapshot()))
.ok_or(AlertReplayRequired)
}
/// Subscribe and atomically snapshot the current buffer.
///
/// The returned snapshot contains alerts emitted before
/// this call; the receiver will deliver alerts emitted
/// after. An alert cannot appear in both.
pub fn subscribe_with_snapshot(&self) -> (Vec<Alert>, broadcast::Receiver<Event>) {
let (snapshot, rx) = self.subscribe_with_deferred_snapshot();
(snapshot(), rx)
}
// Resolve positions after the session snapshot has been captured. Otherwise
// a text commit between the two snapshots can leave only its old block ID
// in replay, although that block is already absent from the session snapshot.
pub(crate) fn subscribe_with_deferred_snapshot(
&self,
) -> (impl FnOnce() -> Vec<Alert>, broadcast::Receiver<Event>) {
let buf = self
.inner
.buffer
.lock()
.expect("alerter buffer mutex poisoned");
let rx = self.inner.working_event_tx.subscribe();
let snapshot: Vec<Alert> = buf.iter().cloned().collect();
(snapshot, rx)
let snapshot: Vec<BufferedAlert> = buf.iter().cloned().collect();
(
move || snapshot.iter().map(BufferedAlert::snapshot).collect(),
rx,
)
}
}
@@ -98,6 +210,266 @@ fn now_ms() -> i64 {
mod tests {
use super::*;
fn positioned_alerter() -> (Alerter, SegmentLogSink, InFlightEvents) {
let (tx, _) = broadcast::channel(1024);
let in_flight = InFlightEvents::new(tx.clone());
let sink = SegmentLogSink::new();
let alerter = Alerter::new(tx);
alerter.attach_position_source(sink.clone(), Some(in_flight.clone()));
(alerter, sink, in_flight)
}
#[test]
fn queued_alert_uses_committed_anchor_when_log_lane_was_delivered_first() {
let (alerter, sink, in_flight) = positioned_alerter();
let text = in_flight.start_text_block();
in_flight.text_delta(text, "finished text".into());
in_flight.text_done(text, "finished text".into());
let (_, mut events) = alerter.subscribe_with_snapshot();
let (_, mut log_entries) = sink.subscribe_with_snapshot();
alerter.alert(AlertLevel::Warn, AlertSource::Worker, "queued".into());
let queued = events.try_recv().unwrap();
let Event::Alert(ref original) = queued else {
panic!("expected alert")
};
assert!(matches!(
original.position,
Some(AlertPosition::AfterInFlight { .. })
));
let original_id = original.id.clone();
let item: session_store::LoggedItem = agen::Item::assistant_message("finished text").into();
let entry = crate::session_history::test_logged_history_entry(item.clone());
let entry_id = entry.metadata.entry_id.0.clone();
in_flight.clear_for_committed_item_then(&item, Some(&entry_id), || {
sink.publish(session_store::LogEntry::AnnotatedAssistantItem { ts: 42, entry });
});
let committed =
crate::ipc::protocol_session::live_log_entry_event(log_entries.try_recv().unwrap())
.unwrap();
assert!(
matches!(committed, Event::SessionEntryCommitted { entry } if entry.entry_id == entry_id)
);
let Event::Alert(forwarded) = alerter.freshen_event(queued).unwrap() else {
panic!("expected alert")
};
assert_eq!(forwarded.id, original_id);
assert_eq!(
forwarded.position,
Some(AlertPosition::AfterEntry { entry_id })
);
let replay = alerter.subscribe_with_snapshot().0;
assert_eq!(
serde_json::to_value(forwarded).unwrap(),
serde_json::to_value(&replay[0]).unwrap()
);
}
#[test]
fn evicted_or_unknown_alert_ids_require_replay_but_legacy_events_pass_through() {
let (alerter, _, _) = positioned_alerter();
let (_, mut events) = alerter.subscribe_with_snapshot();
alerter.alert(AlertLevel::Warn, AlertSource::Worker, "old".into());
let queued = events.try_recv().unwrap();
for _ in 0..MAX_BUFFERED_ALERTS {
alerter.alert(AlertLevel::Warn, AlertSource::Worker, "later".into());
}
assert!(matches!(
alerter.freshen_event(queued),
Err(AlertReplayRequired)
));
let unknown = Event::Alert(Alert {
id: Some("not-retained".into()),
position: Some(AlertPosition::AfterInFlight { block_id: 1 }),
level: AlertLevel::Warn,
source: AlertSource::Worker,
message: "unknown".into(),
timestamp_ms: 42,
});
assert!(matches!(
alerter.freshen_event(unknown),
Err(AlertReplayRequired)
));
let legacy = Event::Alert(Alert {
id: None,
position: None,
level: AlertLevel::Warn,
source: AlertSource::Worker,
message: "legacy".into(),
timestamp_ms: 42,
});
let expected = serde_json::to_value(&legacy).unwrap();
assert_eq!(
serde_json::to_value(alerter.freshen_event(legacy).unwrap()).unwrap(),
expected
);
assert!(matches!(
alerter.freshen_event(Event::Shutdown),
Ok(Event::Shutdown)
));
}
#[test]
fn rowless_run_transitions_preserve_visible_alert_anchor_live_and_on_replay() {
use session_store::LogEntry;
let (alerter, sink, _) = positioned_alerter();
let entry = crate::session_history::test_logged_history_entry(
agen::Item::assistant_message("visible"),
);
let expected = Some(AlertPosition::AfterEntry {
entry_id: entry.metadata.entry_id.0.clone(),
});
sink.publish(LogEntry::AnnotatedAssistantItem { ts: 1, entry });
let (_, mut events) = alerter.subscribe_with_snapshot();
let transitions = [
LogEntry::RunCancelled {
ts: 2,
entry_id: None,
},
LogEntry::RunYielded {
ts: 3,
entry_id: None,
reason: protocol::RunYieldReason::Compaction,
active_run_turn_count: 1,
},
LogEntry::RunResumed {
ts: 4,
entry_id: None,
source: protocol::RunResumeSource::Compaction,
active_run_turn_count: 1,
},
];
for mut transition in transitions {
transition.ensure_session_entry_id();
sink.publish(transition);
alerter.alert(
AlertLevel::Warn,
AlertSource::Worker,
"after transition".into(),
);
let Event::Alert(live) = alerter.freshen_event(events.try_recv().unwrap()).unwrap()
else {
panic!("expected alert")
};
assert_eq!(live.position, expected);
}
let replay = alerter.subscribe_with_snapshot().0;
assert_eq!(replay.len(), 3);
assert!(replay.iter().all(|alert| alert.position == expected));
let (entries, _) = sink.subscribe_with_snapshot();
// Rebuilding the sink cache from retained entries must use the same rule.
sink.replace_silent(entries);
alerter.alert(
AlertLevel::Warn,
AlertSource::Worker,
"after cache rebuild".into(),
);
assert_eq!(
alerter.subscribe_with_snapshot().0.last().unwrap().position,
expected
);
}
#[test]
fn identical_alerts_in_same_millisecond_keep_distinct_ids_across_replay() {
let (alerter, _, _) = positioned_alerter();
let (_, mut rx) = alerter.subscribe_with_snapshot();
for _ in 0..2 {
alerter.alert_at(AlertLevel::Warn, AlertSource::Worker, "same".into(), 42);
}
let (snapshot, _) = alerter.subscribe_with_snapshot();
assert_ne!(snapshot[0].id, snapshot[1].id);
for expected in &snapshot {
let Event::Alert(live) = rx.try_recv().unwrap() else {
panic!("expected alert")
};
assert!(live.id.is_some());
assert_eq!(live.id, expected.id);
assert_eq!(live.timestamp_ms, 42);
assert_eq!(live.position, Some(AlertPosition::Start));
assert_eq!(live.position, expected.position);
}
let (replay, _) = alerter.subscribe_with_snapshot();
assert_eq!(
serde_json::to_value(snapshot).unwrap(),
serde_json::to_value(replay).unwrap()
);
}
#[test]
fn tool_call_anchor_survives_commit_and_new_text() {
let (alerter, sink, in_flight) = positioned_alerter();
in_flight.tool_call_start("call-1".into(), "Read".into());
alerter.alert(AlertLevel::Warn, AlertSource::Worker, "tool alert".into());
let item: session_store::LoggedItem = agen::Item::tool_call("call-1", "Read", "{}").into();
let entry = session_store::LogEntry::AnnotatedAssistantItem {
ts: 1,
entry: crate::session_history::test_logged_history_entry(item.clone()),
};
in_flight.clear_for_committed_item_then(&item, Some("entry-tool"), || sink.publish(entry));
alerter.alert(AlertLevel::Warn, AlertSource::Worker, "after commit".into());
let text = in_flight.start_text_block();
in_flight.text_delta(text, "next".into());
let (snapshot, _) = alerter.subscribe_with_snapshot();
for alert in snapshot {
assert_eq!(
alert.position,
Some(AlertPosition::AfterToolCall {
call_id: "call-1".into()
})
);
}
}
#[test]
fn eviction_drops_position_ownership_and_clear_does_not_retarget_old_alerts() {
let (alerter, _, in_flight) = positioned_alerter();
let text = in_flight.start_text_block();
in_flight.text_delta(text, "discarded".into());
alerter.alert(AlertLevel::Warn, AlertSource::Worker, "old".into());
let weak = {
let buffer = alerter.inner.buffer.lock().unwrap();
Arc::downgrade(buffer[0].position.as_ref().unwrap())
};
for _ in 0..MAX_BUFFERED_ALERTS {
alerter.alert(AlertLevel::Warn, AlertSource::Worker, "new".into());
}
assert!(weak.upgrade().is_none());
let (before, _) = alerter.subscribe_with_snapshot();
in_flight.clear();
let next = in_flight.start_text_block();
in_flight.text_delta(next, "discarded".into());
let item = session_store::LoggedItem::Message {
role: session_store::LoggedRole::Assistant,
content: vec![session_store::LoggedContentPart::Text {
text: "discarded".into(),
}],
};
in_flight.clear_for_committed_item_then(&item, Some("unrelated-entry"), || ());
let (after, _) = alerter.subscribe_with_snapshot();
assert_eq!(before.len(), MAX_BUFFERED_ALERTS);
assert_eq!(
serde_json::to_value(before).unwrap(),
serde_json::to_value(after).unwrap()
);
}
#[test]
fn snapshot_does_not_acquire_in_flight_lock_under_buffer_lock() {
let (alerter, _, in_flight) = positioned_alerter();
alerter.alert(AlertLevel::Warn, AlertSource::Worker, "before".into());
let guard = in_flight.snapshot_guard();
let cloned = alerter.clone();
let (tx, rx) = std::sync::mpsc::channel();
let thread = std::thread::spawn(move || {
let (snapshot, _) = cloned.subscribe_with_snapshot();
tx.send(snapshot).unwrap();
});
let result = rx.recv_timeout(std::time::Duration::from_secs(5));
drop(guard);
thread.join().unwrap();
assert_eq!(result.unwrap().len(), 1);
}
#[test]
fn alert_broadcasts_to_existing_subscriber() {
let (tx, _keep) = broadcast::channel::<Event>(8);
+2 -2
View File
@@ -30,12 +30,12 @@ fn subscribe_before_snapshot<Subscription, Snapshot>(
pub fn subscribe_worker_protocol_session(handle: &WorkerHandle) -> WorkerProtocolSessionStreams {
let ((snapshot_event, log_entries), (alert_snapshot, events)) = subscribe_before_snapshot(
|| handle.alerter.subscribe_with_snapshot(),
|| handle.alerter.subscribe_with_deferred_snapshot(),
|| handle.snapshot_event_with_entry_subscription(),
);
WorkerProtocolSessionStreams {
snapshot_event,
alert_snapshot,
alert_snapshot: alert_snapshot(),
log_entries,
events,
}
+5
View File
@@ -154,6 +154,11 @@ async fn handle_connection(stream: tokio::net::UnixStream, handle: WorkerHandle)
event = streams.events.recv() => {
match event {
Ok(event) => {
let Ok(event) = handle.alerter.freshen_event(event) else {
// The queued Alert was evicted; reconnect must rebuild
// the bounded prefix rather than forward a stale anchor.
break;
};
if writer.write_event(&event).await.is_err() {
break;
}
+129 -10
View File
@@ -24,6 +24,7 @@
use std::sync::{Arc, Mutex as StdMutex};
use protocol::{AlertPosition, SessionSnapshotEntryData};
use session_store::LogEntry;
use tokio::sync::broadcast;
@@ -46,7 +47,7 @@ pub struct SegmentLogSink {
struct SinkInner {
/// Full session log mirror in commit order. Reset on session swap
/// (compaction / fork) via [`SegmentLogSink::reset_with_initial`].
mirror: StdMutex<Vec<LogEntry>>,
mirror: StdMutex<Mirror>,
/// Broadcast channel for live entry updates. The same `Sender`
/// survives session swaps so existing subscribers keep their
/// receiver — they observe the swap as a freshly broadcast
@@ -54,6 +55,46 @@ struct SinkInner {
session_entry_tx: broadcast::Sender<LogEntry>,
}
struct Mirror {
entries: Vec<LogEntry>,
alert_position: AlertPosition,
}
impl Mirror {
fn new(entries: Vec<LogEntry>) -> Self {
let alert_position = latest_alert_position(&entries).unwrap_or(AlertPosition::Start);
Self {
entries,
alert_position,
}
}
}
// Use the public projection so private reasoning/response boundaries cannot become
// anchors. Results/accounting update existing rows, and non-error Run lifecycle
// transitions have no conversation row of their own.
fn latest_alert_position(entries: &[LogEntry]) -> Option<AlertPosition> {
session_store::public_snapshot::project_current_session_snapshot(entries)
.entries
.into_iter()
.rev()
.find_map(|entry| match entry.data {
SessionSnapshotEntryData::ToolCall { call_id, .. } => {
Some(AlertPosition::AfterToolCall { call_id })
}
SessionSnapshotEntryData::ToolResult { .. }
| SessionSnapshotEntryData::Invoke { .. }
| SessionSnapshotEntryData::Usage { .. }
| SessionSnapshotEntryData::RunStats { .. }
| SessionSnapshotEntryData::RunYielded { .. }
| SessionSnapshotEntryData::RunResumed { .. }
| SessionSnapshotEntryData::RunCancelled => None,
_ => Some(AlertPosition::AfterEntry {
entry_id: entry.entry_id,
}),
})
}
impl SegmentLogSink {
/// Create a fresh sink with an empty mirror. Used before any entry
/// has been written (deferred SegmentStart) or as a placeholder in
@@ -62,7 +103,7 @@ impl SegmentLogSink {
let (session_entry_tx, _) = broadcast::channel(BROADCAST_CAPACITY);
Self {
inner: Arc::new(SinkInner {
mirror: StdMutex::new(Vec::new()),
mirror: StdMutex::new(Mirror::new(Vec::new())),
session_entry_tx,
}),
}
@@ -75,7 +116,7 @@ impl SegmentLogSink {
let (session_entry_tx, _) = broadcast::channel(BROADCAST_CAPACITY);
Self {
inner: Arc::new(SinkInner {
mirror: StdMutex::new(entries),
mirror: StdMutex::new(Mirror::new(entries)),
session_entry_tx,
}),
}
@@ -103,7 +144,10 @@ impl SegmentLogSink {
.mirror
.lock()
.expect("session log mirror mutex poisoned");
mirror.push(entry.clone());
if let Some(position) = latest_alert_position(std::slice::from_ref(&entry)) {
mirror.alert_position = position;
}
mirror.entries.push(entry.clone());
if Self::is_live_relevant(&entry) {
// SendError means there are zero subscribers; harmless. The
// mirror lock is held across `send` so subscribers cannot
@@ -147,8 +191,7 @@ impl SegmentLogSink {
.mirror
.lock()
.expect("session log mirror mutex poisoned");
mirror.clear();
mirror.push(initial.clone());
*mirror = Mirror::new(vec![initial.clone()]);
let _ = self.inner.session_entry_tx.send(initial);
}
@@ -164,7 +207,7 @@ impl SegmentLogSink {
.mirror
.lock()
.expect("session log mirror mutex poisoned");
*mirror = entries;
*mirror = Mirror::new(entries);
if let Some(initial) = first {
let _ = self.inner.session_entry_tx.send(initial);
}
@@ -181,7 +224,7 @@ impl SegmentLogSink {
.mirror
.lock()
.expect("session log mirror mutex poisoned");
*mirror = entries;
*mirror = Mirror::new(entries);
}
/// Truncate the mirror without broadcasting.
@@ -191,7 +234,9 @@ impl SegmentLogSink {
.mirror
.lock()
.expect("session log mirror mutex poisoned");
mirror.truncate(entries_len);
mirror.entries.truncate(entries_len);
mirror.alert_position =
latest_alert_position(&mirror.entries).unwrap_or(AlertPosition::Start);
}
/// Atomically read the current mirror and subscribe to subsequent
@@ -203,17 +248,27 @@ impl SegmentLogSink {
.mirror
.lock()
.expect("session log mirror mutex poisoned");
let snapshot = mirror.clone();
let snapshot = mirror.entries.clone();
let rx = self.inner.session_entry_tx.subscribe();
(snapshot, rx)
}
pub(crate) fn alert_position(&self) -> AlertPosition {
self.inner
.mirror
.lock()
.expect("session log mirror mutex poisoned")
.alert_position
.clone()
}
/// Current entry count. Useful for tests / diagnostics.
pub fn len(&self) -> usize {
self.inner
.mirror
.lock()
.expect("session log mirror mutex poisoned")
.entries
.len()
}
@@ -267,6 +322,70 @@ mod tests {
}
}
#[test]
fn alert_anchor_tracks_visible_rows_not_reasoning_results_or_accounting() {
let sink = SegmentLogSink::new();
assert_eq!(sink.alert_position(), AlertPosition::Start);
sink.publish(session_start());
assert_eq!(sink.alert_position(), AlertPosition::Start);
let user = user_input("hello");
sink.publish(user.clone());
let expected = latest_alert_position(&[user.clone()]).unwrap();
assert_eq!(sink.alert_position(), expected);
sink.publish(turn_end(1));
sink.publish(LogEntry::Invoke {
ts: 2,
trigger: protocol::InvokeKind::UserSend,
});
sink.publish(LogEntry::AnnotatedAssistantItem {
ts: 3,
entry: crate::session_history::test_logged_history_entry(
session_store::LoggedItem::Reasoning {
text: "private".into(),
summary: vec![],
encrypted_content: None,
signature: None,
},
),
});
assert_eq!(sink.alert_position(), expected);
let tool = LogEntry::AnnotatedAssistantItem {
ts: 4,
entry: crate::session_history::test_logged_history_entry(agen::Item::tool_call(
"call-1", "Read", "{}",
)),
};
sink.publish(tool.clone());
assert_eq!(
sink.alert_position(),
AlertPosition::AfterToolCall {
call_id: "call-1".into()
}
);
sink.publish(user_input("newer row"));
let newer = sink.alert_position();
sink.publish(LogEntry::AnnotatedToolResult {
ts: 5,
entry: crate::session_history::test_logged_history_entry(agen::Item::tool_result(
"call-1", "done",
)),
});
assert_eq!(sink.alert_position(), newer);
sink.replace_silent(vec![session_start(), user.clone(), tool]);
assert_eq!(
sink.alert_position(),
AlertPosition::AfterToolCall {
call_id: "call-1".into()
}
);
sink.truncate_silent(2);
assert_eq!(sink.alert_position(), expected);
sink.reset_with_initial(session_start());
assert_eq!(sink.alert_position(), AlertPosition::Start);
sink.reset_with_initial_entries(vec![session_start(), user]);
assert_eq!(sink.alert_position(), expected);
}
#[test]
fn publish_then_subscribe_returns_history_in_snapshot() {
let sink = SegmentLogSink::new();
+2
View File
@@ -131,6 +131,8 @@ mod tests {
listener,
vec![
Event::Alert(Alert {
id: None,
position: None,
level: AlertLevel::Warn,
source: AlertSource::Worker,
message: "replayed alert".into(),
+77 -3
View File
@@ -1918,11 +1918,13 @@ where
self.state.increment_entries();
if let Some(in_flight) = &self.in_flight {
let committed_item = match &entry {
LogEntry::AnnotatedAssistantItem { entry, .. } => Some(entry.item.clone()),
LogEntry::AnnotatedAssistantItem { entry, .. } => {
Some((entry.item.clone(), entry.metadata.entry_id.0.clone()))
}
_ => None,
};
if let Some(item_for_clear) = committed_item {
in_flight.clear_for_committed_item_then(&item_for_clear, || {
if let Some((item_for_clear, entry_id)) = committed_item {
in_flight.clear_for_committed_item_then(&item_for_clear, Some(&entry_id), || {
self.sink.publish(entry);
});
return Ok(());
@@ -4253,6 +4255,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// Worker-internal operations (compaction failures, AGENTS.md
/// ingestion warnings) can surface messages to connected clients.
pub fn attach_alerter(&mut self, alerter: Alerter) {
alerter.attach_position_source(self.sink.clone(), self.in_flight.clone());
self.alerter = Some(alerter);
}
@@ -13313,6 +13316,77 @@ mod build_summary_prompt_tests {
}
}
// Real Worker wiring, filesystem Store, log writer and reconnect buffers;
// only the model client is inert because this test never invokes the engine.
#[tokio::test]
async fn alert_pending_text_anchor_resolves_at_durable_commit_before_replay() {
let dir = tempfile::tempdir().unwrap();
let cwd = dir.path().join("workspace");
std::fs::create_dir_all(&cwd).unwrap();
let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap();
let mut worker = Worker::new(
minimal_manifest(),
Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient),
store,
WorkerWorkspaceContext::no_workspace(),
WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()),
Scope::writable(&cwd).unwrap(),
)
.await
.unwrap();
let (tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(tx.clone());
worker.in_flight = Some(in_flight.clone());
let alerter = Alerter::new(tx);
worker.attach_alerter(alerter.clone());
alerter.alert(AlertLevel::Warn, AlertSource::Worker, "before head".into());
assert_eq!(
alerter.subscribe_with_snapshot().0[0].position,
Some(protocol::AlertPosition::Start)
);
worker.materialize_durable_session_head().await.unwrap();
let text = in_flight.start_text_block();
in_flight.text_delta(text, "streamed".into());
alerter.alert(AlertLevel::Warn, AlertSource::Worker, "during text".into());
let (before, _) = alerter.subscribe_with_snapshot();
let snapshot = crate::in_flight::snapshot_from_guard(&in_flight.snapshot_guard());
assert_eq!(
before[1].position,
Some(protocol::AlertPosition::AfterInFlight {
block_id: snapshot.block_ids[0]
})
);
let (deferred, _) = alerter.subscribe_with_deferred_snapshot();
let entry = crate::session_history::test_logged_history_entry(
agen::Item::assistant_message("streamed"),
);
let entry_id = entry.metadata.entry_id.0.clone();
worker
.log_writer_handle()
.append_entry(LogEntry::AnnotatedAssistantItem { ts: 42, entry })
.unwrap();
let (entries, _) = worker.sink.subscribe_with_snapshot();
assert!(entries.iter().any(|entry| matches!(entry, LogEntry::AnnotatedAssistantItem { entry, .. } if entry.metadata.entry_id.0 == entry_id)));
assert!(
crate::in_flight::snapshot_from_guard(&in_flight.snapshot_guard())
.blocks
.is_empty()
);
let replay = deferred();
assert_eq!(replay[1].id, before[1].id);
assert_eq!(
replay[1].position,
Some(protocol::AlertPosition::AfterEntry {
entry_id: entry_id.clone()
})
);
alerter.alert(AlertLevel::Warn, AlertSource::Worker, "after commit".into());
assert_eq!(
alerter.subscribe_with_snapshot().0[2].position,
Some(protocol::AlertPosition::AfterEntry { entry_id })
);
}
#[tokio::test]
async fn materialize_durable_session_head_commits_without_model_run() {
let dir = tempfile::tempdir().unwrap();