Skip to main content

lash_plugin_plan_mode/
plan_mode.rs

1use std::collections::BTreeSet;
2use std::fs;
3use std::path::{Path, PathBuf};
4use std::sync::{Arc, Mutex};
5
6use serde_json::json;
7
8use lash_core::plugin::{
9    PluginCommand, PluginCommandOutcome, PluginDirective, PluginError, PluginFactory,
10    PluginOperation, PluginOperationFailure, PluginRegistrar, PluginSessionContext,
11    PluginSnapshotMeta, SessionParam, SessionPlugin, SnapshotReader, SnapshotWriter,
12    ToolCatalogContribution,
13};
14use lash_core::{
15    JsonSchema, PluginMessage, ToolCall, ToolContext, ToolControl, ToolDefinition, ToolResult,
16    ToolScheduling,
17};
18use lash_lashlang_runtime::{LashlangToolBinding, ToolDefinitionLashlangExt};
19use lash_tool_support::{StaticToolExecute, StaticToolProvider, resolve_under};
20
21mod prompt;
22mod state;
23
24pub use prompt::{
25    PlanModePrompt, PlanModePromptRequest, PlanModePromptResponse, PlanModePromptReview,
26};
27use prompt::{
28    plan_exit_confirmation_display, plan_exit_fresh_context_input, plan_exit_next_turn_input,
29    plan_mode_guidance_message, plan_mode_tool_note,
30};
31#[cfg(test)]
32use state::PLAN_TEMPLATE;
33use state::{
34    PlanModeSnapshot, PlanModeState, PlanReport, effective_run_session_id, plan_display_path,
35    read_plan_report, resolve_plan_path, seed_plan_template,
36};
37
38const PLAN_MODE_STATE_EVENT: &str = "plan_mode.state";
39
40fn default_allowed_tools() -> BTreeSet<String> {
41    [
42        "ask",
43        "fetch_url",
44        "glob",
45        "grep",
46        "read_file",
47        "search_tools",
48        "search_web",
49        "edit",
50        "write",
51        "plan_exit",
52    ]
53    .into_iter()
54    .map(str::to_string)
55    .collect()
56}
57
58fn fresh_context_frame_id() -> String {
59    format!("plan-frame-{}", uuid::Uuid::new_v4().simple())
60}
61
62fn plan_protocol_state_event(
63    session_id: &str,
64    enabled: bool,
65    report: Option<&PlanReport>,
66) -> Result<lash_core::PluginRuntimeEvent, PluginError> {
67    Ok(lash_core::PluginRuntimeEvent::Custom {
68        name: PLAN_MODE_STATE_EVENT.to_string(),
69        payload: serde_json::to_value(plan_mode_payload(session_id, enabled, report)).map_err(
70            |err| PluginError::Session(format!("failed to encode plan mode state: {err}")),
71        )?,
72    })
73}
74
75#[derive(Clone, Debug)]
76pub struct PlanModePluginConfig {
77    pub allowed_tools: BTreeSet<String>,
78}
79
80impl Default for PlanModePluginConfig {
81    fn default() -> Self {
82        Self {
83            allowed_tools: default_allowed_tools(),
84        }
85    }
86}
87
88impl PlanModePluginConfig {
89    pub fn with_allowed_tools<I, S>(mut self, allowed_tools: I) -> Self
90    where
91        I: IntoIterator<Item = S>,
92        S: Into<String>,
93    {
94        self.allowed_tools = allowed_tools.into_iter().map(Into::into).collect();
95        self.allowed_tools.insert("plan_exit".to_string());
96        self
97    }
98}
99
100#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize, JsonSchema)]
101pub struct PlanModeExternalArgs {}
102
103#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize, JsonSchema)]
104pub struct PlanModeExternalStatus {
105    pub session_id: String,
106    pub enabled: bool,
107    pub plan_path: Option<String>,
108}
109
110pub struct PlanModeEnableOp;
111pub struct PlanModeDisableOp;
112pub struct PlanModeToggleOp;
113
114impl PluginOperation for PlanModeEnableOp {
115    const NAME: &'static str = "plan_mode.enable";
116    const DESCRIPTION: &'static str = "Enable plan mode for this session.";
117    const SESSION_PARAM: SessionParam = SessionParam::Required;
118    type Args = PlanModeExternalArgs;
119    type Output = PlanModeExternalStatus;
120}
121
122impl PluginCommand for PlanModeEnableOp {}
123
124impl PluginOperation for PlanModeDisableOp {
125    const NAME: &'static str = "plan_mode.disable";
126    const DESCRIPTION: &'static str = "Disable plan mode for this session.";
127    const SESSION_PARAM: SessionParam = SessionParam::Required;
128    type Args = PlanModeExternalArgs;
129    type Output = PlanModeExternalStatus;
130}
131
132impl PluginCommand for PlanModeDisableOp {}
133
134impl PluginOperation for PlanModeToggleOp {
135    const NAME: &'static str = "plan_mode.toggle";
136    const DESCRIPTION: &'static str = "Toggle plan mode for this session.";
137    const SESSION_PARAM: SessionParam = SessionParam::Required;
138    type Args = PlanModeExternalArgs;
139    type Output = PlanModeExternalStatus;
140}
141
142impl PluginCommand for PlanModeToggleOp {}
143
144async fn ensure_plan_path<H>(
145    state: &Arc<Mutex<PlanModeState>>,
146    session_id: &str,
147    host: &Arc<H>,
148) -> Result<PathBuf, PluginError>
149where
150    H: lash_core::plugin::runtime_host::SessionStateService + ?Sized,
151{
152    if let Some(path) = state
153        .lock()
154        .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
155        .plan_path()
156    {
157        return Ok(path);
158    }
159
160    let snapshot = host.snapshot_session(session_id).await?;
161    let run_session_id =
162        effective_run_session_id(&snapshot.session_id, &snapshot.policy).to_string();
163    let path = resolve_plan_path(&run_session_id).map_err(PluginError::Session)?;
164    state
165        .lock()
166        .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
167        .set_plan_path(path.clone());
168    Ok(path)
169}
170
171fn ensure_plan_path_from_snapshot(
172    state: &Arc<Mutex<PlanModeState>>,
173    snapshot: &lash_core::SessionSnapshot,
174) -> Result<PathBuf, PluginError> {
175    if let Some(path) = state
176        .lock()
177        .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
178        .plan_path()
179    {
180        return Ok(path);
181    }
182    let run_session_id =
183        effective_run_session_id(&snapshot.session_id, &snapshot.policy).to_string();
184    let path = resolve_plan_path(&run_session_id).map_err(PluginError::Session)?;
185    state
186        .lock()
187        .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
188        .set_plan_path(path.clone());
189    Ok(path)
190}
191
192async fn ensure_plan_report_for_tool_context(
193    state: &Arc<Mutex<PlanModeState>>,
194    context: &ToolContext<'_>,
195    seed_if_missing: bool,
196) -> Result<PlanReport, PluginError> {
197    let snapshot = context.sessions().snapshot_current().await?;
198    let path = ensure_plan_path_from_snapshot(state, &snapshot)?;
199    if seed_if_missing {
200        seed_plan_template(&path).map_err(PluginError::Session)?;
201    }
202    read_plan_report(&path).map_err(PluginError::Session)
203}
204
205async fn ensure_plan_report<H>(
206    state: &Arc<Mutex<PlanModeState>>,
207    session_id: &str,
208    host: &Arc<H>,
209    seed_if_missing: bool,
210) -> Result<PlanReport, PluginError>
211where
212    H: lash_core::plugin::runtime_host::SessionStateService + ?Sized,
213{
214    let path = ensure_plan_path(state, session_id, host).await?;
215    if seed_if_missing {
216        seed_plan_template(&path).map_err(PluginError::Session)?;
217    }
218    read_plan_report(&path).map_err(PluginError::Session)
219}
220
221/// Flip plan-mode enablement. Tool Catalog membership (`plan_exit` visibility,
222/// suppression of non-allowed tools) is derived per turn by the plan-mode
223/// catalog contribution keyed on this state, so there is no registry mutation
224/// here.
225fn set_plan_mode_enabled_state(
226    state: &Arc<Mutex<PlanModeState>>,
227    enabled: bool,
228) -> Result<bool, PluginError> {
229    let mut guard = state
230        .lock()
231        .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?;
232    if guard.enabled != enabled {
233        guard.set_enabled(enabled);
234    }
235    Ok(enabled)
236}
237
238fn plan_mode_payload(
239    session_id: &str,
240    enabled: bool,
241    report: Option<&PlanReport>,
242) -> PlanModeExternalStatus {
243    PlanModeExternalStatus {
244        session_id: session_id.to_string(),
245        enabled,
246        plan_path: report.map(|value| value.display_path.clone()),
247    }
248}
249
250fn file_mutation_allowed_for_plan_file(
251    tool_name: &str,
252    args: &serde_json::Value,
253    plan_path: &Path,
254) -> Result<(), String> {
255    let path = args
256        .get("path")
257        .and_then(|value| value.as_str())
258        .ok_or_else(|| format!("plan mode requires `{tool_name}.path`"))?;
259    let cwd = std::env::current_dir().map_err(|err| format!("Failed to determine cwd: {err}"))?;
260    let resolved = resolve_under(&cwd, Path::new(path));
261    if resolved != plan_path {
262        return Err(format!(
263            "plan mode only allows `{tool_name}` to edit `{}`",
264            plan_display_path(plan_path)
265        ));
266    }
267    if tool_name == "edit"
268        && args
269            .get("edits")
270            .and_then(|value| value.as_array())
271            .is_none_or(Vec::is_empty)
272    {
273        return Err(
274            "plan mode requires `edit.edits` to contain at least one replacement".to_string(),
275        );
276    }
277    Ok(())
278}
279
280#[derive(Clone)]
281struct PlanModeTools {
282    state: Arc<Mutex<PlanModeState>>,
283    prompt: Option<Arc<dyn PlanModePrompt>>,
284}
285
286impl PlanModeTools {
287    async fn execute_plan_exit(&self, context: &ToolContext<'_>) -> ToolResult {
288        let enabled = match self.state.lock() {
289            Ok(guard) => guard.enabled,
290            Err(_) => return ToolResult::err(json!("plan mode state poisoned")),
291        };
292        if !enabled {
293            return ToolResult::err(json!("plan mode is not active"));
294        }
295
296        let report = match ensure_plan_report_for_tool_context(&self.state, context, true).await {
297            Ok(report) => report,
298            Err(err) => return ToolResult::err(json!(err.to_string())),
299        };
300
301        let Some(prompt) = &self.prompt else {
302            return ToolResult::err(json!(
303                "plan approval prompts are unavailable in this session"
304            ));
305        };
306        let answer = match prompt
307            .prompt_user(
308                PlanModePromptRequest::single(
309                    format!("Review the plan in `{}`. What next?", report.display_path),
310                    vec![
311                        "Start implementing now".to_string(),
312                        "Keep planning".to_string(),
313                        "Start in fresh context".to_string(),
314                    ],
315                )
316                .with_review("PLAN", report.approval_content())
317                .with_optional_note(),
318            )
319            .await
320        {
321            Ok(answer) => answer,
322            Err(err) => return ToolResult::err(json!(err.to_string())),
323        };
324
325        let selection = match &answer {
326            PlanModePromptResponse::Single { selection, .. } => selection.as_str(),
327        };
328        if selection == "Keep planning" {
329            return ToolResult::ok(json!({
330                "approved": false,
331                "plan_path": report.display_path,
332                "answer": answer,
333            }));
334        }
335
336        let note = match &answer {
337            PlanModePromptResponse::Single { note, .. } => note.clone(),
338        };
339
340        if let Err(err) = set_plan_mode_enabled_state(&self.state, false) {
341            return ToolResult::err(json!(err.to_string()));
342        }
343
344        if selection == "Start in fresh context" {
345            return ToolResult::ok(json!({
346                "approved": true,
347                "plan_path": report.display_path,
348                "execution_mode": "fresh_context",
349            }));
350        }
351
352        ToolResult::ok(json!({
353            "approved": true,
354            "answer": answer,
355            "confirmation_display": plan_exit_confirmation_display(selection, note.as_deref()),
356            "plan_path": report.display_path,
357            "execution_mode": "current_session",
358            "next_turn_input": plan_exit_next_turn_input(&report.display_path, note.as_deref()),
359        }))
360    }
361}
362
363fn plan_mode_provider(
364    state: Arc<Mutex<PlanModeState>>,
365    prompt: Option<Arc<dyn PlanModePrompt>>,
366) -> StaticToolProvider<PlanModeTools> {
367    StaticToolProvider::new(
368        vec![plan_exit_tool_definition()],
369        PlanModeTools { state, prompt },
370    )
371}
372
373#[async_trait::async_trait]
374impl StaticToolExecute for PlanModeTools {
375    async fn execute(&self, call: ToolCall<'_>) -> ToolResult {
376        match call.name {
377            "plan_exit" => self.execute_plan_exit(call.context).await,
378            other => ToolResult::err_fmt(format_args!("Unknown tool: {other}")),
379        }
380    }
381}
382
383fn plan_exit_tool_definition() -> ToolDefinition {
384    ToolDefinition::raw(
385        "tool:plan_exit",
386        "plan_exit",
387        "Ask whether to exit plan mode.",
388        plan_exit_input_schema(),
389        plan_exit_output_schema(),
390    )
391    .with_examples(vec!["await plan.exit({})?".into()])
392    .with_lashlang_binding(LashlangToolBinding::new(["plan"], "exit"))
393    .with_scheduling(ToolScheduling::Parallel)
394}
395
396fn plan_exit_input_schema() -> serde_json::Value {
397    serde_json::json!({
398        "type": "object",
399        "properties": {},
400        "additionalProperties": false
401    })
402}
403
404fn plan_exit_output_schema() -> serde_json::Value {
405    serde_json::json!({
406        "type": "object",
407        "properties": {
408            "approved": { "type": "boolean" },
409            "plan_path": { "type": "string" },
410            "answer": {
411                "type": "object",
412                "properties": {
413                    "kind": { "type": "string", "enum": ["single"] },
414                    "selection": { "type": "string" },
415                    "note": { "type": "string" }
416                },
417                "required": ["kind", "selection"],
418                "additionalProperties": false
419            },
420            "execution_mode": {
421                "type": "string",
422                "enum": ["current_session", "fresh_context"]
423            },
424            "confirmation_display": { "type": "string" },
425            "next_turn_input": { "type": "string" }
426        },
427        "required": ["approved", "plan_path"],
428        "additionalProperties": false
429    })
430}
431
432pub struct PlanModePluginFactory {
433    config: PlanModePluginConfig,
434    prompt: Option<Arc<dyn PlanModePrompt>>,
435}
436
437impl Default for PlanModePluginFactory {
438    fn default() -> Self {
439        Self::new(PlanModePluginConfig::default())
440    }
441}
442
443impl PlanModePluginFactory {
444    pub fn new(config: PlanModePluginConfig) -> Self {
445        Self {
446            config,
447            prompt: None,
448        }
449    }
450
451    pub fn with_prompt(mut self, prompt: Arc<dyn PlanModePrompt>) -> Self {
452        self.prompt = Some(prompt);
453        self
454    }
455}
456
457impl PluginFactory for PlanModePluginFactory {
458    fn id(&self) -> &'static str {
459        "plan_mode"
460    }
461
462    fn build(&self, _ctx: &PluginSessionContext) -> Result<Arc<dyn SessionPlugin>, PluginError> {
463        Ok(Arc::new(PlanModePlugin {
464            state: Arc::new(Mutex::new(PlanModeState::default())),
465            config: self.config.clone(),
466            prompt: self.prompt.clone(),
467        }))
468    }
469}
470
471struct PlanModePlugin {
472    state: Arc<Mutex<PlanModeState>>,
473    config: PlanModePluginConfig,
474    prompt: Option<Arc<dyn PlanModePrompt>>,
475}
476
477impl SessionPlugin for PlanModePlugin {
478    fn id(&self) -> &'static str {
479        "plan_mode"
480    }
481
482    fn register(&self, reg: &mut PluginRegistrar) -> Result<(), PluginError> {
483        reg.tools().provider(Arc::new(plan_mode_provider(
484            Arc::clone(&self.state),
485            self.prompt.clone(),
486        )))?;
487
488        let before_turn_state = Arc::clone(&self.state);
489        reg.turn().before(Arc::new(move |ctx| {
490            let state = Arc::clone(&before_turn_state);
491            Box::pin(async move {
492                let plan_path = {
493                    let mut state = state.lock().map_err(|_| {
494                        PluginError::Session("plan mode state poisoned".to_string())
495                    })?;
496                    let should_inject = state.prepare_turn();
497                    if !should_inject {
498                        return Ok(Vec::new());
499                    }
500                    state.ensure_plan_path_from_state(&ctx.state.to_snapshot())?
501                };
502                seed_plan_template(&plan_path).map_err(PluginError::Session)?;
503                let report = read_plan_report(&plan_path).map_err(PluginError::Session)?;
504                Ok(vec![
505                    PluginDirective::emit_runtime_events(vec![plan_protocol_state_event(
506                        &ctx.session_id,
507                        true,
508                        Some(&report),
509                    )?]),
510                    PluginDirective::EnqueueMessages {
511                        messages: vec![plan_mode_guidance_message(&plan_path)],
512                    },
513                ])
514            })
515        }));
516
517        let checkpoint_state = Arc::clone(&self.state);
518        reg.turn().checkpoint(Arc::new(move |ctx| {
519            let state = Arc::clone(&checkpoint_state);
520            Box::pin(async move {
521                let plan_path = {
522                    let mut state = state.lock().map_err(|_| {
523                        PluginError::Session("plan mode state poisoned".to_string())
524                    })?;
525                    let should_inject = state.checkpoint_injection_needed();
526                    if !should_inject {
527                        return Ok(Vec::new());
528                    }
529                    state.ensure_plan_path_from_state(&ctx.state.to_snapshot())?
530                };
531                seed_plan_template(&plan_path).map_err(PluginError::Session)?;
532                let report = read_plan_report(&plan_path).map_err(PluginError::Session)?;
533                Ok(vec![
534                    PluginDirective::emit_runtime_events(vec![plan_protocol_state_event(
535                        &ctx.session_id,
536                        true,
537                        Some(&report),
538                    )?]),
539                    PluginDirective::EnqueueMessages {
540                        messages: vec![plan_mode_guidance_message(&plan_path)],
541                    },
542                ])
543            })
544        }));
545
546        let after_turn_state = Arc::clone(&self.state);
547        reg.turn().after(Arc::new(move |_ctx| {
548            let state = Arc::clone(&after_turn_state);
549            Box::pin(async move {
550                state
551                    .lock()
552                    .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
553                    .finish_turn();
554                Ok(Vec::new())
555            })
556        }));
557
558        let before_tool_state = Arc::clone(&self.state);
559        let before_tool_config = self.config.clone();
560        reg.tool_calls().before(Arc::new(move |ctx| {
561            let state = Arc::clone(&before_tool_state);
562            let config = before_tool_config.clone();
563            Box::pin(async move {
564                let enabled = state
565                    .lock()
566                    .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
567                    .enabled;
568                if !enabled {
569                    return Ok(Vec::new());
570                }
571
572                if ctx.tool_name != "plan_exit" && !config.allowed_tools.contains(&ctx.tool_name) {
573                    return Ok(vec![PluginDirective::AbortTurn {
574                        code: "plan_mode_tool_blocked".to_string(),
575                        message: format!(
576                            "Plan mode blocks `{}`. Use planning tools or `plan.exit`.",
577                            ctx.tool_name
578                        ),
579                    }]);
580                }
581
582                if matches!(ctx.tool_name.as_str(), "edit" | "write") {
583                    let snapshot = ctx.session_snapshot().await?;
584                    let plan_path = ensure_plan_path_from_snapshot(&state, &snapshot)?;
585                    if let Err(message) =
586                        file_mutation_allowed_for_plan_file(&ctx.tool_name, &ctx.args, &plan_path)
587                    {
588                        return Ok(vec![PluginDirective::AbortTurn {
589                            code: "plan_mode_tool_blocked".to_string(),
590                            message,
591                        }]);
592                    }
593                }
594
595                Ok(Vec::new())
596            })
597        }));
598
599        let after_tool_state = Arc::clone(&self.state);
600        reg.tool_calls().after(Arc::new(move |ctx| {
601            let state = Arc::clone(&after_tool_state);
602            Box::pin(async move {
603                let result_value = ctx.result.value_for_projection();
604                let approved = ctx.tool_name == "plan_exit"
605                    && ctx.result.is_success()
606                    && result_value
607                        .get("approved")
608                        .and_then(|value| value.as_bool())
609                        .unwrap_or(false);
610                if approved {
611                    let mut directives = vec![PluginDirective::emit_runtime_events(vec![
612                        plan_protocol_state_event(&ctx.session_id, false, None)?,
613                    ])];
614                    if result_value
615                        .get("execution_mode")
616                        .and_then(|value| value.as_str())
617                        == Some("fresh_context")
618                    {
619                        let plan_path = result_value
620                            .get("plan_path")
621                            .and_then(|value| value.as_str())
622                            .unwrap_or_default()
623                            .to_string();
624                        let frame_id = fresh_context_frame_id();
625                        let task = plan_exit_fresh_context_input(&plan_path);
626                        directives.push(PluginDirective::short_circuit(
627                            ToolResult::ok(json!({
628                                "approved": true,
629                                "plan_path": plan_path,
630                                "execution_mode": "fresh_context",
631                                "frame_id": frame_id.clone(),
632                            }))
633                            .with_control(
634                                ToolControl::SwitchAgentFrame {
635                                    frame_id,
636                                    initial_nodes: Vec::new(),
637                                    task: Some(task),
638                                },
639                            ),
640                        ));
641                    }
642                    return Ok(directives);
643                }
644
645                let enabled = state
646                    .lock()
647                    .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
648                    .enabled;
649                if !enabled
650                    || !matches!(ctx.tool_name.as_str(), "edit" | "write")
651                    || !ctx.result.is_success()
652                {
653                    return Ok(Vec::new());
654                }
655
656                let snapshot = ctx.session_snapshot().await?;
657                let path = ensure_plan_path_from_snapshot(&state, &snapshot)?;
658                let report = read_plan_report(&path).map_err(PluginError::Session)?;
659                Ok(vec![PluginDirective::emit_runtime_events(vec![
660                    plan_protocol_state_event(&ctx.session_id, true, Some(&report))?,
661                ])])
662            })
663        }));
664
665        // Membership is the execution gate. While plan mode is enabled, remove
666        // every tool that is not allowed (and not `plan_exit`); while it is
667        // disabled, remove `plan_exit` so it is not callable.
668        let tool_catalog_state = Arc::clone(&self.state);
669        let tool_catalog_config = self.config.clone();
670        reg.tool_catalog().contribute(Arc::new(move |ctx| {
671            let enabled = tool_catalog_state
672                .lock()
673                .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
674                .enabled;
675            if !enabled {
676                return Ok(ToolCatalogContribution::remove_tools(["plan_exit"]));
677            }
678
679            let remove = ctx
680                .tools
681                .iter()
682                .filter(|tool| {
683                    tool.name != "plan_exit"
684                        && !tool_catalog_config.allowed_tools.contains(&tool.name)
685                })
686                .map(|tool| tool.name.clone())
687                .collect::<Vec<_>>();
688            Ok(ToolCatalogContribution { remove })
689        }));
690
691        // The plan-mode tool note is an ordinary prompt contribution, gated on
692        // `plan_exit` being a catalog member (i.e. plan mode enabled).
693        let prompt_state = Arc::clone(&self.state);
694        reg.prompt().contribute(Arc::new(move |_ctx| {
695            let state = Arc::clone(&prompt_state);
696            Box::pin(async move {
697                let guard = state
698                    .lock()
699                    .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?;
700                if !guard.enabled {
701                    return Ok(Vec::new());
702                }
703                let note = plan_mode_tool_note(guard.plan_path().as_deref());
704                Ok(vec![
705                    lash_core::PromptContribution::execution("Plan Mode", note)
706                        .requires_tool("plan_exit"),
707                ])
708            })
709        }));
710
711        register_plan_mode_op::<PlanModeEnableOp>(reg, Arc::clone(&self.state))?;
712        register_plan_mode_op::<PlanModeDisableOp>(reg, Arc::clone(&self.state))?;
713        register_plan_mode_op::<PlanModeToggleOp>(reg, Arc::clone(&self.state))?;
714
715        Ok(())
716    }
717
718    fn snapshot(
719        &self,
720        _writer: &mut dyn SnapshotWriter,
721    ) -> Result<PluginSnapshotMeta, PluginError> {
722        let snapshot = self
723            .state
724            .lock()
725            .map_err(|_| PluginError::Snapshot("plan mode state poisoned".to_string()))?
726            .snapshot();
727        Ok(PluginSnapshotMeta {
728            plugin_id: self.id().to_string(),
729            plugin_version: self.version().to_string(),
730            revision: snapshot.generation,
731            state: Some(json!({
732                "enabled": snapshot.enabled,
733                "generation": snapshot.generation,
734                "plan_path": snapshot.plan_path,
735            })),
736        })
737    }
738
739    fn restore(
740        &self,
741        meta: &PluginSnapshotMeta,
742        _reader: &dyn SnapshotReader,
743    ) -> Result<(), PluginError> {
744        let snapshot = meta
745            .state
746            .clone()
747            .map(serde_json::from_value::<PlanModeSnapshot>)
748            .transpose()
749            .map_err(|err| PluginError::Snapshot(err.to_string()))?
750            .unwrap_or_default();
751        self.state
752            .lock()
753            .map_err(|_| PluginError::Snapshot("plan mode state poisoned".to_string()))?
754            .restore_snapshot(snapshot);
755        Ok(())
756    }
757
758    fn snapshot_revision(&self) -> u64 {
759        self.state
760            .lock()
761            .map(|state| state.generation)
762            .unwrap_or_default()
763    }
764}
765
766fn register_plan_mode_op<Op>(
767    reg: &mut PluginRegistrar,
768    state: Arc<Mutex<PlanModeState>>,
769) -> Result<(), PluginError>
770where
771    Op: PluginCommand<Args = PlanModeExternalArgs, Output = PlanModeExternalStatus>,
772{
773    reg.operations()
774        .typed_command::<Op, _, _>(move |ctx, _args| {
775            let state = Arc::clone(&state);
776            async move {
777                let Some(session_id) = ctx.session_id else {
778                    return Err(PluginOperationFailure::new(format!(
779                        "{} requires session_id",
780                        Op::NAME
781                    )));
782                };
783                let target_enabled = match state.lock() {
784                    Ok(guard) => match Op::NAME {
785                        "plan_mode.enable" => true,
786                        "plan_mode.disable" => false,
787                        "plan_mode.toggle" => !guard.enabled,
788                        _ => unreachable!(),
789                    },
790                    Err(_) => return Err(PluginOperationFailure::new("plan mode state poisoned")),
791                };
792                let enabled = set_plan_mode_enabled_state(&state, target_enabled)?;
793                let report =
794                    ensure_plan_report(&state, &session_id, &ctx.sessions, enabled).await?;
795                let status = plan_mode_payload(&session_id, enabled, Some(&report));
796                Ok(
797                    PluginCommandOutcome::new(status).with_events(vec![plan_protocol_state_event(
798                        &session_id,
799                        enabled,
800                        Some(&report),
801                    )?]),
802                )
803            }
804        })
805}
806
807#[cfg(test)]
808mod tests {
809    use super::{
810        PLAN_TEMPLATE, plan_exit_fresh_context_input, plan_exit_next_turn_input,
811        plan_exit_tool_definition, read_plan_report,
812    };
813
814    #[test]
815    fn plan_exit_contract_documents_decision_result() {
816        let definition = plan_exit_tool_definition();
817
818        assert_eq!(
819            definition.contract.input_schema.canonical["additionalProperties"],
820            serde_json::json!(false)
821        );
822        assert_eq!(
823            definition.contract.output_schema.canonical["required"],
824            serde_json::json!(["approved", "plan_path"])
825        );
826        assert!(
827            definition
828                .contract
829                .output_schema
830                .canonical
831                .to_string()
832                .contains("execution_mode")
833        );
834    }
835
836    #[test]
837    fn plan_exit_next_turn_input_appends_user_note() {
838        assert_eq!(
839            plan_exit_next_turn_input(
840                ".lash/plans/run-session.md",
841                Some("start with the safe slice"),
842            ),
843            "The user approved the plan. Execute the plan in `.lash/plans/run-session.md` now — start immediately, do not ask for confirmation.\n\nUser note: start with the safe slice"
844        );
845        assert_eq!(
846            plan_exit_next_turn_input(".lash/plans/run-session.md", Some("   ")),
847            "The user approved the plan. Execute the plan in `.lash/plans/run-session.md` now — start immediately, do not ask for confirmation."
848        );
849    }
850
851    #[test]
852    fn plan_exit_fresh_context_input_is_short_and_direct() {
853        assert_eq!(
854            plan_exit_fresh_context_input(".lash/plans/run-session.md"),
855            "Do a full, faithful implementation of the plan found at: .lash/plans/run-session.md"
856        );
857    }
858
859    #[test]
860    fn seeded_template_is_readable_without_validation() {
861        let dir = tempfile::tempdir().expect("tempdir");
862        let path = dir.path().join("plan.md");
863        std::fs::write(&path, PLAN_TEMPLATE).expect("write template");
864        let report = read_plan_report(&path).expect("report");
865        assert_eq!(report.content.as_deref(), Some(PLAN_TEMPLATE));
866    }
867}