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
66impl AgentMessage {
67    pub fn user_text(s: impl Into<String>) -> AgentMessage {
68        AgentMessage::User {
69            content: vec![ContentBlock::Text { text: s.into() }],
70        }
71    }
72
73    /// A synthetic user message carrying a compaction summary that stands in
74    /// for the summarized prefix of the conversation.
75    pub fn summary(text: impl Into<String>) -> AgentMessage {
76        AgentMessage::User {
77            content: vec![ContentBlock::Text {
78                text: format!(
79                    "Summary of the earlier conversation (auto-generated by compaction):\n\n{}",
80                    text.into()
81                ),
82            }],
83        }
84    }
85
86    pub fn as_text(&self) -> String {
87        let content = match self {
88            AgentMessage::User { content } | AgentMessage::Assistant { content, .. } => content,
89            AgentMessage::ToolResult { .. } => return String::new(),
90        };
91        content
92            .iter()
93            .filter_map(|b| match b {
94                ContentBlock::Text { text } => Some(text.as_str()),
95                _ => None,
96            })
97            .collect::<Vec<_>>()
98            .join("")
99    }
100
101    pub fn tool_calls(&self) -> Vec<&ContentBlock> {
102        match self {
103            AgentMessage::Assistant { content, .. } => content
104                .iter()
105                .filter(|b| matches!(b, ContentBlock::ToolCall { .. }))
106                .collect(),
107            _ => Vec::new(),
108        }
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use serde_json::json;
116
117    #[test]
118    fn content_block_roundtrip_and_tags() {
119        let blocks = vec![
120            ContentBlock::Text { text: "hi".into() },
121            ContentBlock::Thinking { text: "hmm".into() },
122            ContentBlock::ToolCall {
123                id: "t1".into(),
124                name: "run".into(),
125                arguments: json!({"x": 1}),
126            },
127        ];
128        for b in blocks {
129            let v = serde_json::to_value(&b).unwrap();
130            let back: ContentBlock = serde_json::from_value(v.clone()).unwrap();
131            assert_eq!(back, b);
132            let expected = match &b {
133                ContentBlock::Text { .. } => "text",
134                ContentBlock::Thinking { .. } => "thinking",
135                ContentBlock::ToolCall { .. } => "tool_call",
136            };
137            assert_eq!(v["type"], expected);
138        }
139    }
140
141    #[test]
142    fn usage_roundtrip_and_omits_none() {
143        let u = Usage::default();
144        let v = serde_json::to_value(u).unwrap();
145        assert_eq!(v, json!({"input_tokens": 0, "output_tokens": 0}));
146        let back: Usage = serde_json::from_value(v).unwrap();
147        assert_eq!(back, u);
148
149        let u2 = Usage {
150            input_tokens: 10,
151            output_tokens: 20,
152            cache_read_tokens: Some(5),
153            cache_write_tokens: Some(7),
154        };
155        let v2 = serde_json::to_value(u2).unwrap();
156        assert_eq!(
157            v2,
158            json!({
159                "input_tokens": 10,
160                "output_tokens": 20,
161                "cache_read_tokens": 5,
162                "cache_write_tokens": 7
163            })
164        );
165        assert_eq!(serde_json::from_value::<Usage>(v2).unwrap(), u2);
166    }
167
168    #[test]
169    fn stop_reason_roundtrip_snake_case() {
170        for (r, s) in [
171            (StopReason::Stop, "stop"),
172            (StopReason::Length, "length"),
173            (StopReason::ToolUse, "tool_use"),
174            (StopReason::Aborted, "aborted"),
175            (StopReason::Error, "error"),
176            (StopReason::MaxTurns, "max_turns"),
177        ] {
178            let v = serde_json::to_value(r).unwrap();
179            assert_eq!(v, json!(s));
180            assert_eq!(serde_json::from_value::<StopReason>(v).unwrap(), r);
181        }
182    }
183
184    #[test]
185    fn agent_message_roundtrip_and_role_tags() {
186        let msgs = vec![
187            AgentMessage::User {
188                content: vec![ContentBlock::Text { text: "q".into() }],
189            },
190            AgentMessage::Assistant {
191                content: vec![ContentBlock::Text { text: "a".into() }],
192                stop_reason: StopReason::Stop,
193                usage: Some(Usage::default()),
194                model: Some("m".into()),
195            },
196            AgentMessage::ToolResult {
197                tool_call_id: "t1".into(),
198                name: "run".into(),
199                output: "out".into(),
200                is_error: false,
201            },
202        ];
203        for m in msgs {
204            let v = serde_json::to_value(&m).unwrap();
205            let back: AgentMessage = serde_json::from_value(v.clone()).unwrap();
206            assert_eq!(back, m);
207            let expected = match &m {
208                AgentMessage::User { .. } => "user",
209                AgentMessage::Assistant { .. } => "assistant",
210                AgentMessage::ToolResult { .. } => "tool_result",
211            };
212            assert_eq!(v["role"], expected);
213        }
214    }
215
216    #[test]
217    fn assistant_optional_fields_default_and_skip() {
218        let v = json!({
219            "role": "assistant",
220            "content": [{"type": "text", "text": "a"}],
221            "stop_reason": "stop"
222        });
223        let m: AgentMessage = serde_json::from_value(v).unwrap();
224        match &m {
225            AgentMessage::Assistant {
226                content,
227                usage,
228                model,
229                ..
230            } => {
231                assert_eq!(content.len(), 1);
232                assert_eq!(*usage, None);
233                assert_eq!(*model, None);
234            }
235            _ => panic!("wrong variant"),
236        }
237        let out = serde_json::to_value(&m).unwrap();
238        assert!(out.get("usage").is_none());
239        assert!(out.get("model").is_none());
240    }
241
242    #[test]
243    fn user_text_and_as_text() {
244        let m = AgentMessage::user_text("hello");
245        assert_eq!(
246            m,
247            AgentMessage::User {
248                content: vec![ContentBlock::Text {
249                    text: "hello".into()
250                }]
251            }
252        );
253        assert_eq!(m.as_text(), "hello");
254
255        let m2 = AgentMessage::Assistant {
256            content: vec![
257                ContentBlock::Text { text: "a".into() },
258                ContentBlock::Thinking { text: "t".into() },
259                ContentBlock::Text { text: "b".into() },
260            ],
261            stop_reason: StopReason::Stop,
262            usage: None,
263            model: None,
264        };
265        assert_eq!(m2.as_text(), "ab");
266    }
267
268    #[test]
269    fn tool_calls_returns_assistant_tool_call_blocks() {
270        let m = AgentMessage::Assistant {
271            content: vec![
272                ContentBlock::Text { text: "a".into() },
273                ContentBlock::ToolCall {
274                    id: "t1".into(),
275                    name: "run".into(),
276                    arguments: json!({}),
277                },
278                ContentBlock::ToolCall {
279                    id: "t2".into(),
280                    name: "ls".into(),
281                    arguments: json!({}),
282                },
283            ],
284            stop_reason: StopReason::ToolUse,
285            usage: None,
286            model: None,
287        };
288        assert_eq!(m.tool_calls().len(), 2);
289        assert!(matches!(m.tool_calls()[0], ContentBlock::ToolCall { .. }));
290        assert!(
291            m.tool_calls().iter().all(
292                |b| matches!(b, ContentBlock::ToolCall { id, .. } if id == "t1" || id == "t2")
293            )
294        );
295
296        let u = AgentMessage::user_text("x");
297        assert!(u.tool_calls().is_empty());
298    }
299}
wcode_harness/message.rs

wcode_harness/
message.rs
