gemini_adk_rs/text/
timeout.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use async_trait::async_trait;
5
6use super::TextAgent;
7use crate::context::AgentEvent;
8use crate::error::AgentError;
9use crate::middleware::MiddlewareChain;
10use crate::state::State;
11
12/// Wraps an agent with a time limit. Returns `AgentError::Timeout` if exceeded.
13pub struct TimeoutTextAgent {
14    name: String,
15    inner: Arc<dyn TextAgent>,
16    timeout: Duration,
17    middleware: MiddlewareChain,
18}
19
20impl TimeoutTextAgent {
21    /// Create a new timeout agent wrapping an inner agent with a time limit.
22    pub fn new(name: impl Into<String>, inner: Arc<dyn TextAgent>, timeout: Duration) -> Self {
23        Self {
24            name: name.into(),
25            inner,
26            timeout,
27            middleware: MiddlewareChain::new(),
28        }
29    }
30
31    /// Attach a middleware chain. `AgentEvent::Timeout` is emitted through it
32    /// when the inner agent exceeds the limit, so `on_event` observers fire
33    /// before `AgentError::Timeout` is returned.
34    pub fn with_middleware_chain(mut self, chain: MiddlewareChain) -> Self {
35        self.middleware = chain;
36        self
37    }
38}
39
40#[async_trait]
41impl TextAgent for TimeoutTextAgent {
42    fn name(&self) -> &str {
43        &self.name
44    }
45
46    async fn run(&self, state: &State) -> Result<String, AgentError> {
47        match tokio::time::timeout(self.timeout, self.inner.run(state)).await {
48            Ok(result) => result,
49            Err(_) => {
50                let _ = self.middleware.run_on_event(&AgentEvent::Timeout).await;
51                Err(AgentError::Timeout)
52            }
53        }
54    }
55}
gemini_adk_rs/text/timeout.rs

gemini_adk_rs/text/
timeout.rs
