Skip to main content

yapi_types/
rpc.rs

1//! RPC mode: commands read from stdin and the responses written to stdout, as
2//! in `modes/rpc/rpc-types.ts` of pi `v1.0.0`. Events use the JSON mode shapes
3//! in [`crate::event`].
4#![allow(
5    missing_docs,
6    reason = "fields mirror pi's TypeScript types; contracts are noted where they differ from the name"
7)]
8
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11
12use crate::message::{ImageContent, ThinkingLevel};
13use crate::session::FileEntry;
14use crate::settings::QueueMode;
15
16/// How a prompt sent while the agent runs is queued.
17#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
18#[serde(rename_all = "camelCase")]
19pub enum StreamingBehavior {
20    /// Delivered after the current tool calls.
21    Steer,
22    /// Delivered when the run would otherwise end.
23    FollowUp,
24}
25
26/// What became of a prompt once it passed preflight.
27#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
28#[serde(rename_all = "lowercase")]
29pub enum PromptDisposition {
30    /// A run started.
31    Started,
32    /// Queued behind the running agent.
33    Queued,
34    /// Consumed without a run, such as by an extension command.
35    Handled,
36}
37
38/// A command line on stdin, without its `id`.
39#[derive(Clone, Debug, PartialEq, Deserialize)]
40#[serde(
41    tag = "type",
42    rename_all = "snake_case",
43    rename_all_fields = "camelCase"
44)]
45pub enum RpcCommand {
46    Prompt {
47        message: String,
48        #[serde(default)]
49        images: Vec<ImageContent>,
50        #[serde(default)]
51        streaming_behavior: Option<StreamingBehavior>,
52    },
53    Steer {
54        message: String,
55        #[serde(default)]
56        images: Vec<ImageContent>,
57    },
58    FollowUp {
59        message: String,
60        #[serde(default)]
61        images: Vec<ImageContent>,
62    },
63    Abort,
64    ClearQueue,
65    NewSession {
66        #[serde(default)]
67        parent_session: Option<String>,
68    },
69    GetState,
70    SetModel {
71        provider: String,
72        model_id: String,
73    },
74    CycleModel,
75    GetAvailableModels,
76    SetThinkingLevel {
77        level: ThinkingLevel,
78    },
79    CycleThinkingLevel,
80    GetAvailableThinkingLevels,
81    SetSteeringMode {
82        mode: QueueMode,
83    },
84    SetFollowUpMode {
85        mode: QueueMode,
86    },
87    Compact {
88        #[serde(default)]
89        custom_instructions: Option<String>,
90    },
91    SetAutoCompaction {
92        enabled: bool,
93    },
94    SetAutoRetry {
95        enabled: bool,
96    },
97    AbortRetry,
98    Bash {
99        command: String,
100        #[serde(default)]
101        exclude_from_context: Option<bool>,
102    },
103    AbortBash,
104    GetSessionStats,
105    ExportHtml {
106        #[serde(default)]
107        output_path: Option<String>,
108    },
109    SwitchSession {
110        session_path: String,
111    },
112    Fork {
113        entry_id: String,
114    },
115    Clone,
116    GetForkMessages,
117    GetEntries {
118        #[serde(default)]
119        since: Option<String>,
120    },
121    GetTree,
122    GetLastAssistantText,
123    SetSessionName {
124        name: String,
125    },
126    GetMessages,
127    GetCommands,
128    /// Any other `type`.
129    #[serde(other)]
130    Unknown,
131}
132
133/// A response line: `data` is serialized JSON, or absent. A line without a
134/// `type` gets a response without `command`, as `JSON.stringify` drops pi's
135/// `undefined`.
136pub fn response_line(
137    id: Option<&Value>,
138    command: Option<&str>,
139    outcome: Result<Option<&str>, &str>,
140) -> String {
141    let mut line = String::from("{");
142    if let Some(id) = id
143        && let Ok(id) = crate::json::to_string(id)
144    {
145        line.push_str("\"id\":");
146        line.push_str(&id);
147        line.push(',');
148    }
149    line.push_str("\"type\":\"response\"");
150    if let Some(command) = command {
151        line.push_str(",\"command\":");
152        line.push_str(&crate::json::to_string(command).unwrap_or_else(|_| "\"\"".into()));
153    }
154    match outcome {
155        Ok(data) => {
156            line.push_str(",\"success\":true");
157            if let Some(data) = data {
158                line.push_str(",\"data\":");
159                line.push_str(data);
160            }
161        }
162        Err(error) => {
163            line.push_str(",\"success\":false,\"error\":");
164            line.push_str(&crate::json::to_string(error).unwrap_or_else(|_| "\"\"".into()));
165        }
166    }
167    line.push('}');
168    line
169}
170
171/// `get_state` data.
172#[derive(Clone, Debug, PartialEq, Serialize)]
173#[serde(rename_all = "camelCase")]
174pub struct SessionState {
175    #[serde(skip_serializing_if = "Option::is_none")]
176    pub model: Option<Value>,
177    pub thinking_level: ThinkingLevel,
178    pub is_streaming: bool,
179    pub is_compacting: bool,
180    pub steering_mode: QueueMode,
181    pub follow_up_mode: QueueMode,
182    #[serde(skip_serializing_if = "Option::is_none")]
183    pub session_file: Option<String>,
184    pub session_id: String,
185    #[serde(skip_serializing_if = "Option::is_none")]
186    pub session_name: Option<String>,
187    pub auto_compaction_enabled: bool,
188    pub message_count: usize,
189    pub pending_message_count: usize,
190}
191
192/// Token totals in [`SessionStats`].
193#[derive(Clone, Debug, Default, PartialEq, Serialize)]
194#[serde(rename_all = "camelCase")]
195pub struct TokenTotals {
196    pub input: u64,
197    pub output: u64,
198    pub cache_read: u64,
199    pub cache_write: u64,
200    pub total: u64,
201}
202
203/// How full the context window is; `tokens` and `percent` are `null` while
204/// unknown.
205#[derive(Clone, Debug, PartialEq, Serialize)]
206#[serde(rename_all = "camelCase")]
207pub struct ContextUsage {
208    pub tokens: Option<u64>,
209    pub context_window: u64,
210    pub percent: Option<f64>,
211}
212
213/// `get_session_stats` data.
214#[derive(Clone, Debug, PartialEq, Serialize)]
215#[serde(rename_all = "camelCase")]
216pub struct SessionStats {
217    #[serde(skip_serializing_if = "Option::is_none")]
218    pub session_file: Option<String>,
219    pub session_id: String,
220    pub user_messages: usize,
221    pub assistant_messages: usize,
222    pub tool_calls: usize,
223    pub tool_results: usize,
224    pub total_messages: usize,
225    pub tokens: TokenTotals,
226    pub cost: f64,
227    #[serde(skip_serializing_if = "Option::is_none")]
228    pub context_usage: Option<ContextUsage>,
229}
230
231/// What a user `!` command produced: the `bash` response data.
232#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
233#[serde(rename_all = "camelCase")]
234pub struct BashResult {
235    /// Combined, sanitized output; the tail when truncated.
236    pub output: String,
237    /// The exit code; absent when cancelled.
238    #[serde(default, skip_serializing_if = "Option::is_none")]
239    pub exit_code: Option<i32>,
240    /// The command was cancelled.
241    pub cancelled: bool,
242    /// The output was truncated.
243    pub truncated: bool,
244    /// The full output, when it was too long to keep.
245    #[serde(default, skip_serializing_if = "Option::is_none")]
246    pub full_output_path: Option<String>,
247}
248
249/// A fork point in `get_fork_messages` data.
250#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
251#[serde(rename_all = "camelCase")]
252pub struct ForkMessage {
253    pub entry_id: String,
254    pub text: String,
255}
256
257/// Where a resource was found, as pi's `SourceInfo`: `source` is `auto` for
258/// discovered files, `cli` for command-line paths and the package for
259/// package resources; `scope` is `user`, `project` or `temporary`.
260#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
261#[serde(rename_all = "camelCase")]
262pub struct SourceInfo {
263    pub path: String,
264    pub source: String,
265    pub scope: String,
266    pub origin: String,
267    #[serde(default, skip_serializing_if = "Option::is_none")]
268    pub base_dir: Option<String>,
269}
270
271/// Where a [`SlashCommand`] comes from.
272#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
273#[serde(rename_all = "lowercase")]
274pub enum CommandSource {
275    Extension,
276    Prompt,
277    Skill,
278}
279
280/// A command a prompt can invoke, in `get_commands` data.
281#[derive(Clone, Debug, PartialEq, Serialize)]
282#[serde(rename_all = "camelCase")]
283pub struct SlashCommand {
284    pub name: String,
285    #[serde(skip_serializing_if = "Option::is_none")]
286    pub description: Option<String>,
287    pub source: CommandSource,
288    pub source_info: SourceInfo,
289}
290
291/// One node of `get_tree` data, as pi's `SessionTreeNode`.
292pub struct TreeNode<'a> {
293    pub entry: &'a FileEntry,
294    /// Indexes into the same node list.
295    pub children: Vec<usize>,
296    pub label: Option<&'a str>,
297    pub label_timestamp: Option<&'a str>,
298}
299
300/// Serializes the trees under `roots` without recursion, so a long linear
301/// session cannot exhaust the stack.
302pub fn tree_json(nodes: &[TreeNode<'_>], roots: &[usize]) -> serde_json::Result<String> {
303    enum Step {
304        Open(usize),
305        Close(usize),
306        Comma,
307    }
308    let mut out = String::from("[");
309    let mut stack: Vec<Step> = Vec::new();
310    let push_list = |stack: &mut Vec<Step>, list: &[usize]| {
311        for (index, &node) in list.iter().enumerate().rev() {
312            stack.push(Step::Open(node));
313            if index > 0 {
314                stack.push(Step::Comma);
315            }
316        }
317    };
318    push_list(&mut stack, roots);
319    while let Some(step) = stack.pop() {
320        match step {
321            Step::Comma => out.push(','),
322            Step::Open(index) => {
323                let node = &nodes[index];
324                out.push_str("{\"entry\":");
325                out.push_str(&crate::json::to_string(node.entry)?);
326                out.push_str(",\"children\":[");
327                stack.push(Step::Close(index));
328                push_list(&mut stack, &node.children);
329            }
330            Step::Close(index) => {
331                out.push(']');
332                let node = &nodes[index];
333                if let Some(label) = node.label {
334                    out.push_str(",\"label\":");
335                    out.push_str(&crate::json::to_string(label)?);
336                }
337                if let Some(timestamp) = node.label_timestamp {
338                    out.push_str(",\"labelTimestamp\":");
339                    out.push_str(&crate::json::to_string(timestamp)?);
340                }
341                out.push('}');
342            }
343        }
344    }
345    out.push(']');
346    Ok(out)
347}
348
349/// An `extension_ui_response` line on stdin.
350#[derive(Clone, Debug, PartialEq, Deserialize)]
351pub struct ExtensionUiResponse {
352    pub id: String,
353    #[serde(default)]
354    pub value: Option<String>,
355    #[serde(default)]
356    pub confirmed: Option<bool>,
357    #[serde(default)]
358    pub cancelled: Option<bool>,
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364
365    #[test]
366    fn commands_parse_like_pi() {
367        let prompt: RpcCommand = serde_json::from_str(
368            r#"{"type":"prompt","message":"hi","streamingBehavior":"followUp"}"#,
369        )
370        .unwrap();
371        assert_eq!(
372            prompt,
373            RpcCommand::Prompt {
374                message: "hi".into(),
375                images: Vec::new(),
376                streaming_behavior: Some(StreamingBehavior::FollowUp),
377            }
378        );
379        let mode: RpcCommand =
380            serde_json::from_str(r#"{"type":"set_steering_mode","mode":"one-at-a-time"}"#).unwrap();
381        assert_eq!(
382            mode,
383            RpcCommand::SetSteeringMode {
384                mode: QueueMode::OneAtATime
385            }
386        );
387        let unknown: RpcCommand = serde_json::from_str(r#"{"type":"dance"}"#).unwrap();
388        assert_eq!(unknown, RpcCommand::Unknown);
389    }
390
391    #[test]
392    fn responses_match_pi() {
393        let id = Value::from("7");
394        assert_eq!(
395            response_line(Some(&id), Some("abort"), Ok(None)),
396            r#"{"id":"7","type":"response","command":"abort","success":true}"#
397        );
398        assert_eq!(
399            response_line(None, Some("cycle_model"), Ok(Some("null"))),
400            r#"{"type":"response","command":"cycle_model","success":true,"data":null}"#
401        );
402        assert_eq!(
403            response_line(Some(&Value::from(3)), Some("x"), Err("Unknown command: x")),
404            r#"{"id":3,"type":"response","command":"x","success":false,"error":"Unknown command: x"}"#
405        );
406    }
407
408    #[test]
409    fn deep_trees_serialize_without_recursion() {
410        let entry: FileEntry = serde_json::from_str(
411            r#"{"type":"label","id":"a","parentId":null,"timestamp":"t","targetId":"b"}"#,
412        )
413        .unwrap();
414        let depth = 200_000;
415        let nodes: Vec<TreeNode<'_>> = (0..depth)
416            .map(|index| TreeNode {
417                entry: &entry,
418                children: if index + 1 < depth {
419                    vec![index + 1]
420                } else {
421                    Vec::new()
422                },
423                label: (index == 0).then_some("start"),
424                label_timestamp: None,
425            })
426            .collect();
427        let json = tree_json(&nodes, &[0]).unwrap();
428        assert!(json.ends_with(&format!(
429            "{}],\"label\":\"start\"}}]",
430            "]}".repeat(depth - 1)
431        )));
432        let small = tree_json(&nodes[depth - 2..depth - 1], &[]).unwrap();
433        assert_eq!(small, "[]");
434    }
435}