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}