pi-codes 1.4.0

//! Typed models of pi's wire vocabulary: the messages and events shared
//! by `--mode json` (one-shot JSONL event stream) and `--mode rpc`
//! (command protocol; see [`crate::rpc`]).
//!
//! Shapes mirror the documented TypeScript types
//! (`packages/ai/src/types.ts`, `packages/agent/src/types.ts`, and
//! `packages/coding-agent/docs/{json,rpc}.md` in earendil-works/pi).
//! Unknown event types are preserved verbatim in [`PiEvent::Unknown`]
//! rather than erroring, so a newer CLI degrades soft.

use serde::{Deserialize, Serialize};
use serde_json::Value;

/// Per-request/response token accounting on assistant and tool messages.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct Usage {
    #[serde(default)]
    pub input: f64,
    #[serde(default)]
    pub output: f64,
    #[serde(rename = "cacheRead", default)]
    pub cache_read: f64,
    #[serde(rename = "cacheWrite", default)]
    pub cache_write: f64,
    #[serde(
        rename = "totalTokens",
        default,
        skip_serializing_if = "Option::is_none"
    )]
    pub total_tokens: Option<f64>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cost: Option<UsageCost>,
    #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
    pub extra: serde_json::Map<String, Value>,
}

/// Dollar cost breakdown inside [`Usage`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct UsageCost {
    #[serde(default)]
    pub input: f64,
    #[serde(default)]
    pub output: f64,
    #[serde(rename = "cacheRead", default)]
    pub cache_read: f64,
    #[serde(rename = "cacheWrite", default)]
    pub cache_write: f64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub total: Option<f64>,
}

/// A model catalog entry (`get_state`/`set_model`/`get_available_models`).
/// With no provider configured, `get_state` returns a placeholder with
/// `id: "unknown"` and zeroed limits — treat `contextWindow == 0` as
/// "no real model resolved".
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct Model {
    #[serde(default)]
    pub id: String,
    #[serde(default)]
    pub name: String,
    #[serde(default)]
    pub api: String,
    #[serde(default)]
    pub provider: String,
    #[serde(rename = "baseUrl", default)]
    pub base_url: String,
    #[serde(default)]
    pub reasoning: bool,
    #[serde(default)]
    pub input: Vec<String>,
    #[serde(rename = "contextWindow", default)]
    pub context_window: i64,
    #[serde(rename = "maxTokens", default)]
    pub max_tokens: i64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cost: Option<UsageCost>,
    #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
    pub extra: serde_json::Map<String, Value>,
}

/// One block of assistant content: text, thinking, or a tool call.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum ContentBlock {
    Text {
        text: String,
    },
    Thinking {
        thinking: String,
    },
    ToolCall {
        id: String,
        name: String,
        arguments: Value,
        /// Set when the model's full arguments text was not a valid JSON
        /// object (pi 1.3.0+). `arguments` is then `{}`, the message
        /// describes the problem and quotes the text the model sent, and
        /// the call fails tool argument validation instead of running.
        #[serde(
            rename = "argumentsError",
            default,
            skip_serializing_if = "Option::is_none"
        )]
        arguments_error: Option<String>,
    },
    Image {
        data: String,
        #[serde(rename = "mimeType")]
        mime_type: String,
    },
}

