Skip to main content

rpi_agent/
message.rs

1//! Mirrors `packages/agent/src/types.ts` (open enum portion) — the
2//! `AgentMessage = Message | Custom` declaration-merging model.
3//!
4//! TS lets apps add custom message kinds via declaration merging on
5//! `CustomAgentMessages`. Rust has no declaration merging, so we model the same
6//! openness as a four-arm enum: the three base [`Message`] arms plus
7//! [`CustomMessage`]. Apps extend by populating `Custom`, not by adding arms.
8//!
9//! Serialization mirrors the wire shape: `#[serde(tag="kind")]` with
10//! `{kind:"user"|"assistant"|"toolResult"|"custom", ...}`. The three base arms
11//! flatten their `Message` (which itself is tagged on `role`); we serialize
12//! them via a custom impl so the outer `kind` discriminator stays consistent
13//! and the inner role-specific fields survive a round-trip. `Custom` carries a
14//! free-form `data: Value` so each role's structured payload rides along.
15
16use rpi_ai::types::{AssistantMessage, Content, Message, ToolResultMessage, UserMessage};
17use serde::{Deserialize, Serialize};
18
19/// `AgentMessage = Message | Custom`. The open-enum port of TS
20/// `AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages]`.
21///
22/// `PartialEq` is derived so `pi-harness` entry equality (and tests) can compare
23/// persisted messages by value. `rpi_ai`'s base message types already implement
24/// `PartialEq`; `CustomMessage` does too, so the four-arm enum derives cleanly.
25#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
26#[serde(tag = "kind", rename_all = "snake_case")]
27pub enum AgentMessage {
28    User(UserMessage),
29    Assistant(Box<AssistantMessage>),
30    ToolResult(Box<ToolResultMessage>),
31    Custom(CustomMessage),
32}
33
34impl AgentMessage {
35    /// The logical role tag — mirrors `message.role` on the TS side. For base
36    /// messages this is `"user"`/`"assistant"`/`"toolResult"`; for custom it is
37    /// the custom role string. Used by the loop's role checks (e.g. "can't
38    /// continue from an assistant message").
39    pub fn role(&self) -> AgentMessageRole {
40        match self {
41            AgentMessage::User(_) => AgentMessageRole::User,
42            AgentMessage::Assistant(_) => AgentMessageRole::Assistant,
43            AgentMessage::ToolResult(_) => AgentMessageRole::ToolResult,
44            AgentMessage::Custom(c) => AgentMessageRole::Custom(c.role.clone()),
45        }
46    }
47
48    pub fn is_assistant(&self) -> bool {
49        matches!(self, AgentMessage::Assistant(_))
50    }
51
52    /// If this is an assistant message, borrow it.
53    pub fn as_assistant(&self) -> Option<&AssistantMessage> {
54        match self {
55            AgentMessage::Assistant(a) => Some(a),
56            _ => None,
57        }
58    }
59
60    /// If this is an assistant message, own a clone of it.
61    pub fn into_assistant(self) -> Option<AssistantMessage> {
62        match self {
63            AgentMessage::Assistant(a) => Some(*a),
64            _ => None,
65        }
66    }
67}
68
69impl From<UserMessage> for AgentMessage {
70    fn from(m: UserMessage) -> Self {
71        AgentMessage::User(m)
72    }
73}
74impl From<AssistantMessage> for AgentMessage {
75    fn from(m: AssistantMessage) -> Self {
76        AgentMessage::Assistant(Box::new(m))
77    }
78}
79impl From<ToolResultMessage> for AgentMessage {
80    fn from(m: ToolResultMessage) -> Self {
81        AgentMessage::ToolResult(Box::new(m))
82    }
83}
84impl From<Message> for AgentMessage {
85    fn from(m: Message) -> Self {
86        match m {
87            Message::User(u) => AgentMessage::User(u),
88            Message::Assistant(a) => AgentMessage::Assistant(a),
89            Message::ToolResult(t) => AgentMessage::ToolResult(t),
90        }
91    }
92}
93
94/// Typed view of `AgentMessage::role`.
95#[derive(Debug, Clone, PartialEq, Eq)]
96pub enum AgentMessageRole {
97    User,
98    Assistant,
99    ToolResult,
100    Custom(String),
101}
102
103impl AgentMessageRole {
104    pub fn as_str(&self) -> &str {
105        match self {
106            AgentMessageRole::User => "user",
107            AgentMessageRole::Assistant => "assistant",
108            AgentMessageRole::ToolResult => "toolResult",
109            AgentMessageRole::Custom(s) => s.as_str(),
110        }
111    }
112}
113
114/// A custom app message. `role` is the app-defined tag (e.g. `"bashExecution"`,
115/// `"branchSummary"`); `content` is provider-content the renderer emits;
116/// `data` carries the role's structured payload; `timestamp` is ms-since-epoch.
117///
118/// `pi-harness::messages` provides typed constructors per known role;
119/// `pi-agent` stays role-agnostic and only owns this shell.
120///
121/// `PartialEq` is derived so `pi-harness` entry equality (and tests) can compare
122/// persisted custom messages by value. `serde_json::Value` + `Content` both
123/// implement `PartialEq`, so the derive is sound.
124#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
125#[serde(rename_all = "camelCase")]
126pub struct CustomMessage {
127    pub role: String,
128    pub content: Vec<Content>,
129    pub data: serde_json::Value,
130    pub timestamp: i64,
131}
132
133impl CustomMessage {
134    pub fn new(
135        role: impl Into<String>,
136        content: Vec<Content>,
137        data: serde_json::Value,
138        timestamp: i64,
139    ) -> Self {
140        Self {
141            role: role.into(),
142            content,
143            data,
144            timestamp,
145        }
146    }
147}