gemini_adk_fluent_rs/
builder.rs

1//! AgentBuilder — copy-on-write immutable builder for fluent agent construction.
2//!
3//! Every mutation returns a new builder (original unchanged), so builders
4//! are safely shareable as templates.
5
6use std::sync::Arc;
7
8use gemini_adk_rs::error::ConfigError;
9use gemini_adk_rs::llm::BaseLlm;
10use gemini_adk_rs::middleware::Middleware;
11use gemini_adk_rs::text::{LlmTextAgent, TextAgent};
12use gemini_adk_rs::tool::{ToolDispatcher, ToolFunction, ToolKind};
13use gemini_genai_rs::prelude::{Modality, ModelId, Tool, Voice};
14
15use crate::compose::context::ContextComposite;
16use crate::compose::guards::GuardComposite;
17use crate::compose::middleware::MiddlewareComposite;
18use crate::compose::tools::ToolComposite;
19
20type LlmProviderFn = Arc<dyn Fn(&gemini_adk_rs::State) -> Arc<dyn BaseLlm> + Send + Sync>;
21
22/// Inner state of an AgentBuilder — shared via Arc for copy-on-write.
23#[derive(Clone)]
24struct AgentBuilderInner {
25    name: String,
26    model: Option<ModelId>,
27    instruction: Option<String>,
28    instruction_provider: Option<Arc<dyn gemini_adk_rs::instruction::InstructionProvider>>,
29    llm_provider: Option<LlmProviderFn>,
30    voice: Option<Voice>,
31    temperature: Option<f32>,
32    top_p: Option<f32>,
33    top_k: Option<u32>,
34    max_output_tokens: Option<u32>,
35    stop_sequences: Vec<String>,
36    response_modalities: Option<Vec<Modality>>,
37    thinking_budget: Option<u32>,
38    tools: Vec<ToolEntry>,
39    built_in_tools: Vec<Tool>,
40    writes: Vec<String>,
41    reads: Vec<String>,
42    sub_agents: Vec<AgentBuilder>,
43    isolate: bool,
44    stay: bool,
45    description: Option<String>,
46    output_schema: Option<serde_json::Value>,
47    output_key: Option<String>,
48    transfer_to_agent: Option<String>,
49    /// Middleware layers to install on the compiled `LlmTextAgent`.
50    middleware_layers: Vec<Arc<dyn Middleware>>,
51    /// Configuration problems found by setters (which cannot fail), reported
52    /// as one [`ConfigError`] by [`AgentBuilder::build`].
53    config_errors: Vec<String>,
54}
55
56/// An entry in the builder's tool list — either a runtime ToolKind or a declaration.
57#[derive(Clone)]
58pub enum ToolEntry {
59    /// A runtime tool with a handler function.
60    Runtime(Arc<dyn ToolEntryTrait>),
61    /// A wire-level tool declaration (e.g., built-in tools like Google Search).
62    Declaration(Tool),
63}
64
65/// Trait for tool entries that can provide a name (for dedup/inspection).
66pub trait ToolEntryTrait: Send + Sync + 'static {
67    /// The tool's registered name.
68    fn name(&self) -> &str;
69    /// Convert this entry into the runtime `ToolKind` variant for dispatch.
70    fn to_tool_kind(&self) -> ToolKind;
71}
72
73/// Copy-on-write immutable builder for agent construction.
74///
75/// Every setter returns a new `AgentBuilder`, leaving the original unchanged.
76/// This makes builders safe to share as templates.
77///
78/// # Basic Usage
79///
80/// ```rust
81/// use gemini_adk_fluent_rs::builder::AgentBuilder;
82/// use gemini_genai_rs::prelude::ModelId;
83///
84/// let agent = AgentBuilder::new("analyst")
85///     .model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO)
86///     .instruction("Analyze the given topic")
87///     .temperature(0.3);
88///
89/// assert_eq!(agent.name(), "analyst");
90/// assert_eq!(agent.get_temperature(), Some(0.3));
91/// ```
92///
93/// # Copy-on-Write Pattern
94///
95/// Cloning a builder and modifying the clone leaves the original unchanged.
96/// This is useful for creating template builders with shared defaults.
97///
98/// ```rust
99/// use gemini_adk_fluent_rs::builder::AgentBuilder;
100///
101/// let base = AgentBuilder::new("researcher")
102///     .instruction("You are a research assistant.")
103///     .temperature(0.5);
104///
105/// let creative = base.clone().temperature(0.9);
106/// let precise  = base.clone().temperature(0.1);
107///
108/// // Original unchanged
109/// assert_eq!(base.get_temperature(), Some(0.5));
110/// assert_eq!(creative.get_temperature(), Some(0.9));
111/// assert_eq!(precise.get_temperature(), Some(0.1));
112/// ```
113///
114/// # Sampling Parameters
115///
116/// ```rust
117/// use gemini_adk_fluent_rs::builder::AgentBuilder;
118///
119/// let agent = AgentBuilder::new("sampler")
120///     .temperature(0.7)
121///     .top_p(0.95)
122///     .top_k(40)
123///     .max_output_tokens(4096);
124///
125/// assert_eq!(agent.get_top_p(), Some(0.95));
126/// assert_eq!(agent.get_top_k(), Some(40));
127/// assert_eq!(agent.get_max_output_tokens(), Some(4096));
128/// ```
129///
130/// # Built-in Tools
131///
132/// ```rust
133/// use gemini_adk_fluent_rs::builder::AgentBuilder;
134///
135/// let agent = AgentBuilder::new("searcher")
136///     .google_search()
137///     .code_execution()
138///     .url_context();
139///
140/// assert_eq!(agent.tool_count(), 3);
141/// ```
142///
143/// # Thinking Budget
144///
145/// ```rust
146/// use gemini_adk_fluent_rs::builder::AgentBuilder;
147///
148/// let agent = AgentBuilder::new("thinker")
149///     .thinking(2048);
150///
151/// assert_eq!(agent.get_thinking_budget(), Some(2048));
152/// ```
153#[derive(Clone)]
154pub struct AgentBuilder {
155    inner: Arc<AgentBuilderInner>,
156}
157
158impl AgentBuilder {
159    /// Create a new builder with the given agent name.
160    pub fn new(name: impl Into<String>) -> Self {
161        Self {
162            inner: Arc::new(AgentBuilderInner {
163                name: name.into(),
164                model: None,
165                instruction: None,
166                instruction_provider: None,
167                llm_provider: None,
168                voice: None,
169                temperature: None,
170                top_p: None,
171                top_k: None,
172                max_output_tokens: None,
173                stop_sequences: Vec::new(),
174                response_modalities: None,
175                thinking_budget: None,
176                tools: Vec::new(),
177                built_in_tools: Vec::new(),
178                writes: Vec::new(),
179                reads: Vec::new(),
180                sub_agents: Vec::new(),
181                isolate: false,
182                stay: false,
183                description: None,
184                output_schema: None,
185                output_key: None,
186                transfer_to_agent: None,
187                middleware_layers: Vec::new(),
188                config_errors: Vec::new(),
189            }),
190        }
191    }
192
193    // ── Private helper: clone-on-write ──
194
195    fn mutate(&self) -> AgentBuilderInner {
196        (*self.inner).clone()
197    }
198
199    fn with(inner: AgentBuilderInner) -> Self {
200        Self {
201            inner: Arc::new(inner),
202        }
203    }
204
205    // ── Accessors ──
206
207    /// The agent name.
208    pub fn name(&self) -> &str {
209        &self.inner.name
210    }
211
212    /// Configured model, if any.
213    pub fn get_model(&self) -> Option<&ModelId> {
214        self.inner.model.as_ref()
215    }
216
217    /// Configured instruction, if any.
218    pub fn get_instruction(&self) -> Option<&str> {
219        self.inner.instruction.as_deref()
220    }
221
222    /// Configured voice, if any.
223    pub fn get_voice(&self) -> Option<&Voice> {
224        self.inner.voice.as_ref()
225    }
226
227    /// Configured temperature, if any.
228    pub fn get_temperature(&self) -> Option<f32> {
229        self.inner.temperature
230    }
231
232    /// Whether text-only mode is set.
233    pub fn is_text_only(&self) -> bool {
234        self.inner
235            .response_modalities
236            .as_ref()
237            .map(|m| m == &[Modality::Text])
238            .unwrap_or(false)
239    }
240
241    /// Configured thinking budget, if any.
242    pub fn get_thinking_budget(&self) -> Option<u32> {
243        self.inner.thinking_budget
244    }
245
246    /// State keys this agent writes.
247    pub fn get_writes(&self) -> &[String] {
248        &self.inner.writes
249    }
250
251    /// State keys this agent reads.
252    pub fn get_reads(&self) -> &[String] {
253        &self.inner.reads
254    }
255
256    /// Sub-agents registered.
257    pub fn get_sub_agents(&self) -> &[AgentBuilder] {
258        &self.inner.sub_agents
259    }
260
261    /// Whether agent runs in isolated state.
262    pub fn is_isolated(&self) -> bool {
263        self.inner.isolate
264    }
265
266    /// Whether agent stays after transfer.
267    pub fn is_stay(&self) -> bool {
268        self.inner.stay
269    }
270
271    /// Number of tool entries.
272    pub fn tool_count(&self) -> usize {
273        self.inner.tools.len() + self.inner.built_in_tools.len()
274    }
275
276    /// Configured top_p, if any.
277    pub fn get_top_p(&self) -> Option<f32> {
278        self.inner.top_p
279    }
280
281    /// Configured top_k, if any.
282    pub fn get_top_k(&self) -> Option<u32> {
283        self.inner.top_k
284    }
285
286    /// Configured max_output_tokens, if any.
287    pub fn get_max_output_tokens(&self) -> Option<u32> {
288        self.inner.max_output_tokens
289    }
290
291    /// Configured stop sequences.
292    pub fn get_stop_sequences(&self) -> &[String] {
293        &self.inner.stop_sequences
294    }
295
296    /// Configured description, if any.
297    pub fn get_description(&self) -> Option<&str> {
298        self.inner.description.as_deref()
299    }
300
301    /// Configured output schema, if any.
302    pub fn get_output_schema(&self) -> Option<&serde_json::Value> {
303        self.inner.output_schema.as_ref()
304    }
305
306    /// Get the configured output key.
307    pub fn get_output_key(&self) -> Option<&str> {
308        self.inner.output_key.as_deref()
309    }
310
311    /// Configured transfer target agent, if any.
312    pub fn get_transfer_to(&self) -> Option<&str> {
313        self.inner.transfer_to_agent.as_deref()
314    }
315
316    /// Number of registered middleware layers.
317    pub fn middleware_layer_count(&self) -> usize {
318        self.inner.middleware_layers.len()
319    }
320
321    // ── Fluent Setters (copy-on-write) ──
322
323    /// Set the Gemini model.
324    pub fn model(self, model: ModelId) -> Self {
325        let mut inner = self.mutate();
326        inner.model = Some(model);
327        Self::with(inner)
328    }
329
330    /// Set the system instruction.
331    pub fn instruction(self, inst: impl Into<String>) -> Self {
332        let mut inner = self.mutate();
333        inner.instruction = Some(inst.into());
334        Self::with(inner)
335    }
336
337    /// Set a dynamic instruction source — any `Fn(&State) -> String`
338    /// closure or a `TemplateInstruction` (feature `templates`), resolved
339    /// against live session state at the start of every run. Wins over
340    /// [`instruction`](Self::instruction) when both are set.
341    pub fn instruction_provider(
342        self,
343        provider: impl gemini_adk_rs::instruction::InstructionProvider + 'static,
344    ) -> Self {
345        let mut inner = self.mutate();
346        inner.instruction_provider = Some(Arc::new(provider));
347        Self::with(inner)
348    }
349
350    /// Set a dynamic model source, resolved against session state at the
351    /// start of every run — risk-based escalation to a stronger model, cost
352    /// routing to a cheaper one, per-tenant model selection — without
353    /// rebuilding the agent. Wins over the constructor's model when set.
354    pub fn llm_provider(
355        self,
356        provider: impl Fn(&gemini_adk_rs::State) -> Arc<dyn BaseLlm> + Send + Sync + 'static,
357    ) -> Self {
358        let mut inner = self.mutate();
359        inner.llm_provider = Some(Arc::new(provider));
360        Self::with(inner)
361    }
362
363    /// Set the output voice.
364    pub fn voice(self, voice: Voice) -> Self {
365        let mut inner = self.mutate();
366        inner.voice = Some(voice);
367        Self::with(inner)
368    }
369
370    /// Set the temperature.
371    pub fn temperature(self, t: f32) -> Self {
372        let mut inner = self.mutate();
373        inner.temperature = Some(t);
374        Self::with(inner)
375    }
376
377    /// Set text-only mode (no audio output).
378    pub fn text_only(self) -> Self {
379        let mut inner = self.mutate();
380        inner.response_modalities = Some(vec![Modality::Text]);
381        Self::with(inner)
382    }
383
384    /// Set response modalities explicitly.
385    pub fn response_modalities(self, modalities: Vec<Modality>) -> Self {
386        let mut inner = self.mutate();
387        inner.response_modalities = Some(modalities);
388        Self::with(inner)
389    }
390
391    /// Enable thinking with a token budget.
392    pub fn thinking(self, budget: u32) -> Self {
393        let mut inner = self.mutate();
394        inner.thinking_budget = Some(budget);
395        Self::with(inner)
396    }
397
398    /// Add a built-in URL context tool.
399    pub fn url_context(self) -> Self {
400        let mut inner = self.mutate();
401        inner.built_in_tools.push(Tool::url_context());
402        Self::with(inner)
403    }
404
405    /// Add a built-in Google Search tool.
406    pub fn google_search(self) -> Self {
407        let mut inner = self.mutate();
408        inner.built_in_tools.push(Tool::google_search());
409        Self::with(inner)
410    }
411
412    /// Add a built-in code execution tool.
413    pub fn code_execution(self) -> Self {
414        let mut inner = self.mutate();
415        inner.built_in_tools.push(Tool::code_execution());
416        Self::with(inner)
417    }
418
419    /// Declare a state key this agent writes.
420    pub fn writes(self, key: impl Into<String>) -> Self {
421        let mut inner = self.mutate();
422        inner.writes.push(key.into());
423        Self::with(inner)
424    }
425
426    /// Declare a state key this agent reads.
427    pub fn reads(self, key: impl Into<String>) -> Self {
428        let mut inner = self.mutate();
429        inner.reads.push(key.into());
430        Self::with(inner)
431    }
432
433    /// Add a sub-agent for transfer.
434    pub fn sub_agent(self, agent: AgentBuilder) -> Self {
435        let mut inner = self.mutate();
436        inner.sub_agents.push(agent);
437        Self::with(inner)
438    }
439
440    /// Run this agent in isolated state (no shared state).
441    pub fn isolate(self) -> Self {
442        let mut inner = self.mutate();
443        inner.isolate = true;
444        Self::with(inner)
445    }
446
447    /// Keep this agent active after transfer (don't tear down).
448    pub fn stay(self) -> Self {
449        let mut inner = self.mutate();
450        inner.stay = true;
451        Self::with(inner)
452    }
453
454    /// Set top_p (nucleus sampling).
455    pub fn top_p(self, p: f32) -> Self {
456        let mut inner = self.mutate();
457        inner.top_p = Some(p);
458        Self::with(inner)
459    }
460
461    /// Set top_k (top-k sampling).
462    pub fn top_k(self, k: u32) -> Self {
463        let mut inner = self.mutate();
464        inner.top_k = Some(k);
465        Self::with(inner)
466    }
467
468    /// Set maximum output tokens.
469    pub fn max_output_tokens(self, n: u32) -> Self {
470        let mut inner = self.mutate();
471        inner.max_output_tokens = Some(n);
472        Self::with(inner)
473    }
474
475    /// Set stop sequences.
476    pub fn stop_sequences(self, seqs: Vec<String>) -> Self {
477        let mut inner = self.mutate();
478        inner.stop_sequences = seqs;
479        Self::with(inner)
480    }
481
482    /// Set a description for this agent (used in tool/agent metadata).
483    pub fn description(self, desc: impl Into<String>) -> Self {
484        let mut inner = self.mutate();
485        inner.description = Some(desc.into());
486        Self::with(inner)
487    }
488
489    /// Set a JSON schema for structured output.
490    pub fn output_schema(self, schema: serde_json::Value) -> Self {
491        let mut inner = self.mutate();
492        inner.output_schema = Some(schema);
493        Self::with(inner)
494    }
495
496    /// Set the output key — agent's final text response is auto-saved to this state key.
497    pub fn output_key(self, key: impl Into<String>) -> Self {
498        let mut inner = self.mutate();
499        inner.output_key = Some(key.into());
500        Self::with(inner)
501    }
502
503    /// Set a default transfer target agent.
504    pub fn transfer_to(self, agent_name: impl Into<String>) -> Self {
505        let mut inner = self.mutate();
506        inner.transfer_to_agent = Some(agent_name.into());
507        Self::with(inner)
508    }
509
510    // ── Upstream naming aliases ──
511
512    /// Alias for [`instruction`](Self::instruction) — matches upstream Python `Agent.instruct()`.
513    pub fn instruct(self, inst: impl Into<String>) -> Self {
514        self.instruction(inst)
515    }
516
517    /// Alias for [`description`](Self::description) — matches upstream Python `Agent.describe()`.
518    pub fn describe(self, desc: impl Into<String>) -> Self {
519        self.description(desc)
520    }
521
522    /// Register one tool: anything that implements [`ToolFunction`] — a
523    /// `SimpleTool`/`TypedTool`, the value a `#[tool]` function returns, or an
524    /// `Arc<dyn ToolFunction>` you already hold.
525    ///
526    /// ```no_run
527    /// # use gemini_adk_fluent_rs::prelude::*;
528    /// # use std::sync::Arc;
529    /// #[tool("Get the weather for a city")]
530    /// async fn get_weather(city: String) -> Result<serde_json::Value, ToolError> {
531    ///     Ok(serde_json::json!({"city": city, "temp": 22}))
532    /// }
533    /// let agent = AgentBuilder::new("assistant").tool(get_weather());
534    /// ```
535    pub fn tool(self, f: impl ToolFunction + 'static) -> Self {
536        self.tools(ToolComposite::from_function(Arc::new(f)))
537    }
538
539    /// Register tools: a `|`-composed [`ToolComposite`] from the `T`
540    /// namespace, or a single [`ToolFunction`].
541    ///
542    /// `T::mcp(..)` needs an async connection that this synchronous builder
543    /// cannot perform; it is rejected by [`build`](Self::build) with a
544    /// [`ConfigError`] — attach MCP toolsets to a `Live` session instead.
545    ///
546    /// ```no_run
547    /// # use gemini_adk_fluent_rs::prelude::*;
548    /// # use serde_json::json;
549    /// let tools = T::simple("greet", "Greet", |_| async { Ok(json!({})) })
550    ///     | T::google_search();
551    /// AgentBuilder::new("assistant").tools(tools);
552    /// ```
553    pub fn tools(self, tools: impl Into<ToolComposite>) -> Self {
554        use crate::compose::tools::{DeferredTool, ToolResolution};
555        let mut inner = self.mutate();
556        for entry in tools.into().entries {
557            match entry.classify() {
558                ToolResolution::Runtime(f) => {
559                    inner
560                        .tools
561                        .push(ToolEntry::Runtime(Arc::new(ToolFunctionEntry(f))));
562                }
563                ToolResolution::BuiltIn(t) => {
564                    inner.built_in_tools.push(t);
565                }
566                ToolResolution::Agent {
567                    name,
568                    description,
569                    agent,
570                } => {
571                    // Expose the sub-agent as a callable tool over a fresh State.
572                    let tool = gemini_adk_rs::TextAgentTool::from_arc(
573                        name,
574                        description,
575                        agent,
576                        gemini_adk_rs::State::new(),
577                    );
578                    inner
579                        .tools
580                        .push(ToolEntry::Runtime(Arc::new(ToolFunctionEntry(Arc::new(
581                            tool,
582                        )))));
583                }
584                ToolResolution::Deferred(DeferredTool::Mcp { params }) => {
585                    // An MCP toolset needs an async handshake, which the
586                    // synchronous text-agent `build()` cannot perform. It
587                    // belongs on a `Live` session (resolved at connect); make
588                    // `build` fail rather than drop the tool silently — the same
589                    // outcome `Live::connect` gives an unreachable MCP server.
590                    inner.config_errors.push(format!(
591                        "T::mcp({params:?}) cannot be attached to a text AgentBuilder: MCP \
592                         toolsets need an async connection, which only a Live session performs \
593                         (`Live::builder().tools(T::mcp(..))`)"
594                    ));
595                }
596            }
597        }
598        Self::with(inner)
599    }
600
601    /// Attach output guards. Each model response is validated against every
602    /// guard; if any rejects the output the agent run fails with an
603    /// [`AgentError`](gemini_adk_rs::error::AgentError) listing the violations.
604    ///
605    /// Accepts a single guard or a `|`-composed [`GuardComposite`]:
606    ///
607    /// ```no_run
608    /// # use gemini_adk_fluent_rs::prelude::*;
609    /// AgentBuilder::new("writer").guard(G::pii() | G::length(1, 2000));
610    /// ```
611    ///
612    /// The guards are installed as an `after_model` middleware layer, so they
613    /// accumulate with `.middleware(...)` and honor copy-on-write.
614    pub fn guard(self, guard: impl Into<GuardComposite>) -> Self {
615        let mut inner = self.mutate();
616        inner.middleware_layers.push(guard.into().into_middleware());
617        Self::with(inner)
618    }
619
620    /// Attach a context policy that rewrites conversation history before each
621    /// model call (e.g. windowing, role filtering, tool-result exclusion).
622    ///
623    /// Accepts a single policy or a `+`-composed [`ContextComposite`]:
624    ///
625    /// ```no_run
626    /// # use gemini_adk_fluent_rs::prelude::*;
627    /// AgentBuilder::new("chat").context(C::window(10) + C::user_only());
628    /// ```
629    ///
630    /// The policy is installed as a `transform_request` middleware layer.
631    pub fn context(self, policy: impl Into<ContextComposite>) -> Self {
632        let mut inner = self.mutate();
633        inner
634            .middleware_layers
635            .push(policy.into().into_middleware());
636        Self::with(inner)
637    }
638
639    /// Disallow transfer to peer agents.
640    pub fn no_peers(self) -> Self {
641        self.isolate()
642    }
643
644    /// Attach middleware — a `|`-composed [`MiddlewareComposite`] from the
645    /// `M` namespace or a single `Arc<dyn Middleware>`. All layers are
646    /// installed on the compiled `LlmTextAgent` in the order given.
647    ///
648    /// Multiple calls to `.middleware()` accumulate: the new layers are
649    /// appended after any previously registered layers, preserving the
650    /// copy-on-write contract.
651    ///
652    /// ```no_run
653    /// # use gemini_adk_fluent_rs::prelude::*;
654    /// let agent = AgentBuilder::new("analyst")
655    ///     .instruction("Analyze topics")
656    ///     .middleware(M::log() | M::latency());
657    /// ```
658    pub fn middleware(self, middleware: impl Into<MiddlewareComposite>) -> Self {
659        let mut inner = self.mutate();
660        inner.middleware_layers.extend(middleware.into().layers);
661        Self::with(inner)
662    }
663
664    // ── Compilation ──
665
666    /// Compile this builder into an executable `TextAgent`.
667    ///
668    /// The LLM is required because `TextAgent` makes `BaseLlm::generate()` calls.
669    /// Builder configuration (instruction, temperature, tools) is transferred to
670    /// the resulting agent.
671    ///
672    /// Fails with a [`ConfigError`] when the configuration cannot be realized
673    /// by a text agent — today, an MCP toolset (`T::mcp`) in
674    /// [`tools`](Self::tools), which needs the async connect only a `Live`
675    /// session performs.
676    ///
677    /// ```no_run
678    /// # use gemini_adk_fluent_rs::prelude::*;
679    /// # use std::sync::Arc;
680    /// # async fn run() -> Result<(), AgentError> {
681    /// let llm = Arc::new(GeminiLlm::new(GeminiLlmParams::default()));
682    /// let agent = AgentBuilder::new("analyst")
683    ///     .instruction("Analyze the topic")
684    ///     .temperature(0.3)
685    ///     .build(llm)?;
686    ///
687    /// let state = State::new();
688    /// let result = agent.run(&state).await?;
689    /// # let _ = result; Ok(())
690    /// # }
691    /// ```
692    pub fn build(self, llm: Arc<dyn BaseLlm>) -> Result<Arc<dyn TextAgent>, ConfigError> {
693        if !self.inner.config_errors.is_empty() {
694            return Err(ConfigError {
695                issues: self.inner.config_errors.clone(),
696            });
697        }
698        let mut agent = LlmTextAgent::new(&self.inner.name, llm);
699
700        if let Some(inst) = &self.inner.instruction {
701            agent = agent.instruction(inst);
702        }
703        if let Some(provider) = &self.inner.instruction_provider {
704            agent = agent.instruction_provider(provider.clone());
705        }
706        if let Some(provider) = &self.inner.llm_provider {
707            let provider_clone = provider.clone();
708            agent = agent.llm_provider(move |state| provider_clone(state));
709        }
710        if let Some(t) = self.inner.temperature {
711            agent = agent.temperature(t);
712        }
713        if let Some(n) = self.inner.max_output_tokens {
714            agent = agent.max_output_tokens(n);
715        }
716
717        // Build ToolDispatcher from registered tools.
718        if !self.inner.tools.is_empty() {
719            let mut dispatcher = ToolDispatcher::new();
720            for entry in &self.inner.tools {
721                match entry {
722                    ToolEntry::Runtime(t) => {
723                        let kind = t.to_tool_kind();
724                        match kind {
725                            ToolKind::Function(f) => dispatcher.register_function(f),
726                            ToolKind::Streaming(s) => dispatcher.register_streaming(s),
727                            ToolKind::InputStream(i) => dispatcher.register_input_streaming(i),
728                        }
729                    }
730                    ToolEntry::Declaration(_) => {
731                        // Built-in tool declarations (google_search, etc.) are sent
732                        // as-is; they don't have runtime handlers for text dispatch.
733                    }
734                }
735            }
736            if !dispatcher.is_empty() {
737                agent = agent.tools(Arc::new(dispatcher));
738            }
739        }
740
741        // Install middleware layers from the builder.
742        for mw in &self.inner.middleware_layers {
743            agent = agent.add_middleware(mw.clone());
744        }
745
746        Ok(Arc::new(agent))
747    }
748}
749
750/// Adapter that wraps an `Arc<dyn ToolFunction>` as a `ToolEntryTrait`.
751#[derive(Clone)]
752struct ToolFunctionEntry(Arc<dyn ToolFunction>);
753
754impl ToolEntryTrait for ToolFunctionEntry {
755    fn name(&self) -> &str {
756        self.0.name()
757    }
758
759    fn to_tool_kind(&self) -> ToolKind {
760        ToolKind::Function(self.0.clone())
761    }
762}
763
764impl std::fmt::Debug for AgentBuilder {
765    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
766        f.debug_struct("AgentBuilder")
767            .field("name", &self.inner.name)
768            .field("model", &self.inner.model)
769            .field("instruction", &self.inner.instruction)
770            .field("temperature", &self.inner.temperature)
771            .field("text_only", &self.is_text_only())
772            .field("tool_count", &self.tool_count())
773            .field("sub_agents", &self.inner.sub_agents.len())
774            .finish()
775    }
776}
777
778#[cfg(test)]
779mod tests {
780    use super::*;
781    use async_trait::async_trait;
782    use gemini_adk_rs::llm::{LlmError, LlmRequest, LlmResponse};
783    use gemini_genai_rs::prelude::{Content, Part, Role};
784
785    /// A mock LLM for build() tests.
786    struct MockLlm(String);
787
788    #[async_trait]
789    impl BaseLlm for MockLlm {
790        fn model_id(&self) -> &str {
791            "mock"
792        }
793        async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
794            Ok(LlmResponse {
795                content: Content {
796                    role: Some(Role::Model),
797                    parts: vec![Part::Text {
798                        text: self.0.clone(),
799                    }],
800                },
801                finish_reason: Some("STOP".into()),
802                usage: None,
803            })
804        }
805    }
806
807    #[test]
808    fn builder_creates_with_name() {
809        let b = AgentBuilder::new("test-agent");
810        assert_eq!(b.name(), "test-agent");
811    }
812
813    #[test]
814    fn fluent_chaining_works() {
815        let b = AgentBuilder::new("agent")
816            .instruction("Be helpful")
817            .temperature(0.7)
818            .model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO);
819
820        assert_eq!(b.get_instruction(), Some("Be helpful"));
821        assert_eq!(b.get_temperature(), Some(0.7));
822        assert_eq!(b.get_model(), Some(&ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO));
823    }
824
825    #[test]
826    fn copy_on_write_clone_independence() {
827        let base = AgentBuilder::new("base").temperature(0.5);
828        let variant = base.clone().temperature(0.9);
829
830        // Original unchanged
831        assert_eq!(base.get_temperature(), Some(0.5));
832        // Variant has new value
833        assert_eq!(variant.get_temperature(), Some(0.9));
834    }
835
836    #[test]
837    fn text_only_sets_modalities() {
838        let b = AgentBuilder::new("text").text_only();
839        assert!(b.is_text_only());
840    }
841
842    #[test]
843    fn url_context_adds_tool() {
844        let b = AgentBuilder::new("search").url_context();
845        assert_eq!(b.tool_count(), 1);
846    }
847
848    #[test]
849    fn google_search_adds_tool() {
850        let b = AgentBuilder::new("search").google_search();
851        assert_eq!(b.tool_count(), 1);
852    }
853
854    #[test]
855    fn code_execution_adds_tool() {
856        let b = AgentBuilder::new("code").code_execution();
857        assert_eq!(b.tool_count(), 1);
858    }
859
860    #[test]
861    fn thinking_sets_budget() {
862        let b = AgentBuilder::new("thinker").thinking(2048);
863        assert_eq!(b.get_thinking_budget(), Some(2048));
864    }
865
866    #[test]
867    fn writes_and_reads_keys() {
868        let b = AgentBuilder::new("data").writes("output").reads("input");
869        assert_eq!(b.get_writes(), &["output"]);
870        assert_eq!(b.get_reads(), &["input"]);
871    }
872
873    #[test]
874    fn sub_agent_registration() {
875        let child = AgentBuilder::new("child");
876        let parent = AgentBuilder::new("parent").sub_agent(child);
877        assert_eq!(parent.get_sub_agents().len(), 1);
878        assert_eq!(parent.get_sub_agents()[0].name(), "child");
879    }
880
881    #[test]
882    fn isolate_and_stay() {
883        let b = AgentBuilder::new("agent").isolate().stay();
884        assert!(b.is_isolated());
885        assert!(b.is_stay());
886    }
887
888    #[test]
889    fn debug_display() {
890        let b = AgentBuilder::new("debug-test");
891        let debug = format!("{b:?}");
892        assert!(debug.contains("debug-test"));
893    }
894
895    #[test]
896    fn top_p_sets_value() {
897        let b = AgentBuilder::new("agent").top_p(0.95);
898        assert_eq!(b.get_top_p(), Some(0.95));
899    }
900
901    #[test]
902    fn top_k_sets_value() {
903        let b = AgentBuilder::new("agent").top_k(40);
904        assert_eq!(b.get_top_k(), Some(40));
905    }
906
907    #[test]
908    fn max_output_tokens_sets_value() {
909        let b = AgentBuilder::new("agent").max_output_tokens(4096);
910        assert_eq!(b.get_max_output_tokens(), Some(4096));
911    }
912
913    #[test]
914    fn stop_sequences_sets_value() {
915        let b =
916            AgentBuilder::new("agent").stop_sequences(vec!["END".to_string(), "STOP".to_string()]);
917        assert_eq!(b.get_stop_sequences().len(), 2);
918    }
919
920    #[test]
921    fn description_sets_value() {
922        let b = AgentBuilder::new("agent").description("A helpful agent");
923        assert_eq!(b.get_description(), Some("A helpful agent"));
924    }
925
926    #[test]
927    fn output_schema_sets_value() {
928        let schema = serde_json::json!({"type": "object"});
929        let b = AgentBuilder::new("agent").output_schema(schema.clone());
930        assert_eq!(b.get_output_schema(), Some(&schema));
931    }
932
933    #[test]
934    fn transfer_to_sets_value() {
935        let b = AgentBuilder::new("agent").transfer_to("target-agent");
936        assert_eq!(b.get_transfer_to(), Some("target-agent"));
937    }
938
939    #[test]
940    fn full_fluent_chain() {
941        let b = AgentBuilder::new("full-agent")
942            .model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO)
943            .instruction("Be helpful")
944            .temperature(0.7)
945            .top_p(0.95)
946            .top_k(40)
947            .max_output_tokens(4096)
948            .thinking(2048)
949            .description("A fully configured agent")
950            .google_search()
951            .writes("output")
952            .reads("input");
953
954        assert_eq!(b.name(), "full-agent");
955        assert_eq!(b.get_temperature(), Some(0.7));
956        assert_eq!(b.get_top_p(), Some(0.95));
957        assert_eq!(b.get_top_k(), Some(40));
958        assert_eq!(b.get_max_output_tokens(), Some(4096));
959        assert_eq!(b.get_thinking_budget(), Some(2048));
960        assert_eq!(b.get_description(), Some("A fully configured agent"));
961        assert_eq!(b.tool_count(), 1);
962    }
963
964    // ── build() tests ──
965
966    #[tokio::test]
967    async fn build_produces_executable_agent() {
968        let llm: Arc<dyn BaseLlm> = Arc::new(MockLlm("built agent output".into()));
969        let agent = AgentBuilder::new("test")
970            .instruction("Be helpful")
971            .temperature(0.5)
972            .build(llm)
973            .unwrap();
974
975        assert_eq!(agent.name(), "test");
976        let state = gemini_adk_rs::State::new();
977        let result = agent.run(&state).await.unwrap();
978        assert_eq!(result, "built agent output");
979    }
980
981    #[tokio::test]
982    async fn build_stores_output_in_state() {
983        let llm: Arc<dyn BaseLlm> = Arc::new(MockLlm("state output".into()));
984        let agent = AgentBuilder::new("test").build(llm).unwrap();
985        let state = gemini_adk_rs::State::new();
986        agent.run(&state).await.unwrap();
987        assert_eq!(state.get::<String>("output"), Some("state output".into()));
988    }
989
990    #[tokio::test]
991    async fn build_reads_input_from_state() {
992        use gemini_adk_rs::llm::LlmRequest;
993
994        // An LLM that echoes whatever it receives.
995        struct EchoLlm;
996        #[async_trait]
997        impl BaseLlm for EchoLlm {
998            fn model_id(&self) -> &str {
999                "echo"
1000            }
1001            async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
1002                let text: String = req
1003                    .contents
1004                    .iter()
1005                    .flat_map(|c| &c.parts)
1006                    .filter_map(|p| match p {
1007                        Part::Text { text } => Some(text.as_str()),
1008                        _ => None,
1009                    })
1010                    .collect::<Vec<_>>()
1011                    .join("");
1012                Ok(LlmResponse {
1013                    content: Content {
1014                        role: Some(Role::Model),
1015                        parts: vec![Part::Text { text }],
1016                    },
1017                    finish_reason: Some("STOP".into()),
1018                    usage: None,
1019                })
1020            }
1021        }
1022
1023        let agent = AgentBuilder::new("echo").build(Arc::new(EchoLlm)).unwrap();
1024        let state = gemini_adk_rs::State::new();
1025        let _ = state.set("input", "hello from state");
1026        let result = agent.run(&state).await.unwrap();
1027        assert!(result.contains("hello from state"));
1028    }
1029
1030    // ── Middleware end-to-end tests ──
1031
1032    /// A mock LLM that issues one tool call and then returns text.
1033    struct ToolCallingMockLlm {
1034        tool_name: &'static str,
1035        final_text: &'static str,
1036    }
1037
1038    #[async_trait]
1039    impl BaseLlm for ToolCallingMockLlm {
1040        fn model_id(&self) -> &str {
1041            "tool-mock"
1042        }
1043
1044        async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
1045            use gemini_genai_rs::prelude::FunctionCall;
1046
1047            // If any part is a FunctionResponse, we already dispatched — return text.
1048            let already_responded = req
1049                .contents
1050                .iter()
1051                .flat_map(|c| &c.parts)
1052                .any(|p| matches!(p, Part::FunctionResponse { .. }));
1053
1054            if already_responded {
1055                Ok(LlmResponse {
1056                    content: Content {
1057                        role: Some(Role::Model),
1058                        parts: vec![Part::Text {
1059                            text: self.final_text.to_string(),
1060                        }],
1061                    },
1062                    finish_reason: Some("STOP".into()),
1063                    usage: None,
1064                })
1065            } else {
1066                Ok(LlmResponse {
1067                    content: Content {
1068                        role: Some(Role::Model),
1069                        parts: vec![Part::FunctionCall {
1070                            function_call: FunctionCall {
1071                                name: self.tool_name.to_string(),
1072                                args: serde_json::json!({"x": 1}),
1073                                id: Some("call-1".into()),
1074                            },
1075                        }],
1076                    },
1077                    finish_reason: None,
1078                    usage: None,
1079                })
1080            }
1081        }
1082    }
1083
1084    /// Verify that `M::before_model` and `M::after_tool` hooks fire when the agent runs.
1085    #[tokio::test]
1086    async fn middleware_hooks_fire_end_to_end() {
1087        use crate::compose::middleware::M;
1088        use gemini_adk_rs::tool::SimpleTool;
1089        use std::sync::atomic::{AtomicUsize, Ordering};
1090
1091        let before_model_count = Arc::new(AtomicUsize::new(0));
1092        let after_tool_count = Arc::new(AtomicUsize::new(0));
1093
1094        let bm = before_model_count.clone();
1095        let at = after_tool_count.clone();
1096
1097        let mw = M::before_model(move |_req| {
1098            bm.fetch_add(1, Ordering::SeqCst);
1099            Ok(())
1100        }) | M::after_tool(move |_call, _result| {
1101            at.fetch_add(1, Ordering::SeqCst);
1102            Ok(())
1103        });
1104
1105        let llm: Arc<dyn BaseLlm> = Arc::new(ToolCallingMockLlm {
1106            tool_name: "echo_tool",
1107            final_text: "done",
1108        });
1109
1110        let agent = AgentBuilder::new("mw-test")
1111            .middleware(mw)
1112            .tool(SimpleTool::new(
1113                "echo_tool",
1114                "Echo tool",
1115                None,
1116                |_args| async move { Ok(serde_json::json!({"echo": true})) },
1117            ))
1118            .build(llm)
1119            .unwrap();
1120
1121        let state = gemini_adk_rs::State::new();
1122        let result = agent.run(&state).await.unwrap();
1123        assert_eq!(result, "done");
1124
1125        // before_model fires once per LLM call: first call (tool call) + second call (final text).
1126        assert_eq!(
1127            before_model_count.load(Ordering::SeqCst),
1128            2,
1129            "before_model should fire for each generate() call"
1130        );
1131        // after_tool fires once per successful tool dispatch.
1132        assert_eq!(
1133            after_tool_count.load(Ordering::SeqCst),
1134            1,
1135            "after_tool should fire once for the tool dispatch"
1136        );
1137    }
1138
1139    /// Verify copy-on-write: adding middleware to a clone does not affect the original.
1140    #[test]
1141    fn middleware_copy_on_write() {
1142        use crate::compose::middleware::M;
1143
1144        let base = AgentBuilder::new("base").instruction("base");
1145        let with_mw = base.clone().middleware(M::log() | M::latency());
1146
1147        // Original should have no middleware layers.
1148        assert_eq!(base.middleware_layer_count(), 0);
1149        // Clone with middleware should have 2 layers.
1150        assert_eq!(with_mw.middleware_layer_count(), 2);
1151    }
1152
1153    /// Verify `on_error` hook fires when the agent errors.
1154    #[tokio::test]
1155    async fn middleware_on_error_fires_on_failure() {
1156        use crate::compose::middleware::M;
1157        use gemini_adk_rs::llm::LlmError;
1158        use std::sync::atomic::{AtomicUsize, Ordering};
1159
1160        let error_count = Arc::new(AtomicUsize::new(0));
1161        let ec = error_count.clone();
1162
1163        let mw = M::on_error(move |_err| {
1164            ec.fetch_add(1, Ordering::SeqCst);
1165            Ok(())
1166        });
1167
1168        struct FailLlm;
1169        #[async_trait]
1170        impl BaseLlm for FailLlm {
1171            fn model_id(&self) -> &str {
1172                "fail"
1173            }
1174            async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
1175                Err(LlmError::RequestFailed("boom".into()))
1176            }
1177        }
1178
1179        let agent = AgentBuilder::new("error-test")
1180            .middleware(mw)
1181            .build(Arc::new(FailLlm))
1182            .unwrap();
1183
1184        let state = gemini_adk_rs::State::new();
1185        let result = agent.run(&state).await;
1186        assert!(result.is_err(), "agent should fail");
1187        assert_eq!(
1188            error_count.load(Ordering::SeqCst),
1189            1,
1190            "on_error should fire exactly once"
1191        );
1192    }
1193
1194    // ── Guard / context wiring tests ──
1195
1196    /// A mock LLM that echoes a fixed response and records the number of
1197    /// `contents` it was asked to generate from (to observe context rewriting).
1198    struct RecordingLlm {
1199        text: &'static str,
1200        seen_len: Arc<std::sync::atomic::AtomicUsize>,
1201    }
1202
1203    #[async_trait]
1204    impl BaseLlm for RecordingLlm {
1205        fn model_id(&self) -> &str {
1206            "recording-mock"
1207        }
1208
1209        async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
1210            self.seen_len
1211                .store(req.contents.len(), std::sync::atomic::Ordering::SeqCst);
1212            Ok(LlmResponse {
1213                content: Content {
1214                    role: Some(Role::Model),
1215                    parts: vec![Part::Text {
1216                        text: self.text.to_string(),
1217                    }],
1218                },
1219                finish_reason: Some("STOP".into()),
1220                usage: None,
1221            })
1222        }
1223    }
1224
1225    #[tokio::test]
1226    async fn guard_blocks_violating_output() {
1227        use crate::compose::guards::G;
1228
1229        let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1230            text: "you can reach me at agent@example.com",
1231            seen_len: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
1232        });
1233
1234        let agent = AgentBuilder::new("guarded")
1235            .guard(G::pii())
1236            .build(llm)
1237            .unwrap();
1238
1239        let state = gemini_adk_rs::State::new();
1240        let err = agent.run(&state).await.unwrap_err();
1241        assert!(
1242            err.to_string().contains("guard violation"),
1243            "PII guard should veto the response, got: {err}"
1244        );
1245    }
1246
1247    #[tokio::test]
1248    async fn guard_allows_clean_output() {
1249        use crate::compose::guards::G;
1250
1251        let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1252            text: "all clean here",
1253            seen_len: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
1254        });
1255
1256        let agent = AgentBuilder::new("guarded")
1257            .guard(G::pii() | G::length(1, 1000))
1258            .build(llm)
1259            .unwrap();
1260
1261        let state = gemini_adk_rs::State::new();
1262        let result = agent.run(&state).await.unwrap();
1263        assert_eq!(result, "all clean here");
1264    }
1265
1266    #[tokio::test]
1267    async fn context_policy_rewrites_request_history() {
1268        use crate::compose::context::C;
1269
1270        // The agent seeds one user turn; a prepend policy injects a second turn,
1271        // so the LLM should see 2 contents — proving transform_request ran.
1272        let seen = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1273        let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1274            text: "ok",
1275            seen_len: seen.clone(),
1276        });
1277
1278        let agent = AgentBuilder::new("ctx")
1279            .context(C::prepend(Content::user("system preamble")))
1280            .build(llm)
1281            .unwrap();
1282
1283        let state = gemini_adk_rs::State::new();
1284        let _ = state.set("input", "hello");
1285        let _ = agent.run(&state).await.unwrap();
1286        assert_eq!(
1287            seen.load(std::sync::atomic::Ordering::SeqCst),
1288            2,
1289            "context policy should have prepended a turn before the model call"
1290        );
1291    }
1292
1293    #[tokio::test]
1294    async fn context_window_trims_history() {
1295        use crate::compose::context::C;
1296
1297        // window(1) keeps only the last turn. We seed a single input turn and
1298        // prepend two extra turns, then window down to 1 — the model sees 1.
1299        let seen = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1300        let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1301            text: "ok",
1302            seen_len: seen.clone(),
1303        });
1304
1305        let agent = AgentBuilder::new("ctx")
1306            .context(C::prepend(Content::user("a")) + C::prepend(Content::user("b")) + C::window(1))
1307            .build(llm)
1308            .unwrap();
1309
1310        let state = gemini_adk_rs::State::new();
1311        let _ = state.set("input", "hello");
1312        let _ = agent.run(&state).await.unwrap();
1313        assert_eq!(
1314            seen.load(std::sync::atomic::Ordering::SeqCst),
1315            1,
1316            "window(1) should trim history to the last turn"
1317        );
1318    }
1319}
1320
1321#[cfg(test)]
1322mod mcp_rejection_tests {
1323    use super::*;
1324    use gemini_adk_rs::llm::{LlmError, LlmRequest, LlmResponse};
1325
1326    struct NeverLlm;
1327    #[async_trait::async_trait]
1328    impl BaseLlm for NeverLlm {
1329        fn model_id(&self) -> &str {
1330            "never"
1331        }
1332        async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
1333            Err(LlmError::RequestFailed("never".into()))
1334        }
1335    }
1336
1337    /// `T::mcp` on a text agent is a build error naming the tool kind — the
1338    /// same way `Live::connect` fails on it — never a silent drop.
1339    #[test]
1340    fn mcp_toolset_is_a_build_error() {
1341        use crate::compose::tools::T;
1342        let err = AgentBuilder::new("text")
1343            .tools(T::mcp("node ./server.js"))
1344            .build(Arc::new(NeverLlm))
1345            .err()
1346            .expect("build must fail");
1347        assert!(err.to_string().contains("T::mcp"), "{err}");
1348        assert!(err.to_string().contains("Live"), "{err}");
1349    }
1350
1351    /// `tool(..)` takes any `ToolFunction`, including an `Arc<dyn ToolFunction>`.
1352    #[test]
1353    fn tool_accepts_values_and_arcs() {
1354        use gemini_adk_rs::tool::SimpleTool;
1355        let arc: Arc<dyn ToolFunction> = Arc::new(SimpleTool::new("a", "a", None, |_| async {
1356            Ok(serde_json::json!({}))
1357        }));
1358        let b = AgentBuilder::new("t")
1359            .tool(SimpleTool::new("b", "b", None, |_| async {
1360                Ok(serde_json::json!({}))
1361            }))
1362            .tool(arc.clone())
1363            .tools(arc);
1364        assert_eq!(b.tool_count(), 3);
1365    }
1366}
gemini_adk_fluent_rs/builder.rs

gemini_adk_fluent_rs/
builder.rs