/// A conversation message, discriminated on `role`. This is pi's
/// `AgentMessage` union plus the RPC-only `bashExecution` role.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "role", rename_all = "camelCase")]
pub enum PiMessage {
    /// The prompt and tool loadout (pi ≥ 0.86.0). The first request of a
    /// session emits one carrying every prompt section and tool
    /// declaration; later changes emit further system messages that
    /// patch `sections` by name (`null` removes one) and list
    /// `toolsAdded` / `toolsRemoved`. Replaying them in order yields the
    /// current prompt and tools. Streams as `message_start` /
    /// `message_end` ahead of the user message and leads `agent_end`'s
    /// `messages`.
    System {
        /// A plain string or an array of text blocks; empty when the
        /// prompt lives entirely in `sections`.
        content: Value,
        /// Named prompt sections (`preamble`, `tools`, `rules`, `cwd`,
        /// `docs`, `project_context`, ...) — string bodies, or `null`
        /// to remove a section declared earlier.
        #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")]
        sections: serde_json::Map<String, Value>,
        /// Full tool declarations (`name`, `description`, `parameters`)
        /// now available to the model.
        #[serde(rename = "toolsAdded", default, skip_serializing_if = "Vec::is_empty")]
        tools_added: Vec<Value>,
        /// `{ "name": ... }` references for tools withdrawn.
        #[serde(
            rename = "toolsRemoved",
            default,
            skip_serializing_if = "Vec::is_empty"
        )]
        tools_removed: Vec<Value>,
        #[serde(default)]
        timestamp: f64,
        #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
        extra: serde_json::Map<String, Value>,
    },
    User {
        /// A plain string or an array of content blocks — both are legal.
        content: Value,
        #[serde(default)]
        timestamp: f64,
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        attachments: Vec<Value>,
        #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
        extra: serde_json::Map<String, Value>,
    },
    Assistant {
        content: Vec<ContentBlock>,
        #[serde(default)]
        api: String,
        #[serde(default)]
        provider: String,
        #[serde(default)]
        model: String,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        usage: Option<Usage>,
        /// `stop`, `length`, `toolUse`, `error`, `aborted`, or
        /// (harness-v2) `deferred`. Left open for forward compat.
        #[serde(rename = "stopReason", default)]
        stop_reason: String,
        /// The pi thinking level the agent loop requested for this
        /// response (pi ≥ 0.99.0). Absent outside the agent loop and on
        /// messages from older sessions.
        #[serde(
            rename = "thinkingLevel",
            default,
            skip_serializing_if = "Option::is_none"
        )]
        thinking_level: Option<String>,
        #[serde(default)]
        timestamp: f64,
        #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
        extra: serde_json::Map<String, Value>,
    },
    ToolResult {
        #[serde(rename = "toolCallId", default)]
        tool_call_id: String,
        #[serde(rename = "toolName", default)]
        tool_name: String,
        content: Value,
        /// Nested LLM work performed by the tool; contributes to session
        /// totals when present.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        usage: Option<Usage>,
        #[serde(rename = "isError", default)]
        is_error: bool,
        #[serde(default)]
        timestamp: f64,
        #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
        extra: serde_json::Map<String, Value>,
    },
    /// Created by the RPC `bash` command, never by LLM tool calls. On the
    /// next prompt it is rewritten into a `user` message for the model.
    BashExecution {
        #[serde(default)]
        command: String,
        #[serde(default)]
        output: String,
        #[serde(rename = "exitCode", default, skip_serializing_if = "Option::is_none")]
        exit_code: Option<i64>,
        #[serde(default)]
        cancelled: bool,
        #[serde(default)]
        truncated: bool,
        #[serde(
            rename = "fullOutputPath",
            default,
            skip_serializing_if = "Option::is_none"
        )]
        full_output_path: Option<String>,
        #[serde(default)]
        timestamp: f64,
        #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
        extra: serde_json::Map<String, Value>,
    },
}

