store_sqlite/
store.rs

1//! `SqliteMailboxStore`: the SQLite-backed `MailboxStore` implementation.
2
3use std::sync::Mutex;
4
5use chrono::Utc;
6use rusqlite::{params, Connection};
7use uuid::Uuid;
8
9use a2a::message::{Message, MessageKind, MsgState, Part};
10use a2a::task::{Task, TaskState};
11use substrate_core::mailbox_port::{MailboxStore, MailboxTaskState};
12
13use crate::error::StoreError;
14use crate::schema;
15
16/// A `MailboxStore` backed by a SQLite database.
17///
18/// Thread safety: a `Mutex<Connection>` serialises all writes, which is sufficient
19/// for the atomic-claim guarantee (the `UPDATE WHERE state='unread'` rowcount test).
20pub struct SqliteMailboxStore {
21    conn: Mutex<Connection>,
22}
23
24impl SqliteMailboxStore {
25    /// Open (or create) a store at the given file path.
26    pub fn open(path: &str) -> Result<Self, StoreError> {
27        let conn = Connection::open(path)?;
28        schema::init(&conn)?;
29        Ok(Self {
30            conn: Mutex::new(conn),
31        })
32    }
33
34    /// Open a transient in-memory store (useful in tests).
35    pub fn open_in_memory() -> Result<Self, StoreError> {
36        let conn = Connection::open_in_memory()?;
37        schema::init(&conn)?;
38        Ok(Self {
39            conn: Mutex::new(conn),
40        })
41    }
42}
43
44// ── helpers ──────────────────────────────────────────────────────────────────
45
46fn task_state_to_str(s: MailboxTaskState) -> &'static str {
47    match s {
48        MailboxTaskState::Submitted => "submitted",
49        MailboxTaskState::Working => "working",
50        MailboxTaskState::InputRequired => "input_required",
51        MailboxTaskState::Completed => "completed",
52        MailboxTaskState::Failed => "failed",
53        MailboxTaskState::Cancelled => "cancelled",
54    }
55}
56
57fn str_to_task_state(s: &str) -> TaskState {
58    match s {
59        "working" => TaskState::Working,
60        "input_required" => TaskState::InputRequired,
61        "completed" => TaskState::Completed,
62        "failed" => TaskState::Failed,
63        "cancelled" => TaskState::Cancelled,
64        _ => TaskState::Submitted,
65    }
66}
67
68fn msg_kind_to_str(k: &MessageKind) -> &'static str {
69    match k {
70        MessageKind::Task => "task",
71        MessageKind::Reply => "reply",
72        MessageKind::Question => "question",
73        MessageKind::Status => "status",
74        MessageKind::Artifact => "artifact",
75    }
76}
77
78fn str_to_msg_kind(s: &str) -> MessageKind {
79    match s {
80        "reply" => MessageKind::Reply,
81        "question" => MessageKind::Question,
82        "status" => MessageKind::Status,
83        "artifact" => MessageKind::Artifact,
84        _ => MessageKind::Task,
85    }
86}
87
88fn a2a_task_state_to_str(s: TaskState) -> &'static str {
89    match s {
90        TaskState::Submitted => "submitted",
91        TaskState::Working => "working",
92        TaskState::InputRequired => "input_required",
93        TaskState::Completed => "completed",
94        TaskState::Failed => "failed",
95        TaskState::Cancelled => "cancelled",
96    }
97}
98
99// ── MailboxStore impl ─────────────────────────────────────────────────────────
100
101impl MailboxStore for SqliteMailboxStore {
102    type Msg = Message;
103    type Task = Task;
104    type Error = StoreError;
105
106    fn post(&self, msg: &Message) -> Result<(), StoreError> {
107        let conn = self.conn.lock().unwrap();
108        let parts_json = serde_json::to_string(&msg.parts)?;
109        conn.execute(
110            "INSERT INTO mailbox \
111             (id, team_id, task_id, from_agent, to_agent, kind, parts, in_reply_to, state, created_at, consumed_at) \
112             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
113            params![
114                msg.id.to_string(),
115                msg.team_id,
116                msg.task_id.map(|u| u.to_string()),
117                msg.from,
118                msg.to,
119                msg_kind_to_str(&msg.kind),
120                parts_json,
121                msg.in_reply_to.map(|u| u.to_string()),
122                "unread",
123                msg.created_at.to_rfc3339(),
124                msg.consumed_at.map(|dt| dt.to_rfc3339()),
125            ],
126        )?;
127        Ok(())
128    }
129
130    fn inbox(&self, team_id: &str, to: &str) -> Result<Vec<Message>, StoreError> {
131        let conn = self.conn.lock().unwrap();
132        let mut stmt = conn.prepare(
133            "SELECT id, team_id, task_id, from_agent, to_agent, kind, parts, \
134                    in_reply_to, state, created_at, consumed_at \
135             FROM mailbox \
136             WHERE team_id=?1 AND to_agent=?2 AND state='unread' \
137             ORDER BY created_at ASC",
138        )?;
139        let rows: Vec<Message> = stmt
140            .query_map(params![team_id, to], |row| {
141                Ok((
142                    row.get::<_, String>(0)?,
143                    row.get::<_, String>(1)?,
144                    row.get::<_, Option<String>>(2)?,
145                    row.get::<_, String>(3)?,
146                    row.get::<_, String>(4)?,
147                    row.get::<_, String>(5)?,
148                    row.get::<_, String>(6)?,
149                    row.get::<_, Option<String>>(7)?,
150                    row.get::<_, String>(8)?,
151                    row.get::<_, String>(9)?,
152                    row.get::<_, Option<String>>(10)?,
153                ))
154            })?
155            .filter_map(|r| r.ok())
156            .map(
157                |(
158                    id_str,
159                    team_id,
160                    task_id_str,
161                    from,
162                    to,
163                    kind_str,
164                    parts_json,
165                    in_reply_str,
166                    state_str,
167                    created_str,
168                    consumed_str,
169                )| {
170                    let parts: Vec<Part> = serde_json::from_str(&parts_json).unwrap_or_default();
171                    let state = match state_str.as_str() {
172                        "delivered" => MsgState::Delivered,
173                        "consumed" => MsgState::Consumed,
174                        _ => MsgState::Unread,
175                    };
176                    Message {
177                        id: Uuid::parse_str(&id_str).unwrap_or_else(|_| Uuid::new_v4()),
178                        team_id,
179                        task_id: task_id_str.and_then(|s| Uuid::parse_str(&s).ok()),
180                        from,
181                        to,
182                        kind: str_to_msg_kind(&kind_str),
183                        parts,
184                        in_reply_to: in_reply_str.and_then(|s| Uuid::parse_str(&s).ok()),
185                        state,
186                        created_at: created_str.parse().unwrap_or_else(|_| Utc::now()),
187                        consumed_at: consumed_str.and_then(|s| s.parse().ok()),
188                    }
189                },
190            )
191            .collect();
192        Ok(rows)
193    }
194
195    fn claim(&self, message_id: Uuid) -> Result<bool, StoreError> {
196        let conn = self.conn.lock().unwrap();
197        let n = conn.execute(
198            "UPDATE mailbox SET state='delivered' WHERE id=?1 AND state='unread'",
199            params![message_id.to_string()],
200        )?;
201        Ok(n == 1)
202    }
203
204    fn consume(&self, message_id: Uuid) -> Result<(), StoreError> {
205        let conn = self.conn.lock().unwrap();
206        conn.execute(
207            "UPDATE mailbox SET state='consumed', consumed_at=?2 WHERE id=?1",
208            params![message_id.to_string(), Utc::now().to_rfc3339()],
209        )?;
210        Ok(())
211    }
212
213    fn unclaim(&self, message_id: Uuid) -> Result<(), StoreError> {
214        let conn = self.conn.lock().unwrap();
215        conn.execute(
216            "UPDATE mailbox SET state='unread', consumed_at=NULL WHERE id=?1 AND state='delivered'",
217            params![message_id.to_string()],
218        )?;
219        Ok(())
220    }
221
222    fn task_create(&self, task: &Task) -> Result<(), StoreError> {
223        let conn = self.conn.lock().unwrap();
224        conn.execute(
225            "INSERT INTO tasklist \
226             (id, team_id, title, state, owner, parent_task_id, requirement_id, epic_id, created_at, updated_at) \
227             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
228            params![
229                task.id.to_string(),
230                task.team_id,
231                task.title,
232                a2a_task_state_to_str(task.state),
233                task.owner,
234                task.parent_task_id.map(|u| u.to_string()),
235                task.requirement_id,
236                task.epic_id,
237                task.created_at.to_rfc3339(),
238                task.updated_at.to_rfc3339(),
239            ],
240        )?;
241        Ok(())
242    }
243
244    fn task_update(
245        &self,
246        id: Uuid,
247        state: MailboxTaskState,
248        note: Option<&str>,
249    ) -> Result<(), StoreError> {
250        let conn = self.conn.lock().unwrap();
251        conn.execute(
252            "UPDATE tasklist SET state=?2, updated_at=?3, note=?4 WHERE id=?1",
253            params![
254                id.to_string(),
255                task_state_to_str(state),
256                Utc::now().to_rfc3339(),
257                note,
258            ],
259        )?;
260        Ok(())
261    }
262
263    fn task_list(&self, team_id: &str) -> Result<Vec<Task>, StoreError> {
264        let conn = self.conn.lock().unwrap();
265        let mut stmt = conn.prepare(
266            "SELECT id, team_id, title, state, owner, parent_task_id, requirement_id, epic_id, created_at, updated_at \
267             FROM tasklist WHERE team_id=?1 ORDER BY created_at ASC",
268        )?;
269        let tasks: Vec<Task> = stmt
270            .query_map(params![team_id], |row| {
271                Ok((
272                    row.get::<_, String>(0)?,
273                    row.get::<_, String>(1)?,
274                    row.get::<_, String>(2)?,
275                    row.get::<_, String>(3)?,
276                    row.get::<_, String>(4)?,
277                    row.get::<_, Option<String>>(5)?,
278                    row.get::<_, Option<String>>(6)?,
279                    row.get::<_, Option<String>>(7)?,
280                    row.get::<_, String>(8)?,
281                    row.get::<_, String>(9)?,
282                ))
283            })?
284            .filter_map(|r| r.ok())
285            .map(
286                |(
287                    id_str,
288                    team_id,
289                    title,
290                    state_str,
291                    owner,
292                    parent_str,
293                    req_id,
294                    epic_id,
295                    created_str,
296                    updated_str,
297                )| Task {
298                    id: Uuid::parse_str(&id_str).unwrap_or_else(|_| Uuid::new_v4()),
299                    team_id,
300                    title,
301                    state: str_to_task_state(&state_str),
302                    owner,
303                    parent_task_id: parent_str.and_then(|s| Uuid::parse_str(&s).ok()),
304                    requirement_id: req_id,
305                    epic_id,
306                    created_at: created_str.parse().unwrap_or_else(|_| Utc::now()),
307                    updated_at: updated_str.parse().unwrap_or_else(|_| Utc::now()),
308                },
309            )
310            .collect();
311        Ok(tasks)
312    }
313}
store_sqlite/store.rs

store_sqlite/
store.rs
