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#[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}