Coverage for src / osiris_cli / plan_mode.py: 0%

1""" 

2Plan Mode for Osiris CLI 

3 

4Enables systematic planning and execution of complex multi-step tasks. 

5Inspired by Claude Code's plan mode for better task decomposition. 

6""" 

7 

8import json 

9from pathlib import Path 

10from datetime import datetime 

11from typing import List, Dict, Optional, Any 

12from dataclasses import dataclass, asdict 

13from enum import Enum 

14 

15from rich.console import Console 

16from rich.table import Table 

17from rich.panel import Panel 

18from rich.markdown import Markdown 

19 

20from .config import CONFIG_DIR 

21from .logger import get_logger 

22 

23logger = get_logger() 

24console = Console() 

25 

26# Plans directory 

27PLANS_DIR = CONFIG_DIR / "plans" 

28PLANS_DIR.mkdir(parents=True, exist_ok=True) 

29 

30 

31class StepStatus(Enum): 

32 """Status of a plan step""" 

33 PENDING = "pending" 

34 IN_PROGRESS = "in_progress" 

35 COMPLETED = "completed" 

36 FAILED = "failed" 

37 SKIPPED = "skipped" 

38 

39 

40@dataclass 

41class PlanStep: 

42 """A single step in a plan""" 

43 id: int 

44 description: str 

45 status: StepStatus 

46 result: Optional[str] = None 

47 error: Optional[str] = None 

48 started_at: Optional[str] = None 

49 completed_at: Optional[str] = None 

50 

51 def to_dict(self) -> Dict: 

52 data = asdict(self) 

53 data['status'] = self.status.value 

54 return data 

55 

56 @classmethod 

57 def from_dict(cls, data: Dict) -> 'PlanStep': 

58 data['status'] = StepStatus(data['status']) 

59 return cls(**data) 

60 

61 

62@dataclass 

63class Plan: 

64 """A complete execution plan""" 

65 id: str 

66 title: str 

67 description: str 

68 steps: List[PlanStep] 

69 created_at: str 

70 updated_at: str 

71 status: str # "draft", "in_progress", "completed", "failed" 

72 current_step: int = 0 

73 

74 def to_dict(self) -> Dict: 

75 return { 

76 "id": self.id, 

77 "title": self.title, 

78 "description": self.description, 

79 "steps": [step.to_dict() for step in self.steps], 

80 "created_at": self.created_at, 

81 "updated_at": self.updated_at, 

82 "status": self.status, 

83 "current_step": self.current_step 

84 } 

85 

86 @classmethod 

87 def from_dict(cls, data: Dict) -> 'Plan': 

88 data['steps'] = [PlanStep.from_dict(s) for s in data['steps']] 

89 return cls(**data) 

90 

91 

92class PlanMode: 

93 """ 

94 Plan Mode manager for complex task execution. 

95  

96 Features: 

97 - AI-generated step-by-step plans 

98 - Step execution with checkpoints 

99 - Rollback capability 

100 - Progress tracking 

101 - Plan templates 

102 """ 

103 

104 def __init__(self): 

105 self.current_plan: Optional[Plan] = None 

106 logger.info("Plan Mode initialized") 

107 

108 def create_plan( 

109 self, 

110 title: str, 

111 description: str, 

112 steps: List[str] 

113 ) -> Plan: 

114 """ 

115 Create a new plan. 

116  

117 Args: 

118 title: Plan title 

119 description: Plan description 

120 steps: List of step descriptions 

121  

122 Returns: 

123 Created plan 

124 """ 

125 plan_id = f"plan_{datetime.now().strftime('%Y%m%d_%H%M%S')}" 

126 

127 plan_steps = [ 

128 PlanStep( 

129 id=i + 1, 

130 description=desc, 

131 status=StepStatus.PENDING 

132 ) 

133 for i, desc in enumerate(steps) 

134 ] 

135 

136 plan = Plan( 

137 id=plan_id, 

138 title=title, 

139 description=description, 

140 steps=plan_steps, 

141 created_at=datetime.now().isoformat(), 

142 updated_at=datetime.now().isoformat(), 

143 status="draft", 

144 current_step=0 

145 ) 

146 

147 self.current_plan = plan 

148 self.save_plan(plan) 

149 

150 logger.info(f"Created plan: {plan_id} with {len(steps)} steps") 

151 return plan 

152 

153 def save_plan(self, plan: Plan): 

154 """Save plan to disk""" 

155 plan_file = PLANS_DIR / f"{plan.id}.json" 

156 

157 try: 

158 with open(plan_file, 'w') as f: 

159 json.dump(plan.to_dict(), f, indent=2) 

160 logger.debug(f"Saved plan: {plan.id}") 

161 except Exception as e: 

162 logger.error(f"Failed to save plan: {e}") 

