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    /// Decides `T::confirm(..)` tool calls.
52    confirmation_provider: Option<Arc<dyn gemini_adk_rs::confirmation::ConfirmationProvider>>,
53    /// Configuration problems found by setters (which cannot fail), reported
54    /// as one [`ConfigError`] by [`AgentBuilder::build`].
55    config_errors: Vec<String>,
56}
57
58/// An entry in the builder's tool list — either a runtime ToolKind or a declaration.
59#[derive(Clone)]
60pub enum ToolEntry {
61    /// A runtime tool with a handler function.
62    Runtime(Arc<dyn ToolEntryTrait>),
63    /// A wire-level tool declaration (e.g., built-in tools like Google Search).
64    Declaration(Tool),
65}
66
67/// Trait for tool entries that can provide a name (for dedup/inspection).
68pub trait ToolEntryTrait: Send + Sync + 'static {
69    /// The tool's registered name.
70    fn name(&self) -> &str;
71    /// Convert this entry into the runtime `ToolKind` variant for dispatch.
72    fn to_tool_kind(&self) -> ToolKind;
73}
74
75/// Copy-on-write immutable builder for agent construction.
76///
77/// Every setter returns a new `AgentBuilder`, leaving the original unchanged.
78/// This makes builders safe to share as templates.
79///
80/// # Basic Usage
81///
82/// ```rust
83/// use gemini_adk_fluent_rs::builder::AgentBuilder;
84/// use gemini_genai_rs::prelude::ModelId;
85///
86/// let agent = AgentBuilder::new("analyst")
87///     .model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO)
88///     .instruction("Analyze the given topic")
89///     .temperature(0.3);
90///
91/// assert_eq!(agent.name(), "analyst");
92/// assert_eq!(agent.get_temperature(), Some(0.3));
93/// ```
94///
95/// # Copy-on-Write Pattern
96///
97/// Cloning a builder and modifying the clone leaves the original unchanged.
98/// This is useful for creating template builders with shared defaults.
99///
100/// ```rust
101/// use gemini_adk_fluent_rs::builder::AgentBuilder;
102///
103/// let base = AgentBuilder::new("researcher")
104///     .instruction("You are a research assistant.")
105///     .temperature(0.5);
106///
107/// let creative = base.clone().temperature(0.9);
108/// let precise  = base.clone().temperature(0.1);
109///
110/// // Original unchanged
111/// assert_eq!(base.get_temperature(), Some(0.5));
112/// assert_eq!(creative.get_temperature(), Some(0.9));
113/// assert_eq!(precise.get_temperature(), Some(0.1));
114/// ```
115///
116/// # Sampling Parameters
117///
118/// ```rust
119/// use gemini_adk_fluent_rs::builder::AgentBuilder;
120///
121/// let agent = AgentBuilder::new("sampler")
122///     .temperature(0.7)
123///     .top_p(0.95)
124///     .top_k(40)
125///     .max_output_tokens(4096);
126///
127/// assert_eq!(agent.get_top_p(), Some(0.95));
128/// assert_eq!(agent.get_top_k(), Some(40));
129/// assert_eq!(agent.get_max_output_tokens(), Some(4096));
130/// ```
131///
132/// # Built-in Tools
133///
134/// ```rust
135/// use gemini_adk_fluent_rs::builder::AgentBuilder;
136///
137/// let agent = AgentBuilder::new("searcher")
138///     .google_search()
139///     .code_execution()
140///     .url_context();
141///
142/// assert_eq!(agent.tool_count(), 3);
143/// ```
144///
145/// # Thinking Budget
146///
147/// ```rust
148/// use gemini_adk_fluent_rs::builder::AgentBuilder;
149///
150/// let agent = AgentBuilder::new("thinker")
151///     .thinking(2048);
152///
153/// assert_eq!(agent.get_thinking_budget(), Some(2048));
154/// ```
155#[derive(Clone)]
156pub struct AgentBuilder {
157    inner: Arc<AgentBuilderInner>,
158}
159
160impl AgentBuilder {
161    /// Create a new builder with the given agent name.
162    pub fn new(name: impl Into<String>) -> Self {
163        Self {
164            inner: Arc::new(AgentBuilderInner {
165                name: name.into(),
166                model: None,
167                instruction: None,
168                instruction_provider: None,
169                llm_provider: None,
170                voice: None,
171                temperature: None,
172                top_p: None,
173                top_k: None,
174                max_output_tokens: None,
175                stop_sequences: Vec::new(),
176                response_modalities: None,
177                thinking_budget: None,
178                tools: Vec::new(),
179                built_in_tools: Vec::new(),
180                writes: Vec::new(),
181                reads: Vec::new(),
182                sub_agents: Vec::new(),
183                isolate: false,
184                stay: false,
185                description: None,
186                output_schema: None,
187                output_key: None,
188                transfer_to_agent: None,
189                middleware_layers: Vec::new(),
190                confirmation_provider: None,
191                config_errors: Vec::new(),
192            }),
193        }
194    }
195
196    // ── Private helper: clone-on-write ──
197
198    fn mutate(&self) -> AgentBuilderInner {
199        (*self.inner).clone()
200    }
201
202    fn with(inner: AgentBuilderInner) -> Self {
203        Self {
204            inner: Arc::new(inner),
205        }
206    }
207
208    // ── Accessors ──
209
210    /// The agent name.
211    pub fn name(&self) -> &str {
212        &self.inner.name
213    }
214
215    /// Configured model, if any.
216    pub fn get_model(&self) -> Option<&ModelId> {
217        self.inner.model.as_ref()
218    }
219
220    /// Configured instruction, if any.
221    pub fn get_instruction(&self) -> Option<&str> {
222        self.inner.instruction.as_deref()
223    }
224
225    /// Configured voice, if any.
226    pub fn get_voice(&self) -> Option<&Voice> {
227        self.inner.voice.as_ref()
228    }
229
230    /// Configured temperature, if any.
231    pub fn get_temperature(&self) -> Option<f32> {
232        self.inner.temperature
233    }
234
235    /// Whether text-only mode is set.
236    pub fn is_text_only(&self) -> bool {
237        self.inner
238            .response_modalities
239            .as_ref()
240            .map(|m| m == &[Modality::Text])
241            .unwrap_or(false)
242    }
243
244    /// Configured thinking budget, if any.
245    pub fn get_thinking_budget(&self) -> Option<u32> {
246        self.inner.thinking_budget
247    }
248
249    /// State keys this agent writes.
250    pub fn get_writes(&self) -> &[String] {
251        &self.inner.writes
252    }
253
254    /// State keys this agent reads.
255    pub fn get_reads(&self) -> &[String] {
256        &self.inner.reads
257    }
258
259    /// Sub-agents registered.
260    pub fn get_sub_agents(&self) -> &[AgentBuilder] {
261        &self.inner.sub_agents
262    }
263
264    /// Whether agent runs in isolated state.
265    pub fn is_isolated(&self) -> bool {
266        self.inner.isolate
267    }
268
269    /// Whether agent stays after transfer.
270    pub fn is_stay(&self) -> bool {
271        self.inner.stay
272    }
273
274    /// Number of tool entries.
275    pub fn tool_count(&self) -> usize {
276        self.inner.tools.len() + self.inner.built_in_tools.len()
277    }
278
279    /// Configured top_p, if any.
280    pub fn get_top_p(&self) -> Option<f32> {
281        self.inner.top_p
282    }
283
284    /// Configured top_k, if any.
285    pub fn get_top_k(&self) -> Option<u32> {
286        self.inner.top_k
287    }
288
289    /// Configured max_output_tokens, if any.
290    pub fn get_max_output_tokens(&self) -> Option<u32> {
291        self.inner.max_output_tokens
292    }
293
294    /// Configured stop sequences.
295    pub fn get_stop_sequences(&self) -> &[String] {
296        &self.inner.stop_sequences
297    }
298
299    /// Configured description, if any.
300    pub fn get_description(&self) -> Option<&str> {
301        self.inner.description.as_deref()
302    }
303
304    /// Configured output schema, if any.
305    pub fn get_output_schema(&self) -> Option<&serde_json::Value> {
306        self.inner.output_schema.as_ref()
307    }
308
309    /// Get the configured output key.
310    pub fn get_output_key(&self) -> Option<&str> {
311        self.inner.output_key.as_deref()
312    }
313
314    /// Configured transfer target agent, if any.
315    pub fn get_transfer_to(&self) -> Option<&str> {
316        self.inner.transfer_to_agent.as_deref()
317    }
318
319    /// Number of registered middleware layers.
320    pub fn middleware_layer_count(&self) -> usize {
321        self.inner.middleware_layers.len()
322    }
323
324    // ── Fluent Setters (copy-on-write) ──
325
326    /// Set the Gemini model.
327    pub fn model(self, model: ModelId) -> Self {
328        let mut inner = self.mutate();
329        inner.model = Some(model);
330        Self::with(inner)
331    }
332
333    /// Set the system instruction.
334    pub fn instruction(self, inst: impl Into<String>) -> Self {
335        let mut inner = self.mutate();
336        inner.instruction = Some(inst.into());
337        Self::with(inner)
338    }
339
340    /// Set a dynamic instruction source — any `Fn(&State) -> String`
341    /// closure or a `TemplateInstruction` (feature `templates`), resolved
342    /// against live session state at the start of every run. Wins over
343    /// [`instruction`](Self::instruction) when both are set.
344    pub fn instruction_provider(
345        self,
346        provider: impl gemini_adk_rs::instruction::InstructionProvider + 'static,
347    ) -> Self {
348        let mut inner = self.mutate();
349        inner.instruction_provider = Some(Arc::new(provider));
350        Self::with(inner)
351    }
352
353    /// Set a dynamic model source, resolved against session state at the
354    /// start of every run — risk-based escalation to a stronger model, cost
355    /// routing to a cheaper one, per-tenant model selection — without
356    /// rebuilding the agent. Wins over the constructor's model when set.
357    pub fn llm_provider(
358        self,
359        provider: impl Fn(&gemini_adk_rs::State) -> Arc<dyn BaseLlm> + Send + Sync + 'static,
360    ) -> Self {
361        let mut inner = self.mutate();
362        inner.llm_provider = Some(Arc::new(provider));
363        Self::with(inner)
364    }
365
366    /// Set the output voice.
367    pub fn voice(self, voice: Voice) -> Self {
368        let mut inner = self.mutate();
369        inner.voice = Some(voice);
370        Self::with(inner)
371    }
372
373    /// Set the temperature.
374    pub fn temperature(self, t: f32) -> Self {
375        let mut inner = self.mutate();
376        inner.temperature = Some(t);
377        Self::with(inner)
378    }
379
380    /// Set text-only mode (no audio output).
381    pub fn text_only(self) -> Self {
382        let mut inner = self.mutate();
383        inner.response_modalities = Some(vec![Modality::Text]);
384        Self::with(inner)
385    }
386
387    /// Set response modalities explicitly.
388    pub fn response_modalities(self, modalities: Vec<Modality>) -> Self {
389        let mut inner = self.mutate();
390        inner.response_modalities = Some(modalities);
391        Self::with(inner)
392    }
393
394    /// Enable thinking with a token budget.
395    pub fn thinking(self, budget: u32) -> Self {
396        let mut inner = self.mutate();
397        inner.thinking_budget = Some(budget);
398        Self::with(inner)
399    }
400
401    /// Add a built-in URL context tool.
402    pub fn url_context(self) -> Self {
403        let mut inner = self.mutate();
404        inner.built_in_tools.push(Tool::url_context());
405        Self::with(inner)
406    }
407
408    /// Add a built-in Google Search tool.
409    pub fn google_search(self) -> Self {
410        let mut inner = self.mutate();
411        inner.built_in_tools.push(Tool::google_search());
412        Self::with(inner)
413    }
414
415    /// Add a built-in code execution tool.
416    pub fn code_execution(self) -> Self {
417        let mut inner = self.mutate();
418        inner.built_in_tools.push(Tool::code_execution());
419        Self::with(inner)
420    }
421
422    /// Declare a state key this agent writes.
423    pub fn writes(self, key: impl Into<String>) -> Self {
424        let mut inner = self.mutate();
425        inner.writes.push(key.into());
426        Self::with(inner)
427    }
428
429    /// Declare a state key this agent reads.
430    pub fn reads(self, key: impl Into<String>) -> Self {
431        let mut inner = self.mutate();
432        inner.reads.push(key.into());
433        Self::with(inner)
434    }
435
436    /// Declare the artifacts this agent consumes and produces, from the `A`
437    /// namespace: `A::text_input(..) + A::json_output(..)`.
438    ///
439    /// Each input is a read and each output a write of the state key
440    /// `artifact:{name}`, so [`check_contracts`](crate::testing::check_contracts)
441    /// reports an artifact input no agent produces, or one produced twice,
442    /// exactly as it does for state keys.
443    pub fn artifacts(
444        self,
445        artifacts: impl Into<crate::compose::artifacts::ArtifactComposite>,
446    ) -> Self {
447        let artifacts = artifacts.into();
448        let mut builder = self;
449        for input in artifacts.all_inputs() {
450            builder = builder.reads(format!("artifact:{}", input.name));
451        }
452        for output in artifacts.all_outputs() {
453            builder = builder.writes(format!("artifact:{}", output.name));
454        }
455        builder
456    }
457
458    /// Add a sub-agent for transfer.
459    pub fn sub_agent(self, agent: AgentBuilder) -> Self {
460        let mut inner = self.mutate();
461        inner.sub_agents.push(agent);
462        Self::with(inner)
463    }
464
465    /// Run this agent in isolated state (no shared state).
466    pub fn isolate(self) -> Self {
467        let mut inner = self.mutate();
468        inner.isolate = true;
469        Self::with(inner)
470    }
471
472    /// Keep this agent active after transfer (don't tear down).
473    pub fn stay(self) -> Self {
474        let mut inner = self.mutate();
475        inner.stay = true;
476        Self::with(inner)
477    }
478
479    /// Set top_p (nucleus sampling).
480    pub fn top_p(self, p: f32) -> Self {
481        let mut inner = self.mutate();
482        inner.top_p = Some(p);
483        Self::with(inner)
484    }
485
486    /// Set top_k (top-k sampling).
487    pub fn top_k(self, k: u32) -> Self {
488        let mut inner = self.mutate();
489        inner.top_k = Some(k);
490        Self::with(inner)
491    }
492
493    /// Set maximum output tokens.
494    pub fn max_output_tokens(self, n: u32) -> Self {
495        let mut inner = self.mutate();
496        inner.max_output_tokens = Some(n);
497        Self::with(inner)
498    }
499
500    /// Set stop sequences.
501    pub fn stop_sequences(self, seqs: Vec<String>) -> Self {
502        let mut inner = self.mutate();
503        inner.stop_sequences = seqs;
504        Self::with(inner)
505    }
506
507    /// Set a description for this agent (used in tool/agent metadata).
508    pub fn description(self, desc: impl Into<String>) -> Self {
509        let mut inner = self.mutate();
510        inner.description = Some(desc.into());
511        Self::with(inner)
512    }
513
514    /// Set a JSON schema for structured output.
515    pub fn output_schema(self, schema: serde_json::Value) -> Self {
516        let mut inner = self.mutate();
517        inner.output_schema = Some(schema);
518        Self::with(inner)
519    }
520
521    /// Constrain every reply to JSON shaped like `T`, using `T`'s schema.
522    ///
523    /// Read the reply back with [`RunResult::parse`](gemini_adk_rs::text::RunResult::parse),
524    /// or ask for a `T` directly with
525    /// [`TextAgent::ask_as`], which also
526    /// repairs a reply that does not parse.
527    pub fn output<T: schemars::JsonSchema>(self) -> Self {
528        self.output_schema(gemini_adk_rs::tool::wire_schema::<T>())
529    }
530
531    /// Set the output key — agent's final text response is auto-saved to this state key.
532    pub fn output_key(self, key: impl Into<String>) -> Self {
533        let mut inner = self.mutate();
534        inner.output_key = Some(key.into());
535        Self::with(inner)
536    }
537
538    /// Set a default transfer target agent.
539    pub fn transfer_to(self, agent_name: impl Into<String>) -> Self {
540        let mut inner = self.mutate();
541        inner.transfer_to_agent = Some(agent_name.into());
542        Self::with(inner)
543    }
544
545    // ── Upstream naming aliases ──
546
547    /// Alias for [`instruction`](Self::instruction) — matches upstream Python `Agent.instruct()`.
548    #[deprecated(since = "2.1.0", note = "use `instruction`")]
549    pub fn instruct(self, inst: impl Into<String>) -> Self {
550        self.instruction(inst)
551    }
552
553    /// Alias for [`description`](Self::description) — matches upstream Python `Agent.describe()`.
554    #[deprecated(since = "2.1.0", note = "use `description`")]
555    pub fn describe(self, desc: impl Into<String>) -> Self {
556        self.description(desc)
557    }
558
559    /// Register one tool: anything that implements [`ToolFunction`] — a
560    /// `SimpleTool`/`TypedTool`, the value a `#[tool]` function returns, or an
561    /// `Arc<dyn ToolFunction>` you already hold.
562    ///
563    /// ```no_run
564    /// # use gemini_adk_fluent_rs::prelude::*;
565    /// # use std::sync::Arc;
566    /// #[tool("Get the weather for a city")]
567    /// async fn get_weather(city: String) -> Result<serde_json::Value, ToolError> {
568    ///     Ok(serde_json::json!({"city": city, "temp": 22}))
569    /// }
570    /// let agent = AgentBuilder::new("assistant").tool(get_weather());
571    /// ```
572    pub fn tool(self, f: impl ToolFunction + 'static) -> Self {
573        self.tools(ToolComposite::from_function(Arc::new(f)))
574    }
575
576    /// Register tools: a `+`-combined [`ToolComposite`] from the `T`
577    /// namespace, or a single [`ToolFunction`].
578    ///
579    /// `T::mcp(..)` needs an async connection that this synchronous builder
580    /// cannot perform; it is rejected by [`build`](Self::build) with a
581    /// [`ConfigError`] — attach MCP toolsets to a `Live` session instead.
582    ///
583    /// ```no_run
584    /// # use gemini_adk_fluent_rs::prelude::*;
585    /// # use serde_json::json;
586    /// let tools = T::simple("greet", "Greet", |_| async { Ok(json!({})) })
587    ///     + T::google_search();
588    /// AgentBuilder::new("assistant").tools(tools);
589    /// ```
590    pub fn tools(self, tools: impl Into<ToolComposite>) -> Self {
591        use crate::compose::tools::{DeferredTool, ToolResolution};
592        let mut inner = self.mutate();
593        for entry in tools.into().entries {
594            match entry.classify() {
595                ToolResolution::Runtime(f) => {
596                    inner
597                        .tools
598                        .push(ToolEntry::Runtime(Arc::new(ToolFunctionEntry(f))));
599                }
600                ToolResolution::BuiltIn(t) => {
601                    inner.built_in_tools.push(t);
602                }
603                ToolResolution::Agent {
604                    name,
605                    description,
606                    agent,
607                } => {
608                    // Expose the sub-agent as a callable tool over a fresh State.
609                    let tool = gemini_adk_rs::TextAgentTool::from_arc(
610                        name,
611                        description,
612                        agent,
613                        gemini_adk_rs::State::new(),
614                    );
615                    inner
616                        .tools
617                        .push(ToolEntry::Runtime(Arc::new(ToolFunctionEntry(Arc::new(
618                            tool,
619                        )))));
620                }
621                ToolResolution::Deferred(DeferredTool::Mcp { params }) => {
622                    // An MCP toolset needs an async handshake, which the
623                    // synchronous text-agent `build()` cannot perform. It
624                    // belongs on a `Live` session (resolved at connect); make
625                    // `build` fail rather than drop the tool silently — the same
626                    // outcome `Live::connect` gives an unreachable MCP server.
627                    inner.config_errors.push(format!(
628                        "T::mcp({params:?}) cannot be attached to a text AgentBuilder: MCP \
629                         toolsets need an async connection, which only a Live session performs \
630                         (`Live::builder().tools(T::mcp(..))`)"
631                    ));
632                }
633            }
634        }
635        Self::with(inner)
636    }
637
638    /// Attach output guards. Each model response is validated against every
639    /// guard; if any rejects the output the agent run fails with an
640    /// [`AgentError`](gemini_adk_rs::error::AgentError) listing the violations.
641    ///
642    /// Accepts a single guard or a `+`-combined [`GuardComposite`] (every
643    /// guard must pass):
644    ///
645    /// ```no_run
646    /// # use gemini_adk_fluent_rs::prelude::*;
647    /// AgentBuilder::new("writer").guard(G::pii() + G::length(1, 2000));
648    /// ```
649    ///
650    /// The guards are installed as an `after_model` middleware layer, so they
651    /// accumulate with `.middleware(...)` and honor copy-on-write.
652    pub fn guard(self, guard: impl Into<GuardComposite>) -> Self {
653        let mut inner = self.mutate();
654        inner.middleware_layers.push(guard.into().into_middleware());
655        Self::with(inner)
656    }
657
658    /// Attach a context policy that rewrites conversation history before each
659    /// model call (e.g. windowing, role filtering, tool-result exclusion).
660    ///
661    /// Accepts a single policy or a `>>`-chained [`ContextComposite`] (applied
662    /// in order):
663    ///
664    /// ```no_run
665    /// # use gemini_adk_fluent_rs::prelude::*;
666    /// AgentBuilder::new("chat").context(C::window(10) >> C::user_only());
667    /// ```
668    ///
669    /// The policy is installed as a `transform_request` middleware layer.
670    pub fn context(self, policy: impl Into<ContextComposite>) -> Self {
671        let mut inner = self.mutate();
672        inner
673            .middleware_layers
674            .push(policy.into().into_middleware());
675        Self::with(inner)
676    }
677
678    /// Disallow transfer to peer agents.
679    #[deprecated(since = "2.1.0", note = "use `isolate`, which this has always called")]
680    pub fn no_peers(self) -> Self {
681        self.isolate()
682    }
683
684    /// Attach middleware — a `>>`-stacked [`MiddlewareComposite`] from the
685    /// `M` namespace or a single `Arc<dyn Middleware>`. All layers are
686    /// installed on the compiled `LlmTextAgent` in the order given.
687    ///
688    /// Multiple calls to `.middleware()` accumulate: the new layers are
689    /// appended after any previously registered layers, preserving the
690    /// copy-on-write contract.
691    ///
692    /// ```no_run
693    /// # use gemini_adk_fluent_rs::prelude::*;
694    /// let agent = AgentBuilder::new("analyst")
695    ///     .instruction("Analyze topics")
696    ///     .middleware(M::log() >> M::latency());
697    /// ```
698    pub fn middleware(self, middleware: impl Into<MiddlewareComposite>) -> Self {
699        let mut inner = self.mutate();
700        inner.middleware_layers.extend(middleware.into().layers);
701        Self::with(inner)
702    }
703
704    /// Gate `T::confirm(..)` tools behind a confirmation provider, as
705    /// [`Live::confirmation_provider`](crate::live::Live::confirmation_provider)
706    /// does for a session.
707    ///
708    /// A declined call returns [`ToolError::Declined`](gemini_adk_rs::error::ToolError::Declined)
709    /// with the reason to the model instead of running the tool. An agent with
710    /// a confirmation-gated tool and no provider does not build.
711    pub fn confirmation_provider(
712        self,
713        provider: Arc<dyn gemini_adk_rs::confirmation::ConfirmationProvider>,
714    ) -> Self {
715        let mut inner = self.mutate();
716        inner.confirmation_provider = Some(provider);
717        Self::with(inner)
718    }
719
720    // ── Compilation ──
721
722    /// Settings a text agent cannot honour, each with what to do instead.
723    ///
724    /// Voice, audio output and agent transfer belong to Live sessions; a text
725    /// agent that silently ignored them would behave differently from the
726    /// configuration that describes it.
727    fn text_agent_issues(&self) -> Vec<String> {
728        let inner = &self.inner;
729        let mut issues = Vec::new();
730        if inner.voice.is_some() {
731            issues.push(
732                "`voice(..)` has no effect on a text agent, which returns text; \
733                 remove it, or set the voice on `Live::builder()`"
734                    .to_string(),
735            );
736        }
737        if inner
738            .response_modalities
739            .as_ref()
740            .is_some_and(|m| m.iter().any(|m| *m != Modality::Text))
741        {
742            issues.push(
743                "a text agent only produces text; audio and image `response_modalities` \
744                 need a `Live::builder()` session"
745                    .to_string(),
746            );
747        }
748        if let Some(model) = &inner.model
749            && gemini_adk_rs::llm::ModelCapabilities::infer_from_id(model.as_str()).live_bidi
750        {
751            issues.push(format!(
752                "`{model}` is a Live model, which the text API does not serve; use a text \
753                 model such as `gemini-flash-latest`, or run it on `Live::builder()`"
754            ));
755        }
756        if inner.confirmation_provider.is_none() {
757            let gated: Vec<String> = inner
758                .tools
759                .iter()
760                .filter_map(|entry| match entry {
761                    ToolEntry::Runtime(t) => match t.to_tool_kind() {
762                        ToolKind::Function(f) if f.requires_confirmation() => {
763                            Some(f.name().to_string())
764                        }
765                        _ => None,
766                    },
767                    ToolEntry::Declaration(_) => None,
768                })
769                .collect();
770            if !gated.is_empty() {
771                issues.push(format!(
772                    "`{}` must be confirmed before running (`T::confirm`), but nothing can \
773                     confirm it; set `.confirmation_provider(..)`",
774                    gated.join("`, `")
775                ));
776            }
777        }
778        let transfer: Vec<&str> = [
779            (!inner.sub_agents.is_empty(), "sub_agent"),
780            (inner.transfer_to_agent.is_some(), "transfer_to"),
781            (inner.stay, "stay"),
782            (inner.isolate, "isolate/no_peers"),
783        ]
784        .into_iter()
785        .filter_map(|(set, name)| set.then_some(name))
786        .collect();
787        if !transfer.is_empty() {
788            issues.push(format!(
789                "`{}` configure agent transfer, which a text agent does not perform; \
790                 compose text agents with `>>`, `|` and `T::agent(..)` instead",
791                transfer.join("`, `")
792            ));
793        }
794        issues
795    }
796
797    /// Compile this builder into an executable `TextAgent`.
798    ///
799    /// `llm` is any [`BaseLlm`]. Every setting reaches the requests the agent
800    /// sends: model, instruction, sampling (`temperature`, `top_p`, `top_k`,
801    /// `max_output_tokens`, `stop_sequences`), `thinking`, `output_schema`,
802    /// built-in tools and function tools, middleware, and `output_key`.
803    /// `writes`, `reads` and `description` are declarations for
804    /// [`check_contracts`](crate::testing::check_contracts) and tool metadata.
805    ///
806    /// Fails with a [`ConfigError`] listing every setting a text agent cannot
807    /// honour, rather than dropping it: `voice`, audio modalities, a Live
808    /// model, agent-transfer settings (`sub_agent`, `transfer_to`, `stay`,
809    /// `isolate`), an MCP toolset (`T::mcp`), which needs the async connect
810    /// only a `Live` session performs, and a `T::confirm` tool with no
811    /// [`confirmation_provider`](Self::confirmation_provider).
812    ///
813    /// ```no_run
814    /// # use gemini_adk_fluent_rs::prelude::*;
815    /// # use std::sync::Arc;
816    /// # async fn run() -> Result<(), AgentError> {
817    /// let llm = Arc::new(GeminiLlm::new(GeminiLlmParams::default()));
818    /// let agent = AgentBuilder::new("analyst")
819    ///     .instruction("Analyze the topic")
820    ///     .temperature(0.3)
821    ///     .build(llm)?;
822    ///
823    /// let state = State::new();
824    /// let result = agent.run(&state).await?;
825    /// # let _ = result; Ok(())
826    /// # }
827    /// ```
828    pub fn build(self, llm: impl BaseLlm + 'static) -> Result<Arc<dyn TextAgent>, ConfigError> {
829        let mut issues = self.inner.config_errors.clone();
830        issues.extend(self.text_agent_issues());
831        if !issues.is_empty() {
832            return Err(ConfigError { issues });
833        }
834        let inner = &self.inner;
835        let mut agent = LlmTextAgent::new(&inner.name, llm);
836
837        if let Some(inst) = &inner.instruction {
838            agent = agent.instruction(inst);
839        }
840        if let Some(provider) = &inner.instruction_provider {
841            agent = agent.instruction_provider(provider.clone());
842        }
843        if let Some(provider) = &inner.llm_provider {
844            let provider_clone = provider.clone();
845            agent = agent.llm_provider(move |state| provider_clone(state));
846        }
847        if let Some(model) = &inner.model {
848            agent = agent.model(model.as_str());
849        }
850        if let Some(t) = inner.temperature {
851            agent = agent.temperature(t);
852        }
853        if let Some(p) = inner.top_p {
854            agent = agent.top_p(p);
855        }
856        if let Some(k) = inner.top_k {
857            agent = agent.top_k(k);
858        }
859        if let Some(n) = inner.max_output_tokens {
860            agent = agent.max_output_tokens(n);
861        }
862        if !inner.stop_sequences.is_empty() {
863            agent = agent.stop_sequences(inner.stop_sequences.iter().cloned());
864        }
865        if let Some(budget) = inner.thinking_budget {
866            agent = agent.thinking_budget(budget);
867        }
868        if let Some(schema) = &inner.output_schema {
869            agent = agent.response_schema(schema.clone());
870        }
871        if let Some(key) = &inner.output_key {
872            agent = agent.output_key(key);
873        }
874        let declarations = inner.tools.iter().filter_map(|entry| match entry {
875            ToolEntry::Declaration(tool) => Some(tool),
876            ToolEntry::Runtime(_) => None,
877        });
878        for tool in inner.built_in_tools.iter().chain(declarations) {
879            agent = agent.built_in_tool(tool.clone());
880        }
881
882        // Build ToolDispatcher from registered tools.
883        if !self.inner.tools.is_empty() {
884            let mut dispatcher = ToolDispatcher::new();
885            for entry in &self.inner.tools {
886                match entry {
887                    ToolEntry::Runtime(t) => {
888                        let kind = t.to_tool_kind();
889                        match kind {
890                            ToolKind::Function(f) => dispatcher.register_function(f),
891                            ToolKind::Streaming(s) => dispatcher.register_streaming(s),
892                            ToolKind::InputStream(i) => dispatcher.register_input_streaming(i),
893                        }
894                    }
895                    // Sent with the request above; nothing to dispatch.
896                    ToolEntry::Declaration(_) => {}
897                }
898            }
899            if let Some(provider) = &self.inner.confirmation_provider {
900                dispatcher.set_confirmation_provider(provider.clone());
901            }
902            if !dispatcher.is_empty() {
903                agent = agent.tools(Arc::new(dispatcher));
904            }
905        }
906
907        // Install middleware layers from the builder.
908        for mw in &self.inner.middleware_layers {
909            agent = agent.add_middleware(mw.clone());
910        }
911
912        Ok(Arc::new(agent))
913    }
914}
915
916/// Adapter that wraps an `Arc<dyn ToolFunction>` as a `ToolEntryTrait`.
917#[derive(Clone)]
918struct ToolFunctionEntry(Arc<dyn ToolFunction>);
919
920impl ToolEntryTrait for ToolFunctionEntry {
921    fn name(&self) -> &str {
922        self.0.name()
923    }
924
925    fn to_tool_kind(&self) -> ToolKind {
926        ToolKind::Function(self.0.clone())
927    }
928}
929
930impl std::fmt::Debug for AgentBuilder {
931    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
932        f.debug_struct("AgentBuilder")
933            .field("name", &self.inner.name)
934            .field("model", &self.inner.model)
935            .field("instruction", &self.inner.instruction)
936            .field("temperature", &self.inner.temperature)
937            .field("text_only", &self.is_text_only())
938            .field("tool_count", &self.tool_count())
939            .field("sub_agents", &self.inner.sub_agents.len())
940            .finish()
941    }
942}
943
944#[cfg(test)]
945mod tests {
946    use super::*;
947    use async_trait::async_trait;
948    use gemini_adk_rs::llm::{LlmError, LlmRequest, LlmResponse};
949    use gemini_genai_rs::prelude::{Content, Part, Role};
950
951    /// A mock LLM for build() tests.
952    struct MockLlm(String);
953
954    #[async_trait]
955    impl BaseLlm for MockLlm {
956        fn model_id(&self) -> &str {
957            "mock"
958        }
959        async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
960            Ok(LlmResponse {
961                content: Content {
962                    role: Some(Role::Model),
963                    parts: vec![Part::Text {
964                        text: self.0.clone(),
965                    }],
966                },
967                finish_reason: Some("STOP".into()),
968                usage: None,
969            })
970        }
971    }
972
973    #[test]
974    fn builder_creates_with_name() {
975        let b = AgentBuilder::new("test-agent");
976        assert_eq!(b.name(), "test-agent");
977    }
978
979    #[test]
980    fn fluent_chaining_works() {
981        let b = AgentBuilder::new("agent")
982            .instruction("Be helpful")
983            .temperature(0.7)
984            .model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO);
985
986        assert_eq!(b.get_instruction(), Some("Be helpful"));
987        assert_eq!(b.get_temperature(), Some(0.7));
988        assert_eq!(b.get_model(), Some(&ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO));
989    }
990
991    #[test]
992    fn copy_on_write_clone_independence() {
993        let base = AgentBuilder::new("base").temperature(0.5);
994        let variant = base.clone().temperature(0.9);
995
996        // Original unchanged
997        assert_eq!(base.get_temperature(), Some(0.5));
998        // Variant has new value
999        assert_eq!(variant.get_temperature(), Some(0.9));
1000    }
1001
1002    #[test]
1003    fn text_only_sets_modalities() {
1004        let b = AgentBuilder::new("text").text_only();
1005        assert!(b.is_text_only());
1006    }
1007
1008    #[test]
1009    fn url_context_adds_tool() {
1010        let b = AgentBuilder::new("search").url_context();
1011        assert_eq!(b.tool_count(), 1);
1012    }
1013
1014    #[test]
1015    fn google_search_adds_tool() {
1016        let b = AgentBuilder::new("search").google_search();
1017        assert_eq!(b.tool_count(), 1);
1018    }
1019
1020    #[test]
1021    fn code_execution_adds_tool() {
1022        let b = AgentBuilder::new("code").code_execution();
1023        assert_eq!(b.tool_count(), 1);
1024    }
1025
1026    #[test]
1027    fn thinking_sets_budget() {
1028        let b = AgentBuilder::new("thinker").thinking(2048);
1029        assert_eq!(b.get_thinking_budget(), Some(2048));
1030    }
1031
1032    #[test]
1033    fn writes_and_reads_keys() {
1034        let b = AgentBuilder::new("data").writes("output").reads("input");
1035        assert_eq!(b.get_writes(), &["output"]);
1036        assert_eq!(b.get_reads(), &["input"]);
1037    }
1038
1039    #[test]
1040    fn sub_agent_registration() {
1041        let child = AgentBuilder::new("child");
1042        let parent = AgentBuilder::new("parent").sub_agent(child);
1043        assert_eq!(parent.get_sub_agents().len(), 1);
1044        assert_eq!(parent.get_sub_agents()[0].name(), "child");
1045    }
1046
1047    #[test]
1048    fn isolate_and_stay() {
1049        let b = AgentBuilder::new("agent").isolate().stay();
1050        assert!(b.is_isolated());
1051        assert!(b.is_stay());
1052    }
1053
1054    #[test]
1055    fn debug_display() {
1056        let b = AgentBuilder::new("debug-test");
1057        let debug = format!("{b:?}");
1058        assert!(debug.contains("debug-test"));
1059    }
1060
1061    #[test]
1062    fn top_p_sets_value() {
1063        let b = AgentBuilder::new("agent").top_p(0.95);
1064        assert_eq!(b.get_top_p(), Some(0.95));
1065    }
1066
1067    #[test]
1068    fn top_k_sets_value() {
1069        let b = AgentBuilder::new("agent").top_k(40);
1070        assert_eq!(b.get_top_k(), Some(40));
1071    }
1072
1073    #[test]
1074    fn max_output_tokens_sets_value() {
1075        let b = AgentBuilder::new("agent").max_output_tokens(4096);
1076        assert_eq!(b.get_max_output_tokens(), Some(4096));
1077    }
1078
1079    #[test]
1080    fn stop_sequences_sets_value() {
1081        let b =
1082            AgentBuilder::new("agent").stop_sequences(vec!["END".to_string(), "STOP".to_string()]);
1083        assert_eq!(b.get_stop_sequences().len(), 2);
1084    }
1085
1086    #[test]
1087    fn description_sets_value() {
1088        let b = AgentBuilder::new("agent").description("A helpful agent");
1089        assert_eq!(b.get_description(), Some("A helpful agent"));
1090    }
1091
1092    #[test]
1093    fn output_schema_sets_value() {
1094        let schema = serde_json::json!({"type": "object"});
1095        let b = AgentBuilder::new("agent").output_schema(schema.clone());
1096        assert_eq!(b.get_output_schema(), Some(&schema));
1097    }
1098
1099    #[test]
1100    fn transfer_to_sets_value() {
1101        let b = AgentBuilder::new("agent").transfer_to("target-agent");
1102        assert_eq!(b.get_transfer_to(), Some("target-agent"));
1103    }
1104
1105    #[test]
1106    fn full_fluent_chain() {
1107        let b = AgentBuilder::new("full-agent")
1108            .model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO)
1109            .instruction("Be helpful")
1110            .temperature(0.7)
1111            .top_p(0.95)
1112            .top_k(40)
1113            .max_output_tokens(4096)
1114            .thinking(2048)
1115            .description("A fully configured agent")
1116            .google_search()
1117            .writes("output")
1118            .reads("input");
1119
1120        assert_eq!(b.name(), "full-agent");
1121        assert_eq!(b.get_temperature(), Some(0.7));
1122        assert_eq!(b.get_top_p(), Some(0.95));
1123        assert_eq!(b.get_top_k(), Some(40));
1124        assert_eq!(b.get_max_output_tokens(), Some(4096));
1125        assert_eq!(b.get_thinking_budget(), Some(2048));
1126        assert_eq!(b.get_description(), Some("A fully configured agent"));
1127        assert_eq!(b.tool_count(), 1);
1128    }
1129
1130    // ── build() tests ──
1131
1132    #[tokio::test]
1133    async fn build_produces_executable_agent() {
1134        let llm: Arc<dyn BaseLlm> = Arc::new(MockLlm("built agent output".into()));
1135        let agent = AgentBuilder::new("test")
1136            .instruction("Be helpful")
1137            .temperature(0.5)
1138            .build(llm)
1139            .unwrap();
1140
1141        assert_eq!(agent.name(), "test");
1142        let state = gemini_adk_rs::State::new();
1143        let result = agent.run(&state).await.unwrap();
1144        assert_eq!(result, "built agent output");
1145    }
1146
1147    #[tokio::test]
1148    async fn build_stores_output_in_state() {
1149        let llm: Arc<dyn BaseLlm> = Arc::new(MockLlm("state output".into()));
1150        let agent = AgentBuilder::new("test").build(llm).unwrap();
1151        let state = gemini_adk_rs::State::new();
1152        agent.run(&state).await.unwrap();
1153        assert_eq!(state.get::<String>("output"), Some("state output".into()));
1154    }
1155
1156    #[tokio::test]
1157    async fn build_reads_input_from_state() {
1158        use gemini_adk_rs::llm::LlmRequest;
1159
1160        // An LLM that echoes whatever it receives.
1161        struct EchoLlm;
1162        #[async_trait]
1163        impl BaseLlm for EchoLlm {
1164            fn model_id(&self) -> &str {
1165                "echo"
1166            }
1167            async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
1168                let text: String = req
1169                    .contents
1170                    .iter()
1171                    .flat_map(|c| &c.parts)
1172                    .filter_map(|p| match p {
1173                        Part::Text { text } => Some(text.as_str()),
1174                        _ => None,
1175                    })
1176                    .collect::<Vec<_>>()
1177                    .join("");
1178                Ok(LlmResponse {
1179                    content: Content {
1180                        role: Some(Role::Model),
1181                        parts: vec![Part::Text { text }],
1182                    },
1183                    finish_reason: Some("STOP".into()),
1184                    usage: None,
1185                })
1186            }
1187        }
1188
1189        let agent = AgentBuilder::new("echo").build(Arc::new(EchoLlm)).unwrap();
1190        let state = gemini_adk_rs::State::new();
1191        let _ = state.set("input", "hello from state");
1192        let result = agent.run(&state).await.unwrap();
1193        assert!(result.contains("hello from state"));
1194    }
1195
1196    // ── Middleware end-to-end tests ──
1197
1198    /// A mock LLM that issues one tool call and then returns text.
1199    struct ToolCallingMockLlm {
1200        tool_name: &'static str,
1201        final_text: &'static str,
1202    }
1203
1204    #[async_trait]
1205    impl BaseLlm for ToolCallingMockLlm {
1206        fn model_id(&self) -> &str {
1207            "tool-mock"
1208        }
1209
1210        async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
1211            use gemini_genai_rs::prelude::FunctionCall;
1212
1213            // If any part is a FunctionResponse, we already dispatched — return text.
1214            let already_responded = req
1215                .contents
1216                .iter()
1217                .flat_map(|c| &c.parts)
1218                .any(|p| matches!(p, Part::FunctionResponse { .. }));
1219
1220            if already_responded {
1221                Ok(LlmResponse {
1222                    content: Content {
1223                        role: Some(Role::Model),
1224                        parts: vec![Part::Text {
1225                            text: self.final_text.to_string(),
1226                        }],
1227                    },
1228                    finish_reason: Some("STOP".into()),
1229                    usage: None,
1230                })
1231            } else {
1232                Ok(LlmResponse {
1233                    content: Content {
1234                        role: Some(Role::Model),
1235                        parts: vec![Part::FunctionCall {
1236                            function_call: FunctionCall {
1237                                name: self.tool_name.to_string(),
1238                                args: serde_json::json!({"x": 1}),
1239                                id: Some("call-1".into()),
1240                            },
1241                        }],
1242                    },
1243                    finish_reason: None,
1244                    usage: None,
1245                })
1246            }
1247        }
1248    }
1249
1250    /// Verify that `M::before_model` and `M::after_tool` hooks fire when the agent runs.
1251    #[tokio::test]
1252    async fn middleware_hooks_fire_end_to_end() {
1253        use crate::compose::middleware::M;
1254        use gemini_adk_rs::tool::SimpleTool;
1255        use std::sync::atomic::{AtomicUsize, Ordering};
1256
1257        let before_model_count = Arc::new(AtomicUsize::new(0));
1258        let after_tool_count = Arc::new(AtomicUsize::new(0));
1259
1260        let bm = before_model_count.clone();
1261        let at = after_tool_count.clone();
1262
1263        let mw = M::before_model(move |_req| {
1264            bm.fetch_add(1, Ordering::SeqCst);
1265            Ok(())
1266        }) >> M::after_tool(move |_call, _result| {
1267            at.fetch_add(1, Ordering::SeqCst);
1268            Ok(())
1269        });
1270
1271        let llm: Arc<dyn BaseLlm> = Arc::new(ToolCallingMockLlm {
1272            tool_name: "echo_tool",
1273            final_text: "done",
1274        });
1275
1276        let agent = AgentBuilder::new("mw-test")
1277            .middleware(mw)
1278            .tool(SimpleTool::new(
1279                "echo_tool",
1280                "Echo tool",
1281                None,
1282                |_args| async move { Ok(serde_json::json!({"echo": true})) },
1283            ))
1284            .build(llm)
1285            .unwrap();
1286
1287        let state = gemini_adk_rs::State::new();
1288        let result = agent.run(&state).await.unwrap();
1289        assert_eq!(result, "done");
1290
1291        // before_model fires once per LLM call: first call (tool call) + second call (final text).
1292        assert_eq!(
1293            before_model_count.load(Ordering::SeqCst),
1294            2,
1295            "before_model should fire for each generate() call"
1296        );
1297        // after_tool fires once per successful tool dispatch.
1298        assert_eq!(
1299            after_tool_count.load(Ordering::SeqCst),
1300            1,
1301            "after_tool should fire once for the tool dispatch"
1302        );
1303    }
1304
1305    /// Verify copy-on-write: adding middleware to a clone does not affect the original.
1306    #[test]
1307    fn middleware_copy_on_write() {
1308        use crate::compose::middleware::M;
1309
1310        let base = AgentBuilder::new("base").instruction("base");
1311        let with_mw = base.clone().middleware(M::log() >> M::latency());
1312
1313        // Original should have no middleware layers.
1314        assert_eq!(base.middleware_layer_count(), 0);
1315        // Clone with middleware should have 2 layers.
1316        assert_eq!(with_mw.middleware_layer_count(), 2);
1317    }
1318
1319    /// Verify `on_error` hook fires when the agent errors.
1320    #[tokio::test]
1321    async fn middleware_on_error_fires_on_failure() {
1322        use crate::compose::middleware::M;
1323        use gemini_adk_rs::llm::LlmError;
1324        use std::sync::atomic::{AtomicUsize, Ordering};
1325
1326        let error_count = Arc::new(AtomicUsize::new(0));
1327        let ec = error_count.clone();
1328
1329        let mw = M::on_error(move |_err| {
1330            ec.fetch_add(1, Ordering::SeqCst);
1331            Ok(())
1332        });
1333
1334        struct FailLlm;
1335        #[async_trait]
1336        impl BaseLlm for FailLlm {
1337            fn model_id(&self) -> &str {
1338                "fail"
1339            }
1340            async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
1341                Err(LlmError::RequestFailed("boom".into()))
1342            }
1343        }
1344
1345        let agent = AgentBuilder::new("error-test")
1346            .middleware(mw)
1347            .build(Arc::new(FailLlm))
1348            .unwrap();
1349
1350        let state = gemini_adk_rs::State::new();
1351        let result = agent.run(&state).await;
1352        assert!(result.is_err(), "agent should fail");
1353        assert_eq!(
1354            error_count.load(Ordering::SeqCst),
1355            1,
1356            "on_error should fire exactly once"
1357        );
1358    }
1359
1360    // ── Guard / context wiring tests ──
1361
1362    /// A mock LLM that echoes a fixed response and records the number of
1363    /// `contents` it was asked to generate from (to observe context rewriting).
1364    struct RecordingLlm {
1365        text: &'static str,
1366        seen_len: Arc<std::sync::atomic::AtomicUsize>,
1367    }
1368
1369    #[async_trait]
1370    impl BaseLlm for RecordingLlm {
1371        fn model_id(&self) -> &str {
1372            "recording-mock"
1373        }
1374
1375        async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
1376            self.seen_len
1377                .store(req.contents.len(), std::sync::atomic::Ordering::SeqCst);
1378            Ok(LlmResponse {
1379                content: Content {
1380                    role: Some(Role::Model),
1381                    parts: vec![Part::Text {
1382                        text: self.text.to_string(),
1383                    }],
1384                },
1385                finish_reason: Some("STOP".into()),
1386                usage: None,
1387            })
1388        }
1389    }
1390
1391    #[tokio::test]
1392    async fn guard_blocks_violating_output() {
1393        use crate::compose::guards::G;
1394
1395        let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1396            text: "you can reach me at agent@example.com",
1397            seen_len: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
1398        });
1399
1400        let agent = AgentBuilder::new("guarded")
1401            .guard(G::pii())
1402            .build(llm)
1403            .unwrap();
1404
1405        let state = gemini_adk_rs::State::new();
1406        let err = agent.run(&state).await.unwrap_err();
1407        assert!(
1408            err.to_string().contains("guard violation"),
1409            "PII guard should veto the response, got: {err}"
1410        );
1411    }
1412
1413    #[tokio::test]
1414    async fn guard_allows_clean_output() {
1415        use crate::compose::guards::G;
1416
1417        let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1418            text: "all clean here",
1419            seen_len: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
1420        });
1421
1422        let agent = AgentBuilder::new("guarded")
1423            .guard(G::pii() + G::length(1, 1000))
1424            .build(llm)
1425            .unwrap();
1426
1427        let state = gemini_adk_rs::State::new();
1428        let result = agent.run(&state).await.unwrap();
1429        assert_eq!(result, "all clean here");
1430    }
1431
1432    #[tokio::test]
1433    async fn context_policy_rewrites_request_history() {
1434        use crate::compose::context::C;
1435
1436        // The agent seeds one user turn; a prepend policy injects a second turn,
1437        // so the LLM should see 2 contents — proving transform_request ran.
1438        let seen = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1439        let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1440            text: "ok",
1441            seen_len: seen.clone(),
1442        });
1443
1444        let agent = AgentBuilder::new("ctx")
1445            .context(C::prepend(Content::user("system preamble")))
1446            .build(llm)
1447            .unwrap();
1448
1449        let state = gemini_adk_rs::State::new();
1450        let _ = state.set("input", "hello");
1451        let _ = agent.run(&state).await.unwrap();
1452        assert_eq!(
1453            seen.load(std::sync::atomic::Ordering::SeqCst),
1454            2,
1455            "context policy should have prepended a turn before the model call"
1456        );
1457    }
1458
1459    #[tokio::test]
1460    async fn context_window_trims_history() {
1461        use crate::compose::context::C;
1462
1463        // window(1) keeps only the last turn. We seed a single input turn and
1464        // prepend two extra turns, then window down to 1 — the model sees 1.
1465        let seen = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1466        let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1467            text: "ok",
1468            seen_len: seen.clone(),
1469        });
1470
1471        let agent = AgentBuilder::new("ctx")
1472            .context(
1473                C::prepend(Content::user("a")) >> C::prepend(Content::user("b")) >> C::window(1),
1474            )
1475            .build(llm)
1476            .unwrap();
1477
1478        let state = gemini_adk_rs::State::new();
1479        let _ = state.set("input", "hello");
1480        let _ = agent.run(&state).await.unwrap();
1481        assert_eq!(
1482            seen.load(std::sync::atomic::Ordering::SeqCst),
1483            1,
1484            "window(1) should trim history to the last turn"
1485        );
1486    }
1487}
1488
1489#[cfg(test)]
1490mod mcp_rejection_tests {
1491    use super::*;
1492    use gemini_adk_rs::llm::{LlmError, LlmRequest, LlmResponse};
1493
1494    struct NeverLlm;
1495    #[async_trait::async_trait]
1496    impl BaseLlm for NeverLlm {
1497        fn model_id(&self) -> &str {
1498            "never"
1499        }
1500        async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
1501            Err(LlmError::RequestFailed("never".into()))
1502        }
1503    }
1504
1505    /// `T::mcp` on a text agent is a build error naming the tool kind — the
1506    /// same way `Live::connect` fails on it — never a silent drop.
1507    #[test]
1508    fn mcp_toolset_is_a_build_error() {
1509        use crate::compose::tools::T;
1510        let err = AgentBuilder::new("text")
1511            .tools(T::mcp("node ./server.js"))
1512            .build(Arc::new(NeverLlm))
1513            .err()
1514            .expect("build must fail");
1515        assert!(err.to_string().contains("T::mcp"), "{err}");
1516        assert!(err.to_string().contains("Live"), "{err}");
1517    }
1518
1519    /// `tool(..)` takes any `ToolFunction`, including an `Arc<dyn ToolFunction>`.
1520    #[test]
1521    fn tool_accepts_values_and_arcs() {
1522        use gemini_adk_rs::tool::SimpleTool;
1523        let arc: Arc<dyn ToolFunction> = Arc::new(SimpleTool::new("a", "a", None, |_| async {
1524            Ok(serde_json::json!({}))
1525        }));
1526        let b = AgentBuilder::new("t")
1527            .tool(SimpleTool::new("b", "b", None, |_| async {
1528                Ok(serde_json::json!({}))
1529            }))
1530            .tool(arc.clone())
1531            .tools(arc);
1532        assert_eq!(b.tool_count(), 3);
1533    }
1534}
1535
1536#[cfg(test)]
1537mod honest_build_tests {
1538    use super::*;
1539    use gemini_adk_rs::llm::{LlmResponse, MockLlm};
1540
1541    /// Every setting on the builder reaches the request the built agent sends.
1542    /// Before, `build` forwarded only the instruction, temperature, max tokens
1543    /// and function tools, and dropped the rest without a word.
1544    #[tokio::test]
1545    async fn every_setting_reaches_the_request() {
1546        let llm = MockLlm::script([LlmResponse::from_text(r#"{"city":"Paris"}"#)]);
1547        let agent = AgentBuilder::new("configured")
1548            .model(ModelId::new("gemini-2.5-pro"))
1549            .instruction("Answer in JSON.")
1550            .temperature(0.1)
1551            .top_p(0.8)
1552            .top_k(10)
1553            .max_output_tokens(128)
1554            .stop_sequences(vec!["END".into()])
1555            .thinking(256)
1556            .output_schema(serde_json::json!({ "type": "object" }))
1557            .output_key("answer")
1558            .google_search()
1559            .build(llm.clone())
1560            .expect("a text agent can honour all of these");
1561
1562        let state = gemini_adk_rs::State::new();
1563        state.set("input", "Capital of France?").unwrap();
1564        agent.run(&state).await.unwrap();
1565
1566        let sent = llm.last_request().unwrap();
1567        assert_eq!(sent.model.as_deref(), Some("gemini-2.5-pro"));
1568        assert_eq!(sent.system_instruction.as_deref(), Some("Answer in JSON."));
1569        assert_eq!(sent.temperature, Some(0.1));
1570        assert_eq!(sent.top_p, Some(0.8));
1571        assert_eq!(sent.top_k, Some(10));
1572        assert_eq!(sent.max_output_tokens, Some(128));
1573        assert_eq!(sent.stop_sequences, ["END"]);
1574        assert_eq!(sent.thinking_budget, Some(256));
1575        assert_eq!(sent.response_mime_type.as_deref(), Some("application/json"));
1576        assert_eq!(
1577            sent.response_json_schema,
1578            Some(serde_json::json!({ "type": "object" }))
1579        );
1580        assert_eq!(sent.tools, vec![Tool::google_search()]);
1581        assert_eq!(
1582            state.get::<String>("answer").as_deref(),
1583            Some(r#"{"city":"Paris"}"#)
1584        );
1585    }
1586
1587    /// Function tools and built-in tools travel together.
1588    #[tokio::test]
1589    async fn built_in_and_function_tools_are_both_declared() {
1590        use crate::compose::tools::T;
1591        let llm = MockLlm::text("ok");
1592        let agent = AgentBuilder::new("both")
1593            .tools(
1594                T::simple("ping", "Ping", |_| async { Ok(serde_json::json!({})) })
1595                    + T::code_execution(),
1596            )
1597            .url_context()
1598            .build(llm.clone())
1599            .unwrap();
1600        agent.run(&gemini_adk_rs::State::new()).await.unwrap();
1601        let tools = llm.last_request().unwrap().tools;
1602        assert!(tools.contains(&Tool::url_context()), "{tools:?}");
1603        assert!(tools.contains(&Tool::code_execution()), "{tools:?}");
1604        assert!(
1605            tools.iter().any(|t| t
1606                .function_declarations
1607                .as_ref()
1608                .is_some_and(|d| d.iter().any(|f| f.name == "ping"))),
1609            "{tools:?}"
1610        );
1611    }
1612
1613    /// Settings a text agent cannot honour fail the build, all at once, each
1614    /// saying what to do instead.
1615    #[test]
1616    fn live_only_settings_are_build_errors() {
1617        let err = AgentBuilder::new("confused")
1618            .model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO)
1619            .voice(Voice::Kore)
1620            .response_modalities(vec![Modality::Audio])
1621            .sub_agent(AgentBuilder::new("child"))
1622            .transfer_to("child")
1623            .build(MockLlm::text("unused"))
1624            .err()
1625            .expect("build must fail");
1626        assert_eq!(err.issues.len(), 4, "{err}");
1627        let message = err.to_string();
1628        for needle in [
1629            "voice",
1630            "response_modalities",
1631            "Live model",
1632            "sub_agent",
1633            "transfer_to",
1634        ] {
1635            assert!(message.contains(needle), "missing {needle}: {message}");
1636        }
1637    }
1638
1639    fn refund_tool() -> crate::compose::tools::ToolComposite {
1640        use crate::compose::tools::T;
1641        T::confirm(
1642            T::simple("refund", "Refund the order", |_| async {
1643                Ok(serde_json::json!({ "refunded": true }))
1644            }),
1645            "Refunds move money",
1646        )
1647    }
1648
1649    /// A tool marked `T::confirm` used to run unconfirmed on a text agent,
1650    /// because nothing attached a provider. Now it cannot be built that way.
1651    #[test]
1652    fn a_confirm_tool_without_a_provider_is_a_build_error() {
1653        let err = AgentBuilder::new("support")
1654            .tools(refund_tool())
1655            .build(MockLlm::text("x"))
1656            .err()
1657            .expect("build must fail");
1658        assert!(err.to_string().contains("refund"), "{err}");
1659        assert!(err.to_string().contains("confirmation_provider"), "{err}");
1660    }
1661
1662    /// A declined call does not run, and the model is told why.
1663    #[tokio::test]
1664    async fn a_declined_call_tells_the_model_why() {
1665        use gemini_adk_rs::confirmation::StaticConfirmation;
1666        use gemini_genai_rs::prelude::Part;
1667
1668        let llm = MockLlm::script([
1669            LlmResponse::tool_call("refund", serde_json::json!({})),
1670            LlmResponse::from_text("I could not refund that."),
1671        ]);
1672        let agent = AgentBuilder::new("support")
1673            .tools(refund_tool())
1674            .confirmation_provider(StaticConfirmation::deny_all("over the refund limit"))
1675            .build(llm.clone())
1676            .unwrap();
1677        agent.run(&gemini_adk_rs::State::new()).await.unwrap();
1678
1679        let returned = llm
1680            .last_request()
1681            .unwrap()
1682            .contents
1683            .iter()
1684            .flat_map(|c| c.parts.clone())
1685            .find_map(|p| match p {
1686                Part::FunctionResponse { function_response } => Some(function_response.response),
1687                _ => None,
1688            })
1689            .expect("the refusal is sent back");
1690        let error = returned["error"].as_str().unwrap_or_default();
1691        assert!(error.contains("over the refund limit"), "{returned}");
1692        assert!(returned.get("refunded").is_none(), "the tool must not run");
1693    }
1694
1695    /// `output::<T>()` sends `T`'s wire schema, and the reply parses back.
1696    #[tokio::test]
1697    async fn output_type_constrains_the_reply() {
1698        #[derive(serde::Deserialize, schemars::JsonSchema)]
1699        struct City {
1700            name: String,
1701        }
1702        let llm = MockLlm::text(r#"{"name":"Paris"}"#);
1703        let agent = AgentBuilder::new("geo")
1704            .output::<City>()
1705            .build(llm.clone())
1706            .unwrap();
1707        let result = agent
1708            .run_with(
1709                gemini_adk_rs::text::RunRequest::new("Capital of France?"),
1710                &gemini_adk_rs::State::new(),
1711            )
1712            .await
1713            .unwrap();
1714        assert_eq!(result.parse::<City>().unwrap().name, "Paris");
1715        let schema = llm.last_request().unwrap().response_json_schema.unwrap();
1716        assert_eq!(schema["properties"]["name"]["type"], "string");
1717    }
1718
1719    #[test]
1720    fn text_only_is_a_text_agent_setting() {
1721        assert!(
1722            AgentBuilder::new("t")
1723                .text_only()
1724                .build(MockLlm::text("x"))
1725                .is_ok()
1726        );
1727    }
1728}