/// One event from the `--mode json` / `--mode rpc` stdout stream,
/// discriminated on `type`. Known lifecycle events are typed with their
/// primary fields lifted; every variant keeps its full payload reachable
/// (unmodeled fields land in `extra`), and unknown `type`s fall back to
/// [`PiEvent::Unknown`] instead of erroring.
#[derive(Debug, Clone, PartialEq)]
pub enum PiEvent {
    /// `agent_start`
    AgentStart,
    /// `agent_end` — carries the turn's messages.
    AgentEnd { messages: Vec<PiMessage> },
    /// `turn_start`
    TurnStart,
    /// `turn_end`
    TurnEnd {
        message: Box<PiMessage>,
        tool_results: Vec<PiMessage>,
    },
    /// `message_start`
    MessageStart { message: Box<PiMessage> },
    /// `message_update` — streaming delta; in `--mode json` the
    /// cumulative message snapshot is omitted and only the incremental
    /// `assistantMessageEvent` (kept raw here) arrives.
    MessageUpdate {
        message: Option<Box<PiMessage>>,
        usage: Option<Usage>,
        assistant_message_event: Value,
    },
    /// `message_end`
    MessageEnd { message: Box<PiMessage> },
    /// `tool_execution_start`
    ///
    /// On all three tool events, `parent_tool_call_id` is set when another
    /// tool made the call (pi ≥ 0.99.0: `ctx.executeTool()`, e.g. from a
    /// `codemode` script). Such nested calls get a `<parent id>/<n>`
    /// `tool_call_id` and never appear as tool calls or tool results in
    /// the transcript.
    ToolExecutionStart {
        tool_call_id: String,
        tool_name: String,
        args: Value,
        parent_tool_call_id: Option<String>,
    },
    /// `tool_execution_update`
    ToolExecutionUpdate {
        tool_call_id: String,
        tool_name: String,
        args: Value,
        partial_result: Value,
        parent_tool_call_id: Option<String>,
    },
    /// `tool_execution_end`
    ToolExecutionEnd {
        tool_call_id: String,
        tool_name: String,
        result: Value,
        is_error: bool,
        parent_tool_call_id: Option<String>,
    },
    /// Any other `type` (session events like `queue_update`,
    /// `compaction_start`/`compaction_end`, `bash_execution_update`,
    /// extension UI requests, and whatever future CLIs add). The full
    /// object is preserved.
    Unknown { event_type: String, payload: Value },
}

impl PiEvent {
    /// The wire `type` string.
    pub fn event_type(&self) -> &str {
        match self {
            PiEvent::AgentStart => "agent_start",
            PiEvent::AgentEnd { .. } => "agent_end",
            PiEvent::TurnStart => "turn_start",
            PiEvent::TurnEnd { .. } => "turn_end",
            PiEvent::MessageStart { .. } => "message_start",
            PiEvent::MessageUpdate { .. } => "message_update",
            PiEvent::MessageEnd { .. } => "message_end",
            PiEvent::ToolExecutionStart { .. } => "tool_execution_start",
            PiEvent::ToolExecutionUpdate { .. } => "tool_execution_update",
            PiEvent::ToolExecutionEnd { .. } => "tool_execution_end",
            PiEvent::Unknown { event_type, .. } => event_type,
        }
    }

    /// Parse one stdout JSON object into a typed event.
    pub fn from_value(v: Value) -> Result<Self, serde_json::Error> {
        let t = v
            .get("type")
            .and_then(Value::as_str)
            .unwrap_or_default()
            .to_string();
        Ok(match t.as_str() {
            "agent_start" => PiEvent::AgentStart,
            "agent_end" => PiEvent::AgentEnd {
                messages: serde_json::from_value(
                    v.get("messages").cloned().unwrap_or(Value::Array(vec![])),
                )?,
            },
            "turn_start" => PiEvent::TurnStart,
            "turn_end" => PiEvent::TurnEnd {
                message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
                tool_results: serde_json::from_value(
                    v.get("toolResults")
                        .cloned()
                        .unwrap_or(Value::Array(vec![])),
                )?,
            },
            "message_start" => PiEvent::MessageStart {
                message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
            },
            "message_update" => PiEvent::MessageUpdate {
                message: match v.get("message") {
                    Some(m) => Some(serde_json::from_value(m.clone())?),
                    None => None,
                },
                usage: match v.get("usage") {
                    Some(u) => Some(serde_json::from_value(u.clone())?),
                    None => None,
                },
                assistant_message_event: v
                    .get("assistantMessageEvent")
                    .cloned()
                    .unwrap_or(Value::Null),
            },
            "message_end" => PiEvent::MessageEnd {
                message: serde_json::from_value(v.get("message").cloned().unwrap_or(Value::Null))?,
            },
            "tool_execution_start" => PiEvent::ToolExecutionStart {
                tool_call_id: str_field(&v, "toolCallId"),
                tool_name: str_field(&v, "toolName"),
                args: v.get("args").cloned().unwrap_or(Value::Null),
                parent_tool_call_id: opt_str_field(&v, "parentToolCallId"),
            },
            "tool_execution_update" => PiEvent::ToolExecutionUpdate {
                tool_call_id: str_field(&v, "toolCallId"),
                tool_name: str_field(&v, "toolName"),
                args: v.get("args").cloned().unwrap_or(Value::Null),
                partial_result: v.get("partialResult").cloned().unwrap_or(Value::Null),
                parent_tool_call_id: opt_str_field(&v, "parentToolCallId"),
            },
            "tool_execution_end" => PiEvent::ToolExecutionEnd {
                tool_call_id: str_field(&v, "toolCallId"),
                tool_name: str_field(&v, "toolName"),
                result: v.get("result").cloned().unwrap_or(Value::Null),
                is_error: v
                    .get("isError")
                    .and_then(Value::as_bool)
                    .unwrap_or_default(),
                parent_tool_call_id: opt_str_field(&v, "parentToolCallId"),
            },
            _ => PiEvent::Unknown {
                event_type: t,
                payload: v,
            },
        })
    }
}

