1use std::collections::{HashMap, HashSet, VecDeque};
49use std::sync::{Arc, Mutex};
50use std::time::Instant;
51
52#[derive(Debug, Clone)]
54pub struct JournalQueueConfig {
55 pub max_size: usize,
57 pub batch_size: usize,
59 pub flush_interval_ms: u64,
61}
62
63impl Default for JournalQueueConfig {
64 fn default() -> Self {
65 Self {
66 max_size: 5000,
67 batch_size: 100,
68 flush_interval_ms: 50,
69 }
70 }
71}
72
73impl JournalQueueConfig {
74 pub fn from_env() -> Self {
76 let max_size = std::env::var("AGNT5_JOURNAL_QUEUE_SIZE")
77 .ok()
78 .and_then(|s| s.parse().ok())
79 .unwrap_or(5000);
80
81 let batch_size = std::env::var("AGNT5_JOURNAL_BATCH_SIZE")
82 .ok()
83 .and_then(|s| s.parse().ok())
84 .unwrap_or(100);
85
86 let flush_interval_ms = std::env::var("AGNT5_JOURNAL_FLUSH_INTERVAL_MS")
87 .ok()
88 .and_then(|s| s.parse().ok())
89 .unwrap_or(50);
90
91 Self {
92 max_size,
93 batch_size,
94 flush_interval_ms,
95 }
96 }
97}
98
99#[derive(Debug, Clone)]
101pub struct JournalEventMessage {
102 pub run_id: String,
105 pub event_type: String,
107 pub data: Vec<u8>,
109
110 pub correlation_id: String,
113 pub parent_correlation_id: String,
115
116 pub tenant_id: Option<String>,
119 pub source_timestamp_ns: i64,
121 pub metadata: HashMap<String, String>,
123
124 pub queued_at: Instant,
127 pub is_streaming: bool,
129 pub is_sse_only: bool,
131
132 pub content_index: i32,
135 pub sequence: i64,
137}
138
139impl Default for JournalEventMessage {
140 fn default() -> Self {
141 Self {
142 run_id: String::new(),
143 event_type: String::new(),
144 data: Vec::new(),
145 correlation_id: String::new(),
146 parent_correlation_id: String::new(),
147 tenant_id: None,
148 source_timestamp_ns: 0,
149 metadata: HashMap::new(),
150 queued_at: Instant::now(),
151 is_streaming: false,
152 is_sse_only: false,
153 content_index: 0,
154 sequence: 0,
155 }
156 }
157}
158
159impl JournalEventMessage {
160 pub fn is_sse_only_event_type(event_type: &str) -> bool {
162 event_type.starts_with("output.")
165 || event_type.starts_with("lm.stream.")
166 || event_type.starts_with("lm.content_block.")
167 || event_type.starts_with("lm.message.")
168 || event_type.starts_with("lm.thinking.")
169 || event_type.starts_with("lm.tool_call.") || event_type.starts_with("progress.")
171 || event_type.starts_with("log") }
173
174 pub fn is_checkpoint_event_type(event_type: &str) -> bool {
186 !Self::is_sse_only_event_type(event_type)
187 }
188
189 pub fn new(run_id: String, event_type: String, data: Vec<u8>) -> Self {
191 let is_sse_only = Self::is_sse_only_event_type(&event_type);
192 Self {
193 run_id,
194 event_type,
195 data,
196 is_sse_only,
197 queued_at: Instant::now(),
198 ..Default::default()
199 }
200 }
201}
202
203#[derive(Debug, Default, Clone)]
205pub struct JournalQueueMetrics {
206 pub events_queued: u64,
208 pub events_sent: u64,
210 pub events_dropped: u64,
212 pub send_errors: u64,
214 pub boundary_events_sent: u64,
216 pub sse_only_events_sent: u64,
218}
219
220#[derive(Clone)]
226pub struct JournalEventQueue {
227 queue: Arc<Mutex<JournalQueueState>>,
229 config: JournalQueueConfig,
231 metrics: Arc<Mutex<JournalQueueMetrics>>,
233 held_runs: Arc<Mutex<HashSet<String>>>,
237}
238
239#[derive(Default)]
240struct JournalQueueState {
241 order: VecDeque<u64>,
244 events: HashMap<u64, JournalEventMessage>,
245 by_run: HashMap<String, VecDeque<u64>>,
246 next_id: u64,
247}
248
249impl JournalQueueState {
250 fn with_capacity(capacity: usize) -> Self {
251 Self {
252 order: VecDeque::with_capacity(capacity),
253 events: HashMap::with_capacity(capacity),
254 by_run: HashMap::new(),
255 next_id: 0,
256 }
257 }
258
259 fn insert(&mut self, event: JournalEventMessage, front: bool) {
260 let id = self.next_id;
261 self.next_id = self
262 .next_id
263 .checked_add(1)
264 .expect("journal queue event id exhausted");
265 let run_id = event.run_id.clone();
266
267 self.events.insert(id, event);
268 if front {
269 self.order.push_front(id);
270 self.by_run.entry(run_id).or_default().push_front(id);
271 } else {
272 self.order.push_back(id);
273 self.by_run.entry(run_id).or_default().push_back(id);
274 }
275 }
276
277 fn remove_run_reference(&mut self, run_id: &str, id: u64) {
278 let mut remove_run = false;
279 if let Some(ids) = self.by_run.get_mut(run_id) {
280 if ids.front() == Some(&id) {
281 ids.pop_front();
282 } else if let Some(position) = ids.iter().position(|queued_id| *queued_id == id) {
283 ids.remove(position);
284 }
285 remove_run = ids.is_empty();
286 }
287 if remove_run {
288 self.by_run.remove(run_id);
289 }
290 }
291
292 fn remove_event(&mut self, id: u64) -> Option<JournalEventMessage> {
293 let event = self.events.remove(&id)?;
294 self.remove_run_reference(&event.run_id, id);
295 Some(event)
296 }
297
298 fn maybe_compact_order(&mut self) {
299 const MIN_TOMBSTONES_BEFORE_COMPACTION: usize = 1024;
300 let tombstones = self.order.len().saturating_sub(self.events.len());
301 if tombstones >= MIN_TOMBSTONES_BEFORE_COMPACTION
302 && self.order.len() > self.events.len().saturating_mul(2)
303 {
304 self.order.retain(|id| self.events.contains_key(id));
305 }
306 }
307}
308
309impl JournalEventQueue {
310 pub fn new(config: JournalQueueConfig) -> Self {
312 log::info!(
313 "Creating unified journal event queue: max_size={}, batch_size={}, flush_interval_ms={}",
314 config.max_size,
315 config.batch_size,
316 config.flush_interval_ms
317 );
318
319 JournalEventQueue {
320 queue: Arc::new(Mutex::new(JournalQueueState::with_capacity(
321 config.max_size,
322 ))),
323 config,
324 metrics: Arc::new(Mutex::new(JournalQueueMetrics::default())),
325 held_runs: Arc::new(Mutex::new(HashSet::new())),
326 }
327 }
328
329 pub fn push(&self, event: JournalEventMessage) -> Result<(), String> {
336 let mut queue = self
337 .queue
338 .lock()
339 .map_err(|e| format!("Failed to lock journal queue for push: {}", e))?;
340
341 while queue.events.len() >= self.config.max_size {
345 let oldest_telemetry = queue.order.iter().enumerate().find_map(|(index, id)| {
346 queue.events.get(id).and_then(|queued| {
347 JournalEventMessage::is_sse_only_event_type(&queued.event_type)
348 .then_some((index, *id))
349 })
350 });
351
352 if let Some((index, id)) = oldest_telemetry {
353 queue.order.remove(index);
354 let dropped = queue
355 .remove_event(id)
356 .expect("telemetry id came from the same queue");
357 log::warn!(
358 "Journal queue full ({}), dropped oldest telemetry event: type={} run_id={}",
359 self.config.max_size,
360 dropped.event_type,
361 dropped.run_id
362 );
363
364 if let Ok(mut metrics) = self.metrics.lock() {
365 metrics.events_dropped += 1;
366 }
367 } else if JournalEventMessage::is_sse_only_event_type(&event.event_type) {
368 log::warn!(
369 "Journal queue full ({}), dropped incoming telemetry event: type={} run_id={}",
370 self.config.max_size,
371 event.event_type,
372 event.run_id
373 );
374 if let Ok(mut metrics) = self.metrics.lock() {
375 metrics.events_dropped += 1;
376 }
377 return Ok(());
378 } else {
379 return Err(format!(
380 "Journal queue full ({}) with correctness events; rejected event: type={} run_id={}",
381 self.config.max_size, event.event_type, event.run_id
382 ));
383 }
384 }
385
386 log::debug!(
387 "Queued journal event: type={} run_id={} is_sse_only={} queue_size={}",
388 event.event_type,
389 event.run_id,
390 event.is_sse_only,
391 queue.events.len() + 1
392 );
393
394 queue.insert(event, false);
395
396 if let Ok(mut metrics) = self.metrics.lock() {
398 metrics.events_queued += 1;
399 }
400
401 Ok(())
402 }
403
404 pub fn pop(&self) -> Option<JournalEventMessage> {
406 let mut queue = self.queue.lock().ok()?;
407 while let Some(id) = queue.order.pop_front() {
408 if let Some(event) = queue.remove_event(id) {
409 return Some(event);
410 }
411 }
412 None
413 }
414
415 pub fn push_front(&self, event: JournalEventMessage) -> Result<(), String> {
417 let mut queue = self
418 .queue
419 .lock()
420 .map_err(|e| format!("Failed to lock journal queue for push_front: {}", e))?;
421
422 log::debug!(
423 "Re-queuing event at front: type={} run_id={}",
424 event.event_type,
425 event.run_id
426 );
427
428 queue.insert(event, true);
429 Ok(())
430 }
431
432 pub fn drain_batch(&self, max: usize) -> Vec<JournalEventMessage> {
434 let mut queue = match self.queue.lock() {
435 Ok(q) => q,
436 Err(e) => {
437 log::error!("Failed to lock journal queue for drain_batch: {}", e);
438 return Vec::new();
439 }
440 };
441
442 let mut batch = Vec::with_capacity(std::cmp::min(max, queue.events.len()));
443
444 while batch.len() < max {
445 let Some(id) = queue.order.pop_front() else {
446 break;
447 };
448 if let Some(event) = queue.remove_event(id) {
449 batch.push(event);
450 }
451 }
452
453 if !batch.is_empty() {
454 log::debug!(
455 "Drained {} events from journal queue (remaining={})",
456 batch.len(),
457 queue.events.len()
458 );
459 }
460
461 batch
462 }
463
464 pub fn peek_batch_run_ids(&self, max: usize) -> Vec<String> {
468 let queue = match self.queue.lock() {
469 Ok(queue) => queue,
470 Err(error) => {
471 log::error!("Failed to lock journal queue for peek_batch_run_ids: {error}");
472 return Vec::new();
473 }
474 };
475 let held = self.held_runs_snapshot();
476
477 let mut run_ids = Vec::new();
478 let mut seen = HashSet::new();
479 for event in queue
480 .order
481 .iter()
482 .filter_map(|id| queue.events.get(id))
483 .take(max)
484 {
485 if held.contains(&event.run_id) {
486 continue;
487 }
488 if seen.insert(event.run_id.clone()) {
489 run_ids.push(event.run_id.clone());
490 }
491 }
492 run_ids
493 }
494
495 pub fn hold_run(&self, run_id: &str) {
498 self.held_runs
499 .lock()
500 .unwrap_or_else(|poisoned| poisoned.into_inner())
501 .insert(run_id.to_string());
502 }
503
504 pub fn release_run(&self, run_id: &str) -> bool {
506 self.held_runs
507 .lock()
508 .unwrap_or_else(|poisoned| poisoned.into_inner())
509 .remove(run_id)
510 }
511
512 pub fn is_held(&self, run_id: &str) -> bool {
513 self.held_runs
514 .lock()
515 .unwrap_or_else(|poisoned| poisoned.into_inner())
516 .contains(run_id)
517 }
518
519 fn held_runs_snapshot(&self) -> HashSet<String> {
520 self.held_runs
521 .lock()
522 .unwrap_or_else(|poisoned| poisoned.into_inner())
523 .clone()
524 }
525
526 pub fn drain_batch_for_runs(
529 &self,
530 max: usize,
531 allowed_run_ids: &HashSet<String>,
532 ) -> Vec<JournalEventMessage> {
533 let mut queue = match self.queue.lock() {
534 Ok(queue) => queue,
535 Err(error) => {
536 log::error!("Failed to lock journal queue for drain_batch_for_runs: {error}");
537 return Vec::new();
538 }
539 };
540 let mut batch = Vec::with_capacity(std::cmp::min(max, queue.events.len()));
541 let mut skipped = VecDeque::new();
542 let mut examined = 0;
543
544 while examined < max {
545 let Some(id) = queue.order.pop_front() else {
546 break;
547 };
548 let Some(event) = queue.events.get(&id) else {
549 continue;
550 };
551 examined += 1;
552 if allowed_run_ids.contains(&event.run_id) {
553 if let Some(event) = queue.remove_event(id) {
554 batch.push(event);
555 }
556 } else {
557 skipped.push_back(id);
558 }
559 }
560
561 while let Some(id) = skipped.pop_back() {
562 queue.order.push_front(id);
563 }
564
565 if !batch.is_empty() {
566 log::debug!(
567 "Drained {} barrier-protected events from journal queue (remaining={})",
568 batch.len(),
569 queue.events.len()
570 );
571 }
572 batch
573 }
574
575 pub fn drain_all(&self) -> Vec<JournalEventMessage> {
577 let mut queue = match self.queue.lock() {
578 Ok(q) => q,
579 Err(e) => {
580 log::error!("Failed to lock journal queue for drain_all: {}", e);
581 return Vec::new();
582 }
583 };
584
585 let mut events = Vec::with_capacity(queue.events.len());
586 while let Some(id) = queue.order.pop_front() {
587 if let Some(event) = queue.remove_event(id) {
588 events.push(event);
589 }
590 }
591 queue.by_run.clear();
592 queue.events.clear();
593
594 log::debug!("Drained {} events from journal queue", events.len());
595 events
596 }
597
598 pub fn drain_run_events(&self, run_id: &str) -> Vec<JournalEventMessage> {
601 let mut queue = match self.queue.lock() {
602 Ok(q) => q,
603 Err(e) => {
604 log::error!("Failed to lock journal queue for drain_run_events: {}", e);
605 return Vec::new();
606 }
607 };
608
609 let ids = queue.by_run.remove(run_id).unwrap_or_default();
610 let mut matched = Vec::with_capacity(ids.len());
611 for id in ids {
612 if let Some(event) = queue.events.remove(&id) {
613 matched.push(event);
614 }
615 }
616 queue.maybe_compact_order();
617
618 if !matched.is_empty() {
619 log::debug!(
620 "Drained {} events for run_id={} (remaining={})",
621 matched.len(),
622 run_id,
623 queue.events.len()
624 );
625 }
626
627 matched
628 }
629
630 pub fn len(&self) -> usize {
632 self.queue.lock().map(|q| q.events.len()).unwrap_or(0)
633 }
634
635 pub fn is_empty(&self) -> bool {
637 self.len() == 0
638 }
639
640 pub fn contains_run(&self, run_id: &str) -> bool {
645 self.queue
646 .lock()
647 .map(|queue| queue.by_run.contains_key(run_id))
648 .unwrap_or(true)
649 }
650
651 pub fn record_sent(&self, is_sse_only: bool) {
653 if let Ok(mut metrics) = self.metrics.lock() {
654 metrics.events_sent += 1;
655 if is_sse_only {
656 metrics.sse_only_events_sent += 1;
657 } else {
658 metrics.boundary_events_sent += 1;
659 }
660 }
661 }
662
663 pub fn record_sent_batch(&self, count: usize, sse_only_count: usize) {
665 if let Ok(mut metrics) = self.metrics.lock() {
666 metrics.events_sent += count as u64;
667 metrics.sse_only_events_sent += sse_only_count as u64;
668 metrics.boundary_events_sent += (count - sse_only_count) as u64;
669 }
670 }
671
672 pub fn record_error(&self) {
674 if let Ok(mut metrics) = self.metrics.lock() {
675 metrics.send_errors += 1;
676 }
677 }
678
679 pub fn metrics(&self) -> JournalQueueMetrics {
681 self.metrics.lock().map(|m| m.clone()).unwrap_or_default()
682 }
683
684 pub fn get_metrics(&self) -> (u64, u64, u64, u64) {
686 if let Ok(metrics) = self.metrics.lock() {
687 (
688 metrics.events_queued,
689 metrics.events_sent,
690 metrics.events_dropped,
691 metrics.send_errors,
692 )
693 } else {
694 (0, 0, 0, 0)
695 }
696 }
697
698 pub fn oldest_age(&self) -> Option<std::time::Duration> {
700 let queue = self.queue.lock().ok()?;
701 queue
702 .order
703 .iter()
704 .find_map(|id| queue.events.get(id))
705 .map(|event| event.queued_at.elapsed())
706 }
707
708 pub fn batch_size(&self) -> usize {
710 self.config.batch_size
711 }
712
713 pub fn flush_interval_ms(&self) -> u64 {
715 self.config.flush_interval_ms
716 }
717}
718
719#[cfg(test)]
720mod tests {
721 use super::*;
722
723 fn create_test_event(event_type: &str, seq: i64) -> JournalEventMessage {
724 JournalEventMessage {
725 run_id: "test-run".to_string(),
726 event_type: event_type.to_string(),
727 data: format!("{{\"seq\": {}}}", seq).into_bytes(),
728 is_sse_only: JournalEventMessage::is_sse_only_event_type(event_type),
729 sequence: seq,
730 source_timestamp_ns: 1234567890000000000 + seq * 1000000,
731 queued_at: Instant::now(),
732 ..Default::default()
733 }
734 }
735
736 fn create_run_event(run_id: &str, event_type: &str, seq: i64) -> JournalEventMessage {
737 JournalEventMessage {
738 run_id: run_id.to_string(),
739 ..create_test_event(event_type, seq)
740 }
741 }
742
743 #[test]
744 fn test_journal_queue_basic() {
745 let queue = JournalEventQueue::new(JournalQueueConfig {
746 max_size: 10,
747 ..Default::default()
748 });
749
750 queue
752 .push(create_test_event("workflow.started", 1))
753 .unwrap();
754 assert_eq!(queue.len(), 1);
755
756 let event = queue.pop().unwrap();
758 assert_eq!(event.sequence, 1);
759 assert_eq!(queue.len(), 0);
760 }
761
762 #[test]
763 fn test_journal_queue_overflow_drops_oldest_telemetry() {
764 let queue = JournalEventQueue::new(JournalQueueConfig {
765 max_size: 3,
766 ..Default::default()
767 });
768
769 queue.push(create_test_event("output.delta", 1)).unwrap();
771 queue.push(create_test_event("log.info", 2)).unwrap();
772 queue.push(create_test_event("progress.update", 3)).unwrap();
773 assert_eq!(queue.len(), 3);
774
775 queue.push(create_test_event("output.delta", 4)).unwrap();
777 assert_eq!(queue.len(), 3);
778
779 let event = queue.pop().unwrap();
781 assert_eq!(event.sequence, 2);
782 assert_eq!(queue.metrics().events_dropped, 1);
783 }
784
785 #[test]
786 fn test_journal_queue_overflow_never_evicts_correctness_event() {
787 let queue = JournalEventQueue::new(JournalQueueConfig {
788 max_size: 3,
789 ..Default::default()
790 });
791
792 queue
793 .push(create_test_event("workflow.started", 1))
794 .unwrap();
795 queue.push(create_test_event("output.delta", 2)).unwrap();
796 queue
797 .push(create_test_event("workflow.step.started", 3))
798 .unwrap();
799
800 queue
801 .push(create_test_event("workflow.step.completed", 4))
802 .unwrap();
803
804 let sequences: Vec<_> = queue
805 .drain_all()
806 .into_iter()
807 .map(|event| event.sequence)
808 .collect();
809 assert_eq!(sequences, vec![1, 3, 4]);
810 assert_eq!(queue.metrics().events_dropped, 1);
811 }
812
813 #[test]
814 fn test_journal_queue_rejects_correctness_when_only_correctness_is_buffered() {
815 let queue = JournalEventQueue::new(JournalQueueConfig {
816 max_size: 2,
817 ..Default::default()
818 });
819
820 queue
821 .push(create_test_event("workflow.started", 1))
822 .unwrap();
823 queue
824 .push(create_test_event("workflow.step.started", 2))
825 .unwrap();
826
827 let error = queue
828 .push(create_test_event("workflow.step.completed", 3))
829 .unwrap_err();
830
831 assert!(error.contains("rejected event: type=workflow.step.completed"));
832 let sequences: Vec<_> = queue
833 .drain_all()
834 .into_iter()
835 .map(|event| event.sequence)
836 .collect();
837 assert_eq!(sequences, vec![1, 2]);
838 assert_eq!(queue.metrics().events_dropped, 0);
839 }
840
841 #[test]
842 fn test_journal_queue_drops_incoming_telemetry_before_correctness() {
843 let queue = JournalEventQueue::new(JournalQueueConfig {
844 max_size: 2,
845 ..Default::default()
846 });
847
848 queue
849 .push(create_test_event("workflow.started", 1))
850 .unwrap();
851 queue
852 .push(create_test_event("workflow.step.started", 2))
853 .unwrap();
854
855 queue.push(create_test_event("output.delta", 3)).unwrap();
856
857 let sequences: Vec<_> = queue
858 .drain_all()
859 .into_iter()
860 .map(|event| event.sequence)
861 .collect();
862 assert_eq!(sequences, vec![1, 2]);
863 assert_eq!(queue.metrics().events_dropped, 1);
864 }
865
866 #[test]
867 fn test_sse_only_detection() {
868 assert!(!JournalEventMessage::is_sse_only_event_type(
870 "workflow.started"
871 ));
872 assert!(!JournalEventMessage::is_sse_only_event_type(
873 "agent.completed"
874 ));
875 assert!(!JournalEventMessage::is_sse_only_event_type(
876 "lm.call.started"
877 ));
878 assert!(!JournalEventMessage::is_sse_only_event_type(
879 "tool.call.completed"
880 ));
881
882 assert!(JournalEventMessage::is_sse_only_event_type("output.delta"));
884 assert!(JournalEventMessage::is_sse_only_event_type("output.start"));
885 assert!(JournalEventMessage::is_sse_only_event_type("output.stop"));
886 assert!(JournalEventMessage::is_sse_only_event_type(
887 "lm.stream.delta"
888 ));
889 assert!(JournalEventMessage::is_sse_only_event_type(
890 "lm.content_block.started"
891 ));
892 assert!(JournalEventMessage::is_sse_only_event_type(
893 "lm.content_block.delta"
894 ));
895 assert!(JournalEventMessage::is_sse_only_event_type(
896 "lm.content_block.completed"
897 ));
898 assert!(JournalEventMessage::is_sse_only_event_type(
899 "lm.message.delta"
900 ));
901 assert!(JournalEventMessage::is_sse_only_event_type(
902 "lm.thinking.delta"
903 ));
904 assert!(JournalEventMessage::is_sse_only_event_type(
905 "lm.tool_call.start"
906 ));
907 assert!(JournalEventMessage::is_sse_only_event_type(
908 "lm.tool_call.delta"
909 ));
910 assert!(JournalEventMessage::is_sse_only_event_type(
911 "lm.tool_call.stop"
912 ));
913 assert!(JournalEventMessage::is_sse_only_event_type(
914 "progress.update"
915 ));
916 assert!(JournalEventMessage::is_sse_only_event_type("log"));
917 assert!(JournalEventMessage::is_sse_only_event_type("log.info"));
918 assert!(JournalEventMessage::is_sse_only_event_type("log.warn"));
919 assert!(JournalEventMessage::is_sse_only_event_type("log.error"));
920 }
921
922 #[test]
923 fn test_checkpoint_event_detection() {
924 assert!(JournalEventMessage::is_checkpoint_event_type(
926 "workflow.started"
927 ));
928 assert!(JournalEventMessage::is_checkpoint_event_type(
929 "workflow.completed"
930 ));
931 assert!(JournalEventMessage::is_checkpoint_event_type(
932 "workflow.failed"
933 ));
934 assert!(JournalEventMessage::is_checkpoint_event_type(
935 "workflow.paused"
936 ));
937 assert!(JournalEventMessage::is_checkpoint_event_type(
938 "workflow.step.started"
939 ));
940 assert!(JournalEventMessage::is_checkpoint_event_type(
941 "workflow.step.completed"
942 ));
943 assert!(JournalEventMessage::is_checkpoint_event_type(
944 "workflow.step.paused"
945 ));
946 assert!(JournalEventMessage::is_checkpoint_event_type(
947 "agent.started"
948 ));
949 assert!(JournalEventMessage::is_checkpoint_event_type(
950 "agent.completed"
951 ));
952 assert!(JournalEventMessage::is_checkpoint_event_type(
953 "approval.requested"
954 ));
955 assert!(JournalEventMessage::is_checkpoint_event_type(
956 "approval.resolved"
957 ));
958 assert!(JournalEventMessage::is_checkpoint_event_type(
959 "lm.call.started"
960 ));
961 assert!(JournalEventMessage::is_checkpoint_event_type(
962 "lm.call.completed"
963 ));
964 assert!(JournalEventMessage::is_checkpoint_event_type(
965 "lm.completed"
966 ));
967 assert!(JournalEventMessage::is_checkpoint_event_type(
968 "tool.call.started"
969 ));
970 assert!(JournalEventMessage::is_checkpoint_event_type(
971 "tool.call.completed"
972 ));
973 assert!(JournalEventMessage::is_checkpoint_event_type(
974 "tool_call.started"
975 ));
976
977 assert!(!JournalEventMessage::is_checkpoint_event_type(
979 "output.delta"
980 ));
981 assert!(!JournalEventMessage::is_checkpoint_event_type(
982 "lm.stream.delta"
983 ));
984 assert!(!JournalEventMessage::is_checkpoint_event_type(
985 "lm.message.delta"
986 ));
987 assert!(!JournalEventMessage::is_checkpoint_event_type(
988 "lm.thinking.delta"
989 ));
990 assert!(!JournalEventMessage::is_checkpoint_event_type(
991 "progress.update"
992 ));
993 assert!(!JournalEventMessage::is_checkpoint_event_type("log"));
994 }
995
996 #[test]
997 fn test_drain_batch() {
998 let queue = JournalEventQueue::new(JournalQueueConfig {
999 max_size: 10,
1000 batch_size: 3,
1001 ..Default::default()
1002 });
1003
1004 for i in 1..=5 {
1006 queue
1007 .push(create_test_event("workflow.step.completed", i))
1008 .unwrap();
1009 }
1010 assert_eq!(queue.len(), 5);
1011
1012 let batch = queue.drain_batch(3);
1014 assert_eq!(batch.len(), 3);
1015 assert_eq!(queue.len(), 2);
1016
1017 assert_eq!(batch[0].sequence, 1);
1019 assert_eq!(batch[1].sequence, 2);
1020 assert_eq!(batch[2].sequence, 3);
1021
1022 let batch = queue.drain_batch(10);
1024 assert_eq!(batch.len(), 2);
1025 assert!(queue.is_empty());
1026 }
1027
1028 #[test]
1029 fn test_drain_run_events_uses_index_and_preserves_global_fifo() {
1030 let queue = JournalEventQueue::new(JournalQueueConfig {
1031 max_size: 10,
1032 ..Default::default()
1033 });
1034
1035 queue
1036 .push(create_run_event("run-a", "output.delta", 1))
1037 .unwrap();
1038 queue
1039 .push(create_run_event("run-b", "output.delta", 2))
1040 .unwrap();
1041 queue
1042 .push(create_run_event("run-a", "workflow.step.completed", 3))
1043 .unwrap();
1044 queue
1045 .push(create_run_event("run-c", "output.delta", 4))
1046 .unwrap();
1047
1048 let run_a: Vec<_> = queue
1049 .drain_run_events("run-a")
1050 .into_iter()
1051 .map(|event| event.sequence)
1052 .collect();
1053 assert_eq!(run_a, vec![1, 3]);
1054 assert!(!queue.contains_run("run-a"));
1055 assert!(queue.contains_run("run-b"));
1056 assert_eq!(queue.len(), 2);
1057
1058 let remaining: Vec<_> = queue
1059 .drain_batch(10)
1060 .into_iter()
1061 .map(|event| event.sequence)
1062 .collect();
1063 assert_eq!(remaining, vec![2, 4]);
1064 assert!(queue.is_empty());
1065 }
1066
1067 #[test]
1068 fn test_drain_batch_for_runs_does_not_take_unlocked_front_event() {
1069 let queue = JournalEventQueue::new(JournalQueueConfig {
1070 max_size: 10,
1071 ..Default::default()
1072 });
1073 queue
1074 .push(create_run_event("run-a", "output.delta", 1))
1075 .unwrap();
1076 queue
1077 .push(create_run_event("run-b", "output.delta", 2))
1078 .unwrap();
1079 queue
1080 .push(create_run_event("run-a", "output.delta", 3))
1081 .unwrap();
1082
1083 assert_eq!(queue.peek_batch_run_ids(2), vec!["run-a", "run-b"]);
1084
1085 let allowed = HashSet::from(["run-a".to_string()]);
1086 let drained: Vec<_> = queue
1087 .drain_batch_for_runs(2, &allowed)
1088 .into_iter()
1089 .map(|event| event.sequence)
1090 .collect();
1091 assert_eq!(drained, vec![1]);
1092
1093 let remaining: Vec<_> = queue
1094 .drain_all()
1095 .into_iter()
1096 .map(|event| event.sequence)
1097 .collect();
1098 assert_eq!(remaining, vec![2, 3]);
1099 }
1100
1101 #[test]
1102 fn held_runs_are_skipped_by_the_periodic_peek_but_drained_explicitly() {
1103 let queue = JournalEventQueue::new(JournalQueueConfig::default());
1104 queue.hold_run("held");
1105 for (run_id, event_type) in [
1106 ("held", "run.started"),
1107 ("other", "run.started"),
1108 ("held", "function.completed"),
1109 ] {
1110 queue
1111 .push(JournalEventMessage::new(
1112 run_id.to_string(),
1113 event_type.to_string(),
1114 Vec::new(),
1115 ))
1116 .unwrap();
1117 }
1118
1119 assert!(queue.is_held("held"));
1120 assert_eq!(queue.peek_batch_run_ids(10), vec!["other".to_string()]);
1121 let allowed: HashSet<String> = queue.peek_batch_run_ids(10).into_iter().collect();
1122 let periodic = queue.drain_batch_for_runs(10, &allowed);
1123 assert_eq!(periodic.len(), 1);
1124 assert_eq!(periodic[0].run_id, "other");
1125 assert!(queue.contains_run("held"), "held events stay queued");
1126
1127 let held = queue.drain_run_events("held");
1128 assert_eq!(
1129 held.iter()
1130 .map(|event| event.event_type.as_str())
1131 .collect::<Vec<_>>(),
1132 ["run.started", "function.completed"]
1133 );
1134 assert!(queue.release_run("held"));
1135 assert!(!queue.release_run("held"));
1136 assert!(!queue.is_held("held"));
1137 }
1138
1139 #[test]
1140 fn test_push_front() {
1141 let queue = JournalEventQueue::new(JournalQueueConfig::default());
1142
1143 queue
1144 .push(create_test_event("workflow.started", 1))
1145 .unwrap();
1146 queue
1147 .push(create_test_event("workflow.completed", 2))
1148 .unwrap();
1149
1150 let event = queue.pop().unwrap();
1152 assert_eq!(event.sequence, 1);
1153
1154 queue.push_front(event).unwrap();
1156
1157 let event = queue.pop().unwrap();
1159 assert_eq!(event.sequence, 1);
1160 }
1161
1162 #[test]
1163 fn test_metrics() {
1164 let queue = JournalEventQueue::new(JournalQueueConfig::default());
1165
1166 queue
1167 .push(create_test_event("workflow.started", 1))
1168 .unwrap();
1169 queue.push(create_test_event("output.delta", 2)).unwrap();
1170
1171 let (queued, sent, dropped, errors) = queue.get_metrics();
1172 assert_eq!(queued, 2);
1173 assert_eq!(sent, 0);
1174 assert_eq!(dropped, 0);
1175 assert_eq!(errors, 0);
1176
1177 queue.record_sent(false); queue.record_sent(true); let metrics = queue.metrics();
1182 assert_eq!(metrics.events_sent, 2);
1183 assert_eq!(metrics.boundary_events_sent, 1);
1184 assert_eq!(metrics.sse_only_events_sent, 1);
1185
1186 queue.record_error();
1187 let (_, _, _, errors) = queue.get_metrics();
1188 assert_eq!(errors, 1);
1189 }
1190
1191 #[test]
1192 fn test_fifo_order() {
1193 let queue = JournalEventQueue::new(JournalQueueConfig::default());
1194
1195 for i in 1..=5 {
1197 queue
1198 .push(create_test_event("workflow.step.completed", i))
1199 .unwrap();
1200 }
1201
1202 for i in 1..=5 {
1204 let event = queue.pop().unwrap();
1205 assert_eq!(event.sequence, i);
1206 }
1207
1208 assert!(queue.is_empty());
1209 }
1210
1211 #[test]
1212 fn test_new_with_auto_classification() {
1213 let boundary = JournalEventMessage::new(
1214 "run-1".to_string(),
1215 "workflow.started".to_string(),
1216 b"{}".to_vec(),
1217 );
1218 assert!(!boundary.is_sse_only);
1219
1220 let sse_only = JournalEventMessage::new(
1221 "run-1".to_string(),
1222 "output.delta".to_string(),
1223 b"hello".to_vec(),
1224 );
1225 assert!(sse_only.is_sse_only);
1226 }
1227}