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
10pub struct TimeoutTextAgent {
12 name: String,
13 inner: Arc<dyn TextAgent>,
14 timeout: Duration,
15}
16
17impl TimeoutTextAgent {
18 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}