Skip to main content

pi_codes/
rpc.rs

1//! The `pi --mode rpc` command protocol: JSON objects on stdin (one per
2//! line, LF-delimited), `{"type":"response",...}` acknowledgements and
3//! streamed events on stdout.
4//!
5//! Framing contract (from `docs/rpc.md`): split records on `\n` ONLY,
6//! tolerate a trailing `\r` — never use a splitter that also breaks on
7//! U+2028/U+2029, which are legal inside JSON strings.
8
9use crate::io::{Model, PiMessage};
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12
13/// An image attached to a `prompt`/`steer`/`follow_up` command.
14#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
15pub struct ImageContent {
16    /// Always `"image"`.
17    #[serde(rename = "type")]
18    pub content_type: String,
19    /// Base64-encoded bytes.
20    pub data: String,
21    #[serde(rename = "mimeType")]
22    pub mime_type: String,
23}
24
25/// How a `prompt` sent mid-stream should queue.
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
27pub enum StreamingBehavior {
28    /// Deliver after the current assistant turn's tool calls, before the
29    /// next LLM call.
30    #[serde(rename = "steer")]
31    Steer,
32    /// Deliver only once the agent stops.
33    #[serde(rename = "followUp")]
34    FollowUp,
35}
36
37/// A client → agent command. Serializes with the wire's `type` tag; the
38/// optional `id` correlates the matching [`RpcResponse`].
39#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
40#[serde(tag = "type", rename_all = "snake_case")]
41pub enum RpcCommand {
42    Prompt {
43        #[serde(skip_serializing_if = "Option::is_none")]
44        id: Option<String>,
45        message: String,
46        #[serde(skip_serializing_if = "Option::is_none")]
47        images: Option<Vec<ImageContent>>,
48        #[serde(rename = "streamingBehavior", skip_serializing_if = "Option::is_none")]
49        streaming_behavior: Option<StreamingBehavior>,
50    },
51    Steer {
52        #[serde(skip_serializing_if = "Option::is_none")]
53        id: Option<String>,
54        message: String,
55        #[serde(skip_serializing_if = "Option::is_none")]
56        images: Option<Vec<ImageContent>>,
57    },
58    FollowUp {
59        #[serde(skip_serializing_if = "Option::is_none")]
60        id: Option<String>,
61        message: String,
62        #[serde(skip_serializing_if = "Option::is_none")]
63        images: Option<Vec<ImageContent>>,
64    },
65    Abort {
66        #[serde(skip_serializing_if = "Option::is_none")]
67        id: Option<String>,
68    },
69    ClearQueue {
70        #[serde(skip_serializing_if = "Option::is_none")]
71        id: Option<String>,
72    },
73    NewSession {
74        #[serde(skip_serializing_if = "Option::is_none")]
75        id: Option<String>,
76    },
77    GetState {
78        #[serde(skip_serializing_if = "Option::is_none")]
79        id: Option<String>,
80    },
81    GetMessages {
82        #[serde(skip_serializing_if = "Option::is_none")]
83        id: Option<String>,
84    },
85    SetModel {
86        #[serde(skip_serializing_if = "Option::is_none")]
87        id: Option<String>,
88        provider: String,
89        #[serde(rename = "modelId")]
90        model_id: String,
91    },
92    CycleModel {
93        #[serde(skip_serializing_if = "Option::is_none")]
94        id: Option<String>,
95    },
96    GetAvailableModels {
97        #[serde(skip_serializing_if = "Option::is_none")]
98        id: Option<String>,
99    },
100    SetThinkingLevel {
101        #[serde(skip_serializing_if = "Option::is_none")]
102        id: Option<String>,
103        level: String,
104    },
105    CycleThinkingLevel {
106        #[serde(skip_serializing_if = "Option::is_none")]
107        id: Option<String>,
108    },
109    GetAvailableThinkingLevels {
110        #[serde(skip_serializing_if = "Option::is_none")]
111        id: Option<String>,
112    },
113    SetSteeringMode {
114        #[serde(skip_serializing_if = "Option::is_none")]
115        id: Option<String>,
116        mode: String,
117    },
118    SetFollowUpMode {
119        #[serde(skip_serializing_if = "Option::is_none")]
120        id: Option<String>,
121        mode: String,
122    },
123    Compact {
124        #[serde(skip_serializing_if = "Option::is_none")]
125        id: Option<String>,
126        #[serde(skip_serializing_if = "Option::is_none")]
127        instructions: Option<String>,
128    },
129    SetAutoCompaction {
130        #[serde(skip_serializing_if = "Option::is_none")]
131        id: Option<String>,
132        enabled: bool,
133    },
134    SetAutoRetry {
135        #[serde(skip_serializing_if = "Option::is_none")]
136        id: Option<String>,
137        enabled: bool,
138    },
139    AbortRetry {
140        #[serde(skip_serializing_if = "Option::is_none")]
141        id: Option<String>,
142    },
143    /// Run a shell command outside the model; output joins context on the
144    /// NEXT prompt. Streams `bash_execution_update` events carrying this
145    /// command's `id`.
146    Bash {
147        #[serde(skip_serializing_if = "Option::is_none")]
148        id: Option<String>,
149        command: String,
150    },
151    AbortBash {
152        #[serde(skip_serializing_if = "Option::is_none")]
153        id: Option<String>,
154    },
155    GetSessionStats {
156        #[serde(skip_serializing_if = "Option::is_none")]
157        id: Option<String>,
158    },
159    ExportHtml {
160        #[serde(skip_serializing_if = "Option::is_none")]
161        id: Option<String>,
162        #[serde(skip_serializing_if = "Option::is_none")]
163        path: Option<String>,
164    },
165    SetSessionName {
166        #[serde(skip_serializing_if = "Option::is_none")]
167        id: Option<String>,
168        name: String,
169    },
170    /// Anything this crate doesn't model yet — serialized as its raw
171    /// object so new commands work without a crate release.
172    #[serde(untagged)]
173    Raw(Value),
174}
175
176impl RpcCommand {
177    /// Serialize to one LF-terminated wire line.
178    pub fn to_line(&self) -> Result<String, serde_json::Error> {
179        Ok(format!("{}\n", serde_json::to_string(self)?))
180    }
181}
182
183/// The `{"type":"response"}` acknowledgement for one command.
184#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
185pub struct RpcResponse {
186    #[serde(default, skip_serializing_if = "Option::is_none")]
187    pub id: Option<String>,
188    /// Which command this answers (`"parse"` for unparseable input).
189    #[serde(default)]
190    pub command: String,
191    #[serde(default)]
192    pub success: bool,
193    #[serde(default, skip_serializing_if = "Option::is_none")]
194    pub data: Option<Value>,
195    #[serde(default, skip_serializing_if = "Option::is_none")]
196    pub error: Option<String>,
197}
198
199/// Typed view of `get_state` response data.
200#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
201pub struct AgentState {
202    /// `None` when no provider resolved; may also be a placeholder with
203    /// `id: "unknown"` — see [`Model`].
204    #[serde(default, skip_serializing_if = "Option::is_none")]
205    pub model: Option<Model>,
206    #[serde(rename = "thinkingLevel", default)]
207    pub thinking_level: String,
208    #[serde(rename = "isStreaming", default)]
209    pub is_streaming: bool,
210    #[serde(rename = "isCompacting", default)]
211    pub is_compacting: bool,
212    #[serde(rename = "steeringMode", default)]
213    pub steering_mode: String,
214    #[serde(rename = "followUpMode", default)]
215    pub follow_up_mode: String,
216    #[serde(
217        rename = "sessionFile",
218        default,
219        skip_serializing_if = "Option::is_none"
220    )]
221    pub session_file: Option<String>,
222    #[serde(rename = "sessionId", default)]
223    pub session_id: String,
224    #[serde(
225        rename = "sessionName",
226        default,
227        skip_serializing_if = "Option::is_none"
228    )]
229    pub session_name: Option<String>,
230    #[serde(rename = "autoCompactionEnabled", default)]
231    pub auto_compaction_enabled: bool,
232    #[serde(rename = "messageCount", default)]
233    pub message_count: i64,
234    #[serde(rename = "pendingMessageCount", default)]
235    pub pending_message_count: i64,
236    #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
237    pub extra: serde_json::Map<String, Value>,
238}
239
240/// Typed view of `bash` response data.
241#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
242pub struct BashResult {
243    #[serde(default)]
244    pub output: String,
245    #[serde(rename = "exitCode", default, skip_serializing_if = "Option::is_none")]
246    pub exit_code: Option<i64>,
247    #[serde(default)]
248    pub cancelled: bool,
249    #[serde(default)]
250    pub truncated: bool,
251    #[serde(
252        rename = "fullOutputPath",
253        default,
254        skip_serializing_if = "Option::is_none"
255    )]
256    pub full_output_path: Option<String>,
257}
258
259/// Typed view of `get_messages` response data.
260#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
261pub struct Messages {
262    #[serde(default)]
263    pub messages: Vec<PiMessage>,
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    #[test]
271    fn commands_serialize_with_snake_case_type_tags() {
272        let c = RpcCommand::GetState {
273            id: Some("r1".into()),
274        };
275        let v = serde_json::to_value(&c).unwrap();
276        assert_eq!(v["type"], "get_state");
277        assert_eq!(v["id"], "r1");
278
279        let c = RpcCommand::Prompt {
280            id: None,
281            message: "hi".into(),
282            images: None,
283            streaming_behavior: Some(StreamingBehavior::Steer),
284        };
285        let v = serde_json::to_value(&c).unwrap();
286        assert_eq!(v["type"], "prompt");
287        assert_eq!(v["streamingBehavior"], "steer");
288        assert!(v.get("images").is_none());
289    }
290
291    #[test]
292    fn response_envelope_parses_success_and_error() {
293        let r: RpcResponse = serde_json::from_str(
294            r#"{"id":"r1","type":"response","command":"get_state","success":true,"data":{}}"#,
295        )
296        .unwrap();
297        assert!(r.success);
298        assert_eq!(r.command, "get_state");
299
300        let r: RpcResponse = serde_json::from_str(
301            r#"{"type":"response","command":"set_model","success":false,"error":"Model not found: invalid/model"}"#,
302        )
303        .unwrap();
304        assert!(!r.success);
305        assert!(r.error.unwrap().contains("Model not found"));
306    }
307
308    #[test]
309    fn agent_state_parses_credential_free_placeholder() {
310        let s: AgentState = serde_json::from_value(serde_json::json!({
311            "model": {"id":"unknown","name":"unknown","api":"unknown","provider":"unknown",
312                      "baseUrl":"","reasoning":false,"input":[],
313                      "cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0},
314                      "contextWindow":0,"maxTokens":0},
315            "thinkingLevel":"off","isStreaming":false,"isCompacting":false,
316            "steeringMode":"one-at-a-time","followUpMode":"one-at-a-time",
317            "sessionId":"01a05ec0-4e0e-7dd4-a225-2e78cccd0059",
318            "autoCompactionEnabled":true,"messageCount":0,"pendingMessageCount":0
319        }))
320        .unwrap();
321        assert_eq!(s.model.unwrap().context_window, 0);
322        assert_eq!(s.message_count, 0);
323    }
324
325    #[test]
326    fn raw_command_passes_through_unmodeled_types() {
327        let c = RpcCommand::Raw(serde_json::json!({"type": "get_tree", "id": "r9"}));
328        let line = c.to_line().unwrap();
329        assert!(line.ends_with('\n'));
330        assert!(line.contains("\"get_tree\""));
331    }
332}