Skip to main content

deepstrike_sdk/
signals.rs

1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::sync::Mutex;
5use tokio::sync::broadcast;
6use tokio::task::JoinHandle;
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct RuntimeSignal {
10    pub kind: String, // "interrupt" | "scheduled" | "external"
11    pub payload: serde_json::Value,
12    pub priority: u8,
13}
14
15/// Feed signals from any external source (cron, webhook, queue).
16#[async_trait]
17pub trait SignalSource: Send + Sync {
18    async fn next_signal(&self) -> crate::Result<Option<RuntimeSignal>>;
19}
20
21#[derive(Debug, Clone)]
22pub struct ScheduledPrompt {
23    pub goal: String,
24    pub run_at_ms: u64,
25    pub criteria: Vec<String>,
26}
27
28impl ScheduledPrompt {
29    pub fn new(goal: impl Into<String>, run_at_ms: u64) -> Self {
30        Self {
31            goal: goal.into(),
32            run_at_ms,
33            criteria: Vec::new(),
34        }
35    }
36
37    pub fn to_signal(&self) -> RuntimeSignal {
38        RuntimeSignal {
39            kind: "scheduled".into(),
40            payload: serde_json::json!({
41                "goal": self.goal,
42                "criteria": self.criteria,
43                "run_at_ms": self.run_at_ms,
44            }),
45            priority: 0,
46        }
47    }
48}
49
50/// Entry point for all external signals into the agent.
51///
52/// - Cron scheduling: fires [`ScheduledPrompt`]s at the right wall-clock time
53/// - Webhook ingestion: push any [`RuntimeSignal`] directly via [`ingest`](SignalGateway::ingest)
54/// - Subscribe pattern: callers receive signals via [`subscribe`](SignalGateway::subscribe)
55pub struct SignalGateway {
56    tx: broadcast::Sender<RuntimeSignal>,
57    tasks: Mutex<HashMap<String, JoinHandle<()>>>,
58}
59
60impl SignalGateway {
61    pub fn new() -> Self {
62        let (tx, _) = broadcast::channel(1024);
63        Self {
64            tx,
65            tasks: Mutex::new(HashMap::new()),
66        }
67    }
68
69    /// Subscribe to all signals emitted by this gateway.
70    /// Returns a [`GatewayReceiver`] that implements [`SignalSource`].
71    pub fn subscribe(&self) -> GatewayReceiver {
72        GatewayReceiver {
73            rx: tokio::sync::Mutex::new(self.tx.subscribe()),
74        }
75    }
76
77    /// Schedule a [`ScheduledPrompt`] to fire at its `run_at_ms`. Idempotent by goal+time.
78    pub fn schedule(&self, prompt: ScheduledPrompt) {
79        let key = format!("cron:{}:{}", prompt.goal, prompt.run_at_ms);
80        let mut guard = self.tasks.lock().unwrap();
81        if guard.contains_key(&key) {
82            return;
83        }
84
85        let tx = self.tx.clone();
86        let signal = prompt.to_signal();
87        let now_ms = std::time::SystemTime::now()
88            .duration_since(std::time::UNIX_EPOCH)
89            .unwrap_or_default()
90            .as_millis() as u64;
91        let delay_ms = prompt.run_at_ms.saturating_sub(now_ms);
92
93        let handle = tokio::spawn(async move {
94            if delay_ms > 0 {
95                tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
96            }
97            let _ = tx.send(signal);
98        });
99        guard.insert(key, handle);
100    }
101
102    /// Cancel a previously scheduled prompt.
103    pub fn cancel(&self, goal: &str, run_at_ms: u64) {
104        let key = format!("cron:{goal}:{run_at_ms}");
105        if let Some(h) = self.tasks.lock().unwrap().remove(&key) {
106            h.abort();
107        }
108    }
109
110    /// Ingest a raw external signal (e.g. from a webhook handler).
111    pub fn ingest(&self, signal: RuntimeSignal) {
112        let _ = self.tx.send(signal);
113    }
114
115    /// Abort all pending scheduled tasks.
116    pub fn destroy(&self) {
117        for (_, h) in self.tasks.lock().unwrap().drain() {
118            h.abort();
119        }
120    }
121}
122
123impl Default for SignalGateway {
124    fn default() -> Self {
125        Self::new()
126    }
127}
128
129/// A broadcast receiver from a [`SignalGateway`] that implements [`SignalSource`].
130pub struct GatewayReceiver {
131    rx: tokio::sync::Mutex<broadcast::Receiver<RuntimeSignal>>,
132}
133
134#[async_trait]
135impl SignalSource for GatewayReceiver {
136    async fn next_signal(&self) -> crate::Result<Option<RuntimeSignal>> {
137        match self.rx.lock().await.recv().await {
138            Ok(sig) => Ok(Some(sig)),
139            Err(broadcast::error::RecvError::Lagged(_)) => Ok(None),
140            Err(broadcast::error::RecvError::Closed) => Ok(None),
141        }
142    }
143}