Skip to main content

adk_ui/
notifications.rs

1use crate::surface_runtime::SurfaceRef;
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4use tokio::sync::broadcast;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
7#[serde(rename_all = "snake_case")]
8pub enum UiNotificationSeverity {
9    Info,
10    Success,
11    Warning,
12    Error,
13}
14
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
16pub struct UiNotification {
17    pub severity: UiNotificationSeverity,
18    pub title: String,
19    pub body: String,
20    #[serde(skip_serializing_if = "Option::is_none")]
21    pub surface_ref: Option<SurfaceRef>,
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub action_id: Option<String>,
24}
25
26#[derive(Clone)]
27pub struct UiNotificationChannel {
28    sender: broadcast::Sender<UiNotification>,
29}
30
31impl UiNotificationChannel {
32    pub fn new(capacity: usize) -> Self {
33        let (sender, _) = broadcast::channel(capacity.max(1));
34        Self { sender }
35    }
36
37    pub fn publish(&self, notification: UiNotification) -> usize {
38        self.sender.send(notification).unwrap_or(0)
39    }
40
41    pub fn subscribe(&self) -> broadcast::Receiver<UiNotification> {
42        self.sender.subscribe()
43    }
44    pub fn receiver_count(&self) -> usize {
45        self.sender.receiver_count()
46    }
47}
48
49impl Default for UiNotificationChannel {
50    fn default() -> Self {
51        Self::new(128)
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58
59    #[tokio::test]
60    async fn publishes_to_agent_owned_subscribers() {
61        let channel = UiNotificationChannel::new(4);
62        let mut receiver = channel.subscribe();
63        let notification = UiNotification {
64            severity: UiNotificationSeverity::Warning,
65            title: "Inventory".to_string(),
66            body: "Stock is below threshold".to_string(),
67            surface_ref: Some(SurfaceRef::new("inventory", 2)),
68            action_id: Some("open_inventory".to_string()),
69        };
70        assert_eq!(channel.publish(notification.clone()), 1);
71        assert_eq!(receiver.recv().await.unwrap(), notification);
72    }
73}