1use crate::tool::ToolKind;
13use serde::{Deserialize, Serialize};
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
17#[serde(rename_all = "snake_case")]
18pub enum AgentMode {
19 Default,
21 Plan,
23}
24
25impl AgentMode {
26 pub fn as_str(&self) -> &'static str {
27 match self {
28 Self::Default => "default",
29 Self::Plan => "plan",
30 }
31 }
32
33 pub fn restricts_tools(&self) -> bool {
35 matches!(self, Self::Plan)
36 }
37}
38
39impl Default for AgentMode {
40 fn default() -> Self {
41 Self::Default
42 }
43}
44
45impl std::fmt::Display for AgentMode {
46 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47 write!(f, "{}", self.as_str())
48 }
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct ProposedPlan {
54 pub title: String,
56 pub steps: Vec<String>,
58}
59
60impl ProposedPlan {
61 pub fn new(title: String, steps: Vec<String>) -> Self {
62 Self { title, steps }
63 }
64
65 pub fn is_empty(&self) -> bool {
66 self.steps.is_empty() && self.title.is_empty()
67 }
68}
69
70#[derive(Debug, Default)]
84pub struct ProposedPlanParser {
85 buffer: String,
86 in_plan: bool,
87 plan_title: Option<String>,
88 plan_body: String,
89 completed_plans: Vec<ProposedPlan>,
90}
91
92impl ProposedPlanParser {
93 pub fn new() -> Self {
94 Self::default()
95 }
96
97 pub fn push_text(&mut self, text: &str) -> Vec<ProposedPlan> {
100 self.buffer.push_str(text);
101 self.drain_completed()
102 }
103
104 pub fn drain(&mut self) -> Vec<ProposedPlan> {
106 if self.in_plan {
108 self.plan_body.push_str(&self.buffer);
109 self.buffer.clear();
110 if !self.plan_body.is_empty() {
111 let plan = self.finalize_plan();
112 self.completed_plans.push(plan);
113 }
114 }
115 std::mem::take(&mut self.completed_plans)
116 }
117
118 pub fn is_in_plan(&self) -> bool {
120 self.in_plan
121 }
122
123 pub fn partial_body(&self) -> &str {
125 if self.in_plan { &self.plan_body } else { "" }
126 }
127
128 pub fn partial_title(&self) -> Option<&str> {
130 if self.in_plan {
131 self.plan_title.as_deref()
132 } else {
133 None
134 }
135 }
136
137 fn drain_completed(&mut self) -> Vec<ProposedPlan> {
138 loop {
139 if self.in_plan {
140 if let Some(pos) = self.buffer.find("</proposed_plan>") {
142 self.plan_body.push_str(&self.buffer[..pos]);
143 self.buffer = self.buffer[pos + "</proposed_plan>".len()..].to_string();
144 self.in_plan = false;
145 let plan = self.finalize_plan();
146 self.completed_plans.push(plan);
147 } else {
148 let tag = "</proposed_plan>";
151 if self.buffer.len() >= tag.len() {
152 let safe = self.buffer.len() - tag.len() + 1;
153 self.plan_body.push_str(&self.buffer[..safe]);
154 self.buffer = self.buffer[safe..].to_string();
155 } else {
156 self.plan_body.push_str(&self.buffer);
158 self.buffer.clear();
159 }
160 break;
161 }
162 } else {
163 if let Some(tag_info) = self.find_opening_tag() {
165 self.buffer = self.buffer[tag_info.consume_len..].to_string();
166 self.in_plan = true;
167 self.plan_title = tag_info.title;
168 self.plan_body.clear();
169 } else {
170 let tag = "<proposed_plan";
173 let safe = self.buffer.len().saturating_sub(tag.len() - 1);
174 self.buffer = self.buffer[safe..].to_string();
175 break;
176 }
177 }
178 }
179 std::mem::take(&mut self.completed_plans)
180 }
181
182 fn finalize_plan(&mut self) -> ProposedPlan {
183 let title = self.plan_title.take().unwrap_or_default();
184 let body = std::mem::take(&mut self.plan_body);
185 let steps = parse_plan_steps(&body);
186 ProposedPlan::new(title, steps)
187 }
188
189 fn find_opening_tag(&self) -> Option<OpeningTagInfo> {
190 let start = self.buffer.find("<proposed_plan")?;
191 let rest = &self.buffer[start..];
192 let tag_end = rest.find('>')?;
193 let tag_content = &rest[..tag_end];
194 let consume_len = start + tag_end + 1;
195
196 let title = tag_content
197 .find("title=\"")
198 .and_then(|t_pos| {
199 let value_start = t_pos + "title=\"".len();
200 tag_content[value_start..]
201 .find('"')
202 .map(|end| tag_content[value_start..value_start + end].to_string())
203 })
204 .or_else(|| {
205 tag_content.find("title='").and_then(|t_pos| {
206 let value_start = t_pos + "title='".len();
207 tag_content[value_start..]
208 .find('\'')
209 .map(|end| tag_content[value_start..value_start + end].to_string())
210 })
211 });
212
213 Some(OpeningTagInfo { consume_len, title })
214 }
215}
216
217struct OpeningTagInfo {
218 consume_len: usize,
219 title: Option<String>,
220}
221
222fn parse_plan_steps(body: &str) -> Vec<String> {
225 body.lines()
226 .map(|line| line.trim())
227 .filter(|line| !line.is_empty())
228 .map(|line| {
229 let stripped = if let Some(rest) = line.strip_prefix("- ") {
231 rest.to_string()
232 } else if let Some(rest) = line.strip_prefix("* ") {
233 rest.to_string()
234 } else if let Some(rest) = line.strip_prefix("• ") {
235 rest.to_string()
236 } else {
237 let chars: Vec<char> = line.chars().collect();
239 let mut idx = 0;
240 while idx < chars.len() && chars[idx].is_ascii_digit() {
241 idx += 1;
242 }
243 if idx > 0 && idx < chars.len() && (chars[idx] == '.' || chars[idx] == ')') {
244 idx += 1;
245 while idx < chars.len() && chars[idx] == ' ' {
246 idx += 1;
247 }
248 line[idx..].to_string()
249 } else {
250 line.to_string()
251 }
252 };
253 if stripped.is_empty() {
254 line.to_string()
255 } else {
256 stripped
257 }
258 })
259 .collect()
260}
261
262pub fn is_tool_allowed_in_plan_mode(kind: ToolKind) -> bool {
269 is_tool_allowed_in_plan_mode_named("", kind)
270}
271
272pub fn is_tool_allowed_in_plan_mode_named(name: &str, kind: ToolKind) -> bool {
274 match kind {
275 ToolKind::Read => true,
276 ToolKind::Write => matches!(name, "write_file" | "edit" | "multiedit" | "write"),
277 ToolKind::Custom => matches!(name, "plan" | "question"),
278 ToolKind::Command => false,
279 }
280}
281
282#[cfg(test)]
283mod tests {
284 use super::*;
285
286 #[test]
287 fn parser_extracts_simple_plan() {
288 let mut parser = ProposedPlanParser::new();
289 let plans = parser.push_text(
290 "Let me analyze this.\n\
291 <proposed_plan title=\"Fix the bug\">\n\
292 1. Read the file\n\
293 2. Fix the function\n\
294 3. Run tests\n\
295 </proposed_plan>\n\
296 Done.",
297 );
298 assert_eq!(plans.len(), 1);
299 assert_eq!(plans[0].title, "Fix the bug");
300 assert_eq!(plans[0].steps.len(), 3);
301 assert_eq!(plans[0].steps[0], "Read the file");
302 assert_eq!(plans[0].steps[1], "Fix the function");
303 assert_eq!(plans[0].steps[2], "Run tests");
304 }
305
306 #[test]
307 fn parser_extracts_plan_without_title() {
308 let mut parser = ProposedPlanParser::new();
309 let plans = parser.push_text(
310 "<proposed_plan>\n\
311 - Step one\n\
312 - Step two\n\
313 </proposed_plan>",
314 );
315 assert_eq!(plans.len(), 1);
316 assert_eq!(plans[0].title, "");
317 assert_eq!(plans[0].steps, vec!["Step one", "Step two"]);
318 }
319
320 #[test]
321 fn parser_handles_chunked_stream() {
322 let mut parser = ProposedPlanParser::new();
323 let chunks = [
324 "Let me think.\n<propos",
325 "ed_plan title=\"My Plan\">\n",
326 "1. First step\n2. Second ",
327 "step\n</proposed_",
328 "plan>\nDone.",
329 ];
330
331 let mut all_plans = Vec::new();
332 for chunk in &chunks {
333 all_plans.extend(parser.push_text(chunk));
334 }
335
336 assert_eq!(all_plans.len(), 1);
337 assert_eq!(all_plans[0].title, "My Plan");
338 assert_eq!(all_plans[0].steps, vec!["First step", "Second step"]);
339 }
340
341 #[test]
342 fn parser_handles_multiple_plans() {
343 let mut parser = ProposedPlanParser::new();
344 let plans = parser.push_text(
345 "<proposed_plan title=\"Plan A\">\n1. A1\n</proposed_plan>\n\
346 <proposed_plan title=\"Plan B\">\n1. B1\n</proposed_plan>",
347 );
348 assert_eq!(plans.len(), 2);
349 assert_eq!(plans[0].title, "Plan A");
350 assert_eq!(plans[1].title, "Plan B");
351 }
352
353 #[test]
354 fn parser_drain_unclosed_plan() {
355 let mut parser = ProposedPlanParser::new();
356 parser.push_text("<proposed_plan title=\"Unclosed\">\n1. Step\n");
357 let plans = parser.drain();
358 assert_eq!(plans.len(), 1);
359 assert_eq!(plans[0].title, "Unclosed");
360 assert_eq!(plans[0].steps, vec!["Step"]);
361 }
362
363 #[test]
364 fn parser_partial_preview() {
365 let mut parser = ProposedPlanParser::new();
366 parser.push_text("<proposed_plan title=\"Live\">\n1. First");
367 assert!(parser.is_in_plan());
368 assert_eq!(parser.partial_title(), Some("Live"));
369 assert!(parser.partial_body().contains("First"));
370 }
371
372 #[test]
373 fn parser_no_plan_in_regular_text() {
374 let mut parser = ProposedPlanParser::new();
375 let plans = parser.push_text("Just regular text without any plan tags.");
376 assert!(plans.is_empty());
377 assert!(!parser.is_in_plan());
378 }
379
380 #[test]
381 fn parse_plan_steps_strips_numbering() {
382 let steps = parse_plan_steps("1. First\n2. Second\n3. Third");
383 assert_eq!(steps, vec!["First", "Second", "Third"]);
384 }
385
386 #[test]
387 fn parse_plan_steps_strips_bullets() {
388 let steps = parse_plan_steps("- Alpha\n* Beta");
389 assert_eq!(steps, vec!["Alpha", "Beta"]);
390 }
391
392 #[test]
393 fn agent_mode_default_is_default() {
394 assert_eq!(AgentMode::default(), AgentMode::Default);
395 }
396
397 #[test]
398 fn agent_mode_plan_restricts_tools() {
399 assert!(AgentMode::Plan.restricts_tools());
400 assert!(!AgentMode::Default.restricts_tools());
401 }
402
403 #[test]
404 fn plan_mode_allows_read_plan_write_and_question() {
405 assert!(is_tool_allowed_in_plan_mode_named(
406 "read_file",
407 ToolKind::Read
408 ));
409 assert!(is_tool_allowed_in_plan_mode_named("search", ToolKind::Read));
410 assert!(is_tool_allowed_in_plan_mode_named("plan", ToolKind::Read));
411 assert!(is_tool_allowed_in_plan_mode_named(
412 "write_file",
413 ToolKind::Write
414 ));
415 assert!(is_tool_allowed_in_plan_mode_named("edit", ToolKind::Write));
416 assert!(is_tool_allowed_in_plan_mode_named(
417 "question",
418 ToolKind::Custom
419 ));
420 assert!(!is_tool_allowed_in_plan_mode_named(
421 "bash",
422 ToolKind::Command
423 ));
424 assert!(!is_tool_allowed_in_plan_mode_named(
425 "subagent",
426 ToolKind::Command
427 ));
428 assert!(!is_tool_allowed_in_plan_mode_named(
429 "apply_patch",
430 ToolKind::Write
431 ));
432 }
433}