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 System {
118 content: Value,
121 #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")]
125 sections: serde_json::Map<String, Value>,
126 #[serde(rename = "toolsAdded", default, skip_serializing_if = "Vec::is_empty")]
129 tools_added: Vec<Value>,
130 #[serde(
132 rename = "toolsRemoved",
133 default,
134 skip_serializing_if = "Vec::is_empty"
135 )]
136 tools_removed: Vec<Value>,
137 #[serde(default)]
138 timestamp: f64,
139 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
140 extra: serde_json::Map<String, Value>,
141 },
142 User {
143 content: Value,
145 #[serde(default)]
146 timestamp: f64,
147 #[serde(default, skip_serializing_if = "Vec::is_empty")]
148 attachments: Vec<Value>,
149 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
150 extra: serde_json::Map<String, Value>,
151 },
152 Assistant {
153 content: Vec<ContentBlock>,
154 #[serde(default)]
155 api: String,
156 #[serde(default)]
157 provider: String,
158 #[serde(default)]
159 model: String,
160 #[serde(default, skip_serializing_if = "Option::is_none")]
161 usage: Option<Usage>,
162 #[serde(rename = "stopReason", default)]
165 stop_reason: String,
166 #[serde(default)]
167 timestamp: f64,
168 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
169 extra: serde_json::Map<String, Value>,
170 },
171 ToolResult {
172 #[serde(rename = "toolCallId", default)]
173 tool_call_id: String,
174 #[serde(rename = "toolName", default)]
175 tool_name: String,
176 content: Value,
177 #[serde(default, skip_serializing_if = "Option::is_none")]
180 usage: Option<Usage>,
181 #[serde(rename = "isError", default)]
182 is_error: bool,
183 #[serde(default)]
184 timestamp: f64,
185 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
186 extra: serde_json::Map<String, Value>,
187 },
188 BashExecution {
191 #[serde(default)]
192 command: String,
193 #[serde(default)]
194 output: String,
195 #[serde(rename = "exitCode", default, skip_serializing_if = "Option::is_none")]
196 exit_code: Option<i64>,
197 #[serde(default)]
198 cancelled: bool,
199 #[serde(default)]
200 truncated: bool,
201 #[serde(
202 rename = "fullOutputPath",
203 default,
204 skip_serializing_if = "Option::is_none"
205 )]
206 full_output_path: Option<String>,
207 #[serde(default)]
208 timestamp: f64,
209 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
210 extra: serde_json::Map<String, Value>,
211 },
212}
213
214#[derive(Debug, Clone, PartialEq)]
220pub enum PiEvent {
221 AgentStart,
223 AgentEnd { messages: Vec<PiMessage> },
225 TurnStart,
227 TurnEnd {
229 message: Box<PiMessage>,
230 tool_results: Vec<PiMessage>,
231 },
232 MessageStart { message: Box<PiMessage> },
234 MessageUpdate {
238 message: Option<Box<PiMessage>>,
239 usage: Option<Usage>,
240 assistant_message_event: Value,
241 },
242 MessageEnd { message: Box<PiMessage> },
244 ToolExecutionStart {
246 tool_call_id: String,
247 tool_name: String,
248 args: Value,
249 },
250 ToolExecutionUpdate {
252 tool_call_id: String,
253 tool_name: String,
254 args: Value,
255 partial_result: Value,
256 },
257 ToolExecutionEnd {
259 tool_call_id: String,
260 tool_name: String,
261 result: Value,
262 is_error: bool,
263 },
264 Unknown { event_type: String, payload: Value },
269}
270
271impl PiEvent {
272 pub fn event_type(&self) -> &str {
274 match self {
275 PiEvent::AgentStart => "agent_start",
276 PiEvent::AgentEnd { .. } => "agent_end",
277 PiEvent::TurnStart => "turn_start",
278 PiEvent::TurnEnd { .. } => "turn_end",
279 PiEvent::MessageStart { .. } => "message_start",
280 PiEvent::MessageUpdate { .. } => "message_update",
281 PiEvent::MessageEnd { .. } => "message_end",
282 PiEvent::ToolExecutionStart { .. } => "tool_execution_start",
283 PiEvent::ToolExecutionUpdate { .. } => "tool_execution_update",
284 PiEvent::ToolExecutionEnd { .. } => "tool_execution_end",
285 PiEvent::Unknown { event_type, .. } => event_type,
286 }
287 }
288
289 pub fn from_value(v: Value) -> Result<Self, serde_json::Error> {
291 let t = v
292 .get("type")
293 .and_then(Value::as_str)
294 .unwrap_or_default()
295 .to_string();
296 Ok(match t.as_str() {
297 "agent_start" => PiEvent::AgentStart,
298 "agent_end" => PiEvent::AgentEnd {
299 messages: serde_json::from_value(
300 v.get("messages").cloned().unwrap_or(Value::Array(vec![])),
301 )?,
302 },
303 "turn_start" => PiEvent::TurnStart,
304 "turn_end" => PiEvent::TurnEnd {
305 message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
306 tool_results: serde_json::from_value(
307 v.get("toolResults")
308 .cloned()
309 .unwrap_or(Value::Array(vec![])),
310 )?,
311 },
312 "message_start" => PiEvent::MessageStart {
313 message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
314 },
315 "message_update" => PiEvent::MessageUpdate {
316 message: match v.get("message") {
317 Some(m) => Some(serde_json::from_value(m.clone())?),
318 None => None,
319 },
320 usage: match v.get("usage") {
321 Some(u) => Some(serde_json::from_value(u.clone())?),
322 None => None,
323 },
324 assistant_message_event: v
325 .get("assistantMessageEvent")
326 .cloned()
327 .unwrap_or(Value::Null),
328 },
329 "message_end" => PiEvent::MessageEnd {
330 message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
331 },
332 "tool_execution_start" => PiEvent::ToolExecutionStart {
333 tool_call_id: str_field(&v, "toolCallId"),
334 tool_name: str_field(&v, "toolName"),
335 args: v.get("args").cloned().unwrap_or(Value::Null),
336 },
337 "tool_execution_update" => PiEvent::ToolExecutionUpdate {
338 tool_call_id: str_field(&v, "toolCallId"),
339 tool_name: str_field(&v, "toolName"),
340 args: v.get("args").cloned().unwrap_or(Value::Null),
341 partial_result: v.get("partialResult").cloned().unwrap_or(Value::Null),
342 },
343 "tool_execution_end" => PiEvent::ToolExecutionEnd {
344 tool_call_id: str_field(&v, "toolCallId"),
345 tool_name: str_field(&v, "toolName"),
346 result: v.get("result").cloned().unwrap_or(Value::Null),
347 is_error: v
348 .get("isError")
349 .and_then(Value::as_bool)
350 .unwrap_or_default(),
351 },
352 _ => PiEvent::Unknown {
353 event_type: t,
354 payload: v,
355 },
356 })
357 }
358}
359
360fn str_field(v: &Value, k: &str) -> String {
361 v.get(k)
362 .and_then(Value::as_str)
363 .unwrap_or_default()
364 .to_string()
365}
366
367#[cfg(test)]
368mod tests {
369 use super::*;
370
371 #[test]
372 fn assistant_message_from_docs_parses() {
373 let m: PiMessage = serde_json::from_value(serde_json::json!({
374 "role": "assistant",
375 "content": [
376 {"type": "text", "text": "Hello! How can I help?"},
377 {"type": "thinking", "thinking": "User is greeting me..."},
378 {"type": "toolCall", "id": "call_123", "name": "bash", "arguments": {"command": "ls"}}
379 ],
380 "api": "anthropic-messages",
381 "provider": "anthropic",
382 "model": "claude-sonnet-4-20250514",
383 "usage": {
384 "input": 100, "output": 50, "cacheRead": 0, "cacheWrite": 0,
385 "cost": {"input": 0.0003, "output": 0.00075, "cacheRead": 0, "cacheWrite": 0, "total": 0.00105}
386 },
387 "stopReason": "stop",
388 "timestamp": 1733234567890i64
389 }))
390 .unwrap();
391 let PiMessage::Assistant {
392 content,
393 stop_reason,
394 ..
395 } = m
396 else {
397 panic!("wrong role")
398 };
399 assert_eq!(content.len(), 3);
400 assert_eq!(stop_reason, "stop");
401 assert!(matches!(content[2], ContentBlock::ToolCall { .. }));
402 }
403
404 #[test]
405 fn system_message_from_docs_parses() {
406 let m: PiMessage = serde_json::from_value(serde_json::json!({
407 "role": "system",
408 "content": "",
409 "sections": {
410 "preamble": "You are an expert coding assistant...",
411 "tools": "<tools>\n- read: ...\n</tools>",
412 "cwd": "/project"
413 },
414 "toolsAdded": [{"name": "read", "description": "...", "parameters": {}}],
415 "timestamp": 1733234400000i64
416 }))
417 .unwrap();
418 let PiMessage::System {
419 sections,
420 tools_added,
421 tools_removed,
422 ..
423 } = m
424 else {
425 panic!("wrong role")
426 };
427 assert_eq!(sections.len(), 3);
428 assert_eq!(tools_added[0]["name"], "read");
429 assert!(tools_removed.is_empty());
430
431 let patch: PiMessage = serde_json::from_value(serde_json::json!({
432 "role": "system",
433 "content": "",
434 "sections": {"skills": null},
435 "toolsRemoved": [{"name": "write"}],
436 "timestamp": 1733234640000i64
437 }))
438 .unwrap();
439 let PiMessage::System {
440 sections,
441 tools_removed,
442 ..
443 } = patch
444 else {
445 panic!("wrong role")
446 };
447 assert!(sections["skills"].is_null());
448 assert_eq!(tools_removed[0]["name"], "write");
449 }
450
451 #[test]
452 fn bash_execution_message_round_trips() {
453 let v = serde_json::json!({
454 "role": "bashExecution",
455 "command": "ls -la",
456 "output": "total 48",
457 "exitCode": 0,
458 "cancelled": false,
459 "truncated": false,
460 "timestamp": 1733234567890i64
461 });
462 let m: PiMessage = serde_json::from_value(v).unwrap();
463 assert!(matches!(m, PiMessage::BashExecution { .. }));
464 }
465
466 #[test]
467 fn unknown_event_type_is_preserved_not_error() {
468 let e = PiEvent::from_value(serde_json::json!({
469 "type": "queue_update", "steering": [], "followUp": []
470 }))
471 .unwrap();
472 let PiEvent::Unknown {
473 event_type,
474 payload,
475 } = e
476 else {
477 panic!("expected Unknown")
478 };
479 assert_eq!(event_type, "queue_update");
480 assert!(payload.get("steering").is_some());
481 }
482
483 #[test]
484 fn tool_execution_events_lift_fields() {
485 let e = PiEvent::from_value(serde_json::json!({
486 "type": "tool_execution_end",
487 "toolCallId": "call_9", "toolName": "bash",
488 "result": {"ok": true}, "isError": false
489 }))
490 .unwrap();
491 let PiEvent::ToolExecutionEnd {
492 tool_call_id,
493 tool_name,
494 is_error,
495 ..
496 } = e
497 else {
498 panic!("wrong variant")
499 };
500 assert_eq!(
501 (tool_call_id.as_str(), tool_name.as_str()),
502 ("call_9", "bash")
503 );
504 assert!(!is_error);
505 }
506}