1use std::sync::Arc;
10use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
11use std::time::Duration;
12
13use bytes::Bytes;
14use tokio::sync::{broadcast, mpsc};
15use tokio_util::sync::CancellationToken;
16use tracing::Instrument;
17
18use gemini_genai_rs::prelude::{SessionEvent, SessionPhase};
19use gemini_genai_rs::session::SessionWriter;
20
21use crate::state::State;
22use crate::tool::ToolDispatcher;
23
24use super::background_tool::BackgroundToolTracker;
25use super::callbacks::EventCallbacks;
26use super::computed::ComputedRegistry;
27use super::context_writer::PendingContext;
28use super::control_plane::run_control_lane;
29use super::events::LiveEvent;
30use super::extractor::TurnExtractor;
31use super::needs::NeedsFulfillment;
32use super::persistence::SessionPersistence;
33use super::phase::PhaseMachine;
34use super::session_signals::SessionSignals;
35use super::soft_turn::SoftTurnDetector;
36use super::steering::{ContextDelivery, SteeringMode};
37use super::telemetry::SessionTelemetry;
38use super::temporal::TemporalRegistry;
39use super::watcher::WatcherRegistry;
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
52pub enum Delivery {
53 #[default]
59 Lossless,
60 LossyDropNewest,
73}
74
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83pub struct DeliveryConfig {
84 pub audio: Delivery,
86 pub media: Delivery,
88 pub text: Delivery,
90 pub transcript: Delivery,
93 pub thought: Delivery,
95 pub vad: Delivery,
97 pub phase: Delivery,
99}
100
101impl Default for DeliveryConfig {
102 fn default() -> Self {
103 Self {
104 audio: Delivery::Lossless,
105 media: Delivery::Lossless,
106 text: Delivery::Lossless,
107 transcript: Delivery::Lossless,
108 thought: Delivery::Lossless,
109 vad: Delivery::Lossless,
110 phase: Delivery::Lossless,
111 }
112 }
113}
114
115impl DeliveryConfig {
116 pub fn lossless() -> Self {
119 Self::default()
120 }
121
122 pub fn audio(mut self, d: Delivery) -> Self {
124 self.audio = d;
125 self
126 }
127
128 pub fn media(mut self, d: Delivery) -> Self {
130 self.media = d;
131 self
132 }
133
134 pub fn text(mut self, d: Delivery) -> Self {
136 self.text = d;
137 self
138 }
139
140 pub fn transcript(mut self, d: Delivery) -> Self {
142 self.transcript = d;
143 self
144 }
145
146 pub fn thought(mut self, d: Delivery) -> Self {
148 self.thought = d;
149 self
150 }
151
152 pub fn vad(mut self, d: Delivery) -> Self {
154 self.vad = d;
155 self
156 }
157
158 pub fn phase(mut self, d: Delivery) -> Self {
160 self.phase = d;
161 self
162 }
163}
164
165#[derive(Debug, Default)]
171pub(crate) struct DroppedFrames {
172 pub audio: AtomicU64,
173 pub media: AtomicU64,
174 pub text: AtomicU64,
175 pub transcript: AtomicU64,
176 pub thought: AtomicU64,
177 pub vad: AtomicU64,
178 pub phase: AtomicU64,
179}
180
181impl DroppedFrames {
182 #[cfg(test)]
187 pub(crate) fn total(&self) -> u64 {
188 self.audio.load(Ordering::Relaxed)
189 + self.media.load(Ordering::Relaxed)
190 + self.text.load(Ordering::Relaxed)
191 + self.transcript.load(Ordering::Relaxed)
192 + self.thought.load(Ordering::Relaxed)
193 + self.vad.load(Ordering::Relaxed)
194 + self.phase.load(Ordering::Relaxed)
195 }
196}
197
198async fn deliver_fast(
207 tx: &mpsc::Sender<FastEvent>,
208 ev: FastEvent,
209 policy: Delivery,
210 dropped: &AtomicU64,
211) {
212 match policy {
213 Delivery::Lossless => {
214 let _ = tx.send(ev).await;
215 }
216 Delivery::LossyDropNewest => {
217 if let Err(mpsc::error::TrySendError::Full(_)) = tx.try_send(ev) {
218 dropped.fetch_add(1, Ordering::Relaxed);
219 }
220 }
223 }
224}
225
226pub(crate) enum FastEvent {
228 Audio(Bytes),
229 Media(gemini_genai_rs::session::InlineMedia),
230 Text(String),
231 TextComplete(String),
232 InputTranscript(String),
233 OutputTranscript(String),
234 Thought(String),
235 VadStart,
236 VadEnd,
237 Phase(SessionPhase),
238 Interrupted,
240 Barrier(tokio::sync::oneshot::Sender<()>),
242}
243
244pub(crate) enum ControlEvent {
246 Barrier(tokio::sync::oneshot::Sender<()>),
248 ToolCall(Vec<gemini_genai_rs::prelude::FunctionCall>),
249 ToolCallCancelled(Vec<String>),
250 ToolCompleted {
254 call_id: String,
256 name: String,
258 ok: bool,
260 },
261 Interrupted {
264 heard_chars: Option<usize>,
265 },
266 TurnComplete,
267 GenerationComplete,
269 GoAway(Option<std::time::Duration>),
270 Connected,
271 Disconnected(Option<String>),
272 SessionResumeUpdate(gemini_genai_rs::session::ResumeInfo),
273 Error(String),
274 InputTranscript(String),
276 OutputTranscript(String),
277}
278
279pub(crate) struct SharedState {
281 pub interrupted: AtomicBool,
283 pub barge_in: parking_lot::Mutex<CancellationToken>,
291 pub resume_handle: parking_lot::Mutex<Option<String>>,
293 pub last_instruction: parking_lot::Mutex<Option<String>>,
295 pub last_context: parking_lot::Mutex<Vec<String>>,
304 pub pending_context: Option<Arc<PendingContext>>,
306 pub delivery: DeliveryConfig,
308 pub dropped: DroppedFrames,
310 pub redactor: Option<Arc<crate::live::redaction::TranscriptRedactor>>,
313 pub playback: super::playback::PlaybackClock,
315 pub speech: parking_lot::Mutex<super::playback::TurnSpeech>,
318}
319
320pub(crate) struct ControlPlaneConfig {
326 pub base_instruction: Option<String>,
336 pub soft_turn: Option<SoftTurnDetector>,
338 pub steering_mode: SteeringMode,
340 pub context_delivery: ContextDelivery,
344 pub needs_fulfillment: Option<NeedsFulfillment>,
346 pub persistence: Option<Arc<dyn SessionPersistence>>,
348 pub session_id: Option<String>,
350 pub tool_advisory: bool,
352 pub pending_context: Option<Arc<PendingContext>>,
356 pub middleware: Arc<crate::middleware::MiddlewareChain>,
359 pub flow: Option<crate::flow::SharedFlowStack>,
367 pub delivery: DeliveryConfig,
370 pub redactor: Option<Arc<crate::live::redaction::TranscriptRedactor>>,
373 pub lockstep: Option<Arc<Lockstep>>,
376 pub playback: super::playback::PlaybackClock,
379}
380
381#[derive(Debug, Default)]
386pub(crate) struct Lockstep {
387 settled: AtomicU64,
388 notify: tokio::sync::Notify,
389}
390
391impl Lockstep {
392 pub(crate) fn settled(&self) -> u64 {
394 self.settled.load(Ordering::Acquire)
395 }
396
397 fn settle(&self, n: u64) {
398 self.settled.fetch_add(n, Ordering::AcqRel);
399 self.notify.notify_waiters();
400 }
401
402 pub(crate) async fn wait_for(&self, count: u64, timeout: std::time::Duration) {
404 let deadline = tokio::time::Instant::now() + timeout;
405 loop {
406 let notified = self.notify.notified();
407 if self.settled() >= count {
408 return;
409 }
410 if tokio::time::timeout_at(deadline, notified).await.is_err() {
411 return;
412 }
413 }
414 }
415}
416
417impl Default for ControlPlaneConfig {
418 fn default() -> Self {
419 Self {
420 base_instruction: None,
421 soft_turn: None,
422 steering_mode: SteeringMode::default(),
423 context_delivery: ContextDelivery::default(),
424 needs_fulfillment: None,
425 persistence: None,
426 session_id: None,
427 tool_advisory: true,
428 pending_context: None,
429 middleware: Arc::new(crate::middleware::MiddlewareChain::new()),
430 flow: None,
431 delivery: DeliveryConfig::default(),
432 redactor: None,
433 lockstep: None,
434 playback: super::playback::PlaybackClock::new(crate::clock::system_clock()),
435 }
436 }
437}
438
439#[allow(
440 clippy::too_many_arguments,
441 reason = "lane spawn site: parameters are the owned subsystem handles split between the fast and control lanes"
442)]
443pub(crate) fn spawn_event_processor(
444 mut event_rx: broadcast::Receiver<SessionEvent>,
445 callbacks: Arc<EventCallbacks>,
446 dispatcher: Option<Arc<ToolDispatcher>>,
447 writer: Arc<dyn SessionWriter>,
448 extractors: Vec<Arc<dyn TurnExtractor>>,
449 state: State,
450 computed: Option<ComputedRegistry>,
451 phase_machine: Option<tokio::sync::Mutex<PhaseMachine>>,
452 watchers: Option<WatcherRegistry>,
453 temporal: Option<Arc<TemporalRegistry>>,
454 background_tracker: Option<Arc<BackgroundToolTracker>>,
455 execution_modes: std::collections::HashMap<String, super::background_tool::ToolExecutionMode>,
456 control_plane: ControlPlaneConfig,
457 live_event_tx: broadcast::Sender<LiveEvent>,
458) -> (
459 tokio::task::JoinHandle<()>,
460 tokio::task::JoinHandle<()>,
461 mpsc::WeakSender<ControlEvent>,
462) {
463 let shared = Arc::new(SharedState {
464 interrupted: AtomicBool::new(false),
465 barge_in: parking_lot::Mutex::new(CancellationToken::new()),
466 resume_handle: parking_lot::Mutex::new(None),
467 last_instruction: parking_lot::Mutex::new(None),
468 last_context: parking_lot::Mutex::new(Vec::new()),
469 pending_context: control_plane.pending_context.clone(),
470 delivery: control_plane.delivery,
471 dropped: DroppedFrames::default(),
472 redactor: control_plane.redactor.clone(),
473 playback: control_plane.playback.clone(),
474 speech: parking_lot::Mutex::default(),
475 });
476
477 let timer_cancel = CancellationToken::new();
478
479 let (fast_tx, fast_rx) = mpsc::channel::<FastEvent>(512);
490 let (ctrl_tx, ctrl_rx) = mpsc::channel::<ControlEvent>(512);
491
492 let fast_tx_clone = fast_tx.clone();
495 let ctrl_tx_clone = ctrl_tx.clone();
496 let shared_clone = shared.clone();
497 let lockstep = control_plane.lockstep.clone();
498 tokio::spawn(
499 async move {
500 let mut turn = super::turn_trace::TurnTrace::new();
502 loop {
503 match event_rx.recv().await {
504 Ok(event) => {
505 let terminal = matches!(event, SessionEvent::Disconnected(_));
513 let boundary = matches!(event, SessionEvent::TurnComplete);
514 route_event(event, &fast_tx_clone, &ctrl_tx_clone, &shared_clone)
515 .instrument(turn.span())
516 .await;
517 if let Some(lockstep) = &lockstep {
518 let (fast_done, fast_rx) = tokio::sync::oneshot::channel();
519 let (ctrl_done, ctrl_rx) = tokio::sync::oneshot::channel();
520 let _ = fast_tx_clone.send(FastEvent::Barrier(fast_done)).await;
521 let _ = ctrl_tx_clone.send(ControlEvent::Barrier(ctrl_done)).await;
522 let _ = fast_rx.await;
523 let _ = ctrl_rx.await;
524 lockstep.settle(1);
525 }
526 if boundary {
527 turn.advance();
528 tracing::debug!(parent: &turn.span(), "turn started");
529 }
530 if terminal {
531 break;
532 }
533 }
534 Err(broadcast::error::RecvError::Lagged(n)) => {
535 tracing::warn!(skipped = n, "Event processor lagged, skipped events");
536 if let Some(lockstep) = &lockstep {
537 lockstep.settle(n);
538 }
539 }
540 Err(broadcast::error::RecvError::Closed) => break,
541 }
542 }
543 }
544 .in_current_span(),
545 );
546
547 let fast_callbacks = callbacks.clone();
549 let fast_shared = shared.clone();
550 let fast_event_tx = live_event_tx.clone();
551 let fast_handle = tokio::spawn(
552 async move {
553 run_fast_lane(fast_rx, fast_callbacks, fast_shared, fast_event_tx).await;
554 }
555 .in_current_span(),
556 );
557
558 let timer_temporal = temporal.clone();
560 let reprompt = control_plane.flow.as_ref().is_some_and(|stack| {
562 stack
563 .lock()
564 .timing_policies()
565 .values()
566 .any(|t| t.reprompt_after_ms.is_some())
567 });
568 let reprompt_inputs = reprompt.then(|| (state.clone(), writer.clone(), live_event_tx.clone()));
569 let timer_state = state.clone();
570 let timer_writer = writer.clone();
571
572 let ctrl_callbacks = callbacks;
574 let ctrl_shared = shared;
575 let ctrl_timer_cancel = timer_cancel.clone();
576 let ctrl_tx_weak = ctrl_tx.downgrade();
581 let ctrl_handle = tokio::spawn(
582 async move {
583 run_control_lane(
584 ctrl_rx,
585 ctrl_tx_weak,
586 ctrl_callbacks,
587 dispatcher,
588 writer,
589 ctrl_shared,
590 extractors,
591 state,
592 computed,
593 phase_machine,
594 watchers,
595 temporal,
596 background_tracker,
597 execution_modes,
598 control_plane,
599 live_event_tx,
600 )
601 .await;
602 ctrl_timer_cancel.cancel();
603 }
604 .in_current_span(),
605 );
606
607 if let Some(ref temporal_ref) = timer_temporal
609 && temporal_ref.needs_timer()
610 {
611 let t = temporal_ref.clone();
612 let cancel = timer_cancel.clone();
613 tokio::spawn(async move {
614 let mut interval = tokio::time::interval(Duration::from_millis(500));
615 loop {
616 tokio::select! {
617 _ = cancel.cancelled() => break,
618 _ = interval.tick() => {
619 for action in t.check_all(&timer_state, None, &timer_writer) {
620 tokio::spawn(action);
621 }
622 }
623 }
624 }
625 });
626 }
627
628 if let Some((state, writer, events)) = reprompt_inputs {
629 tokio::spawn(super::reprompt::run_reprompt_timer(
630 state,
631 writer,
632 events,
633 timer_cancel.clone(),
634 ));
635 }
636
637 let ctrl_tx_handle = ctrl_tx.downgrade();
642
643 (fast_handle, ctrl_handle, ctrl_tx_handle)
644}
645
646pub(crate) fn spawn_telemetry_lane(
652 mut telem_rx: broadcast::Receiver<SessionEvent>,
653 signals: SessionSignals,
654 telemetry: Arc<SessionTelemetry>,
655 cancel: CancellationToken,
656 on_usage: Option<super::callbacks::UsageCallback>,
657) -> tokio::task::JoinHandle<()> {
658 tokio::spawn(async move {
659 let mut debounce = tokio::time::interval(Duration::from_millis(100));
660 debounce.tick().await;
662 let mut turn = super::turn_trace::TurnTrace::new();
664 let mut turn_usage: Option<gemini_genai_rs::prelude::UsageMetadata> = None;
666 loop {
667 tokio::select! {
668 biased;
669 result = telem_rx.recv() => {
670 match result {
671 Ok(event) => {
672 turn.span().in_scope(|| {
675 match &event {
677 SessionEvent::AudioData(data) => {
678 if let Some(latency) = telemetry.record_audio_out(data.len()) {
679 signals.record_response_latency(latency);
680 tracing::info!(
681 latency_ms = latency.as_millis() as u64,
682 "first model audio after the user's turn"
683 );
684 }
685 }
686 SessionEvent::TextDelta(_) => {
687 if let Some(latency) = telemetry.record_text_out() {
688 signals.record_response_latency(latency);
689 tracing::info!(
690 latency_ms = latency.as_millis() as u64,
691 "first model text after the user's turn"
692 );
693 }
694 }
695 SessionEvent::VoiceActivityEnd => {
696 telemetry.record_vad_end();
697 }
698 SessionEvent::Interrupted => {
699 telemetry.record_interruption();
700 }
701 SessionEvent::TurnComplete => {
702 telemetry.record_turn_complete();
703 if let Some(usage) = turn_usage.take() {
704 telemetry.record_turn_usage(&usage);
705 }
706 }
707 SessionEvent::VoiceActivityStart => {
708 telemetry.mark_turn_start();
709 }
710 SessionEvent::Usage(usage) => {
711 telemetry.record_usage(
712 usage.total_token_count,
713 usage.prompt_token_count,
714 usage.response_token_count,
715 usage.cached_content_token_count,
716 usage.thoughts_token_count,
717 );
718 turn_usage = Some(usage.clone());
719 if let Some(cb) = &on_usage {
720 cb(usage);
721 }
722 }
723 _ => {}
724 }
725 signals.on_event(&event);
727 });
728 if matches!(event, SessionEvent::TurnComplete) {
729 turn.advance();
730 }
731 }
732 Err(broadcast::error::RecvError::Lagged(n)) => {
733 tracing::warn!(skipped = n, "Telemetry lane lagged");
734 }
735 Err(broadcast::error::RecvError::Closed) => break,
736 }
737 }
738 _ = debounce.tick() => {
739 signals.flush_timing();
741 }
742 _ = cancel.cancelled() => break,
743 }
744 }
745 }.in_current_span())
746}
747
748async fn route_event(
750 event: SessionEvent,
751 fast_tx: &mpsc::Sender<FastEvent>,
752 ctrl_tx: &mpsc::Sender<ControlEvent>,
753 shared: &SharedState,
754) {
755 let delivery = &shared.delivery;
756 let dropped = &shared.dropped;
757 match event {
758 SessionEvent::AudioData(data) => {
760 shared.speech.lock().on_audio(data.len(), &shared.playback);
761 deliver_fast(
762 fast_tx,
763 FastEvent::Audio(data),
764 delivery.audio,
765 &dropped.audio,
766 )
767 .await;
768 }
769 SessionEvent::Media(media) => {
770 deliver_fast(
771 fast_tx,
772 FastEvent::Media(media),
773 delivery.media,
774 &dropped.media,
775 )
776 .await;
777 }
778 SessionEvent::TextDelta(text) => {
779 deliver_fast(fast_tx, FastEvent::Text(text), delivery.text, &dropped.text).await;
780 }
781 SessionEvent::TextComplete(text) => {
782 let text = match &shared.redactor {
783 Some(redactor) => redactor.redact(text),
784 None => text,
785 };
786 deliver_fast(
787 fast_tx,
788 FastEvent::TextComplete(text),
789 delivery.text,
790 &dropped.text,
791 )
792 .await;
793 }
794 SessionEvent::InputTranscription(text) => {
797 let text = match &shared.redactor {
798 Some(redactor) => redactor.redact(text),
799 None => text,
800 };
801 deliver_fast(
802 fast_tx,
803 FastEvent::InputTranscript(text.clone()),
804 delivery.transcript,
805 &dropped.transcript,
806 )
807 .await;
808 let _ = ctrl_tx.send(ControlEvent::InputTranscript(text)).await;
809 }
810 SessionEvent::OutputTranscription(text) => {
811 let text = match &shared.redactor {
812 Some(redactor) => redactor.redact(text),
813 None => text,
814 };
815 shared.speech.lock().on_text(&text);
816 deliver_fast(
817 fast_tx,
818 FastEvent::OutputTranscript(text.clone()),
819 delivery.transcript,
820 &dropped.transcript,
821 )
822 .await;
823 let _ = ctrl_tx.send(ControlEvent::OutputTranscript(text)).await;
824 }
825 SessionEvent::Thought(text) => {
826 deliver_fast(
827 fast_tx,
828 FastEvent::Thought(text),
829 delivery.thought,
830 &dropped.thought,
831 )
832 .await;
833 }
834 SessionEvent::VoiceActivityStart => {
835 tracing::debug!("user started speaking");
836 deliver_fast(fast_tx, FastEvent::VadStart, delivery.vad, &dropped.vad).await;
837 }
838 SessionEvent::VoiceActivityEnd => {
839 tracing::debug!("user stopped speaking");
840 deliver_fast(fast_tx, FastEvent::VadEnd, delivery.vad, &dropped.vad).await;
841 }
842 SessionEvent::PhaseChanged(phase) => {
843 deliver_fast(
844 fast_tx,
845 FastEvent::Phase(phase),
846 delivery.phase,
847 &dropped.phase,
848 )
849 .await;
850 }
851 SessionEvent::SessionResumeUpdate(info) => {
852 *shared.resume_handle.lock() = Some(info.handle.clone());
853 let _ = ctrl_tx.send(ControlEvent::SessionResumeUpdate(info)).await;
854 }
855 SessionEvent::GenerationComplete => {
856 let _ = ctrl_tx.send(ControlEvent::GenerationComplete).await;
857 }
858
859 SessionEvent::ToolCall(calls) => {
861 let _ = ctrl_tx.send(ControlEvent::ToolCall(calls)).await;
862 }
863 SessionEvent::ToolCallCancelled(ids) => {
864 let _ = ctrl_tx.send(ControlEvent::ToolCallCancelled(ids)).await;
865 }
866 SessionEvent::Interrupted => {
867 tracing::debug!("user interrupted the model");
868 shared.interrupted.store(true, Ordering::Release);
870 shared.barge_in.lock().cancel();
874 let heard_chars = shared.speech.lock().on_interrupted(&shared.playback);
877 let _ = fast_tx.send(FastEvent::Interrupted).await;
878 let _ = ctrl_tx
879 .send(ControlEvent::Interrupted { heard_chars })
880 .await;
881 }
882 SessionEvent::TurnComplete => {
883 shared.speech.lock().on_turn_complete();
884 let _ = ctrl_tx.send(ControlEvent::TurnComplete).await;
885 }
886 SessionEvent::Usage(_) => {}
888 SessionEvent::GoAway(time_left) => {
889 let _ = ctrl_tx.send(ControlEvent::GoAway(time_left)).await;
890 }
891 SessionEvent::Connected => {
892 let _ = ctrl_tx.send(ControlEvent::Connected).await;
893 }
894 SessionEvent::Disconnected(reason) => {
895 let _ = ctrl_tx.send(ControlEvent::Disconnected(reason)).await;
896 }
897 SessionEvent::Error(err) => {
898 let _ = ctrl_tx.send(ControlEvent::Error(err.to_string())).await;
899 }
900 other => {
904 tracing::debug!(?other, "unhandled SessionEvent variant (newer wire event?)");
905 }
906 }
907}
908
909async fn run_fast_lane(
912 mut rx: mpsc::Receiver<FastEvent>,
913 callbacks: Arc<EventCallbacks>,
914 shared: Arc<SharedState>,
915 event_tx: broadcast::Sender<LiveEvent>,
916) {
917 while let Some(event) = rx.recv().await {
918 match event {
919 FastEvent::Barrier(done) => {
920 let _ = done.send(());
921 }
922 FastEvent::Audio(data) => {
923 if !shared.interrupted.load(Ordering::Acquire) {
925 if let Some(cb) = &callbacks.on_audio {
926 cb(&data);
927 }
928 let _ = event_tx.send(LiveEvent::Audio(data));
929 }
930 }
931 FastEvent::Media(media) => {
932 if !shared.interrupted.load(Ordering::Acquire) {
934 if let Some(cb) = &callbacks.on_media {
935 cb(&media);
936 }
937 let _ = event_tx.send(LiveEvent::Media(media));
938 }
939 }
940 FastEvent::Text(delta) => {
941 if let Some(cb) = &callbacks.on_text {
942 cb(&delta);
943 }
944 let _ = event_tx.send(LiveEvent::TextDelta(delta));
945 }
946 FastEvent::TextComplete(text) => {
947 if let Some(cb) = &callbacks.on_text_complete {
948 cb(&text);
949 }
950 let _ = event_tx.send(LiveEvent::TextComplete(text));
951 }
952 FastEvent::InputTranscript(text) => {
953 if let Some(cb) = &callbacks.on_input_transcript {
955 cb(&text, false);
956 }
957 let _ = event_tx.send(LiveEvent::InputTranscript {
958 text,
959 is_final: false,
960 });
961 }
962 FastEvent::OutputTranscript(text) => {
963 if let Some(cb) = &callbacks.on_output_transcript {
965 cb(&text, false);
966 }
967 let _ = event_tx.send(LiveEvent::OutputTranscript {
968 text,
969 is_final: false,
970 });
971 }
972 FastEvent::Thought(text) => {
973 if let Some(cb) = &callbacks.on_thought {
974 cb(&text);
975 }
976 let _ = event_tx.send(LiveEvent::Thought(text));
977 }
978 FastEvent::VadStart => {
979 if let Some(cb) = &callbacks.on_vad_start {
980 cb();
981 }
982 let _ = event_tx.send(LiveEvent::VadStart);
983 }
984 FastEvent::VadEnd => {
985 if let Some(cb) = &callbacks.on_vad_end {
986 cb();
987 }
988 let _ = event_tx.send(LiveEvent::VadEnd);
989 }
990 FastEvent::Phase(phase) => {
991 if let Some(cb) = &callbacks.on_session_phase {
992 cb(phase);
993 }
994 }
996 FastEvent::Interrupted => {
997 }
1000 }
1001 }
1002}
1003
1004#[cfg(test)]
1005mod tests {
1006 use super::*;
1007 use std::sync::atomic::AtomicUsize;
1008
1009 use crate::live::events::LiveEvent;
1010 use gemini_genai_rs::prelude::FunctionResponse;
1011
1012 fn dummy_event_tx() -> broadcast::Sender<LiveEvent> {
1013 broadcast::channel::<LiveEvent>(16).0
1014 }
1015
1016 #[test]
1017 fn delivery_config_default_is_all_lossless() {
1018 let cfg = DeliveryConfig::default();
1019 assert_eq!(cfg.audio, Delivery::Lossless);
1020 assert_eq!(cfg.text, Delivery::Lossless);
1021 assert_eq!(cfg.transcript, Delivery::Lossless);
1022 assert_eq!(cfg.thought, Delivery::Lossless);
1023 assert_eq!(cfg.vad, Delivery::Lossless);
1024 assert_eq!(cfg.phase, Delivery::Lossless);
1025 assert_eq!(Delivery::default(), Delivery::Lossless);
1027 }
1028
1029 #[tokio::test]
1030 async fn lossy_drop_newest_does_not_block_and_counts_drops() {
1031 let (tx, _rx) = mpsc::channel::<FastEvent>(1);
1034 tx.send(FastEvent::VadStart).await.unwrap(); let dropped = AtomicU64::new(0);
1036
1037 let res = tokio::time::timeout(
1040 Duration::from_millis(100),
1041 deliver_fast(&tx, FastEvent::VadEnd, Delivery::LossyDropNewest, &dropped),
1042 )
1043 .await;
1044 assert!(res.is_ok(), "deliver_fast blocked under LossyDropNewest");
1045 assert_eq!(dropped.load(Ordering::Relaxed), 1);
1046 }
1047
1048 #[tokio::test]
1049 async fn lossless_delivers_on_non_full_channel() {
1050 let (tx, mut rx) = mpsc::channel::<FastEvent>(4);
1051 let dropped = AtomicU64::new(0);
1052
1053 deliver_fast(
1054 &tx,
1055 FastEvent::Text("hello".into()),
1056 Delivery::Lossless,
1057 &dropped,
1058 )
1059 .await;
1060
1061 assert_eq!(dropped.load(Ordering::Relaxed), 0);
1063 match rx.recv().await {
1064 Some(FastEvent::Text(s)) => assert_eq!(s, "hello"),
1065 other => panic!("expected Text frame, got {:?}", other.is_some()),
1066 }
1067 }
1068
1069 #[test]
1070 fn dropped_frames_total_sums_classes() {
1071 let d = DroppedFrames::default();
1072 d.audio.fetch_add(2, Ordering::Relaxed);
1073 d.transcript.fetch_add(3, Ordering::Relaxed);
1074 assert_eq!(d.total(), 5);
1075 }
1076
1077 #[tokio::test]
1078 async fn fast_lane_routes_audio() {
1079 let count = Arc::new(AtomicUsize::new(0));
1080 let count_clone = count.clone();
1081
1082 let callbacks = EventCallbacks {
1083 on_audio: Some(Box::new(move |_| {
1084 count_clone.fetch_add(1, Ordering::SeqCst);
1085 })),
1086 ..Default::default()
1087 };
1088 let callbacks = Arc::new(callbacks);
1089
1090 let (event_tx, _) = broadcast::channel(16);
1091 let event_rx = event_tx.subscribe();
1092
1093 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1094
1095 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1096 event_rx,
1097 callbacks,
1098 None,
1099 writer,
1100 vec![],
1101 State::new(),
1102 None,
1103 None,
1104 None,
1105 None,
1106 None,
1107 std::collections::HashMap::new(),
1108 ControlPlaneConfig::default(),
1109 dummy_event_tx(),
1110 );
1111
1112 let _ = event_tx.send(SessionEvent::AudioData(Bytes::from_static(b"audio1")));
1114 let _ = event_tx.send(SessionEvent::AudioData(Bytes::from_static(b"audio2")));
1115
1116 tokio::time::sleep(Duration::from_millis(50)).await;
1118
1119 assert_eq!(count.load(Ordering::SeqCst), 2);
1120
1121 drop(event_tx);
1123 let _ = fast_handle.await;
1124 let _ = ctrl_handle.await;
1125 }
1126
1127 #[tokio::test]
1128 async fn media_chunks_reach_on_media_not_on_audio() {
1129 let count = Arc::new(AtomicUsize::new(0));
1130 let count_clone = count.clone();
1131
1132 let callbacks = EventCallbacks {
1133 on_media: Some(Box::new(
1134 move |media: &gemini_genai_rs::session::InlineMedia| {
1135 assert!(media.is_video());
1136 count_clone.fetch_add(1, Ordering::SeqCst);
1137 },
1138 )),
1139 ..Default::default()
1140 };
1141 let callbacks = Arc::new(callbacks);
1142
1143 let (event_tx, _) = broadcast::channel(16);
1144 let event_rx = event_tx.subscribe();
1145
1146 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1147
1148 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1149 event_rx,
1150 callbacks,
1151 None,
1152 writer,
1153 vec![],
1154 State::new(),
1155 None,
1156 None,
1157 None,
1158 None,
1159 None,
1160 std::collections::HashMap::new(),
1161 ControlPlaneConfig::default(),
1162 dummy_event_tx(),
1163 );
1164
1165 for _ in 0..2 {
1167 let _ = event_tx.send(SessionEvent::Media(gemini_genai_rs::session::InlineMedia {
1168 mime_type: "video/mp4".into(),
1169 data: Bytes::from_static(b"mp4"),
1170 }));
1171 }
1172
1173 tokio::time::sleep(Duration::from_millis(50)).await;
1175
1176 assert_eq!(count.load(Ordering::SeqCst), 2);
1177
1178 drop(event_tx);
1180 let _ = fast_handle.await;
1181 let _ = ctrl_handle.await;
1182 }
1183
1184 #[tokio::test]
1185 async fn interrupt_suppresses_audio() {
1186 let count = Arc::new(AtomicUsize::new(0));
1187 let count_clone = count.clone();
1188
1189 let callbacks = EventCallbacks {
1190 on_audio: Some(Box::new(move |_| {
1191 count_clone.fetch_add(1, Ordering::SeqCst);
1192 })),
1193 ..Default::default()
1194 };
1195 let callbacks = Arc::new(callbacks);
1196
1197 let (event_tx, _) = broadcast::channel(16);
1198 let event_rx = event_tx.subscribe();
1199
1200 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1201
1202 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1203 event_rx,
1204 callbacks,
1205 None,
1206 writer,
1207 vec![],
1208 State::new(),
1209 None,
1210 None,
1211 None,
1212 None,
1213 None,
1214 std::collections::HashMap::new(),
1215 ControlPlaneConfig::default(),
1216 dummy_event_tx(),
1217 );
1218
1219 let _ = event_tx.send(SessionEvent::AudioData(Bytes::from_static(b"before")));
1221 tokio::time::sleep(Duration::from_millis(20)).await;
1222 let _ = event_tx.send(SessionEvent::Interrupted);
1223 tokio::time::sleep(Duration::from_millis(20)).await;
1224 let _ = event_tx.send(SessionEvent::AudioData(Bytes::from_static(b"during")));
1225 tokio::time::sleep(Duration::from_millis(50)).await;
1226
1227 assert!(count.load(Ordering::SeqCst) >= 1);
1229
1230 drop(event_tx);
1231 let _ = fast_handle.await;
1232 let _ = ctrl_handle.await;
1233 }
1234
1235 #[tokio::test]
1236 async fn control_lane_routes_turn_complete() {
1237 let called = Arc::new(AtomicBool::new(false));
1238 let called_clone = called.clone();
1239
1240 let callbacks = EventCallbacks {
1241 on_turn_complete: Some(Arc::new(move || {
1242 let c = called_clone.clone();
1243 Box::pin(async move {
1244 c.store(true, Ordering::SeqCst);
1245 })
1246 })),
1247 ..Default::default()
1248 };
1249 let callbacks = Arc::new(callbacks);
1250
1251 let (event_tx, _) = broadcast::channel(16);
1252 let event_rx = event_tx.subscribe();
1253
1254 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1255
1256 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1257 event_rx,
1258 callbacks,
1259 None,
1260 writer,
1261 vec![],
1262 State::new(),
1263 None,
1264 None,
1265 None,
1266 None,
1267 None,
1268 std::collections::HashMap::new(),
1269 ControlPlaneConfig::default(),
1270 dummy_event_tx(),
1271 );
1272
1273 let _ = event_tx.send(SessionEvent::TurnComplete);
1274 tokio::time::sleep(Duration::from_millis(50)).await;
1275
1276 assert!(called.load(Ordering::SeqCst));
1277
1278 drop(event_tx);
1279 let _ = fast_handle.await;
1280 let _ = ctrl_handle.await;
1281 }
1282
1283 #[tokio::test]
1284 async fn transcript_accumulates_in_control_lane() {
1285 let callbacks = Arc::new(EventCallbacks::default());
1286
1287 let (event_tx, _) = broadcast::channel(16);
1288 let event_rx = event_tx.subscribe();
1289
1290 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1291
1292 let state = State::new();
1293 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1294 event_rx,
1295 callbacks,
1296 None,
1297 writer,
1298 vec![],
1299 state.clone(),
1300 None,
1301 None,
1302 None,
1303 None,
1304 None,
1305 std::collections::HashMap::new(),
1306 ControlPlaneConfig::default(),
1307 dummy_event_tx(),
1308 );
1309
1310 let _ = event_tx.send(SessionEvent::InputTranscription("Hello ".to_string()));
1312 let _ = event_tx.send(SessionEvent::InputTranscription("world".to_string()));
1313 let _ = event_tx.send(SessionEvent::OutputTranscription("Hi there!".to_string()));
1314 tokio::time::sleep(Duration::from_millis(50)).await;
1315
1316 let _ = event_tx.send(SessionEvent::TurnComplete);
1318 tokio::time::sleep(Duration::from_millis(50)).await;
1319
1320 let tc: u32 = state.session().get("turn_count").unwrap_or(0);
1322 assert_eq!(tc, 1);
1323
1324 drop(event_tx);
1325 let _ = fast_handle.await;
1326 let _ = ctrl_handle.await;
1327 }
1328
1329 #[tokio::test]
1330 async fn an_interrupted_turn_keeps_only_what_was_heard() {
1331 let finals = Arc::new(parking_lot::Mutex::new(Vec::<String>::new()));
1332 let sink = finals.clone();
1333 let callbacks = Arc::new(EventCallbacks {
1334 on_output_transcript: Some(Box::new(move |text: &str, is_final: bool| {
1335 if is_final {
1336 sink.lock().push(text.to_string());
1337 }
1338 })),
1339 ..EventCallbacks::default()
1340 });
1341 let (event_tx, _) = broadcast::channel(64);
1342 let event_rx = event_tx.subscribe();
1343 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1344 let clock = Arc::new(crate::clock::ManualClock::new());
1345 let control_plane = ControlPlaneConfig {
1346 playback: crate::live::playback::PlaybackClock::new(clock.clone()),
1347 ..ControlPlaneConfig::default()
1348 };
1349 let playback = control_plane.playback.clone();
1350 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1351 event_rx,
1352 callbacks,
1353 None,
1354 writer,
1355 vec![],
1356 State::new(),
1357 None,
1358 None,
1359 None,
1360 None,
1361 None,
1362 std::collections::HashMap::new(),
1363 control_plane,
1364 dummy_event_tx(),
1365 );
1366
1367 let _ = event_tx.send(SessionEvent::AudioData(Bytes::from(vec![0u8; 4 * 48_000])));
1370 let _ = event_tx.send(SessionEvent::OutputTranscription(
1371 "Your table for four is booked for Friday at eight, and I have noted the allergy."
1372 .to_string(),
1373 ));
1374 tokio::time::sleep(Duration::from_millis(50)).await;
1375 playback.queued(Duration::from_secs(4));
1376 clock.advance(Duration::from_millis(1_500));
1378 let _ = event_tx.send(SessionEvent::Interrupted);
1379 let _ = event_tx.send(SessionEvent::TurnComplete);
1380 tokio::time::sleep(Duration::from_millis(50)).await;
1381
1382 assert_eq!(*finals.lock(), ["Your table for four is"]);
1383
1384 drop(event_tx);
1385 let _ = fast_handle.await;
1386 let _ = ctrl_handle.await;
1387 }
1388
1389 #[tokio::test]
1390 async fn extractor_runs_on_turn_complete() {
1391 use crate::live::extractor::TurnExtractor;
1392 use crate::live::transcript::TranscriptTurn;
1393 use crate::llm::LlmError;
1394
1395 struct FixedExtractor;
1396
1397 #[async_trait::async_trait]
1398 impl TurnExtractor for FixedExtractor {
1399 fn name(&self) -> &str {
1400 "TestExtractor"
1401 }
1402 fn window_size(&self) -> usize {
1403 3
1404 }
1405 async fn extract(
1406 &self,
1407 _turns: &[TranscriptTurn],
1408 ) -> Result<serde_json::Value, LlmError> {
1409 Ok(serde_json::json!({"score": 0.9, "mood": "happy"}))
1410 }
1411 }
1412
1413 let callbacks = Arc::new(EventCallbacks::default());
1414
1415 let (event_tx, _) = broadcast::channel(16);
1416 let event_rx = event_tx.subscribe();
1417
1418 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1419
1420 let state = State::new();
1421
1422 let extractors: Vec<Arc<dyn TurnExtractor>> = vec![Arc::new(FixedExtractor)];
1423
1424 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1425 event_rx,
1426 callbacks,
1427 None,
1428 writer,
1429 extractors,
1430 state.clone(),
1431 None,
1432 None,
1433 None,
1434 None,
1435 None,
1436 std::collections::HashMap::new(),
1437 ControlPlaneConfig::default(),
1438 dummy_event_tx(),
1439 );
1440
1441 let _ = event_tx.send(SessionEvent::InputTranscription("hi".to_string()));
1443 tokio::time::sleep(Duration::from_millis(20)).await;
1444 let _ = event_tx.send(SessionEvent::TurnComplete);
1445 tokio::time::sleep(Duration::from_millis(100)).await;
1446
1447 let score: Option<f64> = state.get("score");
1449 assert_eq!(score, Some(0.9));
1450 let mood: Option<String> = state.get("mood");
1451 assert_eq!(mood, Some("happy".to_string()));
1452
1453 drop(event_tx);
1454 let _ = fast_handle.await;
1455 let _ = ctrl_handle.await;
1456 }
1457
1458 #[tokio::test]
1467 async fn control_sender_records_text_on_transcript() {
1468 use crate::live::extractor::TurnExtractor;
1469 use crate::live::transcript::TranscriptTurn;
1470 use crate::llm::LlmError;
1471
1472 struct WindowRecorder(Arc<parking_lot::Mutex<Vec<String>>>);
1474
1475 #[async_trait::async_trait]
1476 impl TurnExtractor for WindowRecorder {
1477 fn name(&self) -> &str {
1478 "WindowRecorder"
1479 }
1480 fn window_size(&self) -> usize {
1481 3
1482 }
1483 async fn extract(
1484 &self,
1485 turns: &[TranscriptTurn],
1486 ) -> Result<serde_json::Value, LlmError> {
1487 self.0.lock().extend(turns.iter().map(|t| t.user.clone()));
1488 Ok(serde_json::json!({}))
1489 }
1490 }
1491
1492 let seen = Arc::new(parking_lot::Mutex::new(Vec::new()));
1493
1494 let (event_tx, _) = broadcast::channel(16);
1495 let event_rx = event_tx.subscribe();
1496 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1497 let extractors: Vec<Arc<dyn TurnExtractor>> = vec![Arc::new(WindowRecorder(seen.clone()))];
1498
1499 let (fast_handle, ctrl_handle, ctrl_tx) = spawn_event_processor(
1500 event_rx,
1501 Arc::new(EventCallbacks::default()),
1502 None,
1503 writer,
1504 extractors,
1505 State::new(),
1506 None,
1507 None,
1508 None,
1509 None,
1510 None,
1511 std::collections::HashMap::new(),
1512 ControlPlaneConfig::default(),
1513 dummy_event_tx(),
1514 );
1515
1516 let tx = ctrl_tx.upgrade().expect("control lane is alive");
1518 tx.send(ControlEvent::InputTranscript("I am pescatarian".into()))
1519 .await
1520 .expect("control lane accepts the typed turn");
1521 drop(tx);
1522
1523 tokio::time::sleep(Duration::from_millis(20)).await;
1524 let _ = event_tx.send(SessionEvent::TurnComplete);
1525 tokio::time::sleep(Duration::from_millis(100)).await;
1526
1527 assert_eq!(
1528 seen.lock().as_slice(),
1529 ["I am pescatarian"],
1530 "typed turn must reach the extractor's transcript window"
1531 );
1532
1533 drop(event_tx);
1534 let _ = fast_handle.await;
1535 let _ = ctrl_handle.await;
1536 }
1537
1538 #[tokio::test]
1542 async fn control_sender_does_not_keep_lane_alive() {
1543 let (event_tx, _) = broadcast::channel(16);
1544 let event_rx = event_tx.subscribe();
1545 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1546
1547 let (fast_handle, ctrl_handle, ctrl_tx) = spawn_event_processor(
1548 event_rx,
1549 Arc::new(EventCallbacks::default()),
1550 None,
1551 writer,
1552 vec![],
1553 State::new(),
1554 None,
1555 None,
1556 None,
1557 None,
1558 None,
1559 std::collections::HashMap::new(),
1560 ControlPlaneConfig::default(),
1561 dummy_event_tx(),
1562 );
1563
1564 drop(event_tx);
1566
1567 tokio::time::timeout(Duration::from_secs(2), ctrl_handle)
1568 .await
1569 .expect("control lane must shut down while the handle's sender is held")
1570 .expect("control lane joined cleanly");
1571
1572 assert!(
1573 ctrl_tx.upgrade().is_none(),
1574 "control channel must be closed once the router is gone"
1575 );
1576
1577 let _ = fast_handle.await;
1578 }
1579
1580 #[tokio::test]
1581 async fn telemetry_lane_auto_collects() {
1582 let (event_tx, _) = broadcast::channel(16);
1583 let telem_rx = event_tx.subscribe();
1584
1585 let telemetry = Arc::new(SessionTelemetry::new());
1586 let signals = SessionSignals::new(State::new());
1587 let cancel = CancellationToken::new();
1588
1589 let telem_handle =
1590 spawn_telemetry_lane(telem_rx, signals, telemetry.clone(), cancel.clone(), None);
1591
1592 let _ = event_tx.send(SessionEvent::AudioData(Bytes::from_static(b"chunk1")));
1594 let _ = event_tx.send(SessionEvent::AudioData(Bytes::from_static(b"chunk2")));
1595 let _ = event_tx.send(SessionEvent::VoiceActivityEnd);
1596 tokio::time::sleep(Duration::from_millis(50)).await;
1597 let _ = event_tx.send(SessionEvent::AudioData(Bytes::from_static(b"response")));
1598 tokio::time::sleep(Duration::from_millis(50)).await;
1599
1600 let snap = telemetry.snapshot();
1601 assert_eq!(snap["audio_chunks_out"], 3);
1602 assert!(snap["response_count"].as_u64().unwrap() >= 1);
1603
1604 cancel.cancel();
1605 let _ = telem_handle.await;
1606 }
1607
1608 #[tokio::test]
1609 async fn background_tool_sends_ack_immediately() {
1610 use crate::live::background_tool::{BackgroundToolTracker, ToolExecutionMode};
1611 use crate::tool::{SimpleTool, ToolDispatcher};
1612
1613 let tool = SimpleTool::new(
1615 "slow_search",
1616 "A slow search tool",
1617 Some(serde_json::json!({"type": "object", "properties": {"q": {"type": "string"}}})),
1618 |_args| async move {
1619 tokio::time::sleep(Duration::from_millis(200)).await;
1620 Ok(serde_json::json!({"results": ["found"]}))
1621 },
1622 );
1623
1624 let mut dispatcher = ToolDispatcher::new();
1625 dispatcher.register(tool);
1626
1627 let mut execution_modes = std::collections::HashMap::new();
1628 execution_modes.insert(
1629 "slow_search".to_string(),
1630 ToolExecutionMode::Background {
1631 formatter: None,
1632 scheduling: None,
1633 },
1634 );
1635
1636 let sent = Arc::new(parking_lot::Mutex::new(Vec::<Vec<FunctionResponse>>::new()));
1637 let sent_clone = sent.clone();
1638
1639 struct RecordingWriter {
1641 sent: Arc<parking_lot::Mutex<Vec<Vec<FunctionResponse>>>>,
1642 }
1643
1644 #[async_trait::async_trait]
1645 impl SessionWriter for RecordingWriter {
1646 async fn send_audio(
1647 &self,
1648 _data: bytes::Bytes,
1649 ) -> Result<(), gemini_genai_rs::session::SessionError> {
1650 Ok(())
1651 }
1652 async fn send_text(
1653 &self,
1654 _text: String,
1655 ) -> Result<(), gemini_genai_rs::session::SessionError> {
1656 Ok(())
1657 }
1658 async fn send_video(
1659 &self,
1660 _data: bytes::Bytes,
1661 ) -> Result<(), gemini_genai_rs::session::SessionError> {
1662 Ok(())
1663 }
1664 async fn send_tool_response(
1665 &self,
1666 responses: Vec<FunctionResponse>,
1667 ) -> Result<(), gemini_genai_rs::session::SessionError> {
1668 self.sent.lock().push(responses);
1669 Ok(())
1670 }
1671 async fn update_instruction(
1672 &self,
1673 _instruction: String,
1674 ) -> Result<(), gemini_genai_rs::session::SessionError> {
1675 Ok(())
1676 }
1677 async fn send_client_content(
1678 &self,
1679 _content: Vec<gemini_genai_rs::prelude::Content>,
1680 _turn_complete: bool,
1681 ) -> Result<(), gemini_genai_rs::session::SessionError> {
1682 Ok(())
1683 }
1684 async fn signal_activity_start(
1685 &self,
1686 ) -> Result<(), gemini_genai_rs::session::SessionError> {
1687 Ok(())
1688 }
1689 async fn signal_activity_end(
1690 &self,
1691 ) -> Result<(), gemini_genai_rs::session::SessionError> {
1692 Ok(())
1693 }
1694 async fn disconnect(&self) -> Result<(), gemini_genai_rs::session::SessionError> {
1695 Ok(())
1696 }
1697 }
1698
1699 let writer: Arc<dyn SessionWriter> = Arc::new(RecordingWriter { sent: sent_clone });
1700 let callbacks = Arc::new(EventCallbacks::default());
1701 let tracker = Arc::new(BackgroundToolTracker::new());
1702
1703 let (event_tx, _) = broadcast::channel(16);
1704 let event_rx = event_tx.subscribe();
1705
1706 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1707 event_rx,
1708 callbacks,
1709 Some(Arc::new(dispatcher)),
1710 writer,
1711 vec![],
1712 State::new(),
1713 None,
1714 None,
1715 None,
1716 None,
1717 Some(tracker.clone()),
1718 execution_modes,
1719 ControlPlaneConfig::default(),
1720 dummy_event_tx(),
1721 );
1722
1723 let _ = event_tx.send(SessionEvent::ToolCall(vec![
1725 gemini_genai_rs::prelude::FunctionCall {
1726 name: "slow_search".to_string(),
1727 args: serde_json::json!({"q": "test"}),
1728 id: Some("fc_1".to_string()),
1729 },
1730 ]));
1731
1732 tokio::time::sleep(Duration::from_millis(50)).await;
1734
1735 {
1737 let responses = sent.lock();
1738 assert!(!responses.is_empty(), "Should have sent ack immediately");
1740 assert_eq!(responses[0][0].response["status"], "running");
1741 }
1742
1743 tokio::time::sleep(Duration::from_millis(300)).await;
1745
1746 {
1747 let responses = sent.lock();
1748 assert!(
1750 responses.len() >= 2,
1751 "Should have sent result after completion"
1752 );
1753 assert_eq!(responses[1][0].response["status"], "completed");
1754 }
1755
1756 drop(event_tx);
1757 let _ = fast_handle.await;
1758 let _ = ctrl_handle.await;
1759 }
1760
1761 #[tokio::test]
1762 async fn callback_mode_blocking_awaits_inline() {
1763 use crate::live::ExecutionMode;
1764 use std::sync::atomic::AtomicU32;
1765
1766 let order = Arc::new(AtomicU32::new(0));
1767 let order_clone = order.clone();
1768
1769 let callbacks = EventCallbacks {
1770 on_turn_complete: Some(Arc::new(move || {
1772 let o = order_clone.clone();
1773 Box::pin(async move {
1774 tokio::time::sleep(Duration::from_millis(10)).await;
1776 o.store(1, Ordering::SeqCst);
1777 })
1778 })),
1779 on_turn_complete_mode: ExecutionMode::Blocking,
1780 ..Default::default()
1781 };
1782 let callbacks = Arc::new(callbacks);
1783
1784 let (event_tx, _) = broadcast::channel(16);
1785 let event_rx = event_tx.subscribe();
1786
1787 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1788
1789 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1790 event_rx,
1791 callbacks,
1792 None,
1793 writer,
1794 vec![],
1795 State::new(),
1796 None,
1797 None,
1798 None,
1799 None,
1800 None,
1801 std::collections::HashMap::new(),
1802 ControlPlaneConfig::default(),
1803 dummy_event_tx(),
1804 );
1805
1806 let _ = event_tx.send(SessionEvent::TurnComplete);
1807 tokio::time::sleep(Duration::from_millis(100)).await;
1808
1809 assert_eq!(order.load(Ordering::SeqCst), 1);
1811
1812 drop(event_tx);
1813 let _ = fast_handle.await;
1814 let _ = ctrl_handle.await;
1815 }
1816
1817 #[tokio::test]
1818 async fn interruption_beats_slow_inline_tool() {
1819 use crate::tool::{SimpleTool, ToolDispatcher};
1820
1821 let mut dispatcher = ToolDispatcher::new();
1823 dispatcher.register(SimpleTool::new("slow", "slow", None, |_args| async move {
1824 tokio::time::sleep(Duration::from_secs(5)).await;
1825 Ok(serde_json::json!({"done": true}))
1826 }));
1827
1828 let interrupted_at = Arc::new(parking_lot::Mutex::new(None::<std::time::Instant>));
1829 let flag = interrupted_at.clone();
1830 let callbacks = EventCallbacks {
1831 on_interrupted: Some(Arc::new(move || {
1832 let flag = flag.clone();
1833 Box::pin(async move {
1834 *flag.lock() = Some(std::time::Instant::now());
1835 })
1836 })),
1837 ..Default::default()
1838 };
1839
1840 let (event_tx, _) = broadcast::channel(16);
1841 let event_rx = event_tx.subscribe();
1842 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1843
1844 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1845 event_rx,
1846 Arc::new(callbacks),
1847 Some(Arc::new(dispatcher)),
1848 writer,
1849 vec![],
1850 State::new(),
1851 None,
1852 None,
1853 None,
1854 None,
1855 None,
1856 std::collections::HashMap::new(),
1857 ControlPlaneConfig::default(),
1858 dummy_event_tx(),
1859 );
1860
1861 let start = std::time::Instant::now();
1863 let _ = event_tx.send(SessionEvent::ToolCall(vec![
1864 gemini_genai_rs::prelude::FunctionCall {
1865 name: "slow".to_string(),
1866 args: serde_json::json!({}),
1867 id: Some("fc_slow".to_string()),
1868 },
1869 ]));
1870 tokio::time::sleep(Duration::from_millis(50)).await;
1871 let _ = event_tx.send(SessionEvent::Interrupted);
1872
1873 let mut waited = Duration::ZERO;
1876 while interrupted_at.lock().is_none() && waited < Duration::from_secs(2) {
1877 tokio::time::sleep(Duration::from_millis(25)).await;
1878 waited += Duration::from_millis(25);
1879 }
1880 let fired = (*interrupted_at.lock()).expect("on_interrupted must fire");
1881 assert!(
1882 fired.duration_since(start) < Duration::from_secs(2),
1883 "interruption must not wait for the slow tool"
1884 );
1885
1886 drop(event_tx);
1887 let _ = fast_handle.await;
1888 let _ = ctrl_handle.await;
1889 }
1890
1891 #[tokio::test]
1892 async fn control_lane_exit_persists_final_snapshot_synchronously() {
1893 use crate::live::persistence::{MemoryPersistence, SessionPersistence};
1894
1895 let persistence = Arc::new(MemoryPersistence::new());
1896 let control_plane = ControlPlaneConfig {
1897 persistence: Some(persistence.clone()),
1898 session_id: Some("final-drain".to_string()),
1899 ..Default::default()
1900 };
1901
1902 let (event_tx, _) = broadcast::channel(16);
1903 let event_rx = event_tx.subscribe();
1904 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1905
1906 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1907 event_rx,
1908 Arc::new(EventCallbacks::default()),
1909 None,
1910 writer,
1911 vec![],
1912 State::new(),
1913 None,
1914 None,
1915 None,
1916 None,
1917 None,
1918 std::collections::HashMap::new(),
1919 control_plane,
1920 dummy_event_tx(),
1921 );
1922
1923 let _ = event_tx.send(SessionEvent::InputTranscription("last words".to_string()));
1926 tokio::time::sleep(Duration::from_millis(50)).await;
1927
1928 drop(event_tx);
1931 let _ = fast_handle.await;
1932 let _ = ctrl_handle.await;
1933
1934 let snap = persistence
1935 .load("final-drain")
1936 .await
1937 .unwrap()
1938 .expect("control-lane exit must persist a final snapshot");
1939 assert_eq!(snap.turn_count, 0);
1940 }
1941
1942 #[tokio::test]
1943 async fn lanes_exit_after_terminal_disconnected_event() {
1944 let callbacks = Arc::new(EventCallbacks::default());
1949 let (event_tx, _) = broadcast::channel(16);
1950 let event_rx = event_tx.subscribe();
1951 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
1952
1953 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
1954 event_rx,
1955 callbacks,
1956 None,
1957 writer,
1958 vec![],
1959 State::new(),
1960 None,
1961 None,
1962 None,
1963 None,
1964 None,
1965 std::collections::HashMap::new(),
1966 ControlPlaneConfig::default(),
1967 dummy_event_tx(),
1968 );
1969
1970 let _ = event_tx.send(SessionEvent::Disconnected(None));
1971
1972 let joined = tokio::time::timeout(Duration::from_secs(2), async {
1975 let _ = fast_handle.await;
1976 let _ = ctrl_handle.await;
1977 })
1978 .await;
1979 assert!(
1980 joined.is_ok(),
1981 "lanes must exit after the terminal Disconnected event"
1982 );
1983 drop(event_tx);
1984 }
1985
1986 #[tokio::test]
1987 async fn callback_mode_concurrent_spawns_task() {
1988 use crate::live::ExecutionMode;
1989
1990 let called = Arc::new(AtomicBool::new(false));
1991 let called_clone = called.clone();
1992
1993 let callbacks = EventCallbacks {
1994 on_turn_complete: Some(Arc::new(move || {
1995 let c = called_clone.clone();
1996 Box::pin(async move {
1997 tokio::time::sleep(Duration::from_millis(10)).await;
1998 c.store(true, Ordering::SeqCst);
1999 })
2000 })),
2001 on_turn_complete_mode: ExecutionMode::Concurrent,
2002 ..Default::default()
2003 };
2004 let callbacks = Arc::new(callbacks);
2005
2006 let (event_tx, _) = broadcast::channel(16);
2007 let event_rx = event_tx.subscribe();
2008
2009 let writer: Arc<dyn SessionWriter> = Arc::new(crate::agent_session::NoOpSessionWriter);
2010
2011 let (fast_handle, ctrl_handle, _ctrl_tx) = spawn_event_processor(
2012 event_rx,
2013 callbacks,
2014 None,
2015 writer,
2016 vec![],
2017 State::new(),
2018 None,
2019 None,
2020 None,
2021 None,
2022 None,
2023 std::collections::HashMap::new(),
2024 ControlPlaneConfig::default(),
2025 dummy_event_tx(),
2026 );
2027
2028 let _ = event_tx.send(SessionEvent::TurnComplete);
2029 tokio::time::sleep(Duration::from_millis(100)).await;
2031
2032 assert!(called.load(Ordering::SeqCst));
2034
2035 drop(event_tx);
2036 let _ = fast_handle.await;
2037 let _ = ctrl_handle.await;
2038 }
2039}