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    /// The prompt and tool loadout (pi ≥ 0.86.0). The first request of a
110    /// session emits one carrying every prompt section and tool
111    /// declaration; later changes emit further system messages that
112    /// patch `sections` by name (`null` removes one) and list
113    /// `toolsAdded` / `toolsRemoved`. Replaying them in order yields the
114    /// current prompt and tools. Streams as `message_start` /
115    /// `message_end` ahead of the user message and leads `agent_end`'s
116    /// `messages`.
117    System {
118        /// A plain string or an array of text blocks; empty when the
119        /// prompt lives entirely in `sections`.
120        content: Value,
121        /// Named prompt sections (`preamble`, `tools`, `rules`, `cwd`,
122        /// `docs`, `project_context`, ...) — string bodies, or `null`
123        /// to remove a section declared earlier.
124        #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")]
125        sections: serde_json::Map<String, Value>,
126        /// Full tool declarations (`name`, `description`, `parameters`)
127        /// now available to the model.
128        #[serde(rename = "toolsAdded", default, skip_serializing_if = "Vec::is_empty")]
129        tools_added: Vec<Value>,
130        /// `{ "name": ... }` references for tools withdrawn.
131        #[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        /// A plain string or an array of content blocks — both are legal.
144        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        /// `stop`, `length`, `toolUse`, `error`, `aborted`, or
163        /// (harness-v2) `deferred`. Left open for forward compat.
164        #[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        /// Nested LLM work performed by the tool; contributes to session
178        /// totals when present.
179        #[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    /// Created by the RPC `bash` command, never by LLM tool calls. On the
189    /// next prompt it is rewritten into a `user` message for the model.
190    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/// One event from the `--mode json` / `--mode rpc` stdout stream,
215/// discriminated on `type`. Known lifecycle events are typed with their
216/// primary fields lifted; every variant keeps its full payload reachable
217/// (unmodeled fields land in `extra`), and unknown `type`s fall back to
218/// [`PiEvent::Unknown`] instead of erroring.
219#[derive(Debug, Clone, PartialEq)]
220pub enum PiEvent {
221    /// `agent_start`
222    AgentStart,
223    /// `agent_end` — carries the turn's messages.
224    AgentEnd { messages: Vec<PiMessage> },
225    /// `turn_start`
226    TurnStart,
227    /// `turn_end`
228    TurnEnd {
229        message: Box<PiMessage>,
230        tool_results: Vec<PiMessage>,
231    },
232    /// `message_start`
233    MessageStart { message: Box<PiMessage> },
234    /// `message_update` — streaming delta; in `--mode json` the
235    /// cumulative message snapshot is omitted and only the incremental
236    /// `assistantMessageEvent` (kept raw here) arrives.
237    MessageUpdate {
238        message: Option<Box<PiMessage>>,
239        usage: Option<Usage>,
240        assistant_message_event: Value,
241    },
242    /// `message_end`
243    MessageEnd { message: Box<PiMessage> },
244    /// `tool_execution_start`
245    ToolExecutionStart {
246        tool_call_id: String,
247        tool_name: String,
248        args: Value,
249    },
250    /// `tool_execution_update`
251    ToolExecutionUpdate {
252        tool_call_id: String,
253        tool_name: String,
254        args: Value,
255        partial_result: Value,
256    },
257    /// `tool_execution_end`
258    ToolExecutionEnd {
259        tool_call_id: String,
260        tool_name: String,
261        result: Value,
262        is_error: bool,
263    },
264    /// Any other `type` (session events like `queue_update`,
265    /// `compaction_start`/`compaction_end`, `bash_execution_update`,
266    /// extension UI requests, and whatever future CLIs add). The full
267    /// object is preserved.
268    Unknown { event_type: String, payload: Value },
269}
270
271impl PiEvent {
272    /// The wire `type` string.
273    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    /// Parse one stdout JSON object into a typed event.
290    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}