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