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
20impl Live {
21    /// Connect using a Google AI API key.
22    pub async fn connect_google_ai(
23        mut self,
24        api_key: impl Into<String>,
25    ) -> Result<LiveHandle, gemini_adk_rs::error::AgentError> {
26        self.config.endpoint = ApiEndpoint::google_ai(api_key);
27        self.build_and_connect().await
28    }
29
30    /// Connect using Vertex AI credentials.
31    pub async fn connect_vertex(
32        mut self,
33        project: impl Into<String>,
34        location: impl Into<String>,
35        access_token: impl Into<String>,
36    ) -> Result<LiveHandle, gemini_adk_rs::error::AgentError> {
37        self.config.endpoint = ApiEndpoint::vertex(project, location, access_token);
38        self.build_and_connect().await
39    }
40
41    /// Connect by resolving the platform and credentials from standard
42    /// environment variables — the zero-ceremony entry point.
43    ///
44    /// Resolution (see [`ApiEndpoint::from_env`]):
45    /// - `GOOGLE_GENAI_USE_VERTEXAI=true` → Vertex AI using
46    ///   `GOOGLE_CLOUD_PROJECT`, `GOOGLE_CLOUD_LOCATION` (default
47    ///   `us-central1`), and a token from `GOOGLE_ACCESS_TOKEN`. If that
48    ///   token is unset, this falls back to running
49    ///   `gcloud auth print-access-token`.
50    /// - otherwise → Google AI using `GEMINI_API_KEY` (or
51    ///   `GOOGLE_GENAI_API_KEY` / `GOOGLE_API_KEY`).
52    ///
53    /// ```no_run
54    /// # use gemini_adk_fluent_rs::prelude::*;
55    /// # async fn run() -> Result<(), AgentError> {
56    /// let handle = Live::builder()
57    ///     .model(GeminiModel::Gemini2_0FlashLive)
58    ///     .voice(Voice::Kore)
59    ///     .connect_from_env()
60    ///     .await?;
61    /// # let _ = handle; Ok(())
62    /// # }
63    /// ```
64    pub async fn connect_from_env(
65        mut self,
66    ) -> Result<LiveHandle, gemini_adk_rs::error::AgentError> {
67        self.config.endpoint = resolve_endpoint_from_env()?;
68        self.build_and_connect().await
69    }
70
71    /// Connect using a pre-configured SessionConfig for auth and model.
72    ///
73    /// Merges the provided config's `endpoint` and `model` into the builder's
74    /// config, preserving system instruction, tools, voice, transcription, and
75    /// all other settings configured via the fluent API.
76    pub async fn connect(
77        mut self,
78        config: SessionConfig,
79    ) -> Result<LiveHandle, gemini_adk_rs::error::AgentError> {
80        // Merge auth/model from external config, keep everything else from builder.
81        self.config.endpoint = config.endpoint;
82        self.config.model = config.model;
83        self.build_and_connect().await
84    }
85
86    async fn build_and_connect(mut self) -> Result<LiveHandle, gemini_adk_rs::error::AgentError> {
87        if uses_audio_output(&self.config) {
88            self.config = self.config.voice_realtime_defaults();
89        }
90
91        // Resolve a `.record_wire(path)` request into a FileWireRecorder now
92        // that we are actually connecting.
93        if let Some(path) = self.record_wire_path.take() {
94            let recorder = FileWireRecorder::create(&path).map_err(|e| {
95                gemini_adk_rs::error::AgentError::Config(format!(
96                    "failed to create wire log at {}: {e}",
97                    path.display()
98                ))
99            })?;
100            self.config = self.config.record_wire(std::sync::Arc::new(recorder));
101        }
102
103        // Config-level tool declarations, captured before `config` moves.
104        let builder_config_tools = self.config.tools.clone();
105        let mut builder = LiveSessionBuilder::new(self.config);
106
107        // The session's `State`. A caller-supplied one is used as-is so tools
108        // they already built around it write where the flow monitor and phase
109        // machine read; otherwise agent tools get a fresh one as before.
110        let shared_state = self.state.clone();
111        if let Some(ref state) = shared_state {
112            builder = builder.with_state(state.clone());
113        }
114
115        // Resolve deferred agent tools: register TextAgentTools against it.
116        let mut dispatcher = self.dispatcher;
117        if !self.deferred_agent_tools.is_empty() {
118            let state = shared_state.clone().unwrap_or_default();
119            let d = dispatcher.get_or_insert_with(gemini_adk_rs::tool::ToolDispatcher::new);
120            for deferred in self.deferred_agent_tools {
121                d.register(gemini_adk_rs::TextAgentTool::from_arc(
122                    deferred.name,
123                    deferred.description,
124                    deferred.agent,
125                    state.clone(),
126                ));
127            }
128            builder = builder.with_state(state);
129        }
130
131        // Resolve deferred async tools (MCP connections, etc.).
132        if !self.deferred_tools.is_empty() {
133            let d = dispatcher.get_or_insert_with(gemini_adk_rs::tool::ToolDispatcher::new);
134            for deferred in std::mem::take(&mut self.deferred_tools) {
135                resolve_deferred_tool(deferred, d).await?;
136            }
137        }
138
139        // Attach the confirmation provider so `T::confirm(..)` tools are gated.
140        if let Some(provider) = self.confirmation_provider {
141            dispatcher
142                .get_or_insert_with(gemini_adk_rs::tool::ToolDispatcher::new)
143                .set_confirmation_provider(provider);
144        }
145
146        // Capture the resolved tool names before the dispatcher moves into the
147        // builder. This is the only point where the set is complete: MCP/A2A/
148        // OpenAPI tools exist only after the handshakes above.
149        let resolved_tool_names: Vec<String> = {
150            let mut names = super::introspect::declaration_names(&builder_config_tools);
151            if let Some(d) = &dispatcher {
152                names.extend(super::introspect::declaration_names(
153                    &d.to_tool_declarations(),
154                ));
155            }
156            names.sort();
157            names.dedup();
158            names
159        };
160
161        if let Some(dispatcher) = dispatcher {
162            builder = builder.dispatcher(dispatcher);
163        }
164        if let Some(greeting) = self.greeting {
165            builder = builder.greeting(greeting);
166        }
167        builder = builder.callbacks(self.callbacks);
168        for ext in self.extractors {
169            builder = builder.extractor(ext);
170        }
171
172        // Pass L1 registries
173        if !self.computed.is_empty() {
174            builder = builder.computed(self.computed);
175        }
176        if let Some(initial) = self.initial_phase {
177            let mut pm = PhaseMachine::new(&initial);
178            for phase in self.phases {
179                pm.add_phase(phase);
180            }
181            builder = builder.phase_machine(pm);
182        }
183        if !self.watchers.observed_keys().is_empty() {
184            builder = builder.watchers(self.watchers);
185        }
186        builder = builder.temporal(self.temporal);
187
188        // Pass tool execution modes
189        for (name, mode) in self.tool_execution_modes {
190            builder = builder.tool_execution_mode(name, mode);
191        }
192
193        // Pass control plane configuration
194        if let Some(timeout) = self.soft_turn_timeout {
195            builder = builder.soft_turn_timeout(timeout);
196        }
197        builder = builder.steering_mode(self.steering_mode);
198        builder = builder.context_delivery(self.context_delivery);
199        builder = builder.delivery(self.delivery);
200        if let Some(config) = self.repair_config {
201            builder = builder.repair(config);
202        }
203        if let Some(p) = self.persistence {
204            builder = builder.persistence(p);
205        }
206        if let Some(id) = self.session_id {
207            builder = builder.session_id(id);
208        }
209        for layer in self.middleware_layers {
210            builder = builder.middleware(layer);
211        }
212        if let Some(mut flow) = self.flow {
213            // Merged here rather than in `govern`/`ambient_tools` so the two
214            // compose regardless of the order the caller wrote them in.
215            merge_ambient(&mut flow, &self.ambient_tools);
216            // A flow names tools as strings, and a name that matches nothing is
217            // not inert: an `allow` whitelist containing only a typo denies
218            // every tool for as long as that step is active, silently and for
219            // the rest of the session. The registry to check against only
220            // exists here, after the deferred handshakes above.
221            //
222            // Skipped for `govern_compiled`/`observe_compiled`, whose contract
223            // is that the caller already surfaced these diagnostics via
224            // `Flow::compile`/`compile_with_tools` and connect will not repeat
225            // the work.
226            if !self.flow_precompiled {
227                let registry: Vec<&str> = resolved_tool_names.iter().map(String::as_str).collect();
228                flow.clone()
229                    .compile_with_tools(&registry)
230                    .map_err(|errors| {
231                        gemini_adk_rs::error::AgentError::Config(format!(
232                            "governing flow does not match this session's tools: {errors}. \
233                             Registered tools: [{}]. Use `Flow::compile_with_tools` at load \
234                             time and `govern_compiled` to check this yourself.",
235                            registry.join(", ")
236                        ))
237                    })?;
238            }
239            let mut monitor = gemini_adk_rs::flow::FlowMonitor::new(flow, self.flow_mode);
240            for (step, agent, mode) in self.flow_actions {
241                monitor = monitor.on_enter(step, gemini_adk_rs::flow::run(agent, mode));
242            }
243            builder = builder.flow_monitor(monitor);
244        }
245        builder = builder.tool_advisory(self.tool_advisory);
246        if let Some(interval) = self.telemetry_interval {
247            builder = builder.telemetry_interval(interval);
248        }
249
250        // Spawn fire-and-forget warm-up tasks for OOB LLMs
251        // (pre-establishes TCP+TLS so first extract call is fast)
252        for llm in self.warm_up_llms {
253            tokio::spawn(async move {
254                let _ = llm.warm_up().await;
255            });
256        }
257
258        builder.connect().await
259    }
260}
261
262/// Resolve an [`ApiEndpoint`] from the environment, with a `gcloud` token
263/// fallback for Vertex AI when `GOOGLE_ACCESS_TOKEN` is not set.
264fn resolve_endpoint_from_env() -> Result<ApiEndpoint, gemini_adk_rs::error::AgentError> {
265    use gemini_adk_rs::error::AgentError;
266    use gemini_genai_rs::protocol::types::EndpointEnvError;
267
268    match ApiEndpoint::from_env() {
269        Ok(endpoint) => Ok(endpoint),
270        // Vertex was selected but no token was in the environment — fall back
271        // to Application Default Credentials via the gcloud CLI.
272        Err(EndpointEnvError::Missing("GOOGLE_ACCESS_TOKEN")) => {
273            let project = std::env::var("GOOGLE_CLOUD_PROJECT").map_err(|_| {
274                AgentError::Config("GOOGLE_CLOUD_PROJECT is required for Vertex AI".into())
275            })?;
276            let location = std::env::var("GOOGLE_CLOUD_LOCATION")
277                .unwrap_or_else(|_| "us-central1".to_string());
278            let token = gcloud_access_token()?;
279            Ok(ApiEndpoint::vertex(project, location, token))
280        }
281        Err(e) => Err(AgentError::Config(format!(
282            "connect_from_env: {e}. For Google AI set GEMINI_API_KEY; for Vertex AI set \
283             GOOGLE_GENAI_USE_VERTEXAI=true and GOOGLE_CLOUD_PROJECT (token via \
284             GOOGLE_ACCESS_TOKEN or the gcloud CLI)."
285        ))),
286    }
287}
288
289/// Fetch an OAuth2 access token via `gcloud auth print-access-token`.
290fn gcloud_access_token() -> Result<String, gemini_adk_rs::error::AgentError> {
291    use gemini_adk_rs::error::AgentError;
292
293    let output = std::process::Command::new("gcloud")
294        .args(["auth", "print-access-token"])
295        .output()
296        .map_err(|e| {
297            AgentError::Config(format!(
298                "Vertex AI needs an access token: set GOOGLE_ACCESS_TOKEN, or install the \
299                 gcloud CLI (failed to run `gcloud auth print-access-token`: {e})"
300            ))
301        })?;
302    if !output.status.success() {
303        return Err(AgentError::Config(format!(
304            "`gcloud auth print-access-token` failed: {}",
305            String::from_utf8_lossy(&output.stderr).trim()
306        )));
307    }
308    let token = String::from_utf8_lossy(&output.stdout).trim().to_string();
309    if token.is_empty() {
310        return Err(AgentError::Config(
311            "`gcloud auth print-access-token` returned an empty token".into(),
312        ));
313    }
314    Ok(token)
315}
316
317/// Resolve a single [`DeferredTool`](crate::compose::tools::DeferredTool) into
318/// concrete tool registrations on the dispatcher. Runs at connect time because
319/// these tools require async I/O (a network call or a subprocess handshake).
320async fn resolve_deferred_tool(
321    tool: crate::compose::tools::DeferredTool,
322    dispatcher: &mut gemini_adk_rs::tool::ToolDispatcher,
323) -> Result<(), gemini_adk_rs::error::AgentError> {
324    use crate::compose::tools::DeferredTool;
325    use gemini_adk_rs::error::AgentError;
326    use gemini_adk_rs::tools::mcp::{McpSessionManager, McpTool};
327    use std::sync::Arc;
328
329    match tool {
330        DeferredTool::Mcp { params } => {
331            let manager = Arc::new(McpSessionManager::new(parse_mcp_params(&params)));
332            let infos = manager.list_tools().await.map_err(|e| {
333                AgentError::Config(format!("MCP tool discovery failed for {params:?}: {e}"))
334            })?;
335            for info in infos {
336                dispatcher.register_function(Arc::new(McpTool::new(
337                    info.name,
338                    info.description,
339                    Some(info.input_schema),
340                    manager.clone(),
341                )));
342            }
343            Ok(())
344        }
345        // The following are part of the ADK-parity toolset roadmap; they are
346        // surfaced as explicit connect-time errors rather than silently dropped.
347        DeferredTool::A2a { url, skill } => Err(AgentError::Config(format!(
348            "T::a2a(url={url:?}, skill={skill:?}) is not yet implemented; tracked for ADK parity"
349        ))),
350        DeferredTool::OpenApi { name, spec_url } => Err(AgentError::Config(format!(
351            "T::openapi(name={name:?}, spec_url={spec_url:?}) is not yet implemented; \
352             tracked for ADK parity"
353        ))),
354        DeferredTool::Search { name, .. } => Err(AgentError::Config(format!(
355            "T::search(name={name:?}) is not yet implemented; tracked for ADK parity"
356        ))),
357    }
358}
359
360/// Parse an MCP connection string: an `http(s)://` URL becomes an SSE/HTTP
361/// connection, anything else is treated as a stdio command line.
362fn parse_mcp_params(params: &str) -> gemini_adk_rs::tools::mcp::McpConnectionParams {
363    use gemini_adk_rs::tools::mcp::McpConnectionParams;
364
365    let trimmed = params.trim();
366    if trimmed.starts_with("http://") || trimmed.starts_with("https://") {
367        McpConnectionParams::Sse {
368            url: trimmed.to_string(),
369            headers: None,
370        }
371    } else {
372        let mut parts = trimmed.split_whitespace();
373        let command = parts.next().unwrap_or_default().to_string();
374        let args = parts.map(str::to_string).collect();
375        McpConnectionParams::Stdio {
376            command,
377            args,
378            timeout: Some(std::time::Duration::from_secs(30)),
379        }
380    }
381}
382
383fn uses_audio_output(config: &SessionConfig) -> bool {
384    config
385        .generation_config
386        .response_modalities
387        .as_ref()
388        .map(|modalities| modalities.iter().any(|m| matches!(m, Modality::Audio)))
389        .unwrap_or(true)
390}
391
392#[cfg(test)]
393mod tests {
394    use super::*;
395
396    #[test]
397    fn uses_audio_output_defaults_to_audio() {
398        let config = SessionConfig::new("key");
399        assert!(uses_audio_output(&config));
400    }
401
402    #[test]
403    fn uses_audio_output_respects_text_only() {
404        let config = SessionConfig::new("key").text_only();
405        assert!(!uses_audio_output(&config));
406    }
407
408    // ─── ambient tool merge ─────────────────────────────────────────────────
409
410    use gemini_adk_rs::flow::{Flow, Guard};
411
412    fn whitelisting_flow() -> Flow {
413        Flow::new()
414            .step("book")
415            .allow(["book_table"])
416            .done(Guard::called_ok("book_table"))
417            .build()
418            .expect("flow is structurally valid")
419    }
420
421    #[test]
422    fn merge_ambient_adds_registered_tools() {
423        let mut flow = whitelisting_flow();
424        merge_ambient(&mut flow, &["recall_context".to_string()]);
425        assert_eq!(flow.ambient, ["recall_context"]);
426    }
427
428    #[test]
429    fn merge_ambient_does_not_duplicate() {
430        // The flow already declares it and an extension registers it too.
431        let mut flow = Flow::new()
432            .ambient(["recall_context"])
433            .step("book")
434            .allow(["book_table"])
435            .done(Guard::called_ok("book_table"))
436            .build()
437            .expect("flow is structurally valid");
438        merge_ambient(&mut flow, &["recall_context".to_string()]);
439        merge_ambient(&mut flow, &["recall_context".to_string()]);
440        assert_eq!(
441            flow.ambient,
442            ["recall_context"],
443            "merging is idempotent, so an extension installed twice is harmless"
444        );
445    }
446
447    // ─── connect-time flow validation ───────────────────────────────────────
448    //
449    // These run to a real `connect_google_ai` with a junk key. That is safe and
450    // deliberate: validation happens before `builder.connect()`, so a `Config`
451    // error proves the check fired *and* that it fired before any socket was
452    // opened. A `Session` error would mean the flow was accepted and the
453    // failure came from the network instead.
454
455    fn book_tool() -> crate::compose::tools::ToolComposite {
456        crate::compose::T::simple("book_table", "Book a table", |_| async {
457            Ok(serde_json::json!({"ok": true}))
458        })
459    }
460
461    #[tokio::test]
462    async fn a_flow_naming_an_unregistered_tool_is_refused_at_connect() {
463        let flow = Flow::new()
464            .step("book")
465            .allow(["book_tabel"]) // typo: the registered tool is `book_table`
466            .done(Guard::called_ok("book_tabel"))
467            .build()
468            .expect("structurally valid — the name is the problem, not the shape");
469
470        let err = Live::builder()
471            .with_tools(book_tool())
472            .govern(flow)
473            .connect_google_ai("not-a-real-key")
474            .await
475            .err()
476            .expect("a flow naming a tool that does not exist must not connect");
477
478        let msg = err.to_string();
479        assert!(
480            msg.contains("book_tabel"),
481            "the error must name the tool that does not exist: {msg}"
482        );
483        assert!(
484            msg.contains("book_table"),
485            "and list what is registered, so the typo is visible: {msg}"
486        );
487    }
488
489    #[tokio::test]
490    async fn a_flow_matching_the_registered_tools_passes_validation() {
491        // Reaches the network and fails there — which is the proof that the
492        // flow check passed rather than short-circuiting.
493        let flow = Flow::new()
494            .step("book")
495            .allow(["book_table"])
496            .done(Guard::called_ok("book_table"))
497            .build()
498            .expect("valid");
499
500        let err = Live::builder()
501            .with_tools(book_tool())
502            .govern(flow)
503            .connect_google_ai("not-a-real-key")
504            .await
505            .err()
506            .expect("the key is junk, so this still fails — just not on the flow");
507
508        assert!(
509            !err.to_string().contains("governing flow"),
510            "a flow whose tools all exist must clear validation: {err}"
511        );
512    }
513
514    #[tokio::test]
515    async fn ambient_tools_count_as_registered_for_validation() {
516        // `ambient` joins the tool universe, so an ambient tool that is really
517        // registered must satisfy the check rather than trip it.
518        let flow = Flow::new()
519            .ambient(["book_table"])
520            .step("book")
521            .done(Guard::called_ok("book_table"))
522            .build()
523            .expect("valid");
524
525        let err = Live::builder()
526            .with_tools(book_tool())
527            .govern(flow)
528            .connect_google_ai("not-a-real-key")
529            .await
530            .err()
531            .expect("junk key");
532
533        assert!(
534            !err.to_string().contains("governing flow"),
535            "an ambient tool that exists must not be reported missing: {err}"
536        );
537    }
538
539    #[tokio::test]
540    async fn a_precompiled_flow_is_not_revalidated() {
541        // `govern_compiled` documents that connect does not re-check. Honour it
542        // even when the flow would fail the check, or the documented
543        // compile-once-govern-many path would silently stop working.
544        let compiled = Flow::new()
545            .step("book")
546            .allow(["book_tabel"])
547            .done(Guard::called_ok("book_tabel"))
548            .build()
549            .expect("valid shape")
550            .compile()
551            .expect("compiles without a tool registry");
552
553        let err = Live::builder()
554            .with_tools(book_tool())
555            .govern_compiled(compiled)
556            .connect_google_ai("not-a-real-key")
557            .await
558            .err()
559            .expect("junk key");
560
561        assert!(
562            !err.to_string().contains("governing flow"),
563            "a pre-compiled flow must not be re-validated at connect: {err}"
564        );
565    }
566
567    #[test]
568    fn ambient_tools_registers_regardless_of_govern_order() {
569        // The whole reason the merge happens at connect: an application may
570        // write `.govern(..)` before or after the extension that needs ambient
571        // tools, and neither order may lose the registration.
572        let before = Live::builder()
573            .govern(whitelisting_flow())
574            .ambient_tools(["recall_context"]);
575        let after = Live::builder()
576            .ambient_tools(["recall_context"])
577            .govern(whitelisting_flow());
578        assert_eq!(before.ambient_tool_names(), ["recall_context"]);
579        assert_eq!(after.ambient_tool_names(), ["recall_context"]);
580    }
581}
gemini_adk_fluent_rs/live/connect.rs

gemini_adk_fluent_rs/live/
connect.rs
