gemini_adk_rs/text/
race.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 agents concurrently, returns the first to complete. Cancels the rest.
12pub struct RaceTextAgent {
13    name: String,
14    agents: Vec<Arc<dyn TextAgent>>,
15    middleware: MiddlewareChain,
16}
17
18impl RaceTextAgent {
19    /// Create a new race agent that runs agents concurrently and returns the first result.
20    pub fn new(name: impl Into<String>, agents: Vec<Arc<dyn TextAgent>>) -> Self {
21        Self {
22            name: name.into(),
23            agents,
24            middleware: MiddlewareChain::new(),
25        }
26    }
27
28    /// Attach a middleware chain. `AgentEvent::AgentStarted` is emitted
29    /// through it for every contender as the race starts and
30    /// `AgentEvent::AgentCompleted` once, for the winner; losers are cancelled
31    /// silently.
32    pub fn with_middleware_chain(mut self, chain: MiddlewareChain) -> Self {
33        self.middleware = chain;
34        self
35    }
36}
37
38#[async_trait]
39impl TextAgent for RaceTextAgent {
40    fn name(&self) -> &str {
41        &self.name
42    }
43
44    async fn run(&self, state: &State) -> Result<String, AgentError> {
45        if self.agents.is_empty() {
46            return Err(AgentError::Other("No agents in race".into()));
47        }
48
49        for agent in &self.agents {
50            let _ = self
51                .middleware
52                .run_on_event(&AgentEvent::AgentStarted {
53                    name: agent.name().to_string(),
54                })
55                .await;
56        }
57        let (tx, mut rx) = tokio::sync::mpsc::channel::<(String, Result<String, AgentError>)>(1);
58        let cancel = tokio_util::sync::CancellationToken::new();
59
60        let mut handles = Vec::with_capacity(self.agents.len());
61        for agent in &self.agents {
62            let agent = agent.clone();
63            let state = state.clone();
64            let tx = tx.clone();
65            let cancel = cancel.clone();
66
67            handles.push(tokio::spawn(async move {
68                tokio::select! {
69                    result = agent.run(&state) => {
70                        let _ = tx.send((agent.name().to_string(), result)).await;
71                    }
72                    _ = cancel.cancelled() => {}
73                }
74            }));
75        }
76        drop(tx); // Close our sender so rx completes when all are done.
77
78        let (winner, result) = rx.recv().await.unwrap_or((
79            String::new(),
80            Err(AgentError::Other("All race agents failed".into())),
81        ));
82
83        // Cancel remaining agents.
84        cancel.cancel();
85        for handle in handles {
86            handle.abort();
87        }
88
89        if result.is_ok() {
90            let _ = self
91                .middleware
92                .run_on_event(&AgentEvent::AgentCompleted { name: winner })
93                .await;
94        }
95        result
96    }
97}
gemini_adk_rs/text/race.rs

gemini_adk_rs/text/
race.rs
