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 (including
7//! optional file-upload offloading via [`FileUploadContext`]).
8
9mod parts;
10mod persistence;
11mod queries;
12
13use sqlx::PgPool;
14use std::sync::Arc;
15use systemprompt_database::DbPool;
16use systemprompt_identifiers::{ContextId, TaskId};
17use systemprompt_traits::RepositoryError;
18
19use crate::models::a2a::Message;
20
21pub use parts::{FileUploadContext, PersistPartSqlxParams, get_message_parts};
22pub use persistence::{
23    PersistMessageSqlxParams, PersistMessageWithTxParams, persist_message_sqlx,
24    persist_message_with_tx,
25};
26pub use queries::{
27    get_messages_by_context, get_messages_by_task, get_next_sequence_number,
28    get_next_sequence_number_in_tx, get_next_sequence_number_sqlx, message_exists,
29};
30
31#[derive(Debug, Clone)]
32pub struct MessageRepository {
33    pool: Arc<PgPool>,
34}
35
36impl MessageRepository {
37    pub fn new(db: &DbPool) -> Result<Self, RepositoryError> {
38        let pool = db.pool_arc().map_err(|e| {
39            RepositoryError::InvalidData(format!("PostgreSQL pool not available: {e}"))
40        })?;
41        Ok(Self { pool })
42    }
43
44    pub async fn get_messages_by_task(
45        &self,
46        task_id: &TaskId,
47    ) -> Result<Vec<Message>, RepositoryError> {
48        get_messages_by_task(&self.pool, task_id).await
49    }
50
51    pub async fn get_messages_by_context(
52        &self,
53        context_id: &ContextId,
54    ) -> Result<Vec<Message>, RepositoryError> {
55        get_messages_by_context(&self.pool, context_id).await
56    }
57
58    pub async fn get_next_sequence_number(&self, task_id: &TaskId) -> Result<i32, RepositoryError> {
59        get_next_sequence_number(&self.pool, task_id).await
60    }
61
62    pub async fn persist_message_sqlx(
63        &self,
64        params: PersistMessageSqlxParams<'_>,
65    ) -> Result<(), RepositoryError> {
66        persist_message_sqlx(params).await
67    }
68
69    pub async fn persist_message_with_tx(
70        &self,
71        params: PersistMessageWithTxParams<'_>,
72    ) -> Result<(), RepositoryError> {
73        persist_message_with_tx(params).await
74    }
75}