gemini_adk_rs/orchestration/
mod.rs

1//! Agent orchestration — invoke an agent in a [`AgentMode`].
2//!
3//! An agent is a value ([`TextAgent`] — a local agent, a composed pipeline, or a
4//! remote A2A agent). Orchestration is the single question of *how* you invoke
5//! it; the result always lands in governed `State` under `{name}:result` (or
6//! `{name}:error`), so coordination is reactive and uniform regardless of the
7//! invoker (the model, a `Flow`, an `Extract`, or a watcher).
8//!
9//! | Mode | Sync? | Lowers to |
10//! |------|-------|-----------|
11//! | [`AgentMode::Call`] | sync — caller awaits | [`call_agent`] (agent-as-tool, awaited inline) |
12//! | [`AgentMode::Dispatch`] | async, fire-and-forget | [`BackgroundAgentDispatcher::dispatch`](crate::live::BackgroundAgentDispatcher) |
13//! | [`AgentMode::Background`] | async, model-aware | an agent-tool marked [`ToolExecutionMode::Background`](crate::live::ToolExecutionMode) |
14//!
15//! All three write `{name}:result`, so a `Flow` step can complete on a resolved
16//! result via [`Guard::resolved`](crate::flow::Guard::resolved), and any
17//! consumer reads the value the same way.
18
19use std::future::Future;
20use std::sync::Arc;
21
22use serde_json::Value;
23
24use crate::AsyncSourceFn;
25use crate::error::AgentError;
26use crate::llm::{BaseLlm, LlmRequest};
27use crate::state::State;
28use crate::text::TextAgent;
29
30/// How an agent is invoked.
31#[derive(Clone, Copy, Debug, PartialEq, Eq)]
32pub enum AgentMode {
33    /// Synchronous — the caller awaits the result. Use only for *fast*
34    /// dependencies (a voice session should not block on slow work).
35    Call,
36    /// Asynchronous, fire-and-forget — the conversation does not wait.
37    Dispatch,
38    /// Asynchronous, model-aware — runs detached; the result is delivered back
39    /// to the model via `FunctionResponseScheduling`.
40    Background,
41}
42
43/// State key an agent's successful result is written to.
44pub fn result_key(name: &str) -> String {
45    format!("{name}:result")
46}
47
48/// State key an agent's error is written to.
49pub fn error_key(name: &str) -> String {
50    format!("{name}:error")
51}
52
53/// The provenance source of a value at `key` (e.g. `"agent"`, `"fetch"`,
54/// `"llm"`, or `"extraction"`), if one was recorded under `state_meta:{key}`.
55pub fn provenance(state: &State, key: &str) -> Option<String> {
56    state
57        .get::<serde_json::Value>(&format!("state_meta:{key}"))
58        .and_then(|m| m.get("source").and_then(|s| s.as_str().map(String::from)))
59}
60
61/// Invoke `agent` **synchronously**: run it to completion, write its result to
62/// `{name}:result` (or its error to `{name}:error`), and return the result.
63///
64/// This is the [`AgentMode::Call`] lowering. It uses the same `{name}:result`
65/// convention as [`BackgroundAgentDispatcher::dispatch`](crate::live::BackgroundAgentDispatcher),
66/// so sync and async invocations are observed identically.
67pub async fn call_agent(
68    name: &str,
69    agent: Arc<dyn TextAgent>,
70    state: &State,
71) -> Result<String, AgentError> {
72    let result = agent.run(state).await;
73    match &result {
74        Ok(r) => {
75            let key = result_key(name);
76            let _ = state.set(
77                format!("state_meta:{key}"),
78                serde_json::json!({ "source": "agent", "resolver": name }),
79            );
80            let _ = state.set(key, r);
81        }
82        Err(e) => {
83            let _ = state.set(error_key(name), e.to_string());
84        }
85    }
86    result
87}
88
89enum Source {
90    /// Run a [`TextAgent`] (which reads its inputs from `State`).
91    Agent(Arc<dyn TextAgent>),
92    /// Run an async closure that reads `State` and returns a value — the seam
93    /// for a tool call, an HTTP fetch, or an MCP request.
94    Fetch(AsyncSourceFn),
95    /// One-shot OOB LLM completion over a `State`-interpolated prompt.
96    Llm {
97        /// The out-of-band LLM.
98        llm: Arc<dyn BaseLlm>,
99        /// Prompt template; `{key}` interpolates the `State` value at `key`.
100        prompt: String,
101    },
102}
103
104/// Interpolate `{key}` placeholders in `template` with `State` string values.
105fn interpolate(template: &str, state: &State) -> String {
106    let mut out = String::with_capacity(template.len());
107    let mut rest = template;
108    while let Some(open) = rest.find('{') {
109        out.push_str(&rest[..open]);
110        let after = &rest[open + 1..];
111        let Some(close) = after.find('}') else {
112            out.push_str(&rest[open..]);
113            return out;
114        };
115        let key = after[..close].trim();
116        match state.get::<serde_json::Value>(key) {
117            Some(serde_json::Value::String(s)) => out.push_str(&s),
118            Some(v) => out.push_str(&v.to_string()),
119            None => {}
120        }
121        rest = &after[close + 1..];
122    }
123    out.push_str(rest);
124    out
125}
126
127/// A named async value source whose inputs come from `State` and whose result
128/// lands back in `State` under `{name}:result` (or `{name}:error`).
129///
130/// `Resolver` is the async sibling of the deterministic
131/// [`Recognizer`](crate::extract::Recognizer): both are *inputs from State →
132/// value*. A `Resolver`
133/// generalizes [`call_agent`] from "a sub-agent" to **any** async source — a sub-agent
134/// ([`Resolver::agent`]) or a system fetch / tool call / MCP request
135/// ([`Resolver::fetch`]) — under one result convention, so a `Flow` step can
136/// complete on it via [`Guard::resolved`](crate::flow::Guard::resolved)
137/// regardless of where the value came from.
138pub struct Resolver {
139    name: String,
140    source: Source,
141}
142
143impl Resolver {
144    /// Resolve by running a sub-agent. Its `String` output becomes the result.
145    pub fn agent(name: impl Into<String>, agent: Arc<dyn TextAgent>) -> Self {
146        Self {
147            name: name.into(),
148            source: Source::Agent(agent),
149        }
150    }
151
152    /// Resolve by running an async closure over a clone of `State` — the seam
153    /// for an HTTP fetch, a tool call, or an MCP request. The closure returns
154    /// `Ok(value)` on success or `Err(message)` to record an error.
155    pub fn fetch<F, Fut>(name: impl Into<String>, f: F) -> Self
156    where
157        F: Fn(State) -> Fut + Send + Sync + 'static,
158        Fut: Future<Output = Result<Value, String>> + Send + 'static,
159    {
160        let f = Arc::new(f);
161        Self {
162            name: name.into(),
163            source: Source::Fetch(Arc::new(move |state| {
164                let f = f.clone();
165                Box::pin(async move { f(state).await })
166            })),
167        }
168    }
169
170    /// Resolve by running a one-shot OOB LLM over a `State`-interpolated prompt
171    /// (`{key}` placeholders). The completion text becomes the result.
172    pub fn llm(name: impl Into<String>, llm: Arc<dyn BaseLlm>, prompt: impl Into<String>) -> Self {
173        Self {
174            name: name.into(),
175            source: Source::Llm {
176                llm,
177                prompt: prompt.into(),
178            },
179        }
180    }
181
182    /// The resolver's name (the `{name}:result` prefix it writes).
183    pub fn name(&self) -> &str {
184        &self.name
185    }
186
187    /// The provenance kind of this resolver's source (`agent`/`fetch`/`llm`).
188    fn source_kind(&self) -> &'static str {
189        match &self.source {
190            Source::Agent(_) => "agent",
191            Source::Fetch(_) => "fetch",
192            Source::Llm { .. } => "llm",
193        }
194    }
195
196    /// Resolve **synchronously** ([`AgentMode::Call`]): await the source, write its
197    /// value to `{name}:result` (or its error to `{name}:error`), record its
198    /// provenance under `state_meta:{name}:result`, and return it.
199    pub async fn resolve(&self, state: &State) -> Result<Value, String> {
200        let outcome = match &self.source {
201            Source::Agent(a) => a
202                .run(state)
203                .await
204                .map(Value::from)
205                .map_err(|e| e.to_string()),
206            Source::Fetch(f) => f(state.clone()).await,
207            Source::Llm { llm, prompt } => {
208                let rendered = interpolate(prompt, state);
209                llm.generate(LlmRequest::from_text(rendered))
210                    .await
211                    .map(|r| Value::from(r.text()))
212                    .map_err(|e| e.to_string())
213            }
214        };
215        match &outcome {
216            Ok(v) => {
217                let key = result_key(&self.name);
218                let _ = state.set(
219                    format!("state_meta:{key}"),
220                    serde_json::json!({ "source": self.source_kind(), "resolver": self.name }),
221                );
222                let _ = state.set(key, v.clone());
223            }
224            Err(e) => {
225                let _ = state.set(error_key(&self.name), e);
226            }
227        }
228        outcome
229    }
230
231    /// Resolve **detached** ([`AgentMode::Dispatch`]): spawn the resolution on the
232    /// runtime and return immediately. The conversation does not wait; consumers
233    /// observe completion reactively via `{name}:result`.
234    pub fn dispatch(self, state: State) {
235        tokio::spawn(async move {
236            let _ = self.resolve(&state).await;
237        });
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use async_trait::async_trait;
245    use serde_json::json;
246
247    struct Echo(&'static str);
248    #[async_trait]
249    impl TextAgent for Echo {
250        fn name(&self) -> &str {
251            "echo"
252        }
253        async fn run(&self, _state: &State) -> Result<String, AgentError> {
254            Ok(self.0.to_string())
255        }
256    }
257
258    struct Boom;
259    #[async_trait]
260    impl TextAgent for Boom {
261        fn name(&self) -> &str {
262            "boom"
263        }
264        async fn run(&self, _state: &State) -> Result<String, AgentError> {
265            Err(AgentError::Other("kaboom".into()))
266        }
267    }
268
269    #[tokio::test]
270    async fn call_writes_result_to_state() {
271        let state = State::new();
272        let out = call_agent("verify", Arc::new(Echo("ok-123")), &state)
273            .await
274            .unwrap();
275        assert_eq!(out, "ok-123");
276        assert_eq!(
277            state.get::<String>("verify:result").as_deref(),
278            Some("ok-123")
279        );
280    }
281
282    #[tokio::test]
283    async fn call_writes_error_to_state() {
284        let state = State::new();
285        let r = call_agent("verify", Arc::new(Boom), &state).await;
286        assert!(r.is_err());
287        assert!(state.contains("verify:error"));
288        assert!(!state.contains("verify:result"));
289    }
290
291    #[tokio::test]
292    async fn resolver_fetch_binds_state_and_writes_result() {
293        let state = State::new();
294        let _ = state.set("slot", "afternoon");
295        let r = Resolver::fetch("availability", |s: State| async move {
296            // Inputs come from State; the value is arbitrary JSON.
297            let slot = s.get::<String>("slot").unwrap_or_default();
298            Ok(json!({ "open": slot == "afternoon" }))
299        });
300        let out = r.resolve(&state).await.unwrap();
301        assert_eq!(out, json!({ "open": true }));
302        assert_eq!(
303            state.get::<Value>("availability:result"),
304            Some(json!({ "open": true }))
305        );
306        // Provenance is recorded for the resolved value.
307        assert_eq!(
308            provenance(&state, "availability:result").as_deref(),
309            Some("fetch")
310        );
311    }
312
313    #[tokio::test]
314    async fn resolver_agent_uses_result_convention() {
315        let state = State::new();
316        // An agent resolver shares the `{name}:result` convention with `call`.
317        Resolver::agent("verify", Arc::new(Echo("ok-9")))
318            .resolve(&state)
319            .await
320            .unwrap();
321        assert_eq!(
322            state.get::<String>("verify:result").as_deref(),
323            Some("ok-9")
324        );
325    }
326
327    #[tokio::test]
328    async fn resolver_fetch_records_error() {
329        let state = State::new();
330        let r = Resolver::fetch("lookup", |_s: State| async move {
331            Err::<Value, String>("upstream 503".into())
332        });
333        assert!(r.resolve(&state).await.is_err());
334        assert_eq!(
335            state.get::<String>("lookup:error").as_deref(),
336            Some("upstream 503")
337        );
338        assert!(!state.contains("lookup:result"));
339    }
340
341    #[tokio::test]
342    async fn resolver_llm_interpolates_prompt_and_stores_text() {
343        use crate::llm::{LlmError, LlmResponse};
344        use gemini_genai_rs::prelude::Content;
345
346        struct EchoLlm;
347        #[async_trait]
348        impl BaseLlm for EchoLlm {
349            fn model_id(&self) -> &str {
350                "echo"
351            }
352            async fn generate(&self, request: LlmRequest) -> Result<LlmResponse, LlmError> {
353                // Echo the (interpolated) prompt back as the completion.
354                let prompt = request.contents[0].parts.iter().find_map(|p| match p {
355                    gemini_genai_rs::prelude::Part::Text { text } => Some(text.clone()),
356                    _ => None,
357                });
358                Ok(LlmResponse {
359                    content: Content::model(prompt.unwrap_or_default()),
360                    finish_reason: None,
361                    usage: None,
362                })
363            }
364        }
365
366        let state = State::new();
367        let _ = state.set("topic", "billing");
368        let out = Resolver::llm("summary", Arc::new(EchoLlm), "Summarize the {topic} issue")
369            .resolve(&state)
370            .await
371            .unwrap();
372        assert_eq!(out, json!("Summarize the billing issue"));
373        assert_eq!(
374            state.get::<String>("summary:result").as_deref(),
375            Some("Summarize the billing issue")
376        );
377    }
378
379    #[tokio::test]
380    async fn resolver_dispatch_runs_detached() {
381        let state = State::new();
382        Resolver::fetch("ping", |_s: State| async move { Ok(json!("pong")) })
383            .dispatch(state.clone());
384        // The spawned task writes the result; await it becoming visible.
385        for _ in 0..100 {
386            if state.contains("ping:result") {
387                break;
388            }
389            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
390        }
391        assert_eq!(state.get::<String>("ping:result").as_deref(), Some("pong"));
392    }
393}
gemini_adk_rs/orchestration/mod.rs

gemini_adk_rs/orchestration/
mod.rs
