pi-codes 1.0.2

//! The `pi --mode rpc` command protocol: JSON objects on stdin (one per
//! line, LF-delimited), `{"type":"response",...}` acknowledgements and
//! streamed events on stdout.
//!
//! Framing contract (from `docs/rpc.md`): split records on `\n` ONLY,
//! tolerate a trailing `\r` — never use a splitter that also breaks on
//! U+2028/U+2029, which are legal inside JSON strings.

use crate::io::{Model, PiMessage};
use serde::{Deserialize, Serialize};
use serde_json::Value;

/// An image attached to a `prompt`/`steer`/`follow_up` command.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ImageContent {
    /// Always `"image"`.
    #[serde(rename = "type")]
    pub content_type: String,
    /// Base64-encoded bytes.
    pub data: String,
    #[serde(rename = "mimeType")]
    pub mime_type: String,
}

/// How a `prompt` sent mid-stream should queue.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum StreamingBehavior {
    /// Deliver after the current assistant turn's tool calls, before the
    /// next LLM call.
    #[serde(rename = "steer")]
    Steer,
    /// Deliver only once the agent stops.
    #[serde(rename = "followUp")]
    FollowUp,
}

/// A client → agent command. Serializes with the wire's `type` tag; the
/// optional `id` correlates the matching [`RpcResponse`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum RpcCommand {
    Prompt {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        message: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        images: Option<Vec<ImageContent>>,
        #[serde(rename = "streamingBehavior", skip_serializing_if = "Option::is_none")]
        streaming_behavior: Option<StreamingBehavior>,
    },
    Steer {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        message: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        images: Option<Vec<ImageContent>>,
    },
    FollowUp {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        message: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        images: Option<Vec<ImageContent>>,
    },
    Abort {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    ClearQueue {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    NewSession {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    GetState {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    GetMessages {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    SetModel {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        provider: String,
        #[serde(rename = "modelId")]
        model_id: String,
    },
    CycleModel {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    GetAvailableModels {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    SetThinkingLevel {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        level: String,
    },
    CycleThinkingLevel {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    GetAvailableThinkingLevels {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    SetSteeringMode {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        mode: String,
    },
    SetFollowUpMode {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        mode: String,
    },
    Compact {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        #[serde(skip_serializing_if = "Option::is_none")]
        instructions: Option<String>,
    },
    SetAutoCompaction {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        enabled: bool,
    },
    SetAutoRetry {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        enabled: bool,
    },
    AbortRetry {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    /// Run a shell command outside the model; output joins context on the
    /// NEXT prompt. Streams `bash_execution_update` events carrying this
    /// command's `id`.
    Bash {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        command: String,
    },
    AbortBash {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    GetSessionStats {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
    },
    ExportHtml {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        #[serde(skip_serializing_if = "Option::is_none")]
        path: Option<String>,
    },
    SetSessionName {
        #[serde(skip_serializing_if = "Option::is_none")]
        id: Option<String>,
        name: String,
    },
    /// Anything this crate doesn't model yet — serialized as its raw
    /// object so new commands work without a crate release.
    #[serde(untagged)]
    Raw(Value),
}

impl RpcCommand {
    /// Serialize to one LF-terminated wire line.
    pub fn to_line(&self) -> Result<String, serde_json::Error> {
        Ok(format!("{}\n", serde_json::to_string(self)?))
    }
}

/// The `{"type":"response"}` acknowledgement for one command.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RpcResponse {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    /// Which command this answers (`"parse"` for unparseable input).
    #[serde(default)]
    pub command: String,
    #[serde(default)]
    pub success: bool,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub data: Option<Value>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

impl RpcResponse {
    /// What pi did with the input a successful `prompt`, `steer` or
    /// `follow_up` submitted (`data.disposition`, pi ≥ 0.99.0). `None` on
    /// other commands, on failures, and from older CLIs.
    pub fn disposition(&self) -> Option<InputDisposition> {
        self.data
            .as_ref()?
            .get("disposition")?
            .as_str()
            .map(InputDisposition::from)
    }
}

/// What pi did with a submitted `prompt`/`steer`/`follow_up` input
/// ([`RpcResponse::disposition`], pi ≥ 0.99.0). It describes the submitted
/// input only — not independent work an extension started, and not a
/// guarantee of completion or that a queued message stays queued.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum InputDisposition {
    /// An extension command or input handler consumed the input. No run
    /// started for it, so don't wait for `agent_settled`.
    Handled,
    /// Pi queued the input during a run (including after an input handler
    /// transformed it).
    Queued,
    /// Pi accepted the prompt to start a run. `prompt` only.
    Started,
    /// A disposition not yet known to this version of the crate.
    Unknown(String),
}

impl InputDisposition {
    pub fn as_str(&self) -> &str {
        match self {
            Self::Handled => "handled",
            Self::Queued => "queued",
            Self::Started => "started",
            Self::Unknown(s) => s.as_str(),
        }
    }
}

impl std::fmt::Display for InputDisposition {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.as_str())
    }
}

impl From<&str> for InputDisposition {
    fn from(s: &str) -> Self {
        match s {
            "handled" => Self::Handled,
            "queued" => Self::Queued,
            "started" => Self::Started,
            other => Self::Unknown(other.to_string()),
        }
    }
}

impl Serialize for InputDisposition {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_str(self.as_str())
    }
}

impl<'de> Deserialize<'de> for InputDisposition {
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        let s = String::deserialize(deserializer)?;
        Ok(Self::from(s.as_str()))
    }
}

/// Typed view of `get_state` response data.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct AgentState {
    /// `None` when no provider resolved; may also be a placeholder with
    /// `id: "unknown"` — see [`Model`].
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub model: Option<Model>,
    #[serde(rename = "thinkingLevel", default)]
    pub thinking_level: String,
    #[serde(rename = "isStreaming", default)]
    pub is_streaming: bool,
    #[serde(rename = "isCompacting", default)]
    pub is_compacting: bool,
    #[serde(rename = "steeringMode", default)]
    pub steering_mode: String,
    #[serde(rename = "followUpMode", default)]
    pub follow_up_mode: String,
    #[serde(
        rename = "sessionFile",
        default,
        skip_serializing_if = "Option::is_none"
    )]
    pub session_file: Option<String>,
    #[serde(rename = "sessionId", default)]
    pub session_id: String,
    #[serde(
        rename = "sessionName",
        default,
        skip_serializing_if = "Option::is_none"
    )]
    pub session_name: Option<String>,
    #[serde(rename = "autoCompactionEnabled", default)]
    pub auto_compaction_enabled: bool,
    #[serde(rename = "messageCount", default)]
    pub message_count: i64,
    #[serde(rename = "pendingMessageCount", default)]
    pub pending_message_count: i64,
    #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
    pub extra: serde_json::Map<String, Value>,
}