fn str_field(v: &Value, k: &str) -> String {
    v.get(k)
        .and_then(Value::as_str)
        .unwrap_or_default()
        .to_string()
}

fn opt_str_field(v: &Value, k: &str) -> Option<String> {
    v.get(k).and_then(Value::as_str).map(str::to_string)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn assistant_message_from_docs_parses() {
        let m: PiMessage = serde_json::from_value(serde_json::json!({
            "role": "assistant",
            "content": [
                {"type": "text", "text": "Hello! How can I help?"},
                {"type": "thinking", "thinking": "User is greeting me..."},
                {"type": "toolCall", "id": "call_123", "name": "bash", "arguments": {"command": "ls"}}
            ],
            "api": "anthropic-messages",
            "provider": "anthropic",
            "model": "claude-sonnet-4-20250514",
            "usage": {
                "input": 100, "output": 50, "cacheRead": 0, "cacheWrite": 0,
                "cost": {"input": 0.0003, "output": 0.00075, "cacheRead": 0, "cacheWrite": 0, "total": 0.00105}
            },
            "stopReason": "stop",
            "timestamp": 1733234567890i64
        }))
        .unwrap();
        let PiMessage::Assistant {
            content,
            stop_reason,
            ..
        } = m
        else {
            panic!("wrong role")
        };
        assert_eq!(content.len(), 3);
        assert_eq!(stop_reason, "stop");
        assert!(matches!(
            content[2],
            ContentBlock::ToolCall {
                arguments_error: None,
                ..
            }
        ));
    }

    #[test]
    fn tool_call_arguments_error_round_trips() {
        let raw = serde_json::json!({
            "type": "toolCall",
            "id": "call_7",
            "name": "read",
            "arguments": {},
            "argumentsError": "Tool arguments are not valid JSON: {\"path\": \"a"
        });
        let block: ContentBlock = serde_json::from_value(raw.clone()).unwrap();
        let ContentBlock::ToolCall {
            ref arguments_error,
            ..
        } = block
        else {
            panic!("wrong block")
        };
        assert!(arguments_error
            .as_deref()
            .is_some_and(|e| e.contains("not valid JSON")));
        assert_eq!(serde_json::to_value(&block).unwrap(), raw);
    }

    #[test]
    fn system_message_from_docs_parses() {
        let m: PiMessage = serde_json::from_value(serde_json::json!({
            "role": "system",
            "content": "",
            "sections": {
                "preamble": "You are an expert coding assistant...",
                "tools": "<tools>\n- read: ...\n</tools>",
                "cwd": "/project"
            },
            "toolsAdded": [{"name": "read", "description": "...", "parameters": {}}],
            "timestamp": 1733234400000i64
        }))
        .unwrap();
        let PiMessage::System {
            sections,
            tools_added,
            tools_removed,
            ..
        } = m
        else {
            panic!("wrong role")
        };
        assert_eq!(sections.len(), 3);
        assert_eq!(tools_added[0]["name"], "read");
        assert!(tools_removed.is_empty());

        let patch: PiMessage = serde_json::from_value(serde_json::json!({
            "role": "system",
            "content": "",
            "sections": {"skills": null},
            "toolsRemoved": [{"name": "write"}],
            "timestamp": 1733234640000i64
        }))
        .unwrap();
        let PiMessage::System {
            sections,
            tools_removed,
            ..
        } = patch
        else {
            panic!("wrong role")
        };
        assert!(sections["skills"].is_null());
        assert_eq!(tools_removed[0]["name"], "write");
    }

    #[test]
    fn bash_execution_message_round_trips() {
        let v = serde_json::json!({
            "role": "bashExecution",
            "command": "ls -la",
            "output": "total 48",
            "exitCode": 0,
            "cancelled": false,
            "truncated": false,
            "timestamp": 1733234567890i64
        });
        let m: PiMessage = serde_json::from_value(v).unwrap();
        assert!(matches!(m, PiMessage::BashExecution { .. }));
    }

    #[test]
    fn unknown_event_type_is_preserved_not_error() {
        let e = PiEvent::from_value(serde_json::json!({
            "type": "queue_update", "steering": [], "followUp": []
        }))
        .unwrap();
        let PiEvent::Unknown {
            event_type,
            payload,
        } = e
        else {
            panic!("expected Unknown")
        };
        assert_eq!(event_type, "queue_update");
        assert!(payload.get("steering").is_some());
    }

    #[test]
    fn tool_execution_events_lift_fields() {
        let e = PiEvent::from_value(serde_json::json!({
            "type": "tool_execution_end",
            "toolCallId": "call_9", "toolName": "bash",
            "result": {"ok": true}, "isError": false
        }))
        .unwrap();
        let PiEvent::ToolExecutionEnd {
            tool_call_id,
            tool_name,
            is_error,
            parent_tool_call_id,
            ..
        } = e
        else {
            panic!("wrong variant")
        };
        assert_eq!(
            (tool_call_id.as_str(), tool_name.as_str()),
            ("call_9", "bash")
        );
        assert!(!is_error);
        assert_eq!(parent_tool_call_id, None);
    }

    #[test]
    fn nested_tool_execution_events_carry_parent_id() {
        for (event_type, extra_key) in [
            ("tool_execution_start", "args"),
            ("tool_execution_update", "partialResult"),
            ("tool_execution_end", "result"),
        ] {
            let e = PiEvent::from_value(serde_json::json!({
                "type": event_type,
                "toolCallId": "call_9/1", "toolName": "read",
                extra_key: {}, "parentToolCallId": "call_9"
            }))
            .unwrap();
            let parent = match e {
                PiEvent::ToolExecutionStart {
                    parent_tool_call_id,
                    ..
                }
                | PiEvent::ToolExecutionUpdate {
                    parent_tool_call_id,
                    ..
                }
                | PiEvent::ToolExecutionEnd {
                    parent_tool_call_id,
                    ..
                } => parent_tool_call_id,
                other => panic!("wrong variant: {other:?}"),
            };
            assert_eq!(parent.as_deref(), Some("call_9"), "{event_type}");
        }
    }

    #[test]
    fn assistant_thinking_level_is_typed_and_optional() {
        let with: PiMessage = serde_json::from_value(serde_json::json!({
            "role": "assistant", "content": [], "stopReason": "stop",
            "thinkingLevel": "off", "timestamp": 1790756884219i64
        }))
        .unwrap();
        let PiMessage::Assistant {
            thinking_level,
            extra,
            ..
        } = &with
        else {
            panic!("wrong role")
        };
        assert_eq!(thinking_level.as_deref(), Some("off"));
        assert!(!extra.contains_key("thinkingLevel"));
        assert_eq!(serde_json::to_value(&with).unwrap()["thinkingLevel"], "off");

        let without: PiMessage = serde_json::from_value(serde_json::json!({
            "role": "assistant", "content": [], "stopReason": "stop"
        }))
        .unwrap();
        let PiMessage::Assistant { thinking_level, .. } = &without else {
            panic!("wrong role")
        };
        assert_eq!(*thinking_level, None);
        assert!(serde_json::to_value(&without)
            .unwrap()
            .get("thinkingLevel")
            .is_none());
    }
}