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//!
7//! Copyright (c) systemprompt.io — Business Source License 1.1.
8//! See <https://systemprompt.io> for licensing details.
9
10use std::sync::Arc;
11
12use sqlx::PgPool;
13use systemprompt_database::DbPool;
14use systemprompt_identifiers::{AgentId, ContextId};
15use systemprompt_traits::RepositoryError;
16
17#[derive(Debug, Clone)]
18pub struct ContextNotificationRepository {
19    write_pool: Arc<PgPool>,
20}
21
22impl ContextNotificationRepository {
23    pub fn new(db: &DbPool) -> Result<Self, RepositoryError> {
24        let write_pool = db.write_pool_arc().map_err(|e| {
25            RepositoryError::InvalidData(format!("PostgreSQL write pool not available: {e}"))
26        })?;
27        Ok(Self { write_pool })
28    }
29
30    pub async fn insert(
31        &self,
32        context_id: &ContextId,
33        agent_id: &AgentId,
34        notification_type: &str,
35        notification_data: &serde_json::Value,
36    ) -> Result<i32, RepositoryError> {
37        let row = sqlx::query!(
38            r#"INSERT INTO context_notifications (context_id, agent_id, notification_type, notification_data)
39            VALUES ($1, $2, $3, $4)
40            RETURNING id"#,
41            context_id.as_str(),
42            agent_id.as_str(),
43            notification_type,
44            notification_data,
45        )
46        .fetch_one(self.write_pool.as_ref())
47        .await
48        .map_err(RepositoryError::database)?;
49        Ok(row.id)
50    }
51
52    pub async fn mark_broadcasted(&self, notification_id: i32) -> Result<(), RepositoryError> {
53        sqlx::query!(
54            "UPDATE context_notifications SET broadcasted = true WHERE id = $1",
55            notification_id,
56        )
57        .execute(self.write_pool.as_ref())
58        .await
59        .map_err(RepositoryError::database)?;
60        Ok(())
61    }
62}
systemprompt_agent/repository/context/notifications.rs

systemprompt_agent/repository/context/
notifications.rs
