Skip to main content

pi_codes/
io.rs

1//! Typed models of pi's wire vocabulary: the messages and events shared
2//! by `--mode json` (one-shot JSONL event stream) and `--mode rpc`
3//! (command protocol; see [`crate::rpc`]).
4//!
5//! Shapes mirror the documented TypeScript types
6//! (`packages/ai/src/types.ts`, `packages/agent/src/types.ts`, and
7//! `packages/coding-agent/docs/{json,rpc}.md` in earendil-works/pi).
8//! Unknown event types are preserved verbatim in [`PiEvent::Unknown`]
9//! rather than erroring, so a newer CLI degrades soft.
10
11use serde::{Deserialize, Serialize};
12use serde_json::Value;
13
14/// Per-request/response token accounting on assistant and tool messages.
15#[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/// Dollar cost breakdown inside [`Usage`].
38#[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/// A model catalog entry (`get_state`/`set_model`/`get_available_models`).
53/// With no provider configured, `get_state` returns a placeholder with
54/// `id: "unknown"` and zeroed limits — treat `contextWindow == 0` as
55/// "no real model resolved".
56#[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/// One block of assistant content: text, thinking, or a tool call.
83#[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/// A conversation message, discriminated on `role`. This is pi's
105/// `AgentMessage` union plus the RPC-only `bashExecution` role.
106#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
107#[serde(tag = "role", rename_all = "camelCase")]
108pub enum PiMessage {
109    User {
110        /// A plain string or an array of content blocks — both are legal.
111        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        /// `stop`, `length`, `toolUse`, `error`, `aborted`, or
130        /// (harness-v2) `deferred`. Left open for forward compat.
131        #[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        /// Nested LLM work performed by the tool; contributes to session
145        /// totals when present.
146        #[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    /// Created by the RPC `bash` command, never by LLM tool calls. On the
156    /// next prompt it is rewritten into a `user` message for the model.
157    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/// One event from the `--mode json` / `--mode rpc` stdout stream,
182/// discriminated on `type`. Known lifecycle events are typed with their
183/// primary fields lifted; every variant keeps its full payload reachable
184/// (unmodeled fields land in `extra`), and unknown `type`s fall back to
185/// [`PiEvent::Unknown`] instead of erroring.
186#[derive(Debug, Clone, PartialEq)]
187pub enum PiEvent {
188    /// `agent_start`
189    AgentStart,
190    /// `agent_end` — carries the turn's messages.
191    AgentEnd { messages: Vec<PiMessage> },
192    /// `turn_start`
193    TurnStart,
194    /// `turn_end`
195    TurnEnd {
196        message: Box<PiMessage>,
197        tool_results: Vec<PiMessage>,
198    },
199    /// `message_start`
200    MessageStart { message: Box<PiMessage> },
201    /// `message_update` — streaming delta; in `--mode json` the
202    /// cumulative message snapshot is omitted and only the incremental
203    /// `assistantMessageEvent` (kept raw here) arrives.
204    MessageUpdate {
205        message: Option<Box<PiMessage>>,
206        usage: Option<Usage>,
207        assistant_message_event: Value,
208    },
209    /// `message_end`
210    MessageEnd { message: Box<PiMessage> },
211    /// `tool_execution_start`
212    ToolExecutionStart {
213        tool_call_id: String,
214        tool_name: String,
215        args: Value,
216    },
217    /// `tool_execution_update`
218    ToolExecutionUpdate {
219        tool_call_id: String,
220        tool_name: String,
221        args: Value,
222        partial_result: Value,
223    },
224    /// `tool_execution_end`
225    ToolExecutionEnd {
226        tool_call_id: String,
227        tool_name: String,
228        result: Value,
229        is_error: bool,
230    },
231    /// Any other `type` (session events like `queue_update`,
232    /// `compaction_start`/`compaction_end`, `bash_execution_update`,
233    /// extension UI requests, and whatever future CLIs add). The full
234    /// object is preserved.
235    Unknown { event_type: String, payload: Value },
236}
237
238impl PiEvent {
239    /// The wire `type` string.
240    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    /// Parse one stdout JSON object into a typed event.
257    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}