tokio_prompt_orchestrator/
scheduler.rs

1//! Cron-style scheduler for periodic LLM prompt submissions.
2//!
3//! This module lets you register named prompt templates paired with a
4//! cron-like schedule expression.  A background Tokio task wakes up at the
5//! right wall-clock times and submits the prompt through the pipeline's
6//! standard input channel.
7//!
8//! ## Supported cron expressions
9//!
10//! The parser supports a **simplified** subset of cron:
11//!
12//! | Field | Positions | Accepted values |
13//! |-------|-----------|-----------------|
14//! | minute | 0 | `*`, `*/N`, `0`–`59` |
15//! | hour | 1 | `*`, `*/N`, `0`–`23` |
16//!
17//! Only `minute` and `hour` fields are parsed (a two-field mini-cron).
18//! Second-level precision is not supported.
19//!
20//! ### Expression examples
21//!
22//! | Expression | Meaning |
23//! |------------|---------|
24//! | `* *` | Every minute |
25//! | `*/5 *` | Every 5 minutes |
26//! | `0 *` | At the top of every hour |
27//! | `0 9` | Every day at 09:00 |
28//! | `30 6` | Every day at 06:30 |
29//! | `*/15 8` | Every 15 minutes during the 8 o'clock hour |
30//!
31//! ## Quick start
32//!
33//! ```rust,no_run
34//! use std::sync::Arc;
35//! use tokio::sync::mpsc;
36//! use tokio_prompt_orchestrator::{PromptRequest, SessionId};
37//! use tokio_prompt_orchestrator::scheduler::{Scheduler, ScheduledPrompt};
38//!
39//! #[tokio::main]
40//! async fn main() {
41//!     let (tx, mut rx) = mpsc::channel::<PromptRequest>(256);
42//!
43//!     let scheduler = Arc::new(Scheduler::new(tx));
44//!
45//!     let prompt = ScheduledPrompt::new("health-check", "*/5 *", "Are you operational?")
46//!         .expect("valid cron expression");
47//!     scheduler.add(prompt).await.expect("prompt registered");
48//!
49//!     let handle = scheduler.spawn();
50//!
51//!     // The scheduler runs until the handle is dropped or abort() is called.
52//!     handle.abort();
53//! }
54//! ```
55//!
56//! ## Web API integration
57//!
58//! When the `web-api` feature is enabled, the [`SchedulerState`] type provides
59//! a thread-safe wrapper suitable for use as Axum `State`.
60//! Register routes with [`scheduler_routes`] to expose
61//! `POST /api/v1/schedule`, `GET /api/v1/schedule`, and
62//! `DELETE /api/v1/schedule/:id`.
63
64use crate::{OrchestratorError, PromptRequest, SessionId};
65use chrono::{Local, Timelike};
66use serde::{Deserialize, Serialize};
67use std::collections::HashMap;
68use std::sync::Arc;
69use tokio::sync::mpsc;
70use tokio::sync::RwLock;
71use tokio::task::JoinHandle;
72use tracing::{debug, error, info, warn};
73use uuid::Uuid;
74
75// ============================================================================
76// CronField — parsed single cron field (minute or hour)
77// ============================================================================
78
79/// A parsed value for a single field (minute or hour) in a cron expression.
80#[derive(Debug, Clone, PartialEq, Eq)]
81enum CronField {
82    /// Matches every value (`*`).
83    Any,
84    /// Matches every N-th value starting from 0 (`*/N`).
85    Every(u32),
86    /// Matches exactly this value.
87    Exact(u32),
88}
89
90impl CronField {
91    /// Parse a single cron field token.
92    ///
93    /// # Errors
94    ///
95    /// Returns an error string when `token` cannot be parsed.
96    fn parse(token: &str, max: u32) -> Result<Self, String> {
97        if token == "*" {
98            return Ok(Self::Any);
99        }
100        if let Some(rest) = token.strip_prefix("*/") {
101            let n: u32 = rest
102                .parse()
103                .map_err(|_| format!("invalid step value in '{token}'"))?;
104            if n == 0 || n > max {
105                return Err(format!("step value {n} out of range (1..={max})"));
106            }
107            return Ok(Self::Every(n));
108        }
109        let v: u32 = token
110            .parse()
111            .map_err(|_| format!("invalid numeric value '{token}'"))?;
112        if v > max {
113            return Err(format!("value {v} out of range (0..={max})"));
114        }
115        Ok(Self::Exact(v))
116    }
117
118    /// Return `true` if `value` matches this field.
119    fn matches(&self, value: u32) -> bool {
120        match self {
121            Self::Any => true,
122            Self::Every(n) => value % n == 0,
123            Self::Exact(v) => *v == value,
124        }
125    }
126}
127
128// ============================================================================
129// CronExpression
130// ============================================================================
131
132/// A parsed two-field mini-cron expression `"MINUTE HOUR"`.
133///
134/// See the [module-level documentation](self) for the supported syntax.
135#[derive(Debug, Clone, PartialEq, Eq)]
136pub struct CronExpression {
137    minute: CronField,
138    hour: CronField,
139    /// The original source string, preserved for serialisation.
140    raw: String,
141}
142
143impl CronExpression {
144    /// Parse a cron expression string.
145    ///
146    /// # Errors
147    ///
148    /// Returns [`OrchestratorError::ConfigError`] when the expression cannot
149    /// be parsed.
150    ///
151    /// # Examples
152    ///
153    /// ```
154    /// use tokio_prompt_orchestrator::scheduler::CronExpression;
155    ///
156    /// assert!(CronExpression::parse("*/5 *").is_ok());
157    /// assert!(CronExpression::parse("0 9").is_ok());
158    /// assert!(CronExpression::parse("invalid").is_err());
159    /// ```
160    pub fn parse(expr: &str) -> Result<Self, OrchestratorError> {
161        let parts: Vec<&str> = expr.split_whitespace().collect();
162        if parts.len() != 2 {
163            return Err(OrchestratorError::ConfigError(format!(
164                "cron expression must have exactly 2 fields (minute hour), got '{expr}'"
165            )));
166        }
167        let minute = CronField::parse(parts[0], 59).map_err(|e| {
168            OrchestratorError::ConfigError(format!("minute field: {e}"))
169        })?;
170        let hour = CronField::parse(parts[1], 23).map_err(|e| {
171            OrchestratorError::ConfigError(format!("hour field: {e}"))
172        })?;
173        Ok(Self {
174            minute,
175            hour,
176            raw: expr.to_string(),
177        })
178    }
179
180    /// Return `true` if this expression matches the given `hour` and `minute`.
181    ///
182    /// # Panics
183    ///
184    /// This function does not panic.
185    pub fn matches(&self, hour: u32, minute: u32) -> bool {
186        self.hour.matches(hour) && self.minute.matches(minute)
187    }
188
189    /// Return the raw expression string as originally provided.
190    ///
191    /// # Panics
192    ///
193    /// This function does not panic.
194    pub fn as_str(&self) -> &str {
195        &self.raw
196    }
197}
198
199impl std::fmt::Display for CronExpression {
200    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
201        f.write_str(&self.raw)
202    }
203}
204
205impl Serialize for CronExpression {
206    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
207        s.serialize_str(&self.raw)
208    }
209}
210
211impl<'de> Deserialize<'de> for CronExpression {
212    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
213        let s = String::deserialize(d)?;
214        CronExpression::parse(&s).map_err(serde::de::Error::custom)
215    }
216}
217
218// ============================================================================
219// ScheduledPrompt
220// ============================================================================
221
222/// A prompt template paired with a cron schedule and target pipeline.
223///
224/// Created via [`ScheduledPrompt::new`] and registered with a [`Scheduler`].
225/// Each `ScheduledPrompt` carries a unique `id` (UUID v4) assigned at
226/// construction time.
227#[derive(Debug, Clone, Serialize, Deserialize)]
228pub struct ScheduledPrompt {
229    /// Unique identifier, auto-assigned at construction.
230    pub id: String,
231    /// Human-readable label (e.g. `"health-check"` or `"daily-summary"`).
232    pub name: String,
233    /// The cron expression controlling when this prompt fires.
234    pub schedule: CronExpression,
235    /// The prompt text template submitted to the pipeline.
236    pub prompt_template: String,
237    /// Optional session identifier.  When `None`, a new UUID is used per run.
238    pub session_id: Option<String>,
239    /// Whether this scheduled prompt is currently active.
240    pub enabled: bool,
241    /// Optional arbitrary metadata forwarded with each submitted request.
242    #[serde(default)]
243    pub metadata: HashMap<String, String>,
244}
245
246impl ScheduledPrompt {
247    /// Construct a new `ScheduledPrompt` with a generated UUID and enabled by default.
248    ///
249    /// # Errors
250    ///
251    /// Returns [`OrchestratorError::ConfigError`] when `cron_expr` cannot be
252    /// parsed.
253    ///
254    /// # Examples
255    ///
256    /// ```
257    /// use tokio_prompt_orchestrator::scheduler::ScheduledPrompt;
258    ///
259    /// let sp = ScheduledPrompt::new("health", "*/5 *", "ping").unwrap();
260    /// assert!(sp.enabled);
261    /// assert!(!sp.id.is_empty());
262    /// ```
263    pub fn new(
264        name: impl Into<String>,
265        cron_expr: &str,
266        prompt_template: impl Into<String>,
267    ) -> Result<Self, OrchestratorError> {
268        let schedule = CronExpression::parse(cron_expr)?;
269        Ok(Self {
270            id: Uuid::new_v4().to_string(),
271            name: name.into(),
272            schedule,
273            prompt_template: prompt_template.into(),
274            session_id: None,
275            enabled: true,
276            metadata: HashMap::new(),
277        })
278    }
279
280    /// Builder: set a fixed session identifier for all runs of this prompt.
281    ///
282    /// # Panics
283    ///
284    /// This function does not panic.
285    #[must_use]
286    pub fn with_session(mut self, session_id: impl Into<String>) -> Self {
287        self.session_id = Some(session_id.into());
288        self
289    }
290
291    /// Builder: attach extra metadata to every submission.
292    ///
293    /// # Panics
294    ///
295    /// This function does not panic.
296    #[must_use]
297    pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
298        self.metadata.insert(key.into(), value.into());
299        self
300    }
301
302    /// Builder: start this prompt in a disabled state.
303    ///
304    /// # Panics
305    ///
306    /// This function does not panic.
307    #[must_use]
308    pub fn disabled(mut self) -> Self {
309        self.enabled = false;
310        self
311    }
312}
313
314// ============================================================================
315// Scheduler
316// ============================================================================
317
318/// Tokio-based scheduler that fires prompt submissions on a cron-like interval.
319///
320/// The scheduler holds a collection of [`ScheduledPrompt`]s and a sender end
321/// of the pipeline's input channel.  Call [`Scheduler::spawn`] to start the
322/// background task.
323///
324/// The scheduler checks the current wall-clock time **once per minute** and
325/// fires every matching prompt.
326///
327/// ## Thread Safety
328///
329/// `Scheduler` is `Send + Sync` and is typically shared as `Arc<Scheduler>`.
330/// The internal prompt list is guarded by a [`RwLock`] so that prompts can be
331/// added and removed without stopping the background task.
332pub struct Scheduler {
333    prompts: RwLock<Vec<ScheduledPrompt>>,
334    tx: mpsc::Sender<PromptRequest>,
335}
336
337impl Scheduler {
338    /// Create a new `Scheduler` that submits prompts on `tx`.
339    ///
340    /// # Panics
341    ///
342    /// This function does not panic.
343    pub fn new(tx: mpsc::Sender<PromptRequest>) -> Self {
344        Self {
345            prompts: RwLock::new(Vec::new()),
346            tx,
347        }
348    }
349
350    /// Register a [`ScheduledPrompt`].
351    ///
352    /// Returns the assigned `id` (a UUID v4 string) so callers can reference
353    /// it later with [`Scheduler::remove`] or through the web API.
354    ///
355    /// # Errors
356    ///
357    /// Returns [`OrchestratorError::ConfigError`] if the prompt is disabled
358    /// and thus would never fire.  Disabled prompts *can* be added — they are
359    /// accepted but will not emit requests until re-enabled (e.g. via the web API).
360    ///
361    /// # Panics
362    ///
363    /// This function does not panic.
364    pub async fn add(&self, prompt: ScheduledPrompt) -> Result<String, OrchestratorError> {
365        let id = prompt.id.clone();
366        info!(
367            id = %id,
368            name = %prompt.name,
369            schedule = %prompt.schedule,
370            enabled = prompt.enabled,
371            "registered scheduled prompt"
372        );
373        self.prompts.write().await.push(prompt);
374        Ok(id)
375    }
376
377    /// Remove the scheduled prompt with the given `id`.
378    ///
379    /// Returns `true` if a prompt with that ID existed and was removed, `false`
380    /// if no matching ID was found.
381    ///
382    /// # Panics
383    ///
384    /// This function does not panic.
385    pub async fn remove(&self, id: &str) -> bool {
386        let mut prompts = self.prompts.write().await;
387        let before = prompts.len();
388        prompts.retain(|p| p.id != id);
389        let removed = prompts.len() < before;
390        if removed {
391            info!(id = %id, "removed scheduled prompt");
392        } else {
393            warn!(id = %id, "remove_scheduled_prompt: id not found");
394        }
395        removed
396    }
397
398    /// Return a snapshot of all registered prompts (including disabled ones).
399    ///
400    /// # Panics
401    ///
402    /// This function does not panic.
403    pub async fn list(&self) -> Vec<ScheduledPrompt> {
404        self.prompts.read().await.clone()
405    }
406
407    /// Enable or disable a scheduled prompt by ID.
408    ///
409    /// Returns `true` when the prompt was found and updated.
410    ///
411    /// # Panics
412    ///
413    /// This function does not panic.
414    pub async fn set_enabled(&self, id: &str, enabled: bool) -> bool {
415        let mut prompts = self.prompts.write().await;
416        if let Some(p) = prompts.iter_mut().find(|p| p.id == id) {
417            p.enabled = enabled;
418            info!(id = %id, enabled = enabled, "updated scheduled prompt state");
419            true
420        } else {
421            false
422        }
423    }
424
425    /// Spawn the background scheduler task and return its [`JoinHandle`].
426    ///
427    /// The task loops forever, sleeping until the next minute boundary, then
428    /// submitting all matching prompts.  It stops when the handle is dropped
429    /// or [`JoinHandle::abort`] is called.
430    ///
431    /// # Panics
432    ///
433    /// The spawned task does not panic; any submission errors are logged and
434    /// skipped.
435    pub fn spawn(self: &Arc<Self>) -> JoinHandle<()> {
436        let scheduler = Arc::clone(self);
437        tokio::spawn(async move {
438            info!("Scheduler started");
439            loop {
440                // Sleep until the start of the next minute.
441                let now = Local::now();
442                let secs_into_minute = now.second() as u64;
443                let nanos_remaining = now.nanosecond() as u64;
444                let sleep_ms = (60 - secs_into_minute) * 1000
445                    + if nanos_remaining > 0 { 1 } else { 0 };
446
447                debug!(sleep_ms = sleep_ms, "scheduler sleeping until next minute");
448                tokio::time::sleep(std::time::Duration::from_millis(sleep_ms)).await;
449
450                // Snapshot the current time for matching.
451                let fire_time = Local::now();
452                let hour = fire_time.hour();
453                let minute = fire_time.minute();
454
455                debug!(hour = hour, minute = minute, "scheduler tick");
456
457                let prompts = scheduler.prompts.read().await;
458                for prompt in prompts.iter() {
459                    if !prompt.enabled {
460                        continue;
461                    }
462                    if !prompt.schedule.matches(hour, minute) {
463                        continue;
464                    }
465                    let session = prompt
466                        .session_id
467                        .clone()
468                        .unwrap_or_else(|| Uuid::new_v4().to_string());
469                    let request_id = Uuid::new_v4().to_string();
470                    let req = PromptRequest {
471                        session: SessionId::new(session),
472                        request_id: request_id.clone(),
473                        input: prompt.prompt_template.clone(),
474                        meta: prompt.metadata.clone(),
475                        deadline: None,
476                    };
477                    match scheduler.tx.try_send(req) {
478                        Ok(_) => {
479                            info!(
480                                scheduled_prompt = %prompt.name,
481                                request_id = %request_id,
482                                "submitted scheduled prompt"
483                            );
484                        }
485                        Err(mpsc::error::TrySendError::Full(_)) => {
486                            warn!(
487                                scheduled_prompt = %prompt.name,
488                                "pipeline input channel full; skipping scheduled prompt"
489                            );
490                        }
491                        Err(mpsc::error::TrySendError::Closed(_)) => {
492                            error!("pipeline input channel closed; stopping scheduler");
493                            return;
494                        }
495                    }
496                }
497            }
498        })
499    }
500}
501
502// ============================================================================
503// SchedulerState — web API integration
504// ============================================================================
505
506/// Thread-safe wrapper around a [`Scheduler`] for use as Axum `State`.
507///
508/// Exposes the scheduler over the web API when the `web-api` feature is enabled.
509/// Routes are registered via [`scheduler_routes`].
510///
511/// ## Endpoints (registered by `scheduler_routes`)
512///
513/// | Method | Path | Description |
514/// |--------|------|-------------|
515/// | `POST` | `/api/v1/schedule` | Register a new scheduled prompt |
516/// | `GET` | `/api/v1/schedule` | List all scheduled prompts |
517/// | `DELETE` | `/api/v1/schedule/:id` | Remove a scheduled prompt |
518/// | `PATCH` | `/api/v1/schedule/:id/enable` | Enable a prompt |
519/// | `PATCH` | `/api/v1/schedule/:id/disable` | Disable a prompt |
520#[derive(Clone)]
521pub struct SchedulerState {
522    pub scheduler: Arc<Scheduler>,
523}
524
525impl SchedulerState {
526    /// Wrap an existing [`Scheduler`] in a `SchedulerState`.
527    ///
528    /// # Panics
529    ///
530    /// This function does not panic.
531    pub fn new(scheduler: Arc<Scheduler>) -> Self {
532        Self { scheduler }
533    }
534}
535
536// ============================================================================
537// Web API handlers (feature-gated)
538// ============================================================================
539
540/// Axum request body for `POST /api/v1/schedule`.
541#[cfg(feature = "web-api")]
542#[derive(Debug, Deserialize)]
543pub struct CreateScheduleRequest {
544    /// Human-readable name.
545    pub name: String,
546    /// Cron expression, e.g. `"*/5 *"`.
547    pub schedule: String,
548    /// The prompt template to submit on each firing.
549    pub prompt_template: String,
550    /// Optional session ID.
551    #[serde(default)]
552    pub session_id: Option<String>,
553    /// Optional metadata.
554    #[serde(default)]
555    pub metadata: HashMap<String, String>,
556    /// Whether to start enabled.  Default: `true`.
557    #[serde(default = "default_enabled")]
558    pub enabled: bool,
559}
560
561#[cfg(feature = "web-api")]
562fn default_enabled() -> bool {
563    true
564}
565
566/// Axum response body for schedule endpoints.
567#[cfg(feature = "web-api")]
568#[derive(Debug, Serialize)]
569pub struct ScheduleResponse {
570    pub id: String,
571    pub name: String,
572    pub schedule: String,
573    pub enabled: bool,
574    pub prompt_preview: String,
575}
576
577#[cfg(feature = "web-api")]
578impl From<&ScheduledPrompt> for ScheduleResponse {
579    fn from(p: &ScheduledPrompt) -> Self {
580        Self {
581            id: p.id.clone(),
582            name: p.name.clone(),
583            schedule: p.schedule.to_string(),
584            enabled: p.enabled,
585            prompt_preview: p.prompt_template.chars().take(80).collect(),
586        }
587    }
588}
589
590/// Build an Axum `Router` with all scheduler endpoints.
591///
592/// Merge this router into your application router:
593///
594/// ```rust,no_run
595/// # #[cfg(feature = "web-api")]
596/// # {
597/// use axum::Router;
598/// use tokio_prompt_orchestrator::scheduler::{Scheduler, SchedulerState, scheduler_routes};
599/// use std::sync::Arc;
600/// use tokio::sync::mpsc;
601///
602/// let (tx, _rx) = mpsc::channel(256);
603/// let scheduler = Arc::new(Scheduler::new(tx));
604/// let state = SchedulerState::new(scheduler);
605///
606/// let app: Router = Router::new()
607///     .merge(scheduler_routes(state));
608/// # }
609/// ```
610#[cfg(feature = "web-api")]
611pub fn scheduler_routes(state: SchedulerState) -> axum::Router {
612    use axum::routing::{delete, get, patch, post};
613
614    axum::Router::new()
615        .route("/api/v1/schedule", post(create_schedule_handler))
616        .route("/api/v1/schedule", get(list_schedules_handler))
617        .route("/api/v1/schedule/:id", delete(delete_schedule_handler))
618        .route("/api/v1/schedule/:id/enable", patch(enable_schedule_handler))
619        .route(
620            "/api/v1/schedule/:id/disable",
621            patch(disable_schedule_handler),
622        )
623        .with_state(state)
624}
625
626/// `POST /api/v1/schedule` — register a new scheduled prompt.
627#[cfg(feature = "web-api")]
628async fn create_schedule_handler(
629    axum::extract::State(state): axum::extract::State<SchedulerState>,
630    axum::Json(body): axum::Json<CreateScheduleRequest>,
631) -> axum::response::Response {
632    use axum::{
633        http::StatusCode,
634        response::IntoResponse,
635        Json,
636    };
637
638    let mut prompt = match ScheduledPrompt::new(&body.name, &body.schedule, &body.prompt_template) {
639        Ok(p) => p,
640        Err(e) => {
641            return (
642                StatusCode::UNPROCESSABLE_ENTITY,
643                Json(serde_json::json!({"error": e.to_string()})),
644            )
645                .into_response();
646        }
647    };
648
649    prompt.session_id = body.session_id;
650    prompt.metadata = body.metadata;
651    if !body.enabled {
652        prompt.enabled = false;
653    }
654
655    let resp = ScheduleResponse::from(&prompt);
656
657    match state.scheduler.add(prompt).await {
658        Ok(_) => (StatusCode::CREATED, Json(resp)).into_response(),
659        Err(e) => (
660            StatusCode::INTERNAL_SERVER_ERROR,
661            Json(serde_json::json!({"error": e.to_string()})),
662        )
663            .into_response(),
664    }
665}
666
667/// `GET /api/v1/schedule` — list all scheduled prompts.
668#[cfg(feature = "web-api")]
669async fn list_schedules_handler(
670    axum::extract::State(state): axum::extract::State<SchedulerState>,
671) -> axum::Json<Vec<ScheduleResponse>> {
672    let prompts = state.scheduler.list().await;
673    axum::Json(prompts.iter().map(ScheduleResponse::from).collect())
674}
675
676/// `DELETE /api/v1/schedule/:id` — remove a scheduled prompt.
677#[cfg(feature = "web-api")]
678async fn delete_schedule_handler(
679    axum::extract::State(state): axum::extract::State<SchedulerState>,
680    axum::extract::Path(id): axum::extract::Path<String>,
681) -> axum::response::Response {
682    use axum::{http::StatusCode, response::IntoResponse, Json};
683
684    if state.scheduler.remove(&id).await {
685        (StatusCode::OK, Json(serde_json::json!({"deleted": id}))).into_response()
686    } else {
687        (
688            StatusCode::NOT_FOUND,
689            Json(serde_json::json!({"error": "scheduled prompt not found", "id": id})),
690        )
691            .into_response()
692    }
693}
694
695/// `PATCH /api/v1/schedule/:id/enable` — enable a scheduled prompt.
696#[cfg(feature = "web-api")]
697async fn enable_schedule_handler(
698    axum::extract::State(state): axum::extract::State<SchedulerState>,
699    axum::extract::Path(id): axum::extract::Path<String>,
700) -> axum::response::Response {
701    toggle_enabled(state, id, true).await
702}
703
704/// `PATCH /api/v1/schedule/:id/disable` — disable a scheduled prompt.
705#[cfg(feature = "web-api")]
706async fn disable_schedule_handler(
707    axum::extract::State(state): axum::extract::State<SchedulerState>,
708    axum::extract::Path(id): axum::extract::Path<String>,
709) -> axum::response::Response {
710    toggle_enabled(state, id, false).await
711}
712
713#[cfg(feature = "web-api")]
714async fn toggle_enabled(
715    state: SchedulerState,
716    id: String,
717    enabled: bool,
718) -> axum::response::Response {
719    use axum::{http::StatusCode, response::IntoResponse, Json};
720
721    if state.scheduler.set_enabled(&id, enabled).await {
722        (
723            StatusCode::OK,
724            Json(serde_json::json!({"id": id, "enabled": enabled})),
725        )
726            .into_response()
727    } else {
728        (
729            StatusCode::NOT_FOUND,
730            Json(serde_json::json!({"error": "scheduled prompt not found", "id": id})),
731        )
732            .into_response()
733    }
734}
735
736// ============================================================================
737// Tests
738// ============================================================================
739
740#[cfg(test)]
741mod tests {
742    use super::*;
743
744    #[test]
745    fn cron_field_parse_any() {
746        let f = CronField::parse("*", 59).expect("should parse");
747        assert_eq!(f, CronField::Any);
748        assert!(f.matches(0));
749        assert!(f.matches(30));
750        assert!(f.matches(59));
751    }
752
753    #[test]
754    fn cron_field_parse_every() {
755        let f = CronField::parse("*/5", 59).expect("should parse");
756        assert_eq!(f, CronField::Every(5));
757        assert!(f.matches(0));
758        assert!(f.matches(5));
759        assert!(f.matches(15));
760        assert!(!f.matches(7));
761    }
762
763    #[test]
764    fn cron_field_parse_exact() {
765        let f = CronField::parse("30", 59).expect("should parse");
766        assert_eq!(f, CronField::Exact(30));
767        assert!(f.matches(30));
768        assert!(!f.matches(31));
769    }
770
771    #[test]
772    fn cron_field_invalid_step_zero() {
773        assert!(CronField::parse("*/0", 59).is_err());
774    }
775
776    #[test]
777    fn cron_field_out_of_range() {
778        assert!(CronField::parse("60", 59).is_err());
779    }
780
781    #[test]
782    fn cron_expression_parse_valid() {
783        let expr = CronExpression::parse("*/5 *").expect("valid");
784        assert!(expr.matches(0, 0));
785        assert!(expr.matches(0, 5));
786        assert!(!expr.matches(0, 3));
787    }
788
789    #[test]
790    fn cron_expression_parse_exact_hour_minute() {
791        let expr = CronExpression::parse("30 9").expect("valid");
792        assert!(expr.matches(9, 30));
793        assert!(!expr.matches(9, 31));
794        assert!(!expr.matches(10, 30));
795    }
796
797    #[test]
798    fn cron_expression_every_minute() {
799        let expr = CronExpression::parse("* *").expect("valid");
800        assert!(expr.matches(0, 0));
801        assert!(expr.matches(23, 59));
802    }
803
804    #[test]
805    fn cron_expression_wrong_field_count() {
806        assert!(CronExpression::parse("*").is_err());
807        assert!(CronExpression::parse("* * *").is_err());
808    }
809
810    #[test]
811    fn scheduled_prompt_new() {
812        let sp = ScheduledPrompt::new("test", "*/10 *", "hello").expect("valid");
813        assert!(sp.enabled);
814        assert_eq!(sp.name, "test");
815        assert_eq!(sp.prompt_template, "hello");
816        assert!(!sp.id.is_empty());
817    }
818
819    #[test]
820    fn scheduled_prompt_builders() {
821        let sp = ScheduledPrompt::new("x", "0 9", "ping")
822            .expect("valid")
823            .with_session("my-session")
824            .with_metadata("env", "prod")
825            .disabled();
826        assert_eq!(sp.session_id.as_deref(), Some("my-session"));
827        assert_eq!(sp.metadata.get("env").map(String::as_str), Some("prod"));
828        assert!(!sp.enabled);
829    }
830
831    #[tokio::test]
832    async fn scheduler_add_list_remove() {
833        let (tx, _rx) = mpsc::channel(16);
834        let scheduler = Arc::new(Scheduler::new(tx));
835
836        let sp = ScheduledPrompt::new("test", "* *", "hello").expect("valid");
837        let id = scheduler.add(sp).await.expect("added");
838
839        let list = scheduler.list().await;
840        assert_eq!(list.len(), 1);
841        assert_eq!(list[0].id, id);
842
843        let removed = scheduler.remove(&id).await;
844        assert!(removed);
845
846        assert!(scheduler.list().await.is_empty());
847    }
848
849    #[tokio::test]
850    async fn scheduler_set_enabled() {
851        let (tx, _rx) = mpsc::channel(16);
852        let scheduler = Arc::new(Scheduler::new(tx));
853
854        let sp = ScheduledPrompt::new("test", "* *", "hello").expect("valid");
855        let id = scheduler.add(sp).await.expect("added");
856
857        assert!(scheduler.set_enabled(&id, false).await);
858        let list = scheduler.list().await;
859        assert!(!list[0].enabled);
860
861        assert!(scheduler.set_enabled(&id, true).await);
862        let list = scheduler.list().await;
863        assert!(list[0].enabled);
864    }
865
866    #[test]
867    fn cron_expression_serialises_to_string() {
868        let expr = CronExpression::parse("*/15 8").expect("valid");
869        let json = serde_json::to_string(&expr).expect("serialise");
870        assert_eq!(json, "\"*/15 8\"");
871    }
872
873    #[test]
874    fn cron_expression_deserialises_from_string() {
875        let expr: CronExpression = serde_json::from_str("\"0 9\"").expect("deserialise");
876        assert!(expr.matches(9, 0));
877    }
878}
tokio_prompt_orchestrator/scheduler.rs

tokio_prompt_orchestrator/
scheduler.rs
