systemprompt_agent/repository/context/
notifications.rs

1//! Persistence for inbound A2A notifications that have not yet been
2//! broadcast to subscribed AG-UI clients.
3//!
4//! One row is inserted per notification received. The `broadcasted` flag
5//! flips to `true` once the corresponding fan-out has completed.
6
7use std::sync::Arc;
8
9use sqlx::PgPool;
10use systemprompt_database::DbPool;
11use systemprompt_identifiers::{AgentId, ContextId};
12use systemprompt_traits::RepositoryError;
13
14#[derive(Debug, Clone)]
15pub struct ContextNotificationRepository {
16    write_pool: Arc<PgPool>,
17}
18
19impl ContextNotificationRepository {
20    pub fn new(db: &DbPool) -> Result<Self, RepositoryError> {
21        let write_pool = db.write_pool_arc().map_err(|e| {
22            RepositoryError::InvalidData(format!("PostgreSQL write pool not available: {e}"))
23        })?;
24        Ok(Self { write_pool })
25    }
26
27    pub async fn insert(
28        &self,
29        context_id: &ContextId,
30        agent_id: &AgentId,
31        notification_type: &str,
32        notification_data: &serde_json::Value,
33    ) -> Result<i32, RepositoryError> {
34        let row = sqlx::query!(
35            r#"INSERT INTO context_notifications (context_id, agent_id, notification_type, notification_data)
36            VALUES ($1, $2, $3, $4)
37            RETURNING id"#,
38            context_id.as_str(),
39            agent_id.as_str(),
40            notification_type,
41            notification_data,
42        )
43        .fetch_one(self.write_pool.as_ref())
44        .await
45        .map_err(RepositoryError::database)?;
46        Ok(row.id)
47    }
48
49    pub async fn mark_broadcasted(&self, notification_id: i32) -> Result<(), RepositoryError> {
50        sqlx::query!(
51            "UPDATE context_notifications SET broadcasted = true WHERE id = $1",
52            notification_id,
53        )
54        .execute(self.write_pool.as_ref())
55        .await
56        .map_err(RepositoryError::database)?;
57        Ok(())
58    }
59}
systemprompt_agent/repository/context/notifications.rs

systemprompt_agent/repository/context/
notifications.rs
