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(
170 rename = "thinkingLevel",
171 default,
172 skip_serializing_if = "Option::is_none"
173 )]
174 thinking_level: Option<String>,
175 #[serde(default)]
176 timestamp: f64,
177 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
178 extra: serde_json::Map<String, Value>,
179 },
180 ToolResult {
181 #[serde(rename = "toolCallId", default)]
182 tool_call_id: String,
183 #[serde(rename = "toolName", default)]
184 tool_name: String,
185 content: Value,
186 #[serde(default, skip_serializing_if = "Option::is_none")]
189 usage: Option<Usage>,
190 #[serde(rename = "isError", default)]
191 is_error: bool,
192 #[serde(default)]
193 timestamp: f64,
194 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
195 extra: serde_json::Map<String, Value>,
196 },
197 BashExecution {
200 #[serde(default)]
201 command: String,
202 #[serde(default)]
203 output: String,
204 #[serde(rename = "exitCode", default, skip_serializing_if = "Option::is_none")]
205 exit_code: Option<i64>,
206 #[serde(default)]
207 cancelled: bool,
208 #[serde(default)]
209 truncated: bool,
210 #[serde(
211 rename = "fullOutputPath",
212 default,
213 skip_serializing_if = "Option::is_none"
214 )]
215 full_output_path: Option<String>,
216 #[serde(default)]
217 timestamp: f64,
218 #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
219 extra: serde_json::Map<String, Value>,
220 },
221}
222
223#[derive(Debug, Clone, PartialEq)]
229pub enum PiEvent {
230 AgentStart,
232 AgentEnd { messages: Vec<PiMessage> },
234 TurnStart,
236 TurnEnd {
238 message: Box<PiMessage>,
239 tool_results: Vec<PiMessage>,
240 },
241 MessageStart { message: Box<PiMessage> },
243 MessageUpdate {
247 message: Option<Box<PiMessage>>,
248 usage: Option<Usage>,
249 assistant_message_event: Value,
250 },
251 MessageEnd { message: Box<PiMessage> },
253 ToolExecutionStart {
261 tool_call_id: String,
262 tool_name: String,
263 args: Value,
264 parent_tool_call_id: Option<String>,
265 },
266 ToolExecutionUpdate {
268 tool_call_id: String,
269 tool_name: String,
270 args: Value,
271 partial_result: Value,
272 parent_tool_call_id: Option<String>,
273 },
274 ToolExecutionEnd {
276 tool_call_id: String,
277 tool_name: String,
278 result: Value,
279 is_error: bool,
280 parent_tool_call_id: Option<String>,
281 },
282 Unknown { event_type: String, payload: Value },
287}
288
289impl PiEvent {
290 pub fn event_type(&self) -> &str {
292 match self {
293 PiEvent::AgentStart => "agent_start",
294 PiEvent::AgentEnd { .. } => "agent_end",
295 PiEvent::TurnStart => "turn_start",
296 PiEvent::TurnEnd { .. } => "turn_end",
297 PiEvent::MessageStart { .. } => "message_start",
298 PiEvent::MessageUpdate { .. } => "message_update",
299 PiEvent::MessageEnd { .. } => "message_end",
300 PiEvent::ToolExecutionStart { .. } => "tool_execution_start",
301 PiEvent::ToolExecutionUpdate { .. } => "tool_execution_update",
302 PiEvent::ToolExecutionEnd { .. } => "tool_execution_end",
303 PiEvent::Unknown { event_type, .. } => event_type,
304 }
305 }
306
307 pub fn from_value(v: Value) -> Result<Self, serde_json::Error> {
309 let t = v
310 .get("type")
311 .and_then(Value::as_str)
312 .unwrap_or_default()
313 .to_string();
314 Ok(match t.as_str() {
315 "agent_start" => PiEvent::AgentStart,
316 "agent_end" => PiEvent::AgentEnd {
317 messages: serde_json::from_value(
318 v.get("messages").cloned().unwrap_or(Value::Array(vec![])),
319 )?,
320 },
321 "turn_start" => PiEvent::TurnStart,
322 "turn_end" => PiEvent::TurnEnd {
323 message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
324 tool_results: serde_json::from_value(
325 v.get("toolResults")
326 .cloned()
327 .unwrap_or(Value::Array(vec![])),
328 )?,
329 },
330 "message_start" => PiEvent::MessageStart {
331 message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
332 },
333 "message_update" => PiEvent::MessageUpdate {
334 message: match v.get("message") {
335 Some(m) => Some(serde_json::from_value(m.clone())?),
336 None => None,
337 },
338 usage: match v.get("usage") {
339 Some(u) => Some(serde_json::from_value(u.clone())?),
340 None => None,
341 },
342 assistant_message_event: v
343 .get("assistantMessageEvent")
344 .cloned()
345 .unwrap_or(Value::Null),
346 },
347 "message_end" => PiEvent::MessageEnd {
348 message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
349 },
350 "tool_execution_start" => PiEvent::ToolExecutionStart {
351 tool_call_id: str_field(&v, "toolCallId"),
352 tool_name: str_field(&v, "toolName"),
353 args: v.get("args").cloned().unwrap_or(Value::Null),
354 parent_tool_call_id: opt_str_field(&v, "parentToolCallId"),
355 },
356 "tool_execution_update" => PiEvent::ToolExecutionUpdate {
357 tool_call_id: str_field(&v, "toolCallId"),
358 tool_name: str_field(&v, "toolName"),
359 args: v.get("args").cloned().unwrap_or(Value::Null),
360 partial_result: v.get("partialResult").cloned().unwrap_or(Value::Null),
361 parent_tool_call_id: opt_str_field(&v, "parentToolCallId"),
362 },
363 "tool_execution_end" => PiEvent::ToolExecutionEnd {
364 tool_call_id: str_field(&v, "toolCallId"),
365 tool_name: str_field(&v, "toolName"),
366 result: v.get("result").cloned().unwrap_or(Value::Null),
367 is_error: v
368 .get("isError")
369 .and_then(Value::as_bool)
370 .unwrap_or_default(),
371 parent_tool_call_id: opt_str_field(&v, "parentToolCallId"),
372 },
373 _ => PiEvent::Unknown {
374 event_type: t,
375 payload: v,
376 },
377 })
378 }
379}
380
381fn str_field(v: &Value, k: &str) -> String {
382 v.get(k)
383 .and_then(Value::as_str)
384 .unwrap_or_default()
385 .to_string()
386}
387
388fn opt_str_field(v: &Value, k: &str) -> Option<String> {
389 v.get(k).and_then(Value::as_str).map(str::to_string)
390}
391
392#[cfg(test)]
393mod tests {
394 use super::*;
395
396 #[test]
397 fn assistant_message_from_docs_parses() {
398 let m: PiMessage = serde_json::from_value(serde_json::json!({
399 "role": "assistant",
400 "content": [
401 {"type": "text", "text": "Hello! How can I help?"},
402 {"type": "thinking", "thinking": "User is greeting me..."},
403 {"type": "toolCall", "id": "call_123", "name": "bash", "arguments": {"command": "ls"}}
404 ],
405 "api": "anthropic-messages",
406 "provider": "anthropic",
407 "model": "claude-sonnet-4-20250514",
408 "usage": {
409 "input": 100, "output": 50, "cacheRead": 0, "cacheWrite": 0,
410 "cost": {"input": 0.0003, "output": 0.00075, "cacheRead": 0, "cacheWrite": 0, "total": 0.00105}
411 },
412 "stopReason": "stop",
413 "timestamp": 1733234567890i64
414 }))
415 .unwrap();
416 let PiMessage::Assistant {
417 content,
418 stop_reason,
419 ..
420 } = m
421 else {
422 panic!("wrong role")
423 };
424 assert_eq!(content.len(), 3);
425 assert_eq!(stop_reason, "stop");
426 assert!(matches!(content[2], ContentBlock::ToolCall { .. }));
427 }
428
429 #[test]
430 fn system_message_from_docs_parses() {
431 let m: PiMessage = serde_json::from_value(serde_json::json!({
432 "role": "system",
433 "content": "",
434 "sections": {
435 "preamble": "You are an expert coding assistant...",
436 "tools": "<tools>\n- read: ...\n</tools>",
437 "cwd": "/project"
438 },
439 "toolsAdded": [{"name": "read", "description": "...", "parameters": {}}],
440 "timestamp": 1733234400000i64
441 }))
442 .unwrap();
443 let PiMessage::System {
444 sections,
445 tools_added,
446 tools_removed,
447 ..
448 } = m
449 else {
450 panic!("wrong role")
451 };
452 assert_eq!(sections.len(), 3);
453 assert_eq!(tools_added[0]["name"], "read");
454 assert!(tools_removed.is_empty());
455
456 let patch: PiMessage = serde_json::from_value(serde_json::json!({
457 "role": "system",
458 "content": "",
459 "sections": {"skills": null},
460 "toolsRemoved": [{"name": "write"}],
461 "timestamp": 1733234640000i64
462 }))
463 .unwrap();
464 let PiMessage::System {
465 sections,
466 tools_removed,
467 ..
468 } = patch
469 else {
470 panic!("wrong role")
471 };
472 assert!(sections["skills"].is_null());
473 assert_eq!(tools_removed[0]["name"], "write");
474 }
475
476 #[test]
477 fn bash_execution_message_round_trips() {
478 let v = serde_json::json!({
479 "role": "bashExecution",
480 "command": "ls -la",
481 "output": "total 48",
482 "exitCode": 0,
483 "cancelled": false,
484 "truncated": false,
485 "timestamp": 1733234567890i64
486 });
487 let m: PiMessage = serde_json::from_value(v).unwrap();
488 assert!(matches!(m, PiMessage::BashExecution { .. }));
489 }
490
491 #[test]
492 fn unknown_event_type_is_preserved_not_error() {
493 let e = PiEvent::from_value(serde_json::json!({
494 "type": "queue_update", "steering": [], "followUp": []
495 }))
496 .unwrap();
497 let PiEvent::Unknown {
498 event_type,
499 payload,
500 } = e
501 else {
502 panic!("expected Unknown")
503 };
504 assert_eq!(event_type, "queue_update");
505 assert!(payload.get("steering").is_some());
506 }
507
508 #[test]
509 fn tool_execution_events_lift_fields() {
510 let e = PiEvent::from_value(serde_json::json!({
511 "type": "tool_execution_end",
512 "toolCallId": "call_9", "toolName": "bash",
513 "result": {"ok": true}, "isError": false
514 }))
515 .unwrap();
516 let PiEvent::ToolExecutionEnd {
517 tool_call_id,
518 tool_name,
519 is_error,
520 parent_tool_call_id,
521 ..
522 } = e
523 else {
524 panic!("wrong variant")
525 };
526 assert_eq!(
527 (tool_call_id.as_str(), tool_name.as_str()),
528 ("call_9", "bash")
529 );
530 assert!(!is_error);
531 assert_eq!(parent_tool_call_id, None);
532 }
533
534 #[test]
535 fn nested_tool_execution_events_carry_parent_id() {
536 for (event_type, extra_key) in [
537 ("tool_execution_start", "args"),
538 ("tool_execution_update", "partialResult"),
539 ("tool_execution_end", "result"),
540 ] {
541 let e = PiEvent::from_value(serde_json::json!({
542 "type": event_type,
543 "toolCallId": "call_9/1", "toolName": "read",
544 extra_key: {}, "parentToolCallId": "call_9"
545 }))
546 .unwrap();
547 let parent = match e {
548 PiEvent::ToolExecutionStart {
549 parent_tool_call_id,
550 ..
551 }
552 | PiEvent::ToolExecutionUpdate {
553 parent_tool_call_id,
554 ..
555 }
556 | PiEvent::ToolExecutionEnd {
557 parent_tool_call_id,
558 ..
559 } => parent_tool_call_id,
560 other => panic!("wrong variant: {other:?}"),
561 };
562 assert_eq!(parent.as_deref(), Some("call_9"), "{event_type}");
563 }
564 }
565
566 #[test]
567 fn assistant_thinking_level_is_typed_and_optional() {
568 let with: PiMessage = serde_json::from_value(serde_json::json!({
569 "role": "assistant", "content": [], "stopReason": "stop",
570 "thinkingLevel": "off", "timestamp": 1790756884219i64
571 }))
572 .unwrap();
573 let PiMessage::Assistant {
574 thinking_level,
575 extra,
576 ..
577 } = &with
578 else {
579 panic!("wrong role")
580 };
581 assert_eq!(thinking_level.as_deref(), Some("off"));
582 assert!(!extra.contains_key("thinkingLevel"));
583 assert_eq!(serde_json::to_value(&with).unwrap()["thinkingLevel"], "off");
584
585 let without: PiMessage = serde_json::from_value(serde_json::json!({
586 "role": "assistant", "content": [], "stopReason": "stop"
587 }))
588 .unwrap();
589 let PiMessage::Assistant { thinking_level, .. } = &without else {
590 panic!("wrong role")
591 };
592 assert_eq!(*thinking_level, None);
593 assert!(serde_json::to_value(&without)
594 .unwrap()
595 .get("thinkingLevel")
596 .is_none());
597 }
598}