gemini_adk_fluent_rs/live/
connect.rs

1//! Connection methods for `Live`.
2
3use gemini_adk_rs::live::{LiveHandle, LiveSessionBuilder, PhaseMachine};
4use gemini_genai_rs::prelude::*;
5
6use super::Live;
7
8/// Fold builder-registered ambient tools into the flow's own list.
9///
10/// Idempotent and duplicate-free, because an application may name a tool the
11/// flow already declares — or an extension may be installed twice.
12pub(crate) fn merge_ambient(flow: &mut gemini_adk_rs::flow::Flow, ambient: &[String]) {
13    for tool in ambient {
14        if !flow.ambient.contains(tool) {
15            flow.ambient.push(tool.clone());
16        }
17    }
18}
19
20/// The one governance object a session drives: the main monitor plus every
21/// digression and repair policy the builder holds. Ambient tools are merged
22/// into each digression's flow too, so a cross-cutting tool stays admitted
23/// while the main flow is suspended.
24pub(crate) fn assemble_stack(
25    main: gemini_adk_rs::flow::FlowMonitor,
26    digressions: Vec<gemini_adk_rs::flow::Overlay>,
27    repair: std::collections::BTreeMap<String, gemini_adk_rs::flow::RepairPolicy>,
28    timing: std::collections::BTreeMap<String, gemini_adk_rs::flow::VoiceTiming>,
29    corrections: std::collections::BTreeMap<String, Vec<String>>,
30    verbatim: std::collections::BTreeMap<String, String>,
31    ambient: &[String],
32) -> gemini_adk_rs::flow::FlowStack {
33    let overlays = digressions.into_iter().map(|mut ov| {
34        merge_ambient(ov.flow_mut(), ambient);
35        ov
36    });
37    gemini_adk_rs::flow::FlowStack::from_monitor(main)
38        .with_overlays(overlays)
39        .with_repairs(repair)
40        .with_timings(timing)
41        .with_corrections(corrections)
42        .with_verbatims(verbatim)
43}
44
45impl Live {
46    /// Connect using a Google AI API key.
47    pub async fn connect_google_ai(
48        mut self,
49        api_key: impl Into<String>,
50    ) -> Result<LiveHandle, gemini_adk_rs::error::AgentError> {
51        self.config.endpoint = ApiEndpoint::google_ai(api_key);
52        self.build_and_connect().await
53    }
54
55    /// Connect using Vertex AI credentials.
56    ///
57    /// `access_token` is a fixed token or an [`AccessToken`]; pass
58    /// `AccessToken::from_fn(..)` for a source that is consulted on every
59    /// reconnect, since Vertex tokens expire after about an hour.
60    pub async fn connect_vertex(
61        mut self,
62        project: impl Into<String>,
63        location: impl Into<String>,
64        access_token: impl Into<AccessToken>,
65    ) -> Result<LiveHandle, gemini_adk_rs::error::AgentError> {
66        self.config.endpoint = ApiEndpoint::vertex(project, location, access_token);
67        self.build_and_connect().await
68    }
69
70    /// Connect by resolving the platform and credentials from standard
71    /// environment variables — the zero-ceremony entry point.
72    ///
73    /// Resolution (see [`ApiEndpoint::from_env`]):
74    /// - `GOOGLE_GENAI_USE_VERTEXAI=true` → Vertex AI using
75    ///   `GOOGLE_CLOUD_PROJECT`, `GOOGLE_CLOUD_LOCATION` (default
76    ///   `us-central1`), and a token from `GOOGLE_ACCESS_TOKEN`. If that
77    ///   token is unset, the token comes from the metadata server on Cloud
78    ///   Run, GKE or Compute Engine, else from `gcloud auth
79    ///   print-access-token`, and is refreshed before it expires, so a
80    ///   reconnect late in a long session is still authenticated (see
81    ///   [`GoogleAccessToken`](gemini_genai_rs::transport::auth::GoogleAccessToken)).
82    ///   Builds without the `gemini-llm` or `gcs-store` feature have no HTTP
83    ///   client and use the `gcloud` CLI only.
84    /// - otherwise → Google AI using `GEMINI_API_KEY` (or
85    ///   `GOOGLE_GENAI_API_KEY` / `GOOGLE_API_KEY`).
86    ///
87    /// When no `.model(..)` was set, the wire layer resolves a default the
88    /// target platform actually serves (`ModelId::live_default`): `GEMINI_MODEL`
89    /// from the environment if present, else the platform's current
90    /// native-audio Flash model.
91    ///
92    /// ```no_run
93    /// # use gemini_adk_fluent_rs::prelude::*;
94    /// # async fn run() -> Result<(), AgentError> {
95    /// let handle = Live::builder()
96    ///     .voice(Voice::Kore)
97    ///     .connect_from_env()
98    ///     .await?;
99    /// # let _ = handle; Ok(())
100    /// # }
101    /// ```
102    pub async fn connect_from_env(
103        mut self,
104    ) -> Result<LiveHandle, gemini_adk_rs::error::AgentError> {
105        self.config.endpoint = resolve_endpoint_from_env().await?;
106        self.build_and_connect().await
107    }
108
109    /// Connect using a pre-configured SessionConfig for auth and model.
110    ///
111    /// Merges the provided config's `endpoint` and `model` into the builder's
112    /// config, preserving system instruction, tools, voice, transcription, and
113    /// all other settings configured via the fluent API.
114    pub async fn connect(
115        mut self,
116        config: SessionConfig,
117    ) -> Result<LiveHandle, gemini_adk_rs::error::AgentError> {
118        // Merge auth/model from external config, keep everything else from builder.
119        self.config.endpoint = config.endpoint;
120        self.config.model = config.model;
121        self.build_and_connect().await
122    }
123
124    /// Connect over `transport` instead of a WebSocket to Gemini, with
125    /// everything this builder configured: tools, phases, extractors,
126    /// watchers, governance and callbacks run for real against whatever the
127    /// transport answers.
128    ///
129    /// This is how to test a session offline. Script the server with
130    /// [`ScriptedServer`](crate::testing::ScriptedServer); no network or
131    /// credential is used.
132    pub async fn connect_with_transport<T: gemini_genai_rs::transport::Transport>(
133        self,
134        transport: T,
135    ) -> Result<LiveHandle, gemini_adk_rs::error::AgentError> {
136        self.build_and_connect_via(move |builder| builder.connect_with_transport(transport))
137            .await
138    }
139
140    async fn build_and_connect(self) -> Result<LiveHandle, gemini_adk_rs::error::AgentError> {
141        self.build_and_connect_via(LiveSessionBuilder::connect)
142            .await
143    }
144
145    async fn build_and_connect_via<F, Fut>(
146        mut self,
147        connect: F,
148    ) -> Result<LiveHandle, gemini_adk_rs::error::AgentError>
149    where
150        F: FnOnce(LiveSessionBuilder) -> Fut,
151        Fut: std::future::Future<Output = Result<LiveHandle, gemini_adk_rs::error::AgentError>>,
152    {
153        // Builder setters cannot fail; problems they found are reported here,
154        // with a `T::confirm` tool that nothing can confirm (agent and MCP
155        // tools resolved below are never gated, so this check is complete).
156        let mut issues = std::mem::take(&mut self.config_errors);
157        let unconfirmed = self.unconfirmed_tools();
158        if !unconfirmed.is_empty() {
159            issues.push(
160                crate::testing::LiveViolation::UnconfirmedTools { tools: unconfirmed }.to_string(),
161            );
162        }
163        if !issues.is_empty() {
164            return Err(gemini_adk_rs::error::ConfigError { issues }.into());
165        }
166        if uses_audio_output(&self.config) {
167            self.config = self.config.voice_realtime_defaults();
168        }
169
170        // Client interruption authority: the server must stop deciding, so
171        // its automatic activity detection is disabled in the setup message
172        // (the input VAD's activity marks take over in `send_audio`).
173        if self.input_audio.client_authority {
174            self.config =
175                self.config
176                    .server_vad(gemini_genai_rs::prelude::AutomaticActivityDetection {
177                        disabled: Some(true),
178                        start_of_speech_sensitivity: None,
179                        end_of_speech_sensitivity: None,
180                        prefix_padding_ms: None,
181                        silence_duration_ms: None,
182                    });
183        }
184        let input_audio = std::mem::take(&mut self.input_audio);
185
186        // Resolve a `.record_wire(path)` request into a FileWireRecorder now
187        // that we are actually connecting.
188        if let Some(path) = self.record_wire_path.take() {
189            let recorder =
190                gemini_genai_rs::transport::FileWireRecorder::create(&path).map_err(|e| {
191                    gemini_adk_rs::error::AgentError::Config(format!(
192                        "failed to create wire log at {}: {e}",
193                        path.display()
194                    ))
195                })?;
196            self.config = self.config.record_wire(std::sync::Arc::new(recorder));
197        }
198
199        // Config-level tool declarations, captured before `config` moves.
200        let builder_config_tools = self.config.tools.clone();
201        let mut builder = LiveSessionBuilder::new(self.config);
202
203        // The session's `State`. A caller-supplied one is used as-is so tools
204        // they already built around it write where the flow monitor and phase
205        // machine read; otherwise agent tools get a fresh one as before.
206        // Deferred agent tools and policies need the state before the runtime
207        // exists, so one is made here when the caller gave none.
208        let shared_state = self.state.clone().or_else(|| {
209            (!self.deferred_agent_tools.is_empty() || !self.policies.is_empty())
210                .then(gemini_adk_rs::State::new)
211        });
212        if let Some(ref state) = shared_state {
213            builder = builder.state(state.clone());
214        }
215
216        // Resolve deferred agent tools: register TextAgentTools against it.
217        let mut dispatcher = self.dispatcher;
218        if !self.deferred_agent_tools.is_empty() {
219            let state = shared_state.clone().unwrap_or_default();
220            let d = dispatcher.get_or_insert_with(gemini_adk_rs::tool::ToolDispatcher::new);
221            for deferred in self.deferred_agent_tools {
222                d.register(gemini_adk_rs::TextAgentTool::from_arc(
223                    deferred.name,
224                    deferred.description,
225                    deferred.agent,
226                    state.clone(),
227                ));
228            }
229        }
230
231        // Resolve deferred async tools (MCP connections, etc.).
232        if !self.deferred_tools.is_empty() {
233            let d = dispatcher.get_or_insert_with(gemini_adk_rs::tool::ToolDispatcher::new);
234            for deferred in std::mem::take(&mut self.deferred_tools) {
235                resolve_deferred_tool(deferred, d).await?;
236            }
237        }
238
239        // Enforce redaction and commit-governance policies.
240        if let Some(state) = &shared_state {
241            apply_policies(&self.policies, state, &mut dispatcher)?;
242        }
243
244        // Attach the confirmation provider so `T::confirm(..)` tools are gated.
245        if let Some(provider) = self.confirmation_provider {
246            dispatcher
247                .get_or_insert_with(gemini_adk_rs::tool::ToolDispatcher::new)
248                .set_confirmation_provider(provider);
249        }
250
251        // Capture the resolved tool names before the dispatcher moves into the
252        // builder. This is the only point where the set is complete: MCP tools
253        // exist only after the handshakes above.
254        let resolved_tool_names: Vec<String> = {
255            let mut names = super::introspect::declaration_names(&builder_config_tools);
256            if let Some(d) = &dispatcher {
257                names.extend(super::introspect::declaration_names(
258                    &d.to_tool_declarations(),
259                ));
260            }
261            names.sort();
262            names.dedup();
263            names
264        };
265
266        if let Some(dispatcher) = dispatcher {
267            builder = builder.dispatcher(dispatcher);
268        }
269        if let Some(greeting) = self.greeting {
270            builder = builder.greeting(greeting);
271        }
272        builder = builder.callbacks(self.callbacks);
273        for ext in self.extractors {
274            builder = builder.extractor(ext);
275        }
276
277        // Pass L1 registries
278        if !self.computed.is_empty() {
279            builder = builder.computed(self.computed);
280        }
281        if let Some(initial) = self.initial_phase {
282            let mut pm = PhaseMachine::new(&initial);
283            for phase in self.phases {
284                pm.add_phase(phase);
285            }
286            builder = builder.phase_machine(pm);
287        }
288        if !self.watchers.observed_keys().is_empty() {
289            builder = builder.watchers(self.watchers);
290        }
291        builder = builder.temporal(self.temporal);
292
293        // Pass tool execution modes
294        for (name, mode) in self.tool_execution_modes {
295            builder = builder.tool_execution_mode(name, mode);
296        }
297
298        // Pass control plane configuration
299        if let Some(timeout) = self.soft_turn_timeout {
300            builder = builder.soft_turn_timeout(timeout);
301        }
302        builder = builder.steering_mode(self.steering_mode);
303        builder = builder.context_delivery(self.context_delivery);
304        builder = builder.delivery(self.delivery);
305        if let Some(redactor) = self.redactor {
306            builder = builder.redaction(redactor);
307        }
308        if let Some(config) = self.repair_config {
309            builder = builder.repair(config);
310        }
311        if let Some(p) = self.persistence {
312            builder = builder.persistence(p);
313        }
314        if let Some(id) = self.session_id {
315            builder = builder.session_id(id);
316        }
317        for layer in self.middleware_layers {
318            builder = builder.middleware(layer);
319        }
320        if let Some(mut flow) = self.flow {
321            // Merged here rather than in `govern`/`ambient_tools` so the two
322            // compose regardless of the order the caller wrote them in.
323            merge_ambient(&mut flow, &self.ambient_tools);
324            // A flow names tools as strings, and a name that matches nothing is
325            // not inert: an `allow` whitelist containing only a typo denies
326            // every tool for as long as that step is active, silently and for
327            // the rest of the session. The registry to check against only
328            // exists here, after the deferred handshakes above.
329            //
330            // Skipped for `govern_compiled`/`observe_compiled`, whose contract
331            // is that the caller already surfaced these diagnostics via
332            // `Flow::compile`/`compile_with_tools` and connect will not repeat
333            // the work.
334            if !self.flow_precompiled {
335                let registry: Vec<&str> = resolved_tool_names.iter().map(String::as_str).collect();
336                flow.clone()
337                    .compile_with_tools(&registry)
338                    .map_err(|errors| {
339                        gemini_adk_rs::error::AgentError::Config(format!(
340                            "governing flow does not match this session's tools: {errors}. \
341                             Registered tools: [{}]. Use `Flow::compile_with_tools` at load \
342                             time and `govern_compiled` to check this yourself.",
343                            registry.join(", ")
344                        ))
345                    })?;
346            }
347            let mut monitor = gemini_adk_rs::flow::FlowMonitor::new(flow, self.flow_mode);
348            for (step, agent, mode) in self.flow_actions {
349                monitor = monitor.on_enter(step, gemini_adk_rs::flow::on_enter(agent, mode));
350            }
351            builder = builder.flow_stack(assemble_stack(
352                monitor,
353                self.digressions,
354                self.repair_policies,
355                self.stage_timings,
356                self.corrections,
357                self.verbatims,
358                &self.ambient_tools,
359            ));
360        }
361        builder = builder.tool_advisory(self.tool_advisory);
362        if let Some(clock) = self.clock.take() {
363            builder = builder.clock(clock);
364        }
365        if let Some(interval) = self.telemetry_interval {
366            builder = builder.telemetry_interval(interval);
367        }
368
369        // Spawn fire-and-forget warm-up tasks for OOB LLMs
370        // (pre-establishes TCP+TLS so first extract call is fast)
371        for llm in self.warm_up_llms {
372            tokio::spawn(async move {
373                let _ = llm.warm_up().await;
374            });
375        }
376
377        let handle = connect(builder).await?;
378
379        // Input-audio hardening: materialize the configured stages (in
380        // order) and hand them plus VAD tuning and authority to the handle.
381        if input_audio.is_configured() {
382            let (processors, vad, client_authority, turn_commit) = input_audio.build_processors();
383            handle.configure_input_audio(
384                processors,
385                vad,
386                if client_authority {
387                    gemini_adk_rs::live::ActivityAuthority::Client
388                } else {
389                    gemini_adk_rs::live::ActivityAuthority::Server
390                },
391            );
392            if let Some(config) = turn_commit {
393                handle.set_turn_commit(config);
394            }
395        }
396        Ok(handle)
397    }
398}
399
400/// Enforce `policies` on a session: mark redacted keys on its state and
401/// wrap each commit tool in a [`CommitGuard`](gemini_adk_rs::tool::CommitGuard).
402fn apply_policies(
403    policies: &[crate::policy::Policy],
404    state: &gemini_adk_rs::State,
405    dispatcher: &mut Option<gemini_adk_rs::tool::ToolDispatcher>,
406) -> Result<(), gemini_adk_rs::error::AgentError> {
407    use crate::policy::Policy;
408    use gemini_adk_rs::tool::{CommitGuard, ToolKind};
409
410    let function = |d: &gemini_adk_rs::tool::ToolDispatcher, name: &str| match d.get_tool(name) {
411        Some(ToolKind::Function(f)) => Some(f.clone()),
412        _ => None,
413    };
414    for policy in policies {
415        match policy {
416            Policy::Redact { keys } => state.redact_keys(keys.iter().cloned()),
417            Policy::Commit {
418                tool,
419                idempotency_key,
420                compensate_with,
421            } => {
422                let d = dispatcher.as_mut();
423                let Some(inner) = d.as_ref().and_then(|d| function(d, tool)) else {
424                    return Err(gemini_adk_rs::error::AgentError::Config(format!(
425                        "Policy::commit names `{tool}`, which is not a registered function tool"
426                    )));
427                };
428                let mut guard = CommitGuard::new(inner, state.clone());
429                if let Some(template) = idempotency_key {
430                    guard = guard.idempotency_key(template.clone());
431                }
432                if let Some(undo) = compensate_with {
433                    let Some(undo_fn) = d.as_ref().and_then(|d| function(d, undo)) else {
434                        return Err(gemini_adk_rs::error::AgentError::Config(format!(
435                            "Policy::commit(`{tool}`) compensates with `{undo}`, which is not a \
436                             registered function tool"
437                        )));
438                    };
439                    guard = guard.compensate_with(undo_fn);
440                }
441                if let Some(d) = d {
442                    d.register(guard);
443                }
444            }
445            Policy::SafetyHandoff { .. } => {}
446        }
447    }
448    Ok(())
449}
450
451/// Resolve an [`ApiEndpoint`] from the environment, with a `gcloud` token
452/// fallback for Vertex AI when `GOOGLE_ACCESS_TOKEN` is not set.
453async fn resolve_endpoint_from_env() -> Result<ApiEndpoint, gemini_adk_rs::error::AgentError> {
454    use gemini_adk_rs::error::AgentError;
455    use gemini_genai_rs::protocol::types::EndpointEnvError;
456
457    match ApiEndpoint::from_env() {
458        Ok(endpoint) => Ok(endpoint),
459        // Vertex was selected but no token was in the environment: use the
460        // service's own credentials (metadata server), else the gcloud CLI.
461        Err(EndpointEnvError::Missing("GOOGLE_ACCESS_TOKEN")) => {
462            let project = std::env::var("GOOGLE_CLOUD_PROJECT").map_err(|_| {
463                AgentError::Config("GOOGLE_CLOUD_PROJECT is required for Vertex AI".into())
464            })?;
465            let location = std::env::var("GOOGLE_CLOUD_LOCATION")
466                .unwrap_or_else(|_| "us-central1".to_string());
467            let token = vertex_access_token().await?;
468            Ok(ApiEndpoint::vertex(project, location, token))
469        }
470        Err(e) => Err(AgentError::Config(format!(
471            "connect_from_env: {e}. For Google AI set GEMINI_API_KEY; for Vertex AI set \
472             GOOGLE_GENAI_USE_VERTEXAI=true and GOOGLE_CLOUD_PROJECT (token via \
473             GOOGLE_ACCESS_TOKEN, the metadata server, or the gcloud CLI)."
474        ))),
475    }
476}
477
478/// A Vertex token that keeps itself fresh: from the metadata server when one
479/// answers (Cloud Run, GKE, Compute Engine), else from the gcloud CLI.
480#[cfg(any(feature = "gemini-llm", feature = "gcs-store"))]
481async fn vertex_access_token() -> Result<AccessToken, gemini_adk_rs::error::AgentError> {
482    gemini_genai_rs::transport::auth::GoogleAccessToken::from_env()
483        .into_access_token()
484        .await
485        .map_err(|e| {
486            gemini_adk_rs::error::AgentError::Config(format!(
487                "Vertex AI needs an access token: set GOOGLE_ACCESS_TOKEN, run on Google \
488                 Cloud with a service account, or install the gcloud CLI ({e})"
489            ))
490        })
491}
492
493/// Without an HTTP client, only the gcloud CLI can supply a token.
494#[cfg(not(any(feature = "gemini-llm", feature = "gcs-store")))]
495#[allow(
496    clippy::unused_async,
497    reason = "same signature as the metadata-server variant it replaces"
498)]
499async fn vertex_access_token() -> Result<AccessToken, gemini_adk_rs::error::AgentError> {
500    gcloud_access_token().map(AccessToken::from)
501}
502
503/// Fetch an OAuth2 access token via `gcloud auth print-access-token`.
504#[cfg(not(any(feature = "gemini-llm", feature = "gcs-store")))]
505fn gcloud_access_token() -> Result<String, gemini_adk_rs::error::AgentError> {
506    use gemini_adk_rs::error::AgentError;
507
508    let output = std::process::Command::new("gcloud")
509        .args(["auth", "print-access-token"])
510        .output()
511        .map_err(|e| {
512            AgentError::Config(format!(
513                "Vertex AI needs an access token: set GOOGLE_ACCESS_TOKEN, or install the \
514                 gcloud CLI (failed to run `gcloud auth print-access-token`: {e})"
515            ))
516        })?;
517    if !output.status.success() {
518        return Err(AgentError::Config(format!(
519            "`gcloud auth print-access-token` failed: {}",
520            String::from_utf8_lossy(&output.stderr).trim()
521        )));
522    }
523    let token = String::from_utf8_lossy(&output.stdout).trim().to_string();
524    if token.is_empty() {
525        return Err(AgentError::Config(
526            "`gcloud auth print-access-token` returned an empty token".into(),
527        ));
528    }
529    Ok(token)
530}
531
532/// Resolve a single [`DeferredTool`](crate::compose::tools::DeferredTool) into
533/// concrete tool registrations on the dispatcher. Runs at connect time because
534/// these tools require async I/O (a network call or a subprocess handshake).
535async fn resolve_deferred_tool(
536    tool: crate::compose::tools::DeferredTool,
537    dispatcher: &mut gemini_adk_rs::tool::ToolDispatcher,
538) -> Result<(), gemini_adk_rs::error::AgentError> {
539    use crate::compose::tools::DeferredTool;
540    use gemini_adk_rs::error::AgentError;
541    use gemini_adk_rs::tools::mcp::{McpSessionManager, McpTool};
542    use std::sync::Arc;
543
544    match tool {
545        DeferredTool::Mcp { params } => {
546            let manager = Arc::new(McpSessionManager::new(parse_mcp_params(&params)));
547            let infos = manager.list_tools().await.map_err(|e| {
548                AgentError::Config(format!("MCP tool discovery failed for {params:?}: {e}"))
549            })?;
550            for info in infos {
551                dispatcher.register_function(Arc::new(McpTool::new(
552                    info.name,
553                    info.description,
554                    Some(info.input_schema),
555                    manager.clone(),
556                )));
557            }
558            Ok(())
559        }
560    }
561}
562
563/// Parse an MCP connection string: an `http(s)://` URL becomes an SSE/HTTP
564/// connection, anything else is treated as a stdio command line.
565pub(crate) fn parse_mcp_params(params: &str) -> gemini_adk_rs::tools::mcp::McpConnectionParams {
566    use gemini_adk_rs::tools::mcp::McpConnectionParams;
567
568    let trimmed = params.trim();
569    if trimmed.starts_with("http://") || trimmed.starts_with("https://") {
570        McpConnectionParams::Sse {
571            url: trimmed.to_string(),
572            headers: None,
573        }
574    } else {
575        let mut parts = trimmed.split_whitespace();
576        let command = parts.next().unwrap_or_default().to_string();
577        let args = parts.map(str::to_string).collect();
578        McpConnectionParams::Stdio {
579            command,
580            args,
581            timeout: Some(std::time::Duration::from_secs(30)),
582        }
583    }
584}
585
586fn uses_audio_output(config: &SessionConfig) -> bool {
587    config
588        .generation_config
589        .response_modalities
590        .as_ref()
591        .map(|modalities| modalities.iter().any(|m| matches!(m, Modality::Audio)))
592        .unwrap_or(true)
593}
594
595#[cfg(test)]
596mod tests {
597    use super::*;
598
599    #[test]
600    fn uses_audio_output_defaults_to_audio() {
601        let config = SessionConfig::new("key");
602        assert!(uses_audio_output(&config));
603    }
604
605    #[test]
606    fn uses_audio_output_respects_text_only() {
607        let config = SessionConfig::new("key").text_only();
608        assert!(!uses_audio_output(&config));
609    }
610
611    // ─── default model resolution ───────────────────────────────────────────
612
613    #[test]
614    fn unset_model_is_left_to_the_wire_layer_to_resolve() {
615        let live = Live::builder();
616        assert!(live.config.model.is_none(), "no default is baked in at L2");
617        let live = Live::builder().model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO);
618        assert_eq!(
619            live.config.model.as_ref(),
620            Some(&ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO)
621        );
622    }
623
624    // ─── ambient tool merge ─────────────────────────────────────────────────
625
626    use gemini_adk_rs::flow::{Flow, Guard};
627
628    fn whitelisting_flow() -> Flow {
629        Flow::new()
630            .step("book")
631            .allow(["book_table"])
632            .done(Guard::called_ok("book_table"))
633            .build()
634            .expect("flow is structurally valid")
635    }
636
637    #[test]
638    fn merge_ambient_adds_registered_tools() {
639        let mut flow = whitelisting_flow();
640        merge_ambient(&mut flow, &["recall_context".to_string()]);
641        assert_eq!(flow.ambient, ["recall_context"]);
642    }
643
644    #[test]
645    fn merge_ambient_does_not_duplicate() {
646        // The flow already declares it and an extension registers it too.
647        let mut flow = Flow::new()
648            .ambient(["recall_context"])
649            .step("book")
650            .allow(["book_table"])
651            .done(Guard::called_ok("book_table"))
652            .build()
653            .expect("flow is structurally valid");
654        merge_ambient(&mut flow, &["recall_context".to_string()]);
655        merge_ambient(&mut flow, &["recall_context".to_string()]);
656        assert_eq!(
657            flow.ambient,
658            ["recall_context"],
659            "merging is idempotent, so an extension installed twice is harmless"
660        );
661    }
662
663    // ─── connect-time flow validation ───────────────────────────────────────
664    //
665    // These run to a real `connect_google_ai` with a junk key. That is safe and
666    // deliberate: validation happens before `builder.connect()`, so a `Config`
667    // error proves the check fired *and* that it fired before any socket was
668    // opened. A `Session` error would mean the flow was accepted and the
669    // failure came from the network instead.
670
671    fn book_tool() -> crate::compose::tools::ToolComposite {
672        crate::compose::T::simple("book_table", "Book a table", |_| async {
673            Ok(serde_json::json!({"ok": true}))
674        })
675    }
676
677    #[tokio::test]
678    async fn a_flow_naming_an_unregistered_tool_is_refused_at_connect() {
679        let flow = Flow::new()
680            .step("book")
681            .allow(["book_tabel"]) // typo: the registered tool is `book_table`
682            .done(Guard::called_ok("book_tabel"))
683            .build()
684            .expect("structurally valid — the name is the problem, not the shape");
685
686        let err = Live::builder()
687            .tools(book_tool())
688            .govern(flow)
689            .connect_google_ai("not-a-real-key")
690            .await
691            .err()
692            .expect("a flow naming a tool that does not exist must not connect");
693
694        let msg = err.to_string();
695        assert!(
696            msg.contains("book_tabel"),
697            "the error must name the tool that does not exist: {msg}"
698        );
699        assert!(
700            msg.contains("book_table"),
701            "and list what is registered, so the typo is visible: {msg}"
702        );
703    }
704
705    #[tokio::test]
706    async fn a_flow_matching_the_registered_tools_passes_validation() {
707        // Reaches the network and fails there — which is the proof that the
708        // flow check passed rather than short-circuiting.
709        let flow = Flow::new()
710            .step("book")
711            .allow(["book_table"])
712            .done(Guard::called_ok("book_table"))
713            .build()
714            .expect("valid");
715
716        let err = Live::builder()
717            .tools(book_tool())
718            .govern(flow)
719            .connect_google_ai("not-a-real-key")
720            .await
721            .err()
722            .expect("the key is junk, so this still fails — just not on the flow");
723
724        assert!(
725            !err.to_string().contains("governing flow"),
726            "a flow whose tools all exist must clear validation: {err}"
727        );
728    }
729
730    #[tokio::test]
731    async fn ambient_tools_count_as_registered_for_validation() {
732        // `ambient` joins the tool universe, so an ambient tool that is really
733        // registered must satisfy the check rather than trip it.
734        let flow = Flow::new()
735            .ambient(["book_table"])
736            .step("book")
737            .done(Guard::called_ok("book_table"))
738            .build()
739            .expect("valid");
740
741        let err = Live::builder()
742            .tools(book_tool())
743            .govern(flow)
744            .connect_google_ai("not-a-real-key")
745            .await
746            .err()
747            .expect("junk key");
748
749        assert!(
750            !err.to_string().contains("governing flow"),
751            "an ambient tool that exists must not be reported missing: {err}"
752        );
753    }
754
755    #[tokio::test]
756    async fn a_precompiled_flow_is_not_revalidated() {
757        // `govern_compiled` documents that connect does not re-check. Honour it
758        // even when the flow would fail the check, or the documented
759        // compile-once-govern-many path would silently stop working.
760        let compiled = Flow::new()
761            .step("book")
762            .allow(["book_tabel"])
763            .done(Guard::called_ok("book_tabel"))
764            .build()
765            .expect("valid shape")
766            .compile()
767            .expect("compiles without a tool registry");
768
769        let err = Live::builder()
770            .tools(book_tool())
771            .govern_compiled(compiled)
772            .connect_google_ai("not-a-real-key")
773            .await
774            .err()
775            .expect("junk key");
776
777        assert!(
778            !err.to_string().contains("governing flow"),
779            "a pre-compiled flow must not be re-validated at connect: {err}"
780        );
781    }
782
783    #[test]
784    fn ambient_tools_registers_regardless_of_govern_order() {
785        // The whole reason the merge happens at connect: an application may
786        // write `.govern(..)` before or after the extension that needs ambient
787        // tools, and neither order may lose the registration.
788        let before = Live::builder()
789            .govern(whitelisting_flow())
790            .ambient_tools(["recall_context"]);
791        let after = Live::builder()
792            .ambient_tools(["recall_context"])
793            .govern(whitelisting_flow());
794        assert_eq!(before.ambient_tool_names(), ["recall_context"]);
795        assert_eq!(after.ambient_tool_names(), ["recall_context"]);
796    }
797}
gemini_adk_fluent_rs/live/connect.rs

gemini_adk_fluent_rs/live/
connect.rs
