use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Mutex;
use tokio::sync::broadcast;
use tokio::task::JoinHandle;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RuntimeSignal {
pub kind: String, pub payload: serde_json::Value,
pub priority: u8,
}
#[async_trait]
pub trait SignalSource: Send + Sync {
async fn next_signal(&self) -> crate::Result<Option<RuntimeSignal>>;
}
#[derive(Debug, Clone)]
pub struct ScheduledPrompt {
pub goal: String,
pub run_at_ms: u64,
pub criteria: Vec<String>,
}
impl ScheduledPrompt {
pub fn new(goal: impl Into<String>, run_at_ms: u64) -> Self {
Self {
goal: goal.into(),
run_at_ms,
criteria: Vec::new(),
}
}
pub fn to_signal(&self) -> RuntimeSignal {
RuntimeSignal {
kind: "scheduled".into(),
payload: serde_json::json!({
"goal": self.goal,
"criteria": self.criteria,
"run_at_ms": self.run_at_ms,
}),
priority: 0,
}
}
}
pub struct SignalGateway {
tx: broadcast::Sender<RuntimeSignal>,
tasks: Mutex<HashMap<String, JoinHandle<()>>>,
}
impl SignalGateway {
pub fn new() -> Self {
let (tx, _) = broadcast::channel(1024);
Self {
tx,
tasks: Mutex::new(HashMap::new()),
}
}
pub fn subscribe(&self) -> GatewayReceiver {
GatewayReceiver {
rx: tokio::sync::Mutex::new(self.tx.subscribe()),
}
}
pub fn schedule(&self, prompt: ScheduledPrompt) {
let key = format!("cron:{}:{}", prompt.goal, prompt.run_at_ms);
let mut guard = self.tasks.lock().unwrap();
if guard.contains_key(&key) {
return;
}
let tx = self.tx.clone();
let signal = prompt.to_signal();
let now_ms = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
let delay_ms = prompt.run_at_ms.saturating_sub(now_ms);
let handle = tokio::spawn(async move {
if delay_ms > 0 {
tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
}
let _ = tx.send(signal);
});
guard.insert(key, handle);
}
pub fn cancel(&self, goal: &str, run_at_ms: u64) {
let key = format!("cron:{goal}:{run_at_ms}");
if let Some(h) = self.tasks.lock().unwrap().remove(&key) {
h.abort();
}
}
pub fn ingest(&self, signal: RuntimeSignal) {
let _ = self.tx.send(signal);
}
pub fn destroy(&self) {
for (_, h) in self.tasks.lock().unwrap().drain() {
h.abort();
}
}
}
impl Default for SignalGateway {
fn default() -> Self {
Self::new()
}
}
pub struct GatewayReceiver {
rx: tokio::sync::Mutex<broadcast::Receiver<RuntimeSignal>>,
}
#[async_trait]
impl SignalSource for GatewayReceiver {
async fn next_signal(&self) -> crate::Result<Option<RuntimeSignal>> {
match self.rx.lock().await.recv().await {
Ok(sig) => Ok(Some(sig)),
Err(broadcast::error::RecvError::Lagged(_)) => Ok(None),
Err(broadcast::error::RecvError::Closed) => Ok(None),
}
}
}