Skip to main content

systemprompt_agent/repository/context/message/
mod.rs

1//! Persistence for A2A messages and their constituent parts.
2//!
3//! [`MessageRepository`] reads message history by task or context and writes
4//! messages within a transaction. The submodules split the work: `queries`
5//! handles reads and sequence-number allocation, `persistence` writes the
6//! message row, and `parts` handles the typed text/file/data parts.
7//!
8//! Copyright (c) systemprompt.io — Business Source License 1.1.
9//! See <https://systemprompt.io> for licensing details.
10
11mod parts;
12mod persistence;
13mod queries;
14
15use sqlx::PgPool;
16use std::sync::Arc;
17use systemprompt_database::DbPool;
18use systemprompt_identifiers::{ContextId, TaskId};
19use systemprompt_traits::RepositoryError;
20
21use crate::models::a2a::Message;
22
23pub use parts::{PersistPartSqlxParams, get_message_parts};
24pub use persistence::{
25    PersistMessageSqlxParams, PersistMessageWithTxParams, persist_message_sqlx,
26    persist_message_with_tx,
27};
28pub use queries::{
29    get_messages_by_context, get_messages_by_task, get_next_sequence_number,
30    get_next_sequence_number_in_tx, get_next_sequence_number_sqlx, message_exists,
31};
32
33#[derive(Debug, Clone)]
34pub struct MessageRepository {
35    pool: Arc<PgPool>,
36}
37
38impl MessageRepository {
39    pub fn new(db: &DbPool) -> Result<Self, RepositoryError> {
40        let pool = db.pool_arc().map_err(|e| {
41            RepositoryError::InvalidData(format!("PostgreSQL pool not available: {e}"))
42        })?;
43        Ok(Self { pool })
44    }
45
46    pub async fn get_messages_by_task(
47        &self,
48        task_id: &TaskId,
49    ) -> Result<Vec<Message>, RepositoryError> {
50        get_messages_by_task(&self.pool, task_id).await
51    }
52
53    pub async fn get_messages_by_context(
54        &self,
55        context_id: &ContextId,
56    ) -> Result<Vec<Message>, RepositoryError> {
57        get_messages_by_context(&self.pool, context_id).await
58    }
59
60    pub async fn get_next_sequence_number(&self, task_id: &TaskId) -> Result<i32, RepositoryError> {
61        get_next_sequence_number(&self.pool, task_id).await
62    }
63
64    pub async fn persist_message_sqlx(
65        &self,
66        params: PersistMessageSqlxParams<'_>,
67    ) -> Result<(), RepositoryError> {
68        persist_message_sqlx(params).await
69    }
70
71    pub async fn persist_message_with_tx(
72        &self,
73        params: PersistMessageWithTxParams<'_>,
74    ) -> Result<(), RepositoryError> {
75        persist_message_with_tx(params).await
76    }
77}