wcode_harness/
message.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
4#[serde(tag = "type", rename_all = "snake_case")]
5pub enum ContentBlock {
6    Text {
7        text: String,
8    },
9    Thinking {
10        text: String,
11    },
12    ToolCall {
13        id: String,
14        name: String,
15        arguments: serde_json::Value,
16    },
17}
18
19#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize, PartialEq)]
20#[serde(rename_all = "snake_case")]
21pub struct Usage {
22    pub input_tokens: u64,
23    pub output_tokens: u64,
24    #[serde(default, skip_serializing_if = "Option::is_none")]
25    pub cache_read_tokens: Option<u64>,
26    #[serde(default, skip_serializing_if = "Option::is_none")]
27    pub cache_write_tokens: Option<u64>,
28}
29
30#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
31#[serde(rename_all = "snake_case")]
32pub enum StopReason {
33    Stop,
34    Length,
35    ToolUse,
36    Aborted,
37    Error,
38    /// Run stopped at the per-run turn cap (`LoopConfig::max_turns`) before the
39    /// model signalled completion — the task may be incomplete, but the history
40    /// is intact and the run can be resumed.
41    MaxTurns,
42}
43
44#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
45#[serde(tag = "role", rename_all = "snake_case")]
46pub enum AgentMessage {
47    User {
48        content: Vec<ContentBlock>,
49    },
50    Assistant {
51        content: Vec<ContentBlock>,
52        stop_reason: StopReason,
53        #[serde(default, skip_serializing_if = "Option::is_none")]
54        usage: Option<Usage>,
55        #[serde(default, skip_serializing_if = "Option::is_none")]
56        model: Option<String>,
57    },
58    ToolResult {
59        tool_call_id: String,
60        name: String,
61        output: String,
62        is_error: bool,
63    },
64}
65
66/// The marker [`AgentMessage::summary`] prefixes onto a compaction summary.
67/// A summary is an ordinary `User` message with no distinguishing field, so
68/// this string IS the discriminator: readers that must not mistake a summary
69/// for something the user typed (`crate::title::excerpt`,
70/// `main.rs::first_user_line_of`) match on it.
71///
72/// One literal, one source: [`AgentMessage::summary`] builds its text from it,
73/// and `crate::title` re-exports it for the CLI, which must not depend on
74/// `message` directly.
75pub const COMPACTION_SUMMARY_PREFIX: &str =
76    "Summary of the earlier conversation (auto-generated by compaction):";
77
78impl AgentMessage {
79    pub fn user_text(s: impl Into<String>) -> AgentMessage {
80        AgentMessage::User {
81            content: vec![ContentBlock::Text { text: s.into() }],
82        }
83    }
84
85    /// A synthetic user message carrying a compaction summary that stands in
86    /// for the summarized prefix of the conversation.
87    pub fn summary(text: impl Into<String>) -> AgentMessage {
88        AgentMessage::User {
89            content: vec![ContentBlock::Text {
90                text: format!("{COMPACTION_SUMMARY_PREFIX}\n\n{}", text.into()),
91            }],
92        }
93    }
94
95    pub fn as_text(&self) -> String {
96        let content = match self {
97            AgentMessage::User { content } | AgentMessage::Assistant { content, .. } => content,
98            AgentMessage::ToolResult { .. } => return String::new(),
99        };
100        content
101            .iter()
102            .filter_map(|b| match b {
103                ContentBlock::Text { text } => Some(text.as_str()),
104                _ => None,
105            })
106            .collect::<Vec<_>>()
107            .join("")
108    }
109
110    pub fn tool_calls(&self) -> Vec<&ContentBlock> {
111        match self {
112            AgentMessage::Assistant { content, .. } => content
113                .iter()
114                .filter(|b| matches!(b, ContentBlock::ToolCall { .. }))
115                .collect(),
116            _ => Vec::new(),
117        }
118    }
119}
120
121/// The text of the most recent `Assistant` message, if any — a compact "what it
122/// last said" for a status read (see `crate::protocol::Request::Status`). Joins
123/// the assistant's `Text` blocks with newlines; `None` when there is no assistant
124/// message or its text is blank.
125pub fn last_assistant_text(messages: &[AgentMessage]) -> Option<String> {
126    let message = messages
127        .iter()
128        .rev()
129        .find(|m| matches!(m, AgentMessage::Assistant { .. }))?;
130    let AgentMessage::Assistant { content, .. } = message else {
131        return None;
132    };
133    let text = content
134        .iter()
135        .filter_map(|b| match b {
136            ContentBlock::Text { text } => Some(text.as_str()),
137            _ => None,
138        })
139        .collect::<Vec<_>>()
140        .join("\n");
141    (!text.trim().is_empty()).then_some(text)
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use serde_json::json;
148
149    #[test]
150    fn content_block_roundtrip_and_tags() {
151        let blocks = vec![
152            ContentBlock::Text { text: "hi".into() },
153            ContentBlock::Thinking { text: "hmm".into() },
154            ContentBlock::ToolCall {
155                id: "t1".into(),
156                name: "run".into(),
157                arguments: json!({"x": 1}),
158            },
159        ];
160        for b in blocks {
161            let v = serde_json::to_value(&b).unwrap();
162            let back: ContentBlock = serde_json::from_value(v.clone()).unwrap();
163            assert_eq!(back, b);
164            let expected = match &b {
165                ContentBlock::Text { .. } => "text",
166                ContentBlock::Thinking { .. } => "thinking",
167                ContentBlock::ToolCall { .. } => "tool_call",
168            };
169            assert_eq!(v["type"], expected);
170        }
171    }
172
173    #[test]
174    fn usage_roundtrip_and_omits_none() {
175        let u = Usage::default();
176        let v = serde_json::to_value(u).unwrap();
177        assert_eq!(v, json!({"input_tokens": 0, "output_tokens": 0}));
178        let back: Usage = serde_json::from_value(v).unwrap();
179        assert_eq!(back, u);
180
181        let u2 = Usage {
182            input_tokens: 10,
183            output_tokens: 20,
184            cache_read_tokens: Some(5),
185            cache_write_tokens: Some(7),
186        };
187        let v2 = serde_json::to_value(u2).unwrap();
188        assert_eq!(
189            v2,
190            json!({
191                "input_tokens": 10,
192                "output_tokens": 20,
193                "cache_read_tokens": 5,
194                "cache_write_tokens": 7
195            })
196        );
197        assert_eq!(serde_json::from_value::<Usage>(v2).unwrap(), u2);
198    }
199
200    #[test]
201    fn stop_reason_roundtrip_snake_case() {
202        for (r, s) in [
203            (StopReason::Stop, "stop"),
204            (StopReason::Length, "length"),
205            (StopReason::ToolUse, "tool_use"),
206            (StopReason::Aborted, "aborted"),
207            (StopReason::Error, "error"),
208            (StopReason::MaxTurns, "max_turns"),
209        ] {
210            let v = serde_json::to_value(r).unwrap();
211            assert_eq!(v, json!(s));
212            assert_eq!(serde_json::from_value::<StopReason>(v).unwrap(), r);
213        }
214    }
215
216    #[test]
217    fn agent_message_roundtrip_and_role_tags() {
218        let msgs = vec![
219            AgentMessage::User {
220                content: vec![ContentBlock::Text { text: "q".into() }],
221            },
222            AgentMessage::Assistant {
223                content: vec![ContentBlock::Text { text: "a".into() }],
224                stop_reason: StopReason::Stop,
225                usage: Some(Usage::default()),
226                model: Some("m".into()),
227            },
228            AgentMessage::ToolResult {
229                tool_call_id: "t1".into(),
230                name: "run".into(),
231                output: "out".into(),
232                is_error: false,
233            },
234        ];
235        for m in msgs {
236            let v = serde_json::to_value(&m).unwrap();
237            let back: AgentMessage = serde_json::from_value(v.clone()).unwrap();
238            assert_eq!(back, m);
239            let expected = match &m {
240                AgentMessage::User { .. } => "user",
241                AgentMessage::Assistant { .. } => "assistant",
242                AgentMessage::ToolResult { .. } => "tool_result",
243            };
244            assert_eq!(v["role"], expected);
245        }
246    }
247
248    #[test]
249    fn assistant_optional_fields_default_and_skip() {
250        let v = json!({
251            "role": "assistant",
252            "content": [{"type": "text", "text": "a"}],
253            "stop_reason": "stop"
254        });
255        let m: AgentMessage = serde_json::from_value(v).unwrap();
256        match &m {
257            AgentMessage::Assistant {
258                content,
259                usage,
260                model,
261                ..
262            } => {
263                assert_eq!(content.len(), 1);
264                assert_eq!(*usage, None);
265                assert_eq!(*model, None);
266            }
267            _ => panic!("wrong variant"),
268        }
269        let out = serde_json::to_value(&m).unwrap();
270        assert!(out.get("usage").is_none());
271        assert!(out.get("model").is_none());
272    }
273
274    #[test]
275    fn user_text_and_as_text() {
276        let m = AgentMessage::user_text("hello");
277        assert_eq!(
278            m,
279            AgentMessage::User {
280                content: vec![ContentBlock::Text {
281                    text: "hello".into()
282                }]
283            }
284        );
285        assert_eq!(m.as_text(), "hello");
286
287        let m2 = AgentMessage::Assistant {
288            content: vec![
289                ContentBlock::Text { text: "a".into() },
290                ContentBlock::Thinking { text: "t".into() },
291                ContentBlock::Text { text: "b".into() },
292            ],
293            stop_reason: StopReason::Stop,
294            usage: None,
295            model: None,
296        };
297        assert_eq!(m2.as_text(), "ab");
298    }
299
300    #[test]
301    fn tool_calls_returns_assistant_tool_call_blocks() {
302        let m = AgentMessage::Assistant {
303            content: vec![
304                ContentBlock::Text { text: "a".into() },
305                ContentBlock::ToolCall {
306                    id: "t1".into(),
307                    name: "run".into(),
308                    arguments: json!({}),
309                },
310                ContentBlock::ToolCall {
311                    id: "t2".into(),
312                    name: "ls".into(),
313                    arguments: json!({}),
314                },
315            ],
316            stop_reason: StopReason::ToolUse,
317            usage: None,
318            model: None,
319        };
320        assert_eq!(m.tool_calls().len(), 2);
321        assert!(matches!(m.tool_calls()[0], ContentBlock::ToolCall { .. }));
322        assert!(
323            m.tool_calls().iter().all(
324                |b| matches!(b, ContentBlock::ToolCall { id, .. } if id == "t1" || id == "t2")
325            )
326        );
327
328        let u = AgentMessage::user_text("x");
329        assert!(u.tool_calls().is_empty());
330    }
331}
wcode_harness/message.rs

wcode_harness/
message.rs
