1#![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#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
18#[serde(rename_all = "camelCase")]
19pub enum StreamingBehavior {
20 Steer,
22 FollowUp,
24}
25
26#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
28#[serde(rename_all = "lowercase")]
29pub enum PromptDisposition {
30 Started,
32 Queued,
34 Handled,
36}
37
38#[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 #[serde(other)]
130 Unknown,
131}
132
133pub 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#[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#[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#[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#[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#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
233#[serde(rename_all = "camelCase")]
234pub struct BashResult {
235 pub output: String,
237 #[serde(default, skip_serializing_if = "Option::is_none")]
239 pub exit_code: Option<i32>,
240 pub cancelled: bool,
242 pub truncated: bool,
244 #[serde(default, skip_serializing_if = "Option::is_none")]
246 pub full_output_path: Option<String>,
247}
248
249#[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#[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#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
273#[serde(rename_all = "lowercase")]
274pub enum CommandSource {
275 Extension,
276 Prompt,
277 Skill,
278}
279
280#[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
291pub struct TreeNode<'a> {
293 pub entry: &'a FileEntry,
294 pub children: Vec<usize>,
296 pub label: Option<&'a str>,
297 pub label_timestamp: Option<&'a str>,
298}
299
300pub 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#[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}