1use serde::{Deserialize, Serialize};
12use serde_json::Value;
13
14#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
16pub struct Usage {
17 #[serde(default)]
18 pub input: f64,
19 #[serde(default)]
20 pub output: f64,
21 #[serde(rename = "cacheRead", default)]
22 pub cache_read: f64,
23 #[serde(rename = "cacheWrite", default)]
24 pub cache_write: f64,
25 #[serde(
26 rename = "totalTokens",
27 default,
28 skip_serializing_if = "Option::is_none"
29 )]
30 pub total_tokens: Option<f64>,
31 #[serde(default, skip_serializing_if = "Option::is_none")]
32 pub cost: Option<UsageCost>,
33 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
34 pub extra: serde_json::Map<String, Value>,
35}
36
37#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
39pub struct UsageCost {
40 #[serde(default)]
41 pub input: f64,
42 #[serde(default)]
43 pub output: f64,
44 #[serde(rename = "cacheRead", default)]
45 pub cache_read: f64,
46 #[serde(rename = "cacheWrite", default)]
47 pub cache_write: f64,
48 #[serde(default, skip_serializing_if = "Option::is_none")]
49 pub total: Option<f64>,
50}
51
52#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
57pub struct Model {
58 #[serde(default)]
59 pub id: String,
60 #[serde(default)]
61 pub name: String,
62 #[serde(default)]
63 pub api: String,
64 #[serde(default)]
65 pub provider: String,
66 #[serde(rename = "baseUrl", default)]
67 pub base_url: String,
68 #[serde(default)]
69 pub reasoning: bool,
70 #[serde(default)]
71 pub input: Vec<String>,
72 #[serde(rename = "contextWindow", default)]
73 pub context_window: i64,
74 #[serde(rename = "maxTokens", default)]
75 pub max_tokens: i64,
76 #[serde(default, skip_serializing_if = "Option::is_none")]
77 pub cost: Option<UsageCost>,
78 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
79 pub extra: serde_json::Map<String, Value>,
80}
81
82#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
84#[serde(tag = "type", rename_all = "camelCase")]
85pub enum ContentBlock {
86 Text {
87 text: String,
88 },
89 Thinking {
90 thinking: String,
91 },
92 ToolCall {
93 id: String,
94 name: String,
95 arguments: Value,
96 },
97 Image {
98 data: String,
99 #[serde(rename = "mimeType")]
100 mime_type: String,
101 },
102}
103
104#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
107#[serde(tag = "role", rename_all = "camelCase")]
108pub enum PiMessage {
109 User {
110 content: Value,
112 #[serde(default)]
113 timestamp: f64,
114 #[serde(default, skip_serializing_if = "Vec::is_empty")]
115 attachments: Vec<Value>,
116 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
117 extra: serde_json::Map<String, Value>,
118 },
119 Assistant {
120 content: Vec<ContentBlock>,
121 #[serde(default)]
122 api: String,
123 #[serde(default)]
124 provider: String,
125 #[serde(default)]
126 model: String,
127 #[serde(default, skip_serializing_if = "Option::is_none")]
128 usage: Option<Usage>,
129 #[serde(rename = "stopReason", default)]
132 stop_reason: String,
133 #[serde(default)]
134 timestamp: f64,
135 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
136 extra: serde_json::Map<String, Value>,
137 },
138 ToolResult {
139 #[serde(rename = "toolCallId", default)]
140 tool_call_id: String,
141 #[serde(rename = "toolName", default)]
142 tool_name: String,
143 content: Value,
144 #[serde(default, skip_serializing_if = "Option::is_none")]
147 usage: Option<Usage>,
148 #[serde(rename = "isError", default)]
149 is_error: bool,
150 #[serde(default)]
151 timestamp: f64,
152 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
153 extra: serde_json::Map<String, Value>,
154 },
155 BashExecution {
158 #[serde(default)]
159 command: String,
160 #[serde(default)]
161 output: String,
162 #[serde(rename = "exitCode", default, skip_serializing_if = "Option::is_none")]
163 exit_code: Option<i64>,
164 #[serde(default)]
165 cancelled: bool,
166 #[serde(default)]
167 truncated: bool,
168 #[serde(
169 rename = "fullOutputPath",
170 default,
171 skip_serializing_if = "Option::is_none"
172 )]
173 full_output_path: Option<String>,
174 #[serde(default)]
175 timestamp: f64,
176 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
177 extra: serde_json::Map<String, Value>,
178 },
179}
180
181#[derive(Debug, Clone, PartialEq)]
187pub enum PiEvent {
188 AgentStart,
190 AgentEnd { messages: Vec<PiMessage> },
192 TurnStart,
194 TurnEnd {
196 message: Box<PiMessage>,
197 tool_results: Vec<PiMessage>,
198 },
199 MessageStart { message: Box<PiMessage> },
201 MessageUpdate {
205 message: Option<Box<PiMessage>>,
206 usage: Option<Usage>,
207 assistant_message_event: Value,
208 },
209 MessageEnd { message: Box<PiMessage> },
211 ToolExecutionStart {
213 tool_call_id: String,
214 tool_name: String,
215 args: Value,
216 },
217 ToolExecutionUpdate {
219 tool_call_id: String,
220 tool_name: String,
221 args: Value,
222 partial_result: Value,
223 },
224 ToolExecutionEnd {
226 tool_call_id: String,
227 tool_name: String,
228 result: Value,
229 is_error: bool,
230 },
231 Unknown { event_type: String, payload: Value },
236}
237
238impl PiEvent {
239 pub fn event_type(&self) -> &str {
241 match self {
242 PiEvent::AgentStart => "agent_start",
243 PiEvent::AgentEnd { .. } => "agent_end",
244 PiEvent::TurnStart => "turn_start",
245 PiEvent::TurnEnd { .. } => "turn_end",
246 PiEvent::MessageStart { .. } => "message_start",
247 PiEvent::MessageUpdate { .. } => "message_update",
248 PiEvent::MessageEnd { .. } => "message_end",
249 PiEvent::ToolExecutionStart { .. } => "tool_execution_start",
250 PiEvent::ToolExecutionUpdate { .. } => "tool_execution_update",
251 PiEvent::ToolExecutionEnd { .. } => "tool_execution_end",
252 PiEvent::Unknown { event_type, .. } => event_type,
253 }
254 }
255
256 pub fn from_value(v: Value) -> Result<Self, serde_json::Error> {
258 let t = v
259 .get("type")
260 .and_then(Value::as_str)
261 .unwrap_or_default()
262 .to_string();
263 Ok(match t.as_str() {
264 "agent_start" => PiEvent::AgentStart,
265 "agent_end" => PiEvent::AgentEnd {
266 messages: serde_json::from_value(
267 v.get("messages").cloned().unwrap_or(Value::Array(vec![])),
268 )?,
269 },
270 "turn_start" => PiEvent::TurnStart,
271 "turn_end" => PiEvent::TurnEnd {
272 message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
273 tool_results: serde_json::from_value(
274 v.get("toolResults")
275 .cloned()
276 .unwrap_or(Value::Array(vec![])),
277 )?,
278 },
279 "message_start" => PiEvent::MessageStart {
280 message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
281 },
282 "message_update" => PiEvent::MessageUpdate {
283 message: match v.get("message") {
284 Some(m) => Some(serde_json::from_value(m.clone())?),
285 None => None,
286 },
287 usage: match v.get("usage") {
288 Some(u) => Some(serde_json::from_value(u.clone())?),
289 None => None,
290 },
291 assistant_message_event: v
292 .get("assistantMessageEvent")
293 .cloned()
294 .unwrap_or(Value::Null),
295 },
296 "message_end" => PiEvent::MessageEnd {
297 message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
298 },
299 "tool_execution_start" => PiEvent::ToolExecutionStart {
300 tool_call_id: str_field(&v, "toolCallId"),
301 tool_name: str_field(&v, "toolName"),
302 args: v.get("args").cloned().unwrap_or(Value::Null),
303 },
304 "tool_execution_update" => PiEvent::ToolExecutionUpdate {
305 tool_call_id: str_field(&v, "toolCallId"),
306 tool_name: str_field(&v, "toolName"),
307 args: v.get("args").cloned().unwrap_or(Value::Null),
308 partial_result: v.get("partialResult").cloned().unwrap_or(Value::Null),
309 },
310 "tool_execution_end" => PiEvent::ToolExecutionEnd {
311 tool_call_id: str_field(&v, "toolCallId"),
312 tool_name: str_field(&v, "toolName"),
313 result: v.get("result").cloned().unwrap_or(Value::Null),
314 is_error: v
315 .get("isError")
316 .and_then(Value::as_bool)
317 .unwrap_or_default(),
318 },
319 _ => PiEvent::Unknown {
320 event_type: t,
321 payload: v,
322 },
323 })
324 }
325}
326
327fn str_field(v: &Value, k: &str) -> String {
328 v.get(k)
329 .and_then(Value::as_str)
330 .unwrap_or_default()
331 .to_string()
332}
333
334#[cfg(test)]
335mod tests {
336 use super::*;
337
338 #[test]
339 fn assistant_message_from_docs_parses() {
340 let m: PiMessage = serde_json::from_value(serde_json::json!({
341 "role": "assistant",
342 "content": [
343 {"type": "text", "text": "Hello! How can I help?"},
344 {"type": "thinking", "thinking": "User is greeting me..."},
345 {"type": "toolCall", "id": "call_123", "name": "bash", "arguments": {"command": "ls"}}
346 ],
347 "api": "anthropic-messages",
348 "provider": "anthropic",
349 "model": "claude-sonnet-4-20250514",
350 "usage": {
351 "input": 100, "output": 50, "cacheRead": 0, "cacheWrite": 0,
352 "cost": {"input": 0.0003, "output": 0.00075, "cacheRead": 0, "cacheWrite": 0, "total": 0.00105}
353 },
354 "stopReason": "stop",
355 "timestamp": 1733234567890i64
356 }))
357 .unwrap();
358 let PiMessage::Assistant {
359 content,
360 stop_reason,
361 ..
362 } = m
363 else {
364 panic!("wrong role")
365 };
366 assert_eq!(content.len(), 3);
367 assert_eq!(stop_reason, "stop");
368 assert!(matches!(content[2], ContentBlock::ToolCall { .. }));
369 }
370
371 #[test]
372 fn bash_execution_message_round_trips() {
373 let v = serde_json::json!({
374 "role": "bashExecution",
375 "command": "ls -la",
376 "output": "total 48",
377 "exitCode": 0,
378 "cancelled": false,
379 "truncated": false,
380 "timestamp": 1733234567890i64
381 });
382 let m: PiMessage = serde_json::from_value(v).unwrap();
383 assert!(matches!(m, PiMessage::BashExecution { .. }));
384 }
385
386 #[test]
387 fn unknown_event_type_is_preserved_not_error() {
388 let e = PiEvent::from_value(serde_json::json!({
389 "type": "queue_update", "steering": [], "followUp": []
390 }))
391 .unwrap();
392 let PiEvent::Unknown {
393 event_type,
394 payload,
395 } = e
396 else {
397 panic!("expected Unknown")
398 };
399 assert_eq!(event_type, "queue_update");
400 assert!(payload.get("steering").is_some());
401 }
402
403 #[test]
404 fn tool_execution_events_lift_fields() {
405 let e = PiEvent::from_value(serde_json::json!({
406 "type": "tool_execution_end",
407 "toolCallId": "call_9", "toolName": "bash",
408 "result": {"ok": true}, "isError": false
409 }))
410 .unwrap();
411 let PiEvent::ToolExecutionEnd {
412 tool_call_id,
413 tool_name,
414 is_error,
415 ..
416 } = e
417 else {
418 panic!("wrong variant")
419 };
420 assert_eq!(
421 (tool_call_id.as_str(), tool_name.as_str()),
422 ("call_9", "bash")
423 );
424 assert!(!is_error);
425 }
426}