askit_std_agents/
time.rs

1use std::str::FromStr;
2use std::sync::{Arc, Mutex};
3use std::time::Duration;
4use std::vec;
5
6use agent_stream_kit::{
7    ASKit, Agent, AgentContext, AgentData, AgentError, AgentOutput, AgentSpec, AgentStatus,
8    AgentValue, AsAgent, askit_agent, async_trait,
9};
10use chrono::{DateTime, Local, Utc};
11use cron::Schedule;
12use log;
13use regex::Regex;
14use tokio::task::JoinHandle;
15
16const CATEGORY: &str = "Std/Time";
17
18const PIN_TIME: &str = "time";
19const PIN_VALUE: &str = "value";
20const PIN_UNIT: &str = "unit";
21
22const CONFIG_DELAY: &str = "delay";
23const CONFIG_MAX_NUM_DATA: &str = "max_num_data";
24const CONFIG_INTERVAL: &str = "interval";
25const CONFIG_SCHEDULE: &str = "schedule";
26const CONFIG_TIME: &str = "time";
27
28const DELAY_MS_DEFAULT: i64 = 1000; // 1 second in milliseconds
29const MAX_NUM_DATA_DEFAULT: i64 = 10;
30const INTERVAL_DEFAULT: &str = "10s";
31const TIME_DEFAULT: &str = "1s";
32
33// Delay Agent
34#[askit_agent(
35    title = "Delay",
36    description = "Delays output by a specified time",
37    category = CATEGORY,
38    inputs = [PIN_VALUE],
39    outputs = [PIN_VALUE],
40    integer_config(name = CONFIG_DELAY, default = DELAY_MS_DEFAULT, title = "delay (ms)"),
41    integer_config(name = CONFIG_MAX_NUM_DATA, default = MAX_NUM_DATA_DEFAULT, title = "max num data")
42)]
43struct DelayAgent {
44    data: AgentData,
45    num_waiting_data: Arc<Mutex<i64>>,
46}
47
48#[async_trait]
49impl AsAgent for DelayAgent {
50    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
51        Ok(Self {
52            data: AgentData::new(askit, id, spec),
53            num_waiting_data: Arc::new(Mutex::new(0)),
54        })
55    }
56
57    async fn process(
58        &mut self,
59        ctx: AgentContext,
60        pin: String,
61        value: AgentValue,
62    ) -> Result<(), AgentError> {
63        let config = self.configs()?;
64        let delay_ms = config.get_integer_or(CONFIG_DELAY, DELAY_MS_DEFAULT);
65        let max_num_data = config.get_integer_or(CONFIG_MAX_NUM_DATA, MAX_NUM_DATA_DEFAULT);
66
67        // To avoid generating too many timers
68        {
69            let num_waiting_data = self.num_waiting_data.clone();
70            let mut num_waiting_data = num_waiting_data.lock().unwrap();
71            if *num_waiting_data >= max_num_data {
72                return Ok(());
73            }
74            *num_waiting_data += 1;
75        }
76
77        tokio::time::sleep(Duration::from_millis(delay_ms as u64)).await;
78
79        self.output(ctx.clone(), pin, value.clone()).await?;
80
81        let mut num_waiting_data = self.num_waiting_data.lock().unwrap();
82        *num_waiting_data -= 1;
83
84        Ok(())
85    }
86}
87
88// Interval Timer Agent
89#[askit_agent(
90    title = "Interval Timer",
91    description = "Outputs a unit signal at specified intervals",
92    category = CATEGORY,
93    outputs = [PIN_UNIT],
94    string_config(name = CONFIG_INTERVAL, default = INTERVAL_DEFAULT, description = "(ex. 10s, 5m, 100ms, 1h, 1d)")
95)]
96struct IntervalTimerAgent {
97    data: AgentData,
98    timer_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
99    interval_ms: u64,
100}
101
102impl IntervalTimerAgent {
103    fn start_timer(&mut self) -> Result<(), AgentError> {
104        let timer_handle = self.timer_handle.clone();
105        let interval_ms = self.interval_ms;
106
107        let askit = self.askit().clone();
108        let agent_id = self.id().to_string();
109        let handle = self.runtime().spawn(async move {
110            loop {
111                // Sleep for the configured interval
112                tokio::time::sleep(tokio::time::Duration::from_millis(interval_ms)).await;
113
114                // Check if we've been stopped
115                if let Ok(handle) = timer_handle.lock() {
116                    if handle.is_none() {
117                        break;
118                    }
119                }
120
121                // Create a unit output
122                if let Err(e) = askit.try_send_agent_out(
123                    agent_id.clone(),
124                    AgentContext::new(),
125                    PIN_UNIT.to_string(),
126                    AgentValue::unit(),
127                ) {
128                    log::error!("Failed to send interval timer output: {}", e);
129                }
130            }
131        });
132
133        // Store the timer handle
134        if let Ok(mut timer_handle) = self.timer_handle.lock() {
135            *timer_handle = Some(handle);
136        }
137
138        Ok(())
139    }
140
141    fn stop_timer(&mut self) -> Result<(), AgentError> {
142        // Cancel the timer
143        if let Ok(mut timer_handle) = self.timer_handle.lock() {
144            if let Some(handle) = timer_handle.take() {
145                handle.abort();
146            }
147        }
148        Ok(())
149    }
150}
151
152#[async_trait]
153impl AsAgent for IntervalTimerAgent {
154    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
155        let interval = spec
156            .configs
157            .as_ref()
158            .ok_or(AgentError::NoConfig)?
159            .get_string_or(CONFIG_INTERVAL, INTERVAL_DEFAULT);
160        let interval_ms = parse_duration_to_ms(&interval)?;
161
162        Ok(Self {
163            data: AgentData::new(askit, id, spec),
164            timer_handle: Default::default(),
165            interval_ms,
166        })
167    }
168
169    async fn start(&mut self) -> Result<(), AgentError> {
170        self.start_timer()
171    }
172
173    async fn stop(&mut self) -> Result<(), AgentError> {
174        self.stop_timer()
175    }
176
177    fn configs_changed(&mut self) -> Result<(), AgentError> {
178        // Check if interval has changed
179        let interval = self.configs()?.get_string(CONFIG_INTERVAL)?;
180        let new_interval = parse_duration_to_ms(&interval)?;
181        if new_interval != self.interval_ms {
182            self.interval_ms = new_interval;
183            if *self.status() == AgentStatus::Start {
184                // Restart the timer with the new interval
185                self.stop_timer()?;
186                self.start_timer()?;
187            }
188        }
189        Ok(())
190    }
191}
192
193// OnStart
194#[askit_agent(
195    title = "On Start",
196    category = CATEGORY,
197    outputs = [PIN_UNIT],
198    integer_config(name = CONFIG_DELAY, default = DELAY_MS_DEFAULT, title = "delay (ms)")
199)]
200struct OnStartAgent {
201    data: AgentData,
202}
203
204#[async_trait]
205impl AsAgent for OnStartAgent {
206    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
207        Ok(Self {
208            data: AgentData::new(askit, id, spec),
209        })
210    }
211
212    async fn start(&mut self) -> Result<(), AgentError> {
213        let config = self.configs()?;
214        let delay_ms = config.get_integer_or(CONFIG_DELAY, DELAY_MS_DEFAULT);
215
216        let askit = self.askit().clone();
217        let agent_id = self.id().to_string();
218
219        self.runtime().spawn(async move {
220            tokio::time::sleep(Duration::from_millis(delay_ms as u64)).await;
221
222            if let Err(e) = askit.try_send_agent_out(
223                agent_id,
224                AgentContext::new(),
225                PIN_UNIT.to_string(),
226                AgentValue::unit(),
227            ) {
228                log::error!("Failed to send delayed output: {}", e);
229            }
230        });
231
232        Ok(())
233    }
234}
235
236// Schedule Timer Agent
237#[askit_agent(
238    title = "Schedule Timer",
239    category = CATEGORY,
240    outputs = [PIN_TIME],
241    string_config(name = CONFIG_SCHEDULE, default = "0 0 * * * *", description = "sec min hour day month week year")
242)]
243struct ScheduleTimerAgent {
244    data: AgentData,
245    cron_schedule: Option<Schedule>,
246    timer_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
247}
248
249impl ScheduleTimerAgent {
250    fn start_timer(&mut self) -> Result<(), AgentError> {
251        let Some(schedule) = &self.cron_schedule else {
252            return Err(AgentError::InvalidConfig("No schedule defined".into()));
253        };
254
255        let askit = self.askit().clone();
256        let agent_id = self.id().to_string();
257        let timer_handle = self.timer_handle.clone();
258        let schedule = schedule.clone();
259
260        let handle = self.runtime().spawn(async move {
261            loop {
262                // Calculate the next time this schedule should run
263                let now: DateTime<Utc> = Utc::now();
264                let next = match schedule.upcoming(Utc).next() {
265                    Some(next_time) => next_time,
266                    None => {
267                        log::error!("No upcoming schedule times found");
268                        break;
269                    }
270                };
271
272                // Calculate the duration until the next scheduled time
273                let duration = match (next - now).to_std() {
274                    Ok(duration) => duration,
275                    Err(e) => {
276                        log::error!("Failed to calculate duration until next schedule: {}", e);
277                        // If we can't calculate the duration, sleep for a short time and try again
278                        tokio::time::sleep(Duration::from_secs(60)).await;
279                        continue;
280                    }
281                };
282
283                let next_local = next.with_timezone(&Local);
284                log::debug!(
285                    "Scheduling timer for '{}' to fire at {} (in {:?})",
286                    agent_id,
287                    next_local.format("%Y-%m-%d %H:%M:%S %z"),
288                    duration
289                );
290
291                // Sleep until the next scheduled time
292                tokio::time::sleep(duration).await;
293
294                // Check if we've been stopped
295                if let Ok(handle) = timer_handle.lock() {
296                    if handle.is_none() {
297                        break;
298                    }
299                }
300
301                // Get the current local timestamp (in seconds)
302                let current_local_time = Local::now().timestamp();
303
304                // Output the timestamp as an integer
305                if let Err(e) = askit.try_send_agent_out(
306                    agent_id.clone(),
307                    AgentContext::new(),
308                    PIN_TIME.to_string(),
309                    AgentValue::integer(current_local_time),
310                ) {
311                    log::error!("Failed to send schedule timer output: {}", e);
312                }
313            }
314        });
315
316        // Store the timer handle
317        if let Ok(mut timer_handle) = self.timer_handle.lock() {
318            *timer_handle = Some(handle);
319        }
320
321        Ok(())
322    }
323
324    fn stop_timer(&mut self) -> Result<(), AgentError> {
325        // Cancel the timer
326        if let Ok(mut timer_handle) = self.timer_handle.lock() {
327            if let Some(handle) = timer_handle.take() {
328                handle.abort();
329            }
330        }
331        Ok(())
332    }
333
334    fn parse_schedule(&mut self, schedule_str: &str) -> Result<(), AgentError> {
335        if schedule_str.trim().is_empty() {
336            self.cron_schedule = None;
337            return Ok(());
338        }
339
340        let schedule = Schedule::from_str(schedule_str).map_err(|e| {
341            AgentError::InvalidConfig(format!("Invalid cron schedule '{}': {}", schedule_str, e))
342        })?;
343        self.cron_schedule = Some(schedule);
344        Ok(())
345    }
346}
347
348#[async_trait]
349impl AsAgent for ScheduleTimerAgent {
350    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
351        let schedule_str = spec
352            .configs
353            .as_ref()
354            .map(|cfg| cfg.get_string(CONFIG_SCHEDULE))
355            .transpose()?;
356
357        let mut agent = Self {
358            data: AgentData::new(askit, id, spec),
359            cron_schedule: None,
360            timer_handle: Default::default(),
361        };
362
363        if let Some(schedule_str) = schedule_str {
364            if !schedule_str.is_empty() {
365                agent.parse_schedule(&schedule_str)?;
366            }
367        }
368
369        Ok(agent)
370    }
371
372    async fn start(&mut self) -> Result<(), AgentError> {
373        if self.cron_schedule.is_some() {
374            self.start_timer()?;
375        }
376        Ok(())
377    }
378
379    async fn stop(&mut self) -> Result<(), AgentError> {
380        self.stop_timer()
381    }
382
383    fn configs_changed(&mut self) -> Result<(), AgentError> {
384        // Check if schedule has changed
385        let schedule_str = self.configs()?.get_string(CONFIG_SCHEDULE)?;
386        self.parse_schedule(&schedule_str)?;
387
388        if *self.status() == AgentStatus::Start {
389            // Restart the timer with the new schedule
390            self.stop_timer()?;
391            if self.cron_schedule.is_some() {
392                self.start_timer()?;
393            }
394        }
395        Ok(())
396    }
397}
398
399// Throttle agent
400#[askit_agent(
401    title = "Throttle Time",
402    category = CATEGORY,
403    inputs = [PIN_VALUE],
404    outputs = [PIN_VALUE],
405    string_config(name = CONFIG_TIME, default = TIME_DEFAULT, description = "(ex. 10s, 5m, 100ms, 1h, 1d)"),
406    integer_config(name = CONFIG_MAX_NUM_DATA, title = "max num data", description = "0: no data, -1: all data")
407)]
408struct ThrottleTimeAgent {
409    data: AgentData,
410    timer_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
411    time_ms: u64,
412    max_num_data: i64,
413    waiting_data: Arc<Mutex<Vec<(AgentContext, String, AgentValue)>>>,
414}
415
416impl ThrottleTimeAgent {
417    fn start_timer(&mut self) -> Result<(), AgentError> {
418        let timer_handle = self.timer_handle.clone();
419        let time_ms = self.time_ms;
420
421        let waiting_data = self.waiting_data.clone();
422        let askit = self.askit().clone();
423        let agent_id = self.id().to_string();
424
425        let handle = self.runtime().spawn(async move {
426            loop {
427                // Sleep for the configured interval
428                tokio::time::sleep(tokio::time::Duration::from_millis(time_ms)).await;
429
430                // Check if we've been stopped
431                let mut handle = timer_handle.lock().unwrap();
432                if handle.is_none() {
433                    break;
434                }
435
436                // process the waiting data
437                let mut wd = waiting_data.lock().unwrap();
438                if wd.len() > 0 {
439                    // If there are data waiting, output the first one
440                    let (ctx, pin, data) = wd.remove(0);
441                    askit
442                        .try_send_agent_out(agent_id.clone(), ctx, pin, data)
443                        .unwrap_or_else(|e| {
444                            log::error!("Failed to send delayed output: {}", e);
445                        });
446                }
447
448                // If there are no data waiting, we stop the timer
449                if wd.len() == 0 {
450                    handle.take();
451                    break;
452                }
453            }
454        });
455
456        // Store the timer handle
457        if let Ok(mut timer_handle) = self.timer_handle.lock() {
458            *timer_handle = Some(handle);
459        }
460
461        Ok(())
462    }
463
464    fn stop_timer(&mut self) -> Result<(), AgentError> {
465        // Cancel the timer
466        if let Ok(mut timer_handle) = self.timer_handle.lock() {
467            if let Some(handle) = timer_handle.take() {
468                handle.abort();
469            }
470        }
471        Ok(())
472    }
473}
474
475#[async_trait]
476impl AsAgent for ThrottleTimeAgent {
477    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
478        let time = spec
479            .configs
480            .as_ref()
481            .ok_or(AgentError::NoConfig)?
482            .get_string_or(CONFIG_TIME, TIME_DEFAULT);
483        let time_ms = parse_duration_to_ms(&time)?;
484
485        let max_num_data = spec
486            .configs
487            .as_ref()
488            .ok_or(AgentError::NoConfig)?
489            .get_integer_or(CONFIG_MAX_NUM_DATA, 0);
490
491        Ok(Self {
492            data: AgentData::new(askit, id, spec),
493            timer_handle: Default::default(),
494            time_ms,
495            max_num_data,
496            waiting_data: Arc::new(Mutex::new(vec![])),
497        })
498    }
499
500    async fn stop(&mut self) -> Result<(), AgentError> {
501        self.stop_timer()
502    }
503
504    fn configs_changed(&mut self) -> Result<(), AgentError> {
505        // Check if interval has changed
506        let time = self.configs()?.get_string(CONFIG_TIME)?;
507        let new_time = parse_duration_to_ms(&time)?;
508        if new_time != self.time_ms {
509            self.time_ms = new_time;
510        }
511
512        // Check if max_num_data has changed
513        let max_num_data = self.configs()?.get_integer(CONFIG_MAX_NUM_DATA)?;
514        if self.max_num_data != max_num_data {
515            let mut wd = self.waiting_data.lock().unwrap();
516            let wd_len = wd.len();
517            if max_num_data >= 0 && wd_len > (max_num_data as usize) {
518                // If we have reached the max data to keep, we drop the oldest one
519                wd.drain(0..(wd_len - (max_num_data as usize)));
520            }
521            self.max_num_data = max_num_data;
522        }
523        Ok(())
524    }
525
526    async fn process(
527        &mut self,
528        ctx: AgentContext,
529        pin: String,
530        value: AgentValue,
531    ) -> Result<(), AgentError> {
532        if self.timer_handle.lock().unwrap().is_some() {
533            // If the timer is running, we just add the data to the waiting list
534            let mut wd = self.waiting_data.lock().unwrap();
535
536            // If max_num_data is 0, we don't need to keep any data
537            if self.max_num_data == 0 {
538                return Ok(());
539            }
540
541            wd.push((ctx, pin, value));
542            if self.max_num_data > 0 && wd.len() > self.max_num_data as usize {
543                // If we have reached the max data to keep, we drop the oldest one
544                wd.remove(0);
545            }
546
547            return Ok(());
548        }
549
550        // Start the timer
551        self.start_timer()?;
552
553        // Output the data
554        self.output(ctx, pin, value).await?;
555
556        Ok(())
557    }
558}
559
560// Parse time duration strings like "2s", "10m", "200ms"
561fn parse_duration_to_ms(duration_str: &str) -> Result<u64, AgentError> {
562    const MIN_DURATION: u64 = 10;
563
564    // Regular expression to match number followed by optional unit
565    let re = Regex::new(r"^(\d+)(?:([a-zA-Z]+))?$").expect("Failed to compile regex");
566
567    if let Some(captures) = re.captures(duration_str.trim()) {
568        let value: u64 = captures.get(1).unwrap().as_str().parse().map_err(|e| {
569            AgentError::InvalidConfig(format!(
570                "Invalid number in duration '{}': {}",
571                duration_str, e
572            ))
573        })?;
574
575        // Get the unit if present, default to "s" (seconds)
576        let unit = captures
577            .get(2)
578            .map_or("s".to_string(), |m| m.as_str().to_lowercase());
579
580        // Convert to milliseconds based on unit
581        let milliseconds = match unit.as_str() {
582            "ms" => value,               // already in milliseconds
583            "s" => value * 1000,         // seconds to milliseconds
584            "m" => value * 60 * 1000,    // minutes to milliseconds
585            "h" => value * 3600 * 1000,  // hours to milliseconds
586            "d" => value * 86400 * 1000, // days to milliseconds
587            _ => {
588                return Err(AgentError::InvalidConfig(format!(
589                    "Unknown time unit: {}",
590                    unit
591                )));
592            }
593        };
594
595        // Ensure we don't return less than the minimum duration
596        Ok(std::cmp::max(milliseconds, MIN_DURATION))
597    } else {
598        // If the string doesn't match the pattern, try to parse it as a plain number
599        // and assume it's in seconds
600        let value: u64 = duration_str.parse().map_err(|e| {
601            AgentError::InvalidConfig(format!("Invalid duration format '{}': {}", duration_str, e))
602        })?;
603        Ok(std::cmp::max(value * 1000, MIN_DURATION)) // Convert to ms
604    }
605}
askit_std_agents/time.rs

askit_std_agents/
time.rs