/// Typed view of `bash` response data.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct BashResult {
    #[serde(default)]
    pub output: String,
    #[serde(rename = "exitCode", default, skip_serializing_if = "Option::is_none")]
    pub exit_code: Option<i64>,
    #[serde(default)]
    pub cancelled: bool,
    #[serde(default)]
    pub truncated: bool,
    #[serde(
        rename = "fullOutputPath",
        default,
        skip_serializing_if = "Option::is_none"
    )]
    pub full_output_path: Option<String>,
}

/// Typed view of `get_messages` response data.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct Messages {
    #[serde(default)]
    pub messages: Vec<PiMessage>,
}

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

    #[test]
    fn commands_serialize_with_snake_case_type_tags() {
        let c = RpcCommand::GetState {
            id: Some("r1".into()),
        };
        let v = serde_json::to_value(&c).unwrap();
        assert_eq!(v["type"], "get_state");
        assert_eq!(v["id"], "r1");

        let c = RpcCommand::Prompt {
            id: None,
            message: "hi".into(),
            images: None,
            streaming_behavior: Some(StreamingBehavior::Steer),
        };
        let v = serde_json::to_value(&c).unwrap();
        assert_eq!(v["type"], "prompt");
        assert_eq!(v["streamingBehavior"], "steer");
        assert!(v.get("images").is_none());
    }

    #[test]
    fn response_envelope_parses_success_and_error() {
        let r: RpcResponse = serde_json::from_str(
            r#"{"id":"r1","type":"response","command":"get_state","success":true,"data":{}}"#,
        )
        .unwrap();
        assert!(r.success);
        assert_eq!(r.command, "get_state");

        let r: RpcResponse = serde_json::from_str(
            r#"{"type":"response","command":"set_model","success":false,"error":"Model not found: invalid/model"}"#,
        )
        .unwrap();
        assert!(!r.success);
        assert!(r.error.unwrap().contains("Model not found"));
    }

    #[test]
    fn input_disposition_reads_from_response_data() {
        for (command, wire, expected) in [
            ("prompt", "started", InputDisposition::Started),
            ("prompt", "handled", InputDisposition::Handled),
            ("steer", "queued", InputDisposition::Queued),
            (
                "follow_up",
                "deferred",
                InputDisposition::Unknown("deferred".into()),
            ),
        ] {
            let r: RpcResponse = serde_json::from_value(serde_json::json!({
                "id": "r1", "type": "response", "command": command,
                "success": true, "data": {"disposition": wire}
            }))
            .unwrap();
            let got = r.disposition().expect("disposition present");
            assert_eq!(got, expected);
            assert_eq!(got.to_string(), wire);
            assert_eq!(serde_json::to_value(&got).unwrap(), wire);
        }

        // pi < 0.99.0 acknowledged these commands with no data at all.
        let legacy: RpcResponse = serde_json::from_str(
            r#"{"id":"r1","type":"response","command":"prompt","success":true}"#,
        )
        .unwrap();
        assert_eq!(legacy.disposition(), None);
    }

    #[test]
    fn agent_state_parses_credential_free_placeholder() {
        let s: AgentState = serde_json::from_value(serde_json::json!({
            "model": {"id":"unknown","name":"unknown","api":"unknown","provider":"unknown",
                      "baseUrl":"","reasoning":false,"input":[],
                      "cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0},
                      "contextWindow":0,"maxTokens":0},
            "thinkingLevel":"off","isStreaming":false,"isCompacting":false,
            "steeringMode":"one-at-a-time","followUpMode":"one-at-a-time",
            "sessionId":"01a05ec0-4e0e-7dd4-a225-2e78cccd0059",
            "autoCompactionEnabled":true,"messageCount":0,"pendingMessageCount":0
        }))
        .unwrap();
        assert_eq!(s.model.unwrap().context_window, 0);
        assert_eq!(s.message_count, 0);
    }

    #[test]
    fn raw_command_passes_through_unmodeled_types() {
        let c = RpcCommand::Raw(serde_json::json!({"type": "get_tree", "id": "r9"}));
        let line = c.to_line().unwrap();
        assert!(line.ends_with('\n'));
        assert!(line.contains("\"get_tree\""));
    }
}