163 

164 def load_plan(self, plan_id: str) -> Optional[Plan]: 

165 """Load plan from disk""" 

166 plan_file = PLANS_DIR / f"{plan_id}.json" 

167 

168 if not plan_file.exists(): 

169 logger.warning(f"Plan not found: {plan_id}") 

170 return None 

171 

172 try: 

173 with open(plan_file, 'r') as f: 

174 data = json.load(f) 

175 plan = Plan.from_dict(data) 

176 self.current_plan = plan 

177 logger.info(f"Loaded plan: {plan_id}") 

178 return plan 

179 except Exception as e: 

180 logger.error(f"Failed to load plan: {e}") 

181 return None 

182 

183 def list_plans(self) -> List[Dict[str, str]]: 

184 """List all saved plans""" 

185 plans = [] 

186 

187 for plan_file in PLANS_DIR.glob("*.json"): 

188 try: 

189 with open(plan_file, 'r') as f: 

190 data = json.load(f) 

191 plans.append({ 

192 "id": data["id"], 

193 "title": data["title"], 

194 "status": data["status"], 

195 "created_at": data["created_at"], 

196 "steps": len(data["steps"]) 

197 }) 

198 except Exception as e: 

199 logger.error(f"Failed to read plan {plan_file}: {e}") 

200 

201 return sorted(plans, key=lambda x: x["created_at"], reverse=True) 

202 

203 def display_plan(self, plan: Optional[Plan] = None): 

204 """Display plan in a formatted way""" 

205 plan = plan or self.current_plan 

206 

207 if not plan: 

208 console.print("[yellow]No plan loaded[/yellow]") 

209 return 

210 

211 # Header 

212 console.print(f"\n[bold cyan]📋 Plan: {plan.title}[/bold cyan]") 

213 console.print(f"[dim]{plan.description}[/dim]") 

214 console.print(f"[dim]Status: {plan.status.upper()} | Created: {plan.created_at[:10]}[/dim]\n") 

215 

216 # Steps table 

217 table = Table(show_header=True, header_style="bold") 

218 table.add_column("Step", justify="center", width=6) 

219 table.add_column("Status", width=12) 

220 table.add_column("Description", no_wrap=False) 

221 

222 for step in plan.steps: 

223 # Status emoji 

224 status_map = { 

225 StepStatus.PENDING: "⏸️ Pending", 

226 StepStatus.IN_PROGRESS: "▶️ In Progress", 

227 StepStatus.COMPLETED: "✅ Completed", 

228 StepStatus.FAILED: "❌ Failed", 

229 StepStatus.SKIPPED: "⏭️ Skipped" 

230 } 

231 status_str = status_map[step.status] 

232 

233 # Highlight current step 

234 if step.id == plan.current_step + 1 and step.status != StepStatus.COMPLETED: 

235 table.add_row( 

236 f"[bold cyan]{step.id}[/bold cyan]", 

237 f"[bold cyan]{status_str}[/bold cyan]", 

238 f"[bold cyan]{step.description}[/bold cyan]" 

239 ) 

240 else: 

241 color = { 

242 StepStatus.PENDING: "dim", 

243 StepStatus.IN_PROGRESS: "cyan", 

244 StepStatus.COMPLETED: "green", 

245 StepStatus.FAILED: "red", 

246 StepStatus.SKIPPED: "yellow" 

247 }[step.status] 

248 

249 table.add_row( 

250 f"[{color}]{step.id}[/{color}]", 

251 f"[{color}]{status_str}[/{color}]", 

252 f"[{color}]{step.description}[/{color}]" 

253 ) 

254 

255 console.print(table) 

256 console.print() 

257 

258 # Progress 

259 completed = sum(1 for s in plan.steps if s.status == StepStatus.COMPLETED) 

260 total = len(plan.steps) 

261 progress = (completed / total * 100) if total > 0 else 0 

262 console.print(f"[bold]Progress:[/bold] {completed}/{total} steps ({progress:.1f}%)\n") 

263 

264 def start_step(self, step_number: int) -> bool: 

265 """Mark step as in progress""" 

266 if not self.current_plan: 

267 logger.error("No plan loaded") 

268 return False 

269 

270 if step_number < 1 or step_number > len(self.current_plan.steps): 

271 logger.error(f"Invalid step number: {step_number}") 

272 return False 

273 

274 step = self.current_plan.steps[step_number - 1] 

275 step.status = StepStatus.IN_PROGRESS 

276 step.started_at = datetime.now().isoformat() 

277 

278 self.current_plan.current_step = step_number - 1 

279 self.current_plan.status = "in_progress" 

280 self.current_plan.updated_at = datetime.now().isoformat() 

281 

282 self.save_plan(self.current_plan) 

283 logger.info(f"Started step {step_number}: {step.description}") 

