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};
20use lash_tools::apply_patch::{PatchAction, inspect_patch_ops};
21
22mod prompt;
23mod state;
24
25pub use prompt::{
26    PlanModePrompt, PlanModePromptRequest, PlanModePromptResponse, PlanModePromptReview,
27};
28use prompt::{
29    plan_exit_confirmation_display, plan_exit_fresh_context_input, plan_exit_next_turn_input,
30    plan_mode_guidance_message, plan_mode_tool_note,
31};
32#[cfg(test)]
33use state::PLAN_TEMPLATE;
34use state::{
35    PlanModeSnapshot, PlanModeState, PlanReport, effective_run_session_id, plan_display_path,
36    read_plan_report, resolve_plan_path, seed_plan_template,
37};
38
39const PLAN_MODE_STATE_EVENT: &str = "plan_mode.state";
40
41fn default_allowed_tools() -> BTreeSet<String> {
42    [
43        "ask",
44        "fetch_url",
45        "glob",
46        "grep",
47        "read_file",
48        "search_tools",
49        "search_web",
50        "apply_patch",
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 patch_allowed_for_plan_file(args: &serde_json::Value, plan_path: &Path) -> Result<(), String> {
251    let input = args
252        .get("input")
253        .and_then(|value| value.as_str())
254        .ok_or_else(|| "plan mode requires `apply_patch.input`".to_string())?;
255    let workdir = args
256        .get("workdir")
257        .and_then(|value| value.as_str())
258        .filter(|value| !value.is_empty());
259    let ops = inspect_patch_ops(input, workdir)?;
260    if ops.is_empty() {
261        return Err("plan mode requires a non-empty patch".to_string());
262    }
263    for op in ops {
264        match op.action {
265            PatchAction::Add | PatchAction::Update
266                if op.path == plan_path && op.move_path.is_none() => {}
267            PatchAction::Add | PatchAction::Update => {
268                return Err(format!(
269                    "plan mode only allows `apply_patch` to edit `{}`",
270                    plan_display_path(plan_path)
271                ));
272            }
273            PatchAction::Delete => {
274                return Err(format!(
275                    "plan mode does not allow deleting `{}`",
276                    plan_display_path(plan_path)
277                ));
278            }
279        }
280    }
281    Ok(())
282}
283
284#[derive(Clone)]
285struct PlanModeTools {
286    state: Arc<Mutex<PlanModeState>>,
287    prompt: Option<Arc<dyn PlanModePrompt>>,
288}
289
290impl PlanModeTools {
291    async fn execute_plan_exit(&self, context: &ToolContext<'_>) -> ToolResult {
292        let enabled = match self.state.lock() {
293            Ok(guard) => guard.enabled,
294            Err(_) => return ToolResult::err(json!("plan mode state poisoned")),
295        };
296        if !enabled {
297            return ToolResult::err(json!("plan mode is not active"));
298        }
299
300        let report = match ensure_plan_report_for_tool_context(&self.state, context, true).await {
301            Ok(report) => report,
302            Err(err) => return ToolResult::err(json!(err.to_string())),
303        };
304
305        let Some(prompt) = &self.prompt else {
306            return ToolResult::err(json!(
307                "plan approval prompts are unavailable in this session"
308            ));
309        };
310        let answer = match prompt
311            .prompt_user(
312                PlanModePromptRequest::single(
313                    format!("Review the plan in `{}`. What next?", report.display_path),
314                    vec![
315                        "Start implementing now".to_string(),
316                        "Keep planning".to_string(),
317                        "Start in fresh context".to_string(),
318                    ],
319                )
320                .with_review("PLAN", report.approval_content())
321                .with_optional_note(),
322            )
323            .await
324        {
325            Ok(answer) => answer,
326            Err(err) => return ToolResult::err(json!(err.to_string())),
327        };
328
329        let selection = match &answer {
330            PlanModePromptResponse::Single { selection, .. } => selection.as_str(),
331        };
332        if selection == "Keep planning" {
333            return ToolResult::ok(json!({
334                "approved": false,
335                "plan_path": report.display_path,
336                "answer": answer,
337            }));
338        }
339
340        let note = match &answer {
341            PlanModePromptResponse::Single { note, .. } => note.clone(),
342        };
343
344        if let Err(err) = set_plan_mode_enabled_state(&self.state, false) {
345            return ToolResult::err(json!(err.to_string()));
346        }
347
348        if selection == "Start in fresh context" {
349            return ToolResult::ok(json!({
350                "approved": true,
351                "plan_path": report.display_path,
352                "execution_mode": "fresh_context",
353            }));
354        }
355
356        ToolResult::ok(json!({
357            "approved": true,
358            "answer": answer,
359            "confirmation_display": plan_exit_confirmation_display(selection, note.as_deref()),
360            "plan_path": report.display_path,
361            "execution_mode": "current_session",
362            "next_turn_input": plan_exit_next_turn_input(&report.display_path, note.as_deref()),
363        }))
364    }
365}
366
367fn plan_mode_provider(
368    state: Arc<Mutex<PlanModeState>>,
369    prompt: Option<Arc<dyn PlanModePrompt>>,
370) -> StaticToolProvider<PlanModeTools> {
371    StaticToolProvider::new(
372        vec![plan_exit_tool_definition()],
373        PlanModeTools { state, prompt },
374    )
375}
376
377#[async_trait::async_trait]
378impl StaticToolExecute for PlanModeTools {
379    async fn execute(&self, call: ToolCall<'_>) -> ToolResult {
380        match call.name {
381            "plan_exit" => self.execute_plan_exit(call.context).await,
382            other => ToolResult::err_fmt(format_args!("Unknown tool: {other}")),
383        }
384    }
385}
386
387fn plan_exit_tool_definition() -> ToolDefinition {
388    ToolDefinition::raw(
389        "tool:plan_exit",
390        "plan_exit",
391        "Ask whether to exit plan mode.",
392        plan_exit_input_schema(),
393        plan_exit_output_schema(),
394    )
395    .with_examples(vec!["await plan.exit({})?".into()])
396    .with_lashlang_binding(LashlangToolBinding::new(["plan"], "exit"))
397    .with_scheduling(ToolScheduling::Parallel)
398}
399
400fn plan_exit_input_schema() -> serde_json::Value {
401    serde_json::json!({
402        "type": "object",
403        "properties": {},
404        "additionalProperties": false
405    })
406}
407
408fn plan_exit_output_schema() -> serde_json::Value {
409    serde_json::json!({
410        "type": "object",
411        "properties": {
412            "approved": { "type": "boolean" },
413            "plan_path": { "type": "string" },
414            "answer": {
415                "type": "object",
416                "properties": {
417                    "kind": { "type": "string", "enum": ["single"] },
418                    "selection": { "type": "string" },
419                    "note": { "type": "string" }
420                },
421                "required": ["kind", "selection"],
422                "additionalProperties": false
423            },
424            "execution_mode": {
425                "type": "string",
426                "enum": ["current_session", "fresh_context"]
427            },
428            "confirmation_display": { "type": "string" },
429            "next_turn_input": { "type": "string" }
430        },
431        "required": ["approved", "plan_path"],
432        "additionalProperties": false
433    })
434}
435
436pub struct PlanModePluginFactory {
437    config: PlanModePluginConfig,
438    prompt: Option<Arc<dyn PlanModePrompt>>,
439}
440
441impl Default for PlanModePluginFactory {
442    fn default() -> Self {
443        Self::new(PlanModePluginConfig::default())
444    }
445}
446
447impl PlanModePluginFactory {
448    pub fn new(config: PlanModePluginConfig) -> Self {
449        Self {
450            config,
451            prompt: None,
452        }
453    }
454
455    pub fn with_prompt(mut self, prompt: Arc<dyn PlanModePrompt>) -> Self {
456        self.prompt = Some(prompt);
457        self
458    }
459}
460
461impl PluginFactory for PlanModePluginFactory {
462    fn id(&self) -> &'static str {
463        "plan_mode"
464    }
465
466    fn build(&self, _ctx: &PluginSessionContext) -> Result<Arc<dyn SessionPlugin>, PluginError> {
467        Ok(Arc::new(PlanModePlugin {
468            state: Arc::new(Mutex::new(PlanModeState::default())),
469            config: self.config.clone(),
470            prompt: self.prompt.clone(),
471        }))
472    }
473}
474
475struct PlanModePlugin {
476    state: Arc<Mutex<PlanModeState>>,
477    config: PlanModePluginConfig,
478    prompt: Option<Arc<dyn PlanModePrompt>>,
479}
480
481impl SessionPlugin for PlanModePlugin {
482    fn id(&self) -> &'static str {
483        "plan_mode"
484    }
485
486    fn register(&self, reg: &mut PluginRegistrar) -> Result<(), PluginError> {
487        reg.tools().provider(Arc::new(plan_mode_provider(
488            Arc::clone(&self.state),
489            self.prompt.clone(),
490        )))?;
491
492        let before_turn_state = Arc::clone(&self.state);
493        reg.turn().before(Arc::new(move |ctx| {
494            let state = Arc::clone(&before_turn_state);
495            Box::pin(async move {
496                let plan_path = {
497                    let mut state = state.lock().map_err(|_| {
498                        PluginError::Session("plan mode state poisoned".to_string())
499                    })?;
500                    let should_inject = state.prepare_turn();
501                    if !should_inject {
502                        return Ok(Vec::new());
503                    }
504                    state.ensure_plan_path_from_state(&ctx.state.to_snapshot())?
505                };
506                seed_plan_template(&plan_path).map_err(PluginError::Session)?;
507                let report = read_plan_report(&plan_path).map_err(PluginError::Session)?;
508                Ok(vec![
509                    PluginDirective::emit_runtime_events(vec![plan_protocol_state_event(
510                        &ctx.session_id,
511                        true,
512                        Some(&report),
513                    )?]),
514                    PluginDirective::EnqueueMessages {
515                        messages: vec![plan_mode_guidance_message(&plan_path)],
516                    },
517                ])
518            })
519        }));
520
521        let checkpoint_state = Arc::clone(&self.state);
522        reg.turn().checkpoint(Arc::new(move |ctx| {
523            let state = Arc::clone(&checkpoint_state);
524            Box::pin(async move {
525                let plan_path = {
526                    let mut state = state.lock().map_err(|_| {
527                        PluginError::Session("plan mode state poisoned".to_string())
528                    })?;
529                    let should_inject = state.checkpoint_injection_needed();
530                    if !should_inject {
531                        return Ok(Vec::new());
532                    }
533                    state.ensure_plan_path_from_state(&ctx.state.to_snapshot())?
534                };
535                seed_plan_template(&plan_path).map_err(PluginError::Session)?;
536                let report = read_plan_report(&plan_path).map_err(PluginError::Session)?;
537                Ok(vec![
538                    PluginDirective::emit_runtime_events(vec![plan_protocol_state_event(
539                        &ctx.session_id,
540                        true,
541                        Some(&report),
542                    )?]),
543                    PluginDirective::EnqueueMessages {
544                        messages: vec![plan_mode_guidance_message(&plan_path)],
545                    },
546                ])
547            })
548        }));
549
550        let after_turn_state = Arc::clone(&self.state);
551        reg.turn().after(Arc::new(move |_ctx| {
552            let state = Arc::clone(&after_turn_state);
553            Box::pin(async move {
554                state
555                    .lock()
556                    .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
557                    .finish_turn();
558                Ok(Vec::new())
559            })
560        }));
561
562        let before_tool_state = Arc::clone(&self.state);
563        let before_tool_config = self.config.clone();
564        reg.tool_calls().before(Arc::new(move |ctx| {
565            let state = Arc::clone(&before_tool_state);
566            let config = before_tool_config.clone();
567            Box::pin(async move {
568                let enabled = state
569                    .lock()
570                    .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
571                    .enabled;
572                if !enabled {
573                    return Ok(Vec::new());
574                }
575
576                if ctx.tool_name != "plan_exit" && !config.allowed_tools.contains(&ctx.tool_name) {
577                    return Ok(vec![PluginDirective::AbortTurn {
578                        code: "plan_mode_tool_blocked".to_string(),
579                        message: format!(
580                            "Plan mode blocks `{}`. Use planning tools or `plan.exit`.",
581                            ctx.tool_name
582                        ),
583                    }]);
584                }
585
586                if ctx.tool_name == "apply_patch" {
587                    let snapshot = ctx.session_snapshot().await?;
588                    let plan_path = ensure_plan_path_from_snapshot(&state, &snapshot)?;
589                    if let Err(message) = patch_allowed_for_plan_file(&ctx.args, &plan_path) {
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 || ctx.tool_name != "apply_patch" || !ctx.result.is_success() {
652                    return Ok(Vec::new());
653                }
654
655                let snapshot = ctx.session_snapshot().await?;
656                let path = ensure_plan_path_from_snapshot(&state, &snapshot)?;
657                let report = read_plan_report(&path).map_err(PluginError::Session)?;
658                Ok(vec![PluginDirective::emit_runtime_events(vec![
659                    plan_protocol_state_event(&ctx.session_id, true, Some(&report))?,
660                ])])
661            })
662        }));
663
664        // Membership is the execution gate. While plan mode is enabled, remove
665        // every tool that is not allowed (and not `plan_exit`); while it is
666        // disabled, remove `plan_exit` so it is not callable.
667        let tool_catalog_state = Arc::clone(&self.state);
668        let tool_catalog_config = self.config.clone();
669        reg.tool_catalog().contribute(Arc::new(move |ctx| {
670            let enabled = tool_catalog_state
671                .lock()
672                .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
673                .enabled;
674            if !enabled {
675                return Ok(ToolCatalogContribution::remove_tools(["plan_exit"]));
676            }
677
678            let remove = ctx
679                .tools
680                .iter()
681                .filter(|tool| {
682                    tool.name != "plan_exit"
683                        && !tool_catalog_config.allowed_tools.contains(&tool.name)
684                })
685                .map(|tool| tool.name.clone())
686                .collect::<Vec<_>>();
687            Ok(ToolCatalogContribution { remove })
688        }));
689
690        // The plan-mode tool note is an ordinary prompt contribution, gated on
691        // `plan_exit` being a catalog member (i.e. plan mode enabled).
692        let prompt_state = Arc::clone(&self.state);
693        reg.prompt().contribute(Arc::new(move |_ctx| {
694            let state = Arc::clone(&prompt_state);
695            Box::pin(async move {
696                let guard = state
697                    .lock()
698                    .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?;
699                if !guard.enabled {
700                    return Ok(Vec::new());
701                }
702                let note = plan_mode_tool_note(guard.plan_path().as_deref());
703                Ok(vec![
704                    lash_core::PromptContribution::execution("Plan Mode", note)
705                        .requires_tool("plan_exit"),
706                ])
707            })
708        }));
709
710        register_plan_mode_op::<PlanModeEnableOp>(reg, Arc::clone(&self.state))?;
711        register_plan_mode_op::<PlanModeDisableOp>(reg, Arc::clone(&self.state))?;
712        register_plan_mode_op::<PlanModeToggleOp>(reg, Arc::clone(&self.state))?;
713
714        Ok(())
715    }
716
717    fn snapshot(
718        &self,
719        _writer: &mut dyn SnapshotWriter,
720    ) -> Result<PluginSnapshotMeta, PluginError> {
721        let snapshot = self
722            .state
723            .lock()
724            .map_err(|_| PluginError::Snapshot("plan mode state poisoned".to_string()))?
725            .snapshot();
726        Ok(PluginSnapshotMeta {
727            plugin_id: self.id().to_string(),
728            plugin_version: self.version().to_string(),
729            revision: snapshot.generation,
730            state: Some(json!({
731                "enabled": snapshot.enabled,
732                "generation": snapshot.generation,
733                "plan_path": snapshot.plan_path,
734            })),
735        })
736    }
737
738    fn restore(
739        &self,
740        meta: &PluginSnapshotMeta,
741        _reader: &dyn SnapshotReader,
742    ) -> Result<(), PluginError> {
743        let snapshot = meta
744            .state
745            .clone()
746            .map(serde_json::from_value::<PlanModeSnapshot>)
747            .transpose()
748            .map_err(|err| PluginError::Snapshot(err.to_string()))?
749            .unwrap_or_default();
750        self.state
751            .lock()
752            .map_err(|_| PluginError::Snapshot("plan mode state poisoned".to_string()))?
753            .restore_snapshot(snapshot);
754        Ok(())
755    }
756
757    fn snapshot_revision(&self) -> u64 {
758        self.state
759            .lock()
760            .map(|state| state.generation)
761            .unwrap_or_default()
762    }
763}
764
765fn register_plan_mode_op<Op>(
766    reg: &mut PluginRegistrar,
767    state: Arc<Mutex<PlanModeState>>,
768) -> Result<(), PluginError>
769where
770    Op: PluginCommand<Args = PlanModeExternalArgs, Output = PlanModeExternalStatus>,
771{
772    reg.operations()
773        .typed_command::<Op, _, _>(move |ctx, _args| {
774            let state = Arc::clone(&state);
775            async move {
776                let Some(session_id) = ctx.session_id else {
777                    return Err(PluginOperationFailure::new(format!(
778                        "{} requires session_id",
779                        Op::NAME
780                    )));
781                };
782                let target_enabled = match state.lock() {
783                    Ok(guard) => match Op::NAME {
784                        "plan_mode.enable" => true,
785                        "plan_mode.disable" => false,
786                        "plan_mode.toggle" => !guard.enabled,
787                        _ => unreachable!(),
788                    },
789                    Err(_) => return Err(PluginOperationFailure::new("plan mode state poisoned")),
790                };
791                let enabled = set_plan_mode_enabled_state(&state, target_enabled)?;
792                let report =
793                    ensure_plan_report(&state, &session_id, &ctx.sessions, enabled).await?;
794                let status = plan_mode_payload(&session_id, enabled, Some(&report));
795                Ok(
796                    PluginCommandOutcome::new(status).with_events(vec![plan_protocol_state_event(
797                        &session_id,
798                        enabled,
799                        Some(&report),
800                    )?]),
801                )
802            }
803        })
804}
805
806#[cfg(test)]
807mod tests {
808    use super::{
809        PLAN_TEMPLATE, plan_exit_fresh_context_input, plan_exit_next_turn_input,
810        plan_exit_tool_definition, read_plan_report,
811    };
812
813    #[test]
814    fn plan_exit_contract_documents_decision_result() {
815        let definition = plan_exit_tool_definition();
816
817        assert_eq!(
818            definition.contract.input_schema["additionalProperties"],
819            serde_json::json!(false)
820        );
821        assert_eq!(
822            definition.contract.output_schema["required"],
823            serde_json::json!(["approved", "plan_path"])
824        );
825        assert!(
826            definition
827                .contract
828                .output_schema
829                .to_string()
830                .contains("execution_mode")
831        );
832    }
833
834    #[test]
835    fn plan_exit_next_turn_input_appends_user_note() {
836        assert_eq!(
837            plan_exit_next_turn_input(
838                ".lash/plans/run-session.md",
839                Some("start with the safe slice"),
840            ),
841            "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"
842        );
843        assert_eq!(
844            plan_exit_next_turn_input(".lash/plans/run-session.md", Some("   ")),
845            "The user approved the plan. Execute the plan in `.lash/plans/run-session.md` now — start immediately, do not ask for confirmation."
846        );
847    }
848
849    #[test]
850    fn plan_exit_fresh_context_input_is_short_and_direct() {
851        assert_eq!(
852            plan_exit_fresh_context_input(".lash/plans/run-session.md"),
853            "Do a full, faithful implementation of the plan found at: .lash/plans/run-session.md"
854        );
855    }
856
857    #[test]
858    fn seeded_template_is_readable_without_validation() {
859        let dir = tempfile::tempdir().expect("tempdir");
860        let path = dir.path().join("plan.md");
861        std::fs::write(&path, PLAN_TEMPLATE).expect("write template");
862        let report = read_plan_report(&path).expect("report");
863        assert_eq!(report.content.as_deref(), Some(PLAN_TEMPLATE));
864    }
865}