gemini_adk_rs/text/
parallel.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4
5use super::TextAgent;
6use crate::context::AgentEvent;
7use crate::error::AgentError;
8use crate::middleware::MiddlewareChain;
9use crate::state::State;
10
11/// Runs text agents concurrently. All branches share state. Results are
12/// collected and joined with newlines.
13pub struct ParallelTextAgent {
14    name: String,
15    branches: Vec<Arc<dyn TextAgent>>,
16    middleware: MiddlewareChain,
17}
18
19impl ParallelTextAgent {
20    /// Create a new parallel agent that runs branches concurrently.
21    pub fn new(name: impl Into<String>, branches: Vec<Arc<dyn TextAgent>>) -> Self {
22        Self {
23            name: name.into(),
24            branches,
25            middleware: MiddlewareChain::new(),
26        }
27    }
28
29    /// Attach a middleware chain. `AgentEvent::AgentStarted` is emitted
30    /// through it as each branch is spawned and `AgentEvent::AgentCompleted`
31    /// as each branch is joined (in branch order), so `on_event` observers see
32    /// the fan-out and fan-in.
33    pub fn with_middleware_chain(mut self, chain: MiddlewareChain) -> Self {
34        self.middleware = chain;
35        self
36    }
37}
38
39#[async_trait]
40impl TextAgent for ParallelTextAgent {
41    fn name(&self) -> &str {
42        &self.name
43    }
44
45    async fn run(&self, state: &State) -> Result<String, AgentError> {
46        let mut handles = Vec::with_capacity(self.branches.len());
47
48        for branch in &self.branches {
49            let _ = self
50                .middleware
51                .run_on_event(&AgentEvent::AgentStarted {
52                    name: branch.name().to_string(),
53                })
54                .await;
55            let branch = branch.clone();
56            let state = state.clone();
57            handles.push(tokio::spawn(async move { branch.run(&state).await }));
58        }
59
60        let mut results = Vec::with_capacity(handles.len());
61        for (branch, handle) in self.branches.iter().zip(handles) {
62            let result = handle
63                .await
64                .map_err(|e| AgentError::Other(format!("Join error: {e}")))?;
65            results.push(result?);
66            let _ = self
67                .middleware
68                .run_on_event(&AgentEvent::AgentCompleted {
69                    name: branch.name().to_string(),
70                })
71                .await;
72        }
73
74        let combined = results.join("\n");
75        let _ = state.set("output", &combined);
76        Ok(combined)
77    }
78}