284 

285 return True 

286 

287 def complete_step( 

288 self, 

289 step_number: int, 

290 result: Optional[str] = None 

291 ) -> bool: 

292 """Mark step as completed""" 

293 if not self.current_plan: 

294 return False 

295 

296 if step_number < 1 or step_number > len(self.current_plan.steps): 

297 return False 

298 

299 step = self.current_plan.steps[step_number - 1] 

300 step.status = StepStatus.COMPLETED 

301 step.completed_at = datetime.now().isoformat() 

302 step.result = result 

303 

304 self.current_plan.updated_at = datetime.now().isoformat() 

305 

306 # Check if all steps completed 

307 if all(s.status == StepStatus.COMPLETED for s in self.current_plan.steps): 

308 self.current_plan.status = "completed" 

309 logger.info(f"Plan completed: {self.current_plan.id}") 

310 

311 self.save_plan(self.current_plan) 

312 logger.info(f"Completed step {step_number}") 

313 

314 return True 

315 

316 def fail_step( 

317 self, 

318 step_number: int, 

319 error: str 

320 ) -> bool: 

321 """Mark step as failed""" 

322 if not self.current_plan: 

323 return False 

324 

325 if step_number < 1 or step_number > len(self.current_plan.steps): 

326 return False 

327 

328 step = self.current_plan.steps[step_number - 1] 

329 step.status = StepStatus.FAILED 

330 step.completed_at = datetime.now().isoformat() 

331 step.error = error 

332 

333 self.current_plan.status = "failed" 

334 self.current_plan.updated_at = datetime.now().isoformat() 

335 

336 self.save_plan(self.current_plan) 

337 logger.warning(f"Step {step_number} failed: {error}") 

338 

339 return True 

340 

341 def get_next_step(self) -> Optional[PlanStep]: 

342 """Get next pending step""" 

343 if not self.current_plan: 

344 return None 

345 

346 for step in self.current_plan.steps: 

347 if step.status == StepStatus.PENDING: 

348 return step 

349 

350 return None 

351 

352 def cancel_plan(self) -> bool: 

353 if not self.current_plan: 

354 return False 

355 

356 plan = self.current_plan 

357 logger.info(f"Cancelling plan {plan.id}") 

358 plan.status = "cancelled" 

359 plan.updated_at = datetime.now().isoformat() 

360 plan.current_step = 0 

361 for step in plan.steps: 

362 if step.status not in {StepStatus.COMPLETED, StepStatus.FAILED}: 

363 step.status = StepStatus.SKIPPED 

364 step.completed_at = datetime.now().isoformat() 

365 self.save_plan(plan) 

366 self.current_plan = None 

367 return True 

368 

369 async def generate_plan_with_ai( 

370 self, 

371 task_description: str, 

372 ai_client 

373 ) -> Optional[Plan]: 

374 """ 

375 Use AI to generate a plan for a complex task. 

376  

377 Args: 

378 task_description: Description of the task 

379 ai_client: AI client to use for generation 

380  

381 Returns: 

382 Generated plan or None 

383 """ 

384 prompt = f"""Break down this task into clear, actionable steps: 

385 

386Task: {task_description} 

387 

388Generate a step-by-step plan. Each step should be: 

389- Specific and actionable 

390- Testable/verifiable 

391- In logical order 

392- Independent where possible 

393 

394Format your response as JSON: 

395{{ 

396 "title": "Brief plan title", 

397 "description": "Overview of the plan", 

398 "steps": [ 

399 "Step 1 description", 

400 "Step 2 description", 

401 ... 

402 ] 

403}}""" 

404 

405 try: 

406 logger.info("Generating plan with AI") 

407 

408 # Call AI 

409 messages = [{"role": "user", "content": prompt}] 

410 response = await ai_client.create_completion(messages, stream=False) 

411 

412 content = response.choices[0].message.content 

413 

414 # Parse JSON from response 

415 # Try to extract JSON from markdown code blocks if present 

416 if "```json" in content: 

417 content = content.split("```json")[1].split("```")[0].strip() 

418 elif "```" in content: 

419 content = content.split("```")[1].split("```")[0].strip() 

420 

421 plan_data = json.loads(content) 

422 

423 # Create plan 

424 plan = self.create_plan( 

425 title=plan_data["title"], 

426 description=plan_data["description"], 

427 steps=plan_data["steps"] 

428 ) 

429 

430 logger.info(f"Generated plan with {len(plan.steps)} steps") 

431 return plan 

432 

433 except Exception as e: 

434 logger.error(f"Failed to generate plan: {e}", exc_info=True) 

435 console.print(f"[red]Failed to generate plan: {e}[/red]") 

436 return None 

437 

438 

439# Global plan mode instance 

440plan_mode = PlanMode()