Skip to main content

atm_storage/
contract.rs

1use serde::{Deserialize, Serialize};
2use serde_json::{Map, Value};
3use std::fmt;
4use std::ops::Deref;
5use std::str::FromStr;
6
7use crate::error::AtmError;
8use crate::schema::{AtmMessageId, MessageEnvelope};
9use crate::types::{AgentName, IsoTimestamp, ModelName, PaneId, TaskId, TeamName};
10
11#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash)]
12#[serde(transparent)]
13pub struct MessageKey(String);
14
15impl MessageKey {
16    pub fn new(value: impl Into<String>) -> Result<Self, AtmError> {
17        let value = value.into();
18        if value.trim().is_empty() {
19            return Err(
20                AtmError::validation("message key must not be blank").with_recovery(
21                    "Populate a stable ATM message key before calling the storage contract.",
22                ),
23            );
24        }
25        Ok(Self(value))
26    }
27
28    pub fn into_inner(self) -> String {
29        self.0
30    }
31
32    pub fn as_str(&self) -> &str {
33        &self.0
34    }
35
36    pub fn as_atm_message_id(&self) -> Result<AtmMessageId, AtmError> {
37        let raw = self.as_str().strip_prefix("atm:").unwrap_or(self.as_str());
38        raw.parse::<AtmMessageId>()
39            .map_err(|error| AtmError::validation(format!("message key parse failed: {error}")))
40    }
41}
42
43impl From<AtmMessageId> for MessageKey {
44    fn from(value: AtmMessageId) -> Self {
45        Self(format!("atm:{value}"))
46    }
47}
48
49impl FromStr for MessageKey {
50    type Err = AtmError;
51
52    fn from_str(value: &str) -> Result<Self, Self::Err> {
53        Self::new(value)
54    }
55}
56
57impl fmt::Display for MessageKey {
58    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59        f.write_str(self.as_str())
60    }
61}
62
63impl AsRef<str> for MessageKey {
64    fn as_ref(&self) -> &str {
65        self.as_str()
66    }
67}
68
69#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash)]
70#[serde(transparent)]
71pub struct TaskState(String);
72
73impl TaskState {
74    pub fn new(value: impl Into<String>) -> Result<Self, AtmError> {
75        let value = value.into();
76        if value.trim().is_empty() {
77            return Err(
78                AtmError::validation("task state must not be blank").with_recovery(
79                    "Populate a non-empty task state before calling the storage contract.",
80                ),
81            );
82        }
83        Ok(Self(value))
84    }
85}
86
87impl AsRef<str> for TaskState {
88    fn as_ref(&self) -> &str {
89        &self.0
90    }
91}
92
93impl Deref for TaskState {
94    type Target = str;
95
96    fn deref(&self) -> &Self::Target {
97        &self.0
98    }
99}
100
101impl fmt::Display for TaskState {
102    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
103        f.write_str(self.as_ref())
104    }
105}
106
107impl FromStr for TaskState {
108    type Err = AtmError;
109
110    fn from_str(value: &str) -> Result<Self, Self::Err> {
111        Self::new(value)
112    }
113}
114
115impl PartialEq<&str> for TaskState {
116    fn eq(&self, other: &&str) -> bool {
117        self.as_ref() == *other
118    }
119}
120
121#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash)]
122#[serde(transparent)]
123pub struct AckTransition(String);
124
125impl AckTransition {
126    pub fn new(value: impl Into<String>) -> Result<Self, AtmError> {
127        let value = value.into();
128        if value.trim().is_empty() {
129            return Err(
130                AtmError::validation("ack transition must not be blank").with_recovery(
131                    "Populate a non-empty ack transition before calling the storage contract.",
132                ),
133            );
134        }
135        Ok(Self(value))
136    }
137}
138
139impl AsRef<str> for AckTransition {
140    fn as_ref(&self) -> &str {
141        &self.0
142    }
143}
144
145impl Deref for AckTransition {
146    type Target = str;
147
148    fn deref(&self) -> &Self::Target {
149        &self.0
150    }
151}
152
153impl fmt::Display for AckTransition {
154    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
155        f.write_str(self.as_ref())
156    }
157}
158
159impl FromStr for AckTransition {
160    type Err = AtmError;
161
162    fn from_str(value: &str) -> Result<Self, Self::Err> {
163        Self::new(value)
164    }
165}
166
167#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
168pub struct Message {
169    pub team: TeamName,
170    pub agent: AgentName,
171    pub message_key: MessageKey,
172    pub envelope: MessageEnvelope,
173}
174
175#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
176pub struct MailMessageState {
177    pub team: TeamName,
178    pub agent: AgentName,
179    pub actor: AgentName,
180    pub message_key: MessageKey,
181    pub read: bool,
182    #[serde(default, skip_serializing_if = "Option::is_none")]
183    pub pending_ack_at: Option<IsoTimestamp>,
184    #[serde(default, skip_serializing_if = "Option::is_none")]
185    pub acknowledged_at: Option<IsoTimestamp>,
186    #[serde(default, skip_serializing_if = "Option::is_none")]
187    pub expires_at: Option<IsoTimestamp>,
188    #[serde(default, skip_serializing_if = "Option::is_none")]
189    pub deleted_at: Option<IsoTimestamp>,
190    #[serde(default, skip_serializing_if = "Option::is_none")]
191    pub updated_at: Option<IsoTimestamp>,
192}
193
194/// Opaque hash or content-addressable identifier that marks the last
195/// successfully ingested message boundary for a replay source.
196#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
197#[serde(transparent)]
198pub struct MessageFingerprint(String);
199
200impl MessageFingerprint {
201    pub fn new(value: impl Into<String>) -> Self {
202        Self(value.into())
203    }
204
205    pub fn as_str(&self) -> &str {
206        &self.0
207    }
208}
209
210impl fmt::Display for MessageFingerprint {
211    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
212        f.write_str(self.as_str())
213    }
214}
215
216impl From<String> for MessageFingerprint {
217    fn from(value: String) -> Self {
218        Self::new(value)
219    }
220}
221
222impl AsRef<str> for MessageFingerprint {
223    fn as_ref(&self) -> &str {
224        self.as_str()
225    }
226}
227
228#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
229pub struct MessageQuery {
230    pub team: TeamName,
231    pub agent: AgentName,
232    #[serde(default, skip_serializing_if = "Option::is_none")]
233    pub sender: Option<AgentName>,
234    #[serde(default, skip_serializing_if = "Option::is_none")]
235    pub task_id: Option<TaskId>,
236    #[serde(default, skip_serializing_if = "Option::is_none")]
237    pub limit: Option<usize>,
238}
239
240#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
241#[serde(rename_all = "kebab-case")]
242pub enum RosterMemberKind {
243    Permanent,
244    Ephemeral,
245}
246
247#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
248#[serde(rename_all = "kebab-case")]
249pub enum RosterHarness {
250    ClaudeCode,
251    CodexCli,
252    GeminiCli,
253    Opencode,
254}
255
256#[derive(Debug, Clone, PartialEq, Eq)]
257pub enum AgentType {
258    GeneralPurpose,
259    Plan,
260    Lead,
261    Qa,
262    Worker,
263    Unknown(String),
264}
265
266impl Default for AgentType {
267    fn default() -> Self {
268        Self::Unknown(String::new())
269    }
270}
271
272impl From<String> for AgentType {
273    fn from(value: String) -> Self {
274        match value.as_str() {
275            "general-purpose" => Self::GeneralPurpose,
276            "plan" => Self::Plan,
277            "lead" => Self::Lead,
278            "qa" => Self::Qa,
279            "worker" => Self::Worker,
280            _ => Self::Unknown(value),
281        }
282    }
283}
284
285impl AgentType {
286    pub fn as_str(&self) -> &str {
287        match self {
288            Self::GeneralPurpose => "general-purpose",
289            Self::Plan => "plan",
290            Self::Lead => "lead",
291            Self::Qa => "qa",
292            Self::Worker => "worker",
293            Self::Unknown(value) => value,
294        }
295    }
296}
297
298impl From<AgentType> for String {
299    fn from(value: AgentType) -> Self {
300        value.as_str().to_string()
301    }
302}
303
304impl Serialize for AgentType {
305    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
306    where
307        S: serde::Serializer,
308    {
309        serializer.serialize_str(self.as_str())
310    }
311}
312
313impl fmt::Display for AgentType {
314    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
315        f.write_str(self.as_str())
316    }
317}
318
319impl<'de> Deserialize<'de> for AgentType {
320    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
321    where
322        D: serde::Deserializer<'de>,
323    {
324        Ok(Self::from(String::deserialize(deserializer)?))
325    }
326}
327
328#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
329pub struct RosterMember {
330    pub team_name: TeamName,
331    pub agent_name: AgentName,
332    pub member_kind: RosterMemberKind,
333    pub harness: RosterHarness,
334    #[serde(default)]
335    pub agent_type: AgentType,
336    #[serde(default)]
337    pub model: ModelName,
338    #[serde(default, skip_serializing_if = "Option::is_none")]
339    pub recipient_pane_id: Option<PaneId>,
340    #[serde(default)]
341    pub metadata_json: Map<String, Value>,
342}
343
344#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
345pub struct RosterSnapshot {
346    pub team_name: TeamName,
347    pub members: Vec<RosterMember>,
348    #[serde(default, skip_serializing_if = "Option::is_none")]
349    pub refreshed_at: Option<IsoTimestamp>,
350}
351
352#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
353pub struct MessageReceivedEvent {
354    pub team: TeamName,
355    pub agent: AgentName,
356    pub message_key: MessageKey,
357    pub timestamp: IsoTimestamp,
358}
359
360#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
361pub struct RosterChangedEvent {
362    pub team: TeamName,
363    pub member_count: usize,
364    pub timestamp: IsoTimestamp,
365}
366
367pub trait MessageStore {
368    fn save_message(&self, message: &Message) -> Result<(), AtmError>;
369    fn load_message(&self, key: &MessageKey) -> Result<Option<Message>, AtmError>;
370    fn list_messages(&self, query: &MessageQuery) -> Result<Vec<Message>, AtmError>;
371    fn delete_message(&self, key: &MessageKey) -> Result<(), AtmError>;
372}
373
374pub trait RosterStore {
375    fn load_roster(&self, team: &TeamName) -> Result<RosterSnapshot, AtmError>;
376    fn save_roster(&self, roster: &RosterSnapshot) -> Result<(), AtmError>;
377    fn list_teams(&self) -> Result<Vec<TeamName>, AtmError>;
378}
379
380pub trait StorageNotifier {
381    fn message_received(&self, event: &MessageReceivedEvent) -> Result<(), AtmError>;
382    fn roster_changed(&self, event: &RosterChangedEvent) -> Result<(), AtmError>;
383}
384
385#[cfg(test)]
386mod tests {
387    use super::{
388        Message, MessageKey, MessageQuery, MessageReceivedEvent, MessageStore, RosterChangedEvent,
389        RosterHarness, RosterMember, RosterMemberKind, RosterSnapshot, RosterStore,
390        StorageNotifier,
391    };
392    use crate::ROLE_WORKER;
393    use crate::error::AtmError;
394    use crate::schema::MessageEnvelope;
395    use crate::types::{AgentName, IsoTimestamp, ModelName, TeamName};
396    use chrono::Utc;
397    use serde_json::Map;
398
399    #[derive(Default)]
400    struct DummyStore;
401
402    impl MessageStore for DummyStore {
403        fn save_message(&self, _message: &Message) -> Result<(), AtmError> {
404            Ok(())
405        }
406
407        fn load_message(&self, _key: &MessageKey) -> Result<Option<Message>, AtmError> {
408            Ok(None)
409        }
410
411        fn list_messages(&self, _query: &MessageQuery) -> Result<Vec<Message>, AtmError> {
412            Ok(Vec::new())
413        }
414
415        fn delete_message(&self, _key: &MessageKey) -> Result<(), AtmError> {
416            Ok(())
417        }
418    }
419
420    impl RosterStore for DummyStore {
421        fn load_roster(&self, team: &TeamName) -> Result<RosterSnapshot, AtmError> {
422            Ok(RosterSnapshot {
423                team_name: team.clone(),
424                members: Vec::new(),
425                refreshed_at: None,
426            })
427        }
428
429        fn save_roster(&self, _roster: &RosterSnapshot) -> Result<(), AtmError> {
430            Ok(())
431        }
432
433        fn list_teams(&self) -> Result<Vec<TeamName>, AtmError> {
434            Ok(Vec::new())
435        }
436    }
437
438    impl StorageNotifier for DummyStore {
439        fn message_received(&self, _event: &MessageReceivedEvent) -> Result<(), AtmError> {
440            Ok(())
441        }
442
443        fn roster_changed(&self, _event: &RosterChangedEvent) -> Result<(), AtmError> {
444            Ok(())
445        }
446    }
447
448    #[test]
449    fn storage_traits_are_object_safe() {
450        let store = DummyStore;
451        let message_store: &dyn MessageStore = &store;
452        let roster_store: &dyn RosterStore = &store;
453        let notifier: &dyn StorageNotifier = &store;
454
455        let team: TeamName = "test-team".parse().expect("team");
456        let agent: AgentName = ROLE_WORKER.parse().expect("agent");
457        let key = MessageKey::new("atm:test-1").expect("key");
458
459        let message = Message {
460            team: team.clone(),
461            agent: agent.clone(),
462            message_key: key.clone(),
463            envelope: MessageEnvelope {
464                from: agent.clone(),
465                text: "hello".to_string(),
466                timestamp: IsoTimestamp::from_datetime(Utc::now()),
467                read: false,
468                source_team: Some(team.clone()),
469                summary: None,
470                message_id: None,
471                pending_ack_at: None,
472                acknowledged_at: None,
473                acknowledges_message_id: None,
474                parent_message_id: None,
475                thread_mode: None,
476                expires_at: None,
477                task_id: None,
478                extra: Map::new(),
479            },
480        };
481        let roster = RosterSnapshot {
482            team_name: team.clone(),
483            members: vec![RosterMember {
484                team_name: team.clone(),
485                agent_name: agent.clone(),
486                member_kind: RosterMemberKind::Permanent,
487                harness: RosterHarness::ClaudeCode,
488                agent_type: super::AgentType::Worker,
489                model: ModelName::default(),
490                recipient_pane_id: None,
491                metadata_json: Map::new(),
492            }],
493            refreshed_at: None,
494        };
495
496        message_store.save_message(&message).expect("save message");
497        assert!(message_store.load_message(&key).expect("load").is_none());
498        assert!(
499            message_store
500                .list_messages(&MessageQuery {
501                    team: team.clone(),
502                    agent: agent.clone(),
503                    sender: None,
504                    task_id: None,
505                    limit: Some(5),
506                })
507                .expect("list")
508                .is_empty()
509        );
510        message_store.delete_message(&key).expect("delete");
511
512        roster_store.save_roster(&roster).expect("save roster");
513        assert_eq!(
514            roster_store
515                .load_roster(&team)
516                .expect("load roster")
517                .team_name,
518            team
519        );
520        assert!(roster_store.list_teams().expect("list teams").is_empty());
521
522        notifier
523            .message_received(&MessageReceivedEvent {
524                team: team.clone(),
525                agent: agent.clone(),
526                message_key: key,
527                timestamp: IsoTimestamp::from_datetime(Utc::now()),
528            })
529            .expect("message notification");
530        notifier
531            .roster_changed(&RosterChangedEvent {
532                team,
533                member_count: 1,
534                timestamp: IsoTimestamp::from_datetime(Utc::now()),
535            })
536            .expect("roster notification");
537    }
538}