1use std::collections::HashMap;
4use std::sync::Arc;
5
6use tokio_util::sync::CancellationToken;
7
8use gemini_genai_rs::prelude::{ConnectBuilder, SessionConfig, SessionEvent, SessionPhase};
9use gemini_genai_rs::session::{SessionError, SessionHandle, SessionWriter};
10
11use crate::error::AgentError;
12use crate::state::State;
13use crate::tool::ToolDispatcher;
14
15use super::background_tool::{BackgroundToolTracker, ToolExecutionMode};
16use super::callbacks::EventCallbacks;
17use super::computed::ComputedRegistry;
18use super::context_writer::{DeferredWriter, PendingContext};
19use super::extractor::TurnExtractor;
20use super::handle::LiveHandle;
21use super::needs::{NeedsFulfillment, RepairConfig};
22use super::persistence::SessionPersistence;
23use super::phase::PhaseMachine;
24use super::processor::{spawn_event_processor, spawn_telemetry_lane, ControlPlaneConfig};
25use super::session_signals::SessionSignals;
26use super::soft_turn::SoftTurnDetector;
27use super::steering::{ContextDelivery, SteeringMode};
28use super::telemetry::SessionTelemetry;
29use super::temporal::TemporalRegistry;
30use super::watcher::WatcherRegistry;
31
32pub struct LiveSessionBuilder {
42 config: SessionConfig,
43 callbacks: EventCallbacks,
44 dispatcher: Option<Arc<ToolDispatcher>>,
45 extractors: Vec<Arc<dyn TurnExtractor>>,
46 computed: Option<ComputedRegistry>,
47 phase_machine: Option<PhaseMachine>,
48 watchers: Option<WatcherRegistry>,
49 temporal: Option<TemporalRegistry>,
50 greeting: Option<String>,
51 state: Option<State>,
52 execution_modes: HashMap<String, ToolExecutionMode>,
53 soft_turn_timeout: Option<std::time::Duration>,
55 steering_mode: SteeringMode,
56 context_delivery: ContextDelivery,
57 delivery: super::processor::DeliveryConfig,
58 repair_config: Option<RepairConfig>,
59 persistence: Option<Arc<dyn SessionPersistence>>,
60 session_id: Option<String>,
61 tool_advisory: bool,
62 telemetry_interval: Option<std::time::Duration>,
63 middleware: Vec<Arc<dyn crate::middleware::Middleware>>,
64 flow: Option<crate::flow::FlowMonitor>,
65}
66
67impl LiveSessionBuilder {
68 pub fn new(config: SessionConfig) -> Self {
70 Self {
71 config,
72 callbacks: EventCallbacks::default(),
73 dispatcher: None,
74 extractors: Vec::new(),
75 computed: None,
76 phase_machine: None,
77 watchers: None,
78 temporal: None,
79 greeting: None,
80 state: None,
81 execution_modes: HashMap::new(),
82 soft_turn_timeout: None,
83 steering_mode: SteeringMode::default(),
84 context_delivery: ContextDelivery::default(),
85 delivery: super::processor::DeliveryConfig::default(),
86 repair_config: None,
87 persistence: None,
88 session_id: None,
89 tool_advisory: true,
90 telemetry_interval: None,
91 middleware: Vec::new(),
92 flow: None,
93 }
94 }
95
96 pub fn middleware(mut self, layer: Arc<dyn crate::middleware::Middleware>) -> Self {
102 self.middleware.push(layer);
103 self
104 }
105
106 pub fn flow_monitor(mut self, monitor: crate::flow::FlowMonitor) -> Self {
108 self.flow = Some(monitor);
109 self
110 }
111
112 pub fn with_state(mut self, state: State) -> Self {
118 self.state = Some(state);
119 self
120 }
121
122 pub fn greeting(mut self, prompt: impl Into<String>) -> Self {
124 self.greeting = Some(prompt.into());
125 self
126 }
127
128 pub fn dispatcher(mut self, dispatcher: ToolDispatcher) -> Self {
130 for tool in dispatcher.to_tool_declarations() {
132 self.config = self.config.add_tool(tool);
133 }
134 self.dispatcher = Some(Arc::new(dispatcher));
135 self
136 }
137
138 pub fn callbacks(mut self, callbacks: EventCallbacks) -> Self {
140 self.callbacks = callbacks;
141 self
142 }
143
144 pub fn extractor(mut self, extractor: Arc<dyn TurnExtractor>) -> Self {
146 self.extractors.push(extractor);
147 self
148 }
149
150 pub fn computed(mut self, registry: ComputedRegistry) -> Self {
152 self.computed = Some(registry);
153 self
154 }
155
156 pub fn phase_machine(mut self, machine: PhaseMachine) -> Self {
158 self.phase_machine = Some(machine);
159 self
160 }
161
162 pub fn watchers(mut self, registry: WatcherRegistry) -> Self {
164 self.watchers = Some(registry);
165 self
166 }
167
168 pub fn temporal(mut self, registry: TemporalRegistry) -> Self {
170 self.temporal = Some(registry);
171 self
172 }
173
174 pub fn tool_execution_mode(
179 mut self,
180 tool_name: impl Into<String>,
181 mode: ToolExecutionMode,
182 ) -> Self {
183 self.execution_modes.insert(tool_name.into(), mode);
184 self
185 }
186
187 pub fn soft_turn_timeout(mut self, timeout: std::time::Duration) -> Self {
194 self.soft_turn_timeout = Some(timeout);
195 self
196 }
197
198 pub fn steering_mode(mut self, mode: SteeringMode) -> Self {
200 self.steering_mode = mode;
201 self
202 }
203
204 pub fn context_delivery(mut self, mode: ContextDelivery) -> Self {
209 self.context_delivery = mode;
210 self
211 }
212
213 pub fn delivery(mut self, delivery: super::processor::DeliveryConfig) -> Self {
220 self.delivery = delivery;
221 self
222 }
223
224 pub fn repair(mut self, config: RepairConfig) -> Self {
229 self.repair_config = Some(config);
230 self
231 }
232
233 pub fn persistence(mut self, backend: Arc<dyn SessionPersistence>) -> Self {
235 self.persistence = Some(backend);
236 self
237 }
238
239 pub fn session_id(mut self, id: impl Into<String>) -> Self {
241 self.session_id = Some(id.into());
242 self
243 }
244
245 pub fn tool_advisory(mut self, enabled: bool) -> Self {
247 self.tool_advisory = enabled;
248 self
249 }
250
251 pub fn telemetry_interval(mut self, interval: std::time::Duration) -> Self {
256 self.telemetry_interval = Some(interval);
257 self
258 }
259
260 pub async fn connect(self) -> Result<LiveHandle, AgentError> {
273 let mut plan = self.into_plan()?;
274
275 let config = plan.config.take().expect("plan always carries a config");
278 let session = ConnectBuilder::new(config)
279 .build()
280 .await
281 .map_err(AgentError::Session)?;
282
283 let mut events = session.subscribe();
293 tokio::select! {
294 () = session.wait_for_phase(SessionPhase::Active) => {}
295 failure = wait_for_connect_failure(&mut events) => {
296 return Err(AgentError::Session(SessionError::SetupFailed(
297 gemini_genai_rs::session::SetupError::ServerRejected {
298 code: None,
299 message: failure,
300 },
301 )));
302 }
303 }
304
305 let runtime = build_runtime(plan, session);
306 spawn_lanes(runtime).await
307 }
308
309 pub(crate) fn into_plan(self) -> Result<SessionPlan, AgentError> {
316 if let Some(ref pm) = self.phase_machine {
318 pm.validate().map_err(AgentError::Config)?;
319 }
320 if let Some(ref computed) = self.computed {
321 computed.validate().map_err(AgentError::Config)?;
322 }
323
324 let mut config = self.config;
326 for (tool_name, mode) in &self.execution_modes {
327 if matches!(
328 mode,
329 super::background_tool::ToolExecutionMode::Background { .. }
330 ) {
331 for tool in &mut config.tools {
332 if let Some(ref mut decls) = tool.function_declarations {
333 for decl in decls {
334 if decl.name == *tool_name {
335 decl.behavior = Some(
336 gemini_genai_rs::prelude::FunctionCallingBehavior::NonBlocking,
337 );
338 }
339 }
340 }
341 }
342 }
343 }
344
345 let base_instruction = config.system_instruction.as_ref().map(|content| {
350 content
351 .parts
352 .iter()
353 .filter_map(|part| match part {
354 gemini_genai_rs::prelude::Part::Text { text } => Some(text.as_str()),
355 _ => None,
356 })
357 .collect::<Vec<_>>()
358 .join("\n")
359 });
360
361 Ok(SessionPlan {
362 config: Some(config),
363 base_instruction,
364 callbacks: self.callbacks,
365 dispatcher: self.dispatcher,
366 extractors: self.extractors,
367 computed: self.computed,
368 phase_machine: self.phase_machine,
369 watchers: self.watchers,
370 temporal: self.temporal,
371 greeting: self.greeting,
372 state: self.state,
373 execution_modes: self.execution_modes,
374 soft_turn_timeout: self.soft_turn_timeout,
375 steering_mode: self.steering_mode,
376 context_delivery: self.context_delivery,
377 delivery: self.delivery,
378 repair_config: self.repair_config,
379 persistence: self.persistence,
380 session_id: self.session_id,
381 tool_advisory: self.tool_advisory,
382 telemetry_interval: self.telemetry_interval,
383 middleware: self.middleware,
384 flow: self.flow,
385 })
386 }
387}
388
389pub(crate) struct SessionPlan {
397 config: Option<SessionConfig>,
400 base_instruction: Option<String>,
403 callbacks: EventCallbacks,
404 dispatcher: Option<Arc<ToolDispatcher>>,
405 extractors: Vec<Arc<dyn TurnExtractor>>,
406 computed: Option<ComputedRegistry>,
407 phase_machine: Option<PhaseMachine>,
408 watchers: Option<WatcherRegistry>,
409 temporal: Option<TemporalRegistry>,
410 greeting: Option<String>,
411 state: Option<State>,
412 execution_modes: HashMap<String, ToolExecutionMode>,
413 soft_turn_timeout: Option<std::time::Duration>,
414 steering_mode: SteeringMode,
415 context_delivery: ContextDelivery,
416 delivery: super::processor::DeliveryConfig,
417 repair_config: Option<RepairConfig>,
418 persistence: Option<Arc<dyn SessionPersistence>>,
419 session_id: Option<String>,
420 tool_advisory: bool,
421 telemetry_interval: Option<std::time::Duration>,
422 middleware: Vec<Arc<dyn crate::middleware::Middleware>>,
423 flow: Option<crate::flow::FlowMonitor>,
424}
425
426pub(crate) struct SessionRuntime {
434 session: SessionHandle,
435 callbacks: Arc<EventCallbacks>,
436 dispatcher: Option<Arc<ToolDispatcher>>,
437 extractors: Vec<Arc<dyn TurnExtractor>>,
438 computed: Option<ComputedRegistry>,
439 phase_machine: Option<tokio::sync::Mutex<PhaseMachine>>,
440 watchers: Option<WatcherRegistry>,
441 temporal: Option<Arc<TemporalRegistry>>,
442 greeting: Option<String>,
443 state: State,
444 execution_modes: HashMap<String, ToolExecutionMode>,
445 background_tracker: Arc<BackgroundToolTracker>,
446 telemetry: Arc<SessionTelemetry>,
447 telemetry_interval: Option<std::time::Duration>,
448 control_plane: ControlPlaneConfig,
449 pending_context: Option<Arc<PendingContext>>,
450 writer: Arc<dyn SessionWriter>,
452 user_writer: Arc<dyn SessionWriter>,
454 event_rx: tokio::sync::broadcast::Receiver<gemini_genai_rs::prelude::SessionEvent>,
455 telem_rx: tokio::sync::broadcast::Receiver<gemini_genai_rs::prelude::SessionEvent>,
456 on_usage_cb: Option<super::callbacks::UsageCallback>,
457 live_event_tx: tokio::sync::broadcast::Sender<super::events::LiveEvent>,
458 telem_cancel: CancellationToken,
459 flow_monitor: Option<crate::flow::SharedFlowMonitor>,
460}
461
462pub(crate) fn build_runtime(plan: SessionPlan, session: SessionHandle) -> SessionRuntime {
467 let flow_monitor = plan.flow.map(crate::flow::FlowMonitor::into_shared);
470 let mut callbacks = plan.callbacks;
471 let on_usage_cb = callbacks.on_usage.take();
472 let callbacks = Arc::new(callbacks);
473 let raw_writer: Arc<dyn SessionWriter> = Arc::new(session.clone());
474 let state = plan.state.unwrap_or_default();
475
476 let event_rx = session.subscribe();
478 let telem_rx = session.subscribe();
479
480 if let Some(ref pm) = plan.phase_machine {
482 let _ = state.session().set("phase", pm.current());
483 if let Some(phase) = pm.current_phase() {
484 if !phase.needs.is_empty() {
485 let _ = state.set("session:phase_needs", phase.needs.clone());
486 }
487 }
488 }
489
490 let phase_machine_mutex = plan.phase_machine.map(tokio::sync::Mutex::new);
491 let temporal_arc = plan.temporal.map(Arc::new);
492 let background_tracker = Arc::new(BackgroundToolTracker::new());
493
494 let telemetry = Arc::new(SessionTelemetry::new());
496 let telem_cancel = CancellationToken::new();
497
498 let mut control_plane = ControlPlaneConfig {
500 soft_turn: plan.soft_turn_timeout.map(SoftTurnDetector::new),
501 steering_mode: plan.steering_mode,
502 context_delivery: plan.context_delivery,
503 delivery: plan.delivery,
504 needs_fulfillment: plan.repair_config.map(NeedsFulfillment::new),
505 persistence: plan.persistence,
506 session_id: plan.session_id,
507 tool_advisory: plan.tool_advisory,
508 base_instruction: plan.base_instruction,
509 pending_context: None, middleware: {
511 let mut chain = crate::middleware::MiddlewareChain::new();
512 for layer in plan.middleware {
513 chain.add(layer);
514 }
515 Arc::new(chain)
516 },
517 flow: flow_monitor.clone(),
518 };
519
520 let pending_context = if plan.context_delivery == ContextDelivery::Deferred {
525 Some(Arc::new(PendingContext::new()))
526 } else {
527 None
528 };
529
530 let (writer, user_writer) = if let Some(ref pending) = pending_context {
532 let deferred: Arc<dyn SessionWriter> =
533 Arc::new(DeferredWriter::new(raw_writer.clone(), pending.clone()));
534 (raw_writer, deferred)
538 } else {
539 (raw_writer.clone(), raw_writer)
540 };
541
542 control_plane.pending_context = pending_context.clone();
544
545 use super::events::LiveEvent;
547 use tokio::sync::broadcast;
548 let (live_event_tx, _) = broadcast::channel::<LiveEvent>(4096);
549
550 SessionRuntime {
551 session,
552 callbacks,
553 dispatcher: plan.dispatcher,
554 extractors: plan.extractors,
555 computed: plan.computed,
556 phase_machine: phase_machine_mutex,
557 watchers: plan.watchers,
558 temporal: temporal_arc,
559 greeting: plan.greeting,
560 state,
561 execution_modes: plan.execution_modes,
562 background_tracker,
563 telemetry,
564 telemetry_interval: plan.telemetry_interval,
565 control_plane,
566 pending_context,
567 writer,
568 user_writer,
569 event_rx,
570 telem_rx,
571 on_usage_cb,
572 live_event_tx,
573 telem_cancel,
574 flow_monitor,
575 }
576}
577
578async fn wait_for_connect_failure(
592 events: &mut tokio::sync::broadcast::Receiver<SessionEvent>,
593) -> String {
594 let mut last_error: Option<String> = None;
595 loop {
596 match events.recv().await {
597 Ok(SessionEvent::Error(msg)) => last_error = Some(msg),
598 Ok(SessionEvent::Disconnected(reason)) => {
599 let reason = reason.unwrap_or_else(|| "connection closed".to_string());
600 return match last_error {
601 Some(err) => format!("{reason} ({err})"),
602 None => reason,
603 };
604 }
605 Ok(_) => {}
606 Err(tokio::sync::broadcast::error::RecvError::Closed) => {
608 return last_error.unwrap_or_else(|| "session ended during setup".to_string());
609 }
610 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {}
611 }
612 }
613}
614
615pub(crate) async fn spawn_lanes(rt: SessionRuntime) -> Result<LiveHandle, AgentError> {
616 use super::events::LiveEvent;
617
618 let session_signals = SessionSignals::new(rt.state.clone());
620 let _telem_handle = spawn_telemetry_lane(
621 rt.telem_rx,
622 session_signals,
623 rt.telemetry.clone(),
624 rt.telem_cancel.clone(),
625 rt.on_usage_cb,
626 );
627
628 let greeting_writer = rt.user_writer.clone();
630 let (fast_handle, ctrl_handle, ctrl_tx) = spawn_event_processor(
631 rt.event_rx,
632 rt.callbacks,
633 rt.dispatcher,
634 rt.writer,
635 rt.extractors,
636 rt.state.clone(),
637 rt.computed,
638 rt.phase_machine,
639 rt.watchers,
640 rt.temporal,
641 Some(rt.background_tracker.clone()),
642 rt.execution_modes,
643 rt.control_plane,
644 rt.live_event_tx.clone(),
645 );
646
647 if let Some(interval) = rt.telemetry_interval {
649 let telem_tx = rt.live_event_tx.clone();
650 let telem_ref = rt.telemetry.clone();
651 tokio::spawn(async move {
652 let mut tick = tokio::time::interval(interval);
653 let mut prev_turns = 0u64;
654 loop {
655 tick.tick().await;
656 let snap = telem_ref.snapshot();
657 if let Some(obj) = snap.as_object() {
658 let tc = obj
659 .get("turn_count")
660 .or_else(|| obj.get("response_count"))
661 .and_then(|v| v.as_u64())
662 .unwrap_or(0);
663 if tc > prev_turns {
664 let latency = obj
665 .get("last_response_latency_ms")
666 .and_then(|v| v.as_u64())
667 .unwrap_or(0) as u32;
668 let prompt = obj
669 .get("prompt_token_count")
670 .and_then(|v| v.as_u64())
671 .unwrap_or(0) as u32;
672 let response = obj
673 .get("response_token_count")
674 .and_then(|v| v.as_u64())
675 .unwrap_or(0) as u32;
676 let _ = telem_tx.send(LiveEvent::TurnMetrics {
677 turn: tc as u32,
678 latency_ms: latency,
679 prompt_tokens: prompt,
680 response_tokens: response,
681 });
682 prev_turns = tc;
683 }
684 }
685 if telem_tx.send(LiveEvent::Telemetry(snap)).is_err() {
686 break;
687 }
688 }
689 });
690 }
691
692 if let Some(greeting) = rt.greeting {
694 greeting_writer
695 .send_text(greeting)
696 .await
697 .map_err(AgentError::Session)?;
698 }
699
700 Ok(LiveHandle::new(
701 rt.session,
702 rt.user_writer,
703 fast_handle,
704 ctrl_handle,
705 rt.state,
706 rt.telemetry,
707 rt.live_event_tx,
708 rt.pending_context,
709 rt.flow_monitor,
710 rt.background_tracker,
711 rt.telem_cancel,
712 )
713 .with_control_sender(ctrl_tx))
714}
715
716#[cfg(test)]
717mod tests {
718 use super::*;
719
720 #[test]
721 fn builder_creates_with_defaults() {
722 let config = SessionConfig::new("test-key");
723 let builder = LiveSessionBuilder::new(config);
724 assert!(builder.dispatcher.is_none());
725 assert!(builder.computed.is_none());
726 assert!(builder.phase_machine.is_none());
727 assert!(builder.watchers.is_none());
728 assert!(builder.temporal.is_none());
729 }
730
731 #[test]
732 fn into_plan_derives_defaults() {
733 let config = SessionConfig::new("test-key");
734 let plan = LiveSessionBuilder::new(config)
735 .into_plan()
736 .expect("default builder should produce a plan");
737
738 assert!(plan.config.is_some());
740 assert!(plan.dispatcher.is_none());
742 assert!(plan.phase_machine.is_none());
743 assert!(plan.persistence.is_none());
744 assert!(plan.session_id.is_none());
745 assert!(plan.greeting.is_none());
746 assert!(plan.soft_turn_timeout.is_none());
747 assert!(plan.telemetry_interval.is_none());
748 assert!(plan.repair_config.is_none());
749 assert!(plan.flow.is_none());
750 assert!(plan.execution_modes.is_empty());
751 assert!(plan.middleware.is_empty());
752 assert_eq!(plan.steering_mode, SteeringMode::default());
753 assert_eq!(plan.context_delivery, ContextDelivery::default());
754 assert!(plan.tool_advisory);
756 }
757
758 #[test]
759 fn into_plan_carries_persistence_and_session_id() {
760 let config = SessionConfig::new("test-key");
761 let plan = LiveSessionBuilder::new(config)
762 .session_id("user-123-session-456")
763 .into_plan()
764 .expect("plan derivation should succeed");
765
766 assert_eq!(plan.session_id.as_deref(), Some("user-123-session-456"));
767 }
768
769 #[test]
770 fn into_plan_carries_steering_and_context_delivery() {
771 let config = SessionConfig::new("test-key");
772 let plan = LiveSessionBuilder::new(config)
773 .steering_mode(SteeringMode::ContextInjection)
774 .context_delivery(ContextDelivery::Deferred)
775 .tool_advisory(false)
776 .into_plan()
777 .expect("plan derivation should succeed");
778
779 assert_eq!(plan.steering_mode, SteeringMode::ContextInjection);
780 assert_eq!(plan.context_delivery, ContextDelivery::Deferred);
781 assert!(!plan.tool_advisory);
782 }
783
784 #[test]
785 fn into_plan_carries_greeting_and_telemetry_interval() {
786 let config = SessionConfig::new("test-key");
787 let plan = LiveSessionBuilder::new(config)
788 .greeting("Hello there")
789 .telemetry_interval(std::time::Duration::from_secs(5))
790 .soft_turn_timeout(std::time::Duration::from_secs(2))
791 .into_plan()
792 .expect("plan derivation should succeed");
793
794 assert_eq!(plan.greeting.as_deref(), Some("Hello there"));
795 assert_eq!(
796 plan.telemetry_interval,
797 Some(std::time::Duration::from_secs(5))
798 );
799 assert_eq!(
800 plan.soft_turn_timeout,
801 Some(std::time::Duration::from_secs(2))
802 );
803 }
804
805 #[test]
806 fn into_plan_validates_phase_machine() {
807 let config = SessionConfig::new("test-key");
810 let pm = PhaseMachine::new("nonexistent");
811 let result = LiveSessionBuilder::new(config)
812 .phase_machine(pm)
813 .into_plan();
814 assert!(result.is_err(), "invalid phase machine should fail to plan");
815 }
816
817 #[test]
818 fn into_plan_carries_valid_phase_machine_and_seeds_nothing() {
819 let config = SessionConfig::new("test-key");
822 let mut pm = PhaseMachine::new("start");
823 pm.add_phase(crate::live::phase::Phase::new("start", "Start phase"));
824 let plan = LiveSessionBuilder::new(config)
825 .phase_machine(pm)
826 .into_plan()
827 .expect("valid phase machine should plan");
828
829 assert!(plan.phase_machine.is_some());
830 }
831
832 #[test]
833 fn into_plan_applies_non_blocking_to_background_tools() {
834 use gemini_genai_rs::prelude::{FunctionCallingBehavior, FunctionDeclaration, Tool};
835
836 let decl = FunctionDeclaration {
837 name: "search_kb".into(),
838 description: "Search".into(),
839 parameters: None,
840 behavior: None,
841 };
842 let config = SessionConfig::new("test-key").add_tool(Tool::functions(vec![decl]));
843
844 let plan = LiveSessionBuilder::new(config)
845 .tool_execution_mode(
846 "search_kb",
847 ToolExecutionMode::Background {
848 formatter: None,
849 scheduling: None,
850 },
851 )
852 .into_plan()
853 .expect("plan derivation should succeed");
854
855 let cfg = plan.config.expect("config carried");
856 let decl = cfg.tools[0]
857 .function_declarations
858 .as_ref()
859 .unwrap()
860 .iter()
861 .find(|d| d.name == "search_kb")
862 .unwrap();
863 assert_eq!(decl.behavior, Some(FunctionCallingBehavior::NonBlocking));
864 }
865}