1use crate::io::{Model, PiMessage};
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12
13#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
15pub struct ImageContent {
16 #[serde(rename = "type")]
18 pub content_type: String,
19 pub data: String,
21 #[serde(rename = "mimeType")]
22 pub mime_type: String,
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
27pub enum StreamingBehavior {
28 #[serde(rename = "steer")]
31 Steer,
32 #[serde(rename = "followUp")]
34 FollowUp,
35}
36
37#[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 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 #[serde(untagged)]
173 Raw(Value),
174}
175
176impl RpcCommand {
177 pub fn to_line(&self) -> Result<String, serde_json::Error> {
179 Ok(format!("{}\n", serde_json::to_string(self)?))
180 }
181}
182
183#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
185pub struct RpcResponse {
186 #[serde(default, skip_serializing_if = "Option::is_none")]
187 pub id: Option<String>,
188 #[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#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
201pub struct AgentState {
202 #[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#[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#[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}