Skip to main content

rs_adk/text/
timeout.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use async_trait::async_trait;
5
6use super::TextAgent;
7use crate::error::AgentError;
8use crate::state::State;
9
10/// Wraps an agent with a time limit. Returns `AgentError::Timeout` if exceeded.
11pub struct TimeoutTextAgent {
12    name: String,
13    inner: Arc<dyn TextAgent>,
14    timeout: Duration,
15}
16
17impl TimeoutTextAgent {
18    /// Create a new timeout agent wrapping an inner agent with a time limit.
19    pub fn new(name: impl Into<String>, inner: Arc<dyn TextAgent>, timeout: Duration) -> Self {
20        Self {
21            name: name.into(),
22            inner,
23            timeout,
24        }
25    }
26}
27
28#[async_trait]
29impl TextAgent for TimeoutTextAgent {
30    fn name(&self) -> &str {
31        &self.name
32    }
33
34    async fn run(&self, state: &State) -> Result<String, AgentError> {
35        match tokio::time::timeout(self.timeout, self.inner.run(state)).await {
36            Ok(result) => result,
37            Err(_) => Err(AgentError::Timeout),
38        }
39    }
40}