1use crate::error::AgentError;
8use crate::types::*;
9use std::sync::{
10 OnceLock,
11 atomic::{AtomicBool, Ordering},
12};
13
14pub const ENTER_PLAN_MODE_TOOL_NAME: &str = "EnterPlanMode";
15pub const EXIT_PLAN_MODE_TOOL_NAME: &str = "ExitPlanModeV2";
16
17static IN_PLAN_MODE: OnceLock<AtomicBool> = OnceLock::new();
19
20fn is_in_plan_mode() -> bool {
21 IN_PLAN_MODE
22 .get_or_init(|| AtomicBool::new(false))
23 .load(Ordering::SeqCst)
24}
25
26fn set_plan_mode(val: bool) {
27 IN_PLAN_MODE
28 .get_or_init(|| AtomicBool::new(false))
29 .store(val, Ordering::SeqCst);
30}
31
32static CURRENT_PLAN: OnceLock<std::sync::Mutex<String>> = OnceLock::new();
34
35fn get_plan() -> String {
36 CURRENT_PLAN
37 .get_or_init(|| std::sync::Mutex::new(String::new()))
38 .lock()
39 .unwrap()
40 .clone()
41}
42
43fn set_plan(plan: String) {
44 *CURRENT_PLAN
45 .get_or_init(|| std::sync::Mutex::new(String::new()))
46 .lock()
47 .unwrap() = plan;
48}
49
50pub struct EnterPlanModeTool;
52
53impl EnterPlanModeTool {
54 pub fn new() -> Self {
55 Self
56 }
57
58 pub fn name(&self) -> &str {
59 ENTER_PLAN_MODE_TOOL_NAME
60 }
61
62 pub fn description(&self) -> &str {
63 "Enter structured planning mode. Switches from implementation to planning workflow where you can explore the codebase and design an implementation approach."
64 }
65
66 pub fn user_facing_name(&self, _input: Option<&serde_json::Value>) -> String {
67 "EnterPlanMode".to_string()
68 }
69
70 pub fn get_tool_use_summary(&self, _input: Option<&serde_json::Value>) -> Option<String> {
71 None
72 }
73
74 pub fn render_tool_result_message(
75 &self,
76 content: &serde_json::Value,
77 ) -> Option<String> {
78 content["content"].as_str().map(|s| s.to_string())
79 }
80
81 pub fn input_schema(&self) -> ToolInputSchema {
82 ToolInputSchema {
83 schema_type: "object".to_string(),
84 properties: serde_json::json!({
85 "allowedPrompts": {
86 "type": "array",
87 "items": { "type": "string" },
88 "description": "Prompt-based permissions needed to implement the plan. These are shell command patterns that will be allowed during plan execution."
89 }
90 }),
91 required: None,
92 }
93 }
94
95 pub async fn execute(
96 &self,
97 input: serde_json::Value,
98 _context: &ToolContext,
99 ) -> Result<ToolResult, AgentError> {
100 let allowed = input["allowedPrompts"]
101 .as_array()
102 .map(|arr| arr.iter().filter_map(|v| v.as_str()).collect::<Vec<_>>())
103 .unwrap_or_default();
104
105 set_plan_mode(true);
106
107 let response = if allowed.is_empty() {
108 "Switched to plan mode. You can now explore the codebase and design an implementation plan. \
109 When ready, use ExitPlanMode to present the plan for user approval."
110 .to_string()
111 } else {
112 format!(
113 "Switched to plan mode with permissions: {}.\n\
114 You can now explore the codebase and design an implementation plan.\n\
115 The following shell command patterns will be allowed during plan execution:\n\
116 - {}\n\
117 When ready, use ExitPlanMode to present the plan for user approval.",
118 allowed.len(),
119 allowed.join("\n- ")
120 )
121 };
122
123 Ok(ToolResult {
124 result_type: "text".to_string(),
125 tool_use_id: "enter_plan_mode".to_string(),
126 content: response,
127 is_error: Some(false),
128 was_persisted: None,
129 })
130 }
131}
132
133impl Default for EnterPlanModeTool {
134 fn default() -> Self {
135 Self::new()
136 }
137}
138
139pub struct ExitPlanModeTool;
141
142impl ExitPlanModeTool {
143 pub fn new() -> Self {
144 Self
145 }
146
147 pub fn name(&self) -> &str {
148 EXIT_PLAN_MODE_TOOL_NAME
149 }
150
151 pub fn description(&self) -> &str {
152 "Exit plan mode and present the plan for user approval. Call this when you have finished designing the implementation approach."
153 }
154
155 pub fn user_facing_name(&self, _input: Option<&serde_json::Value>) -> String {
156 "ExitPlanMode".to_string()
157 }
158
159 pub fn get_tool_use_summary(&self, _input: Option<&serde_json::Value>) -> Option<String> {
160 None
161 }
162
163 pub fn render_tool_result_message(
164 &self,
165 content: &serde_json::Value,
166 ) -> Option<String> {
167 content["content"].as_str().map(|s| s.to_string())
168 }
169
170 pub fn input_schema(&self) -> ToolInputSchema {
171 ToolInputSchema {
172 schema_type: "object".to_string(),
173 properties: serde_json::json!({}),
174 required: None,
175 }
176 }
177
178 pub async fn execute(
179 &self,
180 _input: serde_json::Value,
181 _context: &ToolContext,
182 ) -> Result<ToolResult, AgentError> {
183 if !is_in_plan_mode() {
184 return Ok(ToolResult {
185 result_type: "text".to_string(),
186 tool_use_id: "".to_string(),
187 content: "Error: Not currently in plan mode. Use EnterPlanMode first.".to_string(),
188 is_error: Some(true),
189 was_persisted: None,
190 });
191 }
192
193 set_plan_mode(false);
194
195 let plan = get_plan();
196 let response = if plan.is_empty() {
197 "Exiting plan mode. No plan has been created yet.\n\
198 You should first explore the codebase and design an implementation approach\n\
199 before exiting plan mode."
200 .to_string()
201 } else {
202 format!(
203 "Plan submitted for user approval.\n\
204 The plan will be presented to the user for review and approval.\n\
205 Once approved, you can proceed with implementation.\n\n\
206 Plan summary:\n{}",
207 plan
208 )
209 };
210
211 Ok(ToolResult {
212 result_type: "text".to_string(),
213 tool_use_id: "exit_plan_mode".to_string(),
214 content: response,
215 is_error: Some(false),
216 was_persisted: None,
217 })
218 }
219}
220
221impl Default for ExitPlanModeTool {
222 fn default() -> Self {
223 Self::new()
224 }
225}
226
227pub fn reset_plan_for_testing() {
229 IN_PLAN_MODE
231 .get_or_init(|| AtomicBool::new(false))
232 .store(false, Ordering::SeqCst);
233 if let Some(plan_mutex) = CURRENT_PLAN.get() {
235 let mut plan = plan_mutex.lock().unwrap();
236 plan.clear();
237 }
238}
239
240#[cfg(test)]
241mod tests {
242 use super::*;
243
244 #[test]
245 fn test_enter_plan_mode_name() {
246 let tool = EnterPlanModeTool::new();
247 assert_eq!(tool.name(), ENTER_PLAN_MODE_TOOL_NAME);
248 }
249
250 #[test]
251 fn test_exit_plan_mode_name() {
252 let tool = ExitPlanModeTool::new();
253 assert_eq!(tool.name(), EXIT_PLAN_MODE_TOOL_NAME);
254 }
255
256 #[test]
257 fn test_enter_plan_mode_schema() {
258 let tool = EnterPlanModeTool::new();
259 let schema = tool.input_schema();
260 assert!(schema.properties.get("allowedPrompts").is_some());
261 }
262
263 #[tokio::test]
264 async fn test_enter_plan_mode_sets_flag() {
265 let tool = EnterPlanModeTool::new();
266 let input = serde_json::json!({});
267 let context = ToolContext::default();
268 let result = tool.execute(input, &context).await;
269 assert!(result.is_ok());
270 assert!(is_in_plan_mode());
271 }
272
273 #[tokio::test]
274 async fn test_exit_plan_mode_clears_flag() {
275 set_plan_mode(true);
276 let tool = ExitPlanModeTool::new();
277 let input = serde_json::json!({});
278 let context = ToolContext::default();
279 let result = tool.execute(input, &context).await;
280 assert!(result.is_ok());
281 assert!(!is_in_plan_mode());
282 }
283
284 #[tokio::test]
285 async fn test_exit_plan_mode_not_in_mode() {
286 set_plan_mode(false);
287 let tool = ExitPlanModeTool::new();
288 let input = serde_json::json!({});
289 let context = ToolContext::default();
290 let result = tool.execute(input, &context).await;
291 assert!(result.is_ok());
292 let content = result.unwrap().content;
293 assert!(content.contains("Not currently in plan mode"));
294 }
295
296 #[tokio::test]
297 async fn test_enter_plan_mode_with_permissions() {
298 let tool = EnterPlanModeTool::new();
299 let input = serde_json::json!({
300 "allowedPrompts": ["npm run build", "git commit"]
301 });
302 let context = ToolContext::default();
303 let result = tool.execute(input, &context).await;
304 assert!(result.is_ok());
305 let content = result.unwrap().content;
306 assert!(content.contains("permissions"));
307 assert!(content.contains("npm run build"));
308 }
309}