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};
17use lash_tool_support::{
18    LashlangToolBinding, StaticToolExecute, StaticToolProvider, ToolDefinitionLashlangExt,
19    resolve_under,
20};
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        "edit",
51        "write",
52        "plan_exit",
53    ]
54    .into_iter()
55    .map(str::to_string)
56    .collect()
57}
58
59fn fresh_context_frame_id() -> String {
60    format!("plan-frame-{}", uuid::Uuid::new_v4().simple())
61}
62
63fn plan_protocol_state_event(
64    session_id: &str,
65    enabled: bool,
66    report: Option<&PlanReport>,
67) -> Result<lash_core::PluginRuntimeEvent, PluginError> {
68    Ok(lash_core::PluginRuntimeEvent::Custom {
69        name: PLAN_MODE_STATE_EVENT.to_string(),
70        payload: serde_json::to_value(plan_mode_payload(session_id, enabled, report)).map_err(
71            |err| PluginError::Session(format!("failed to encode plan mode state: {err}")),
72        )?,
73    })
74}
75
76#[derive(Clone, Debug)]
77pub struct PlanModePluginConfig {
78    pub allowed_tools: BTreeSet<String>,
79}
80
81impl Default for PlanModePluginConfig {
82    fn default() -> Self {
83        Self {
84            allowed_tools: default_allowed_tools(),
85        }
86    }
87}
88
89impl PlanModePluginConfig {
90    pub fn with_allowed_tools<I, S>(mut self, allowed_tools: I) -> Self
91    where
92        I: IntoIterator<Item = S>,
93        S: Into<String>,
94    {
95        self.allowed_tools = allowed_tools.into_iter().map(Into::into).collect();
96        self.allowed_tools.insert("plan_exit".to_string());
97        self
98    }
99}
100
101#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize, JsonSchema)]
102pub struct PlanModeExternalArgs {}
103
104#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize, JsonSchema)]
105pub struct PlanModeExternalStatus {
106    pub session_id: String,
107    pub enabled: bool,
108    pub plan_path: Option<String>,
109}
110
111pub struct PlanModeEnableOp;
112pub struct PlanModeDisableOp;
113pub struct PlanModeToggleOp;
114
115impl PluginOperation for PlanModeEnableOp {
116    const NAME: &'static str = "plan_mode.enable";
117    const DESCRIPTION: &'static str = "Enable plan mode for this session.";
118    const SESSION_PARAM: SessionParam = SessionParam::Required;
119    type Args = PlanModeExternalArgs;
120    type Output = PlanModeExternalStatus;
121}
122
123impl PluginCommand for PlanModeEnableOp {}
124
125impl PluginOperation for PlanModeDisableOp {
126    const NAME: &'static str = "plan_mode.disable";
127    const DESCRIPTION: &'static str = "Disable plan mode for this session.";
128    const SESSION_PARAM: SessionParam = SessionParam::Required;
129    type Args = PlanModeExternalArgs;
130    type Output = PlanModeExternalStatus;
131}
132
133impl PluginCommand for PlanModeDisableOp {}
134
135impl PluginOperation for PlanModeToggleOp {
136    const NAME: &'static str = "plan_mode.toggle";
137    const DESCRIPTION: &'static str = "Toggle plan mode for this session.";
138    const SESSION_PARAM: SessionParam = SessionParam::Required;
139    type Args = PlanModeExternalArgs;
140    type Output = PlanModeExternalStatus;
141}
142
143impl PluginCommand for PlanModeToggleOp {}
144
145async fn ensure_plan_path<H>(
146    state: &Arc<Mutex<PlanModeState>>,
147    session_id: &str,
148    host: &Arc<H>,
149) -> Result<PathBuf, PluginError>
150where
151    H: lash_core::plugin::runtime_host::SessionStateService + ?Sized,
152{
153    if let Some(path) = state
154        .lock()
155        .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
156        .plan_path()
157    {
158        return Ok(path);
159    }
160
161    let snapshot = host.snapshot_session(session_id).await?;
162    let run_session_id =
163        effective_run_session_id(&snapshot.session_id, &snapshot.policy).to_string();
164    let path = resolve_plan_path(&run_session_id).map_err(PluginError::Session)?;
165    state
166        .lock()
167        .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
168        .set_plan_path(path.clone());
169    Ok(path)
170}
171
172fn ensure_plan_path_from_snapshot(
173    state: &Arc<Mutex<PlanModeState>>,
174    snapshot: &lash_core::SessionSnapshot,
175) -> Result<PathBuf, PluginError> {
176    if let Some(path) = state
177        .lock()
178        .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
179        .plan_path()
180    {
181        return Ok(path);
182    }
183    let run_session_id =
184        effective_run_session_id(&snapshot.session_id, &snapshot.policy).to_string();
185    let path = resolve_plan_path(&run_session_id).map_err(PluginError::Session)?;
186    state
187        .lock()
188        .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?
189        .set_plan_path(path.clone());
190    Ok(path)
191}
192
193async fn ensure_plan_report_for_tool_context(
194    state: &Arc<Mutex<PlanModeState>>,
195    context: &ToolContext<'_>,
196    seed_if_missing: bool,
197) -> Result<PlanReport, PluginError> {
198    let snapshot = context.sessions().snapshot_current().await?;
199    let path = ensure_plan_path_from_snapshot(state, &snapshot)?;
200    if seed_if_missing {
201        seed_plan_template(&path).map_err(PluginError::Session)?;
202    }
203    read_plan_report(&path).map_err(PluginError::Session)
204}
205
206async fn ensure_plan_report<H>(
207    state: &Arc<Mutex<PlanModeState>>,
208    session_id: &str,
209    host: &Arc<H>,
210    seed_if_missing: bool,
211) -> Result<PlanReport, PluginError>
212where
213    H: lash_core::plugin::runtime_host::SessionStateService + ?Sized,
214{
215    let path = ensure_plan_path(state, session_id, host).await?;
216    if seed_if_missing {
217        seed_plan_template(&path).map_err(PluginError::Session)?;
218    }
219    read_plan_report(&path).map_err(PluginError::Session)
220}
221
222/// Flip plan-mode enablement. Tool Catalog membership (`plan_exit` visibility,
223/// suppression of non-allowed tools) is derived per turn by the plan-mode
224/// catalog contribution keyed on this state, so there is no registry mutation
225/// here.
226fn set_plan_mode_enabled_state(
227    state: &Arc<Mutex<PlanModeState>>,
228    enabled: bool,
229) -> Result<bool, PluginError> {
230    let mut guard = state
231        .lock()
232        .map_err(|_| PluginError::Session("plan mode state poisoned".to_string()))?;
233    if guard.enabled != enabled {
234        guard.set_enabled(enabled);
235    }
236    Ok(enabled)
237}
238
239fn plan_mode_payload(
240    session_id: &str,
241    enabled: bool,
242    report: Option<&PlanReport>,
243) -> PlanModeExternalStatus {
244    PlanModeExternalStatus {
245        session_id: session_id.to_string(),
246        enabled,
247        plan_path: report.map(|value| value.display_path.clone()),
248    }
249}
250
251fn file_mutation_allowed_for_plan_file(
252    tool_name: &str,
253    args: &serde_json::Value,
254    plan_path: &Path,
255) -> Result<(), String> {
256    let path = args
257        .get("path")
258        .and_then(|value| value.as_str())
259        .ok_or_else(|| format!("plan mode requires `{tool_name}.path`"))?;
260    let cwd = std::env::current_dir().map_err(|err| format!("Failed to determine cwd: {err}"))?;
261    let resolved = resolve_under(&cwd, Path::new(path));
262    if resolved != plan_path {
263        return Err(format!(
264            "plan mode only allows `{tool_name}` to edit `{}`",
265            plan_display_path(plan_path)
266        ));
267    }
268    if tool_name == "edit"
269        && args
270            .get("edits")
271            .and_then(|value| value.as_array())
272            .is_none_or(Vec::is_empty)
273    {
274        return Err(
275            "plan mode requires `edit.edits` to contain at least one replacement".to_string(),
276        );
277    }
278    Ok(())
279}
280
281#[derive(Clone)]
282struct PlanModeTools {
283    state: Arc<Mutex<PlanModeState>>,
284    prompt: Option<Arc<dyn PlanModePrompt>>,
285}
286
287impl PlanModeTools {
288    async fn execute_plan_exit(&self, context: &ToolContext<'_>) -> ToolResult {
289        let enabled = match self.state.lock() {
290            Ok(guard) => guard.enabled,
291            Err(_) => return ToolResult::err(json!("plan mode state poisoned")),
292        };
293        if !enabled {
294            return ToolResult::err(json!("plan mode is not active"));
295        }
296
297        let report = match ensure_plan_report_for_tool_context(&self.state, context, true).await {
298            Ok(report) => report,
299            Err(err) => return ToolResult::err(json!(err.to_string())),
300        };
301
302        let Some(prompt) = &self.prompt else {
303            return ToolResult::err(json!(
304                "plan approval prompts are unavailable in this session"
305            ));
306        };
307        let answer = match prompt
308            .prompt_user(
309                PlanModePromptRequest::single(
310                    format!("Review the plan in `{}`. What next?", report.display_path),
311                    vec![
312                        "Start implementing now".to_string(),
313                        "Keep planning".to_string(),
314                        "Start in fresh context".to_string(),
315                    ],
316                )
317                .with_review("PLAN", report.approval_content())
318                .with_optional_note(),
319            )
320            .await
321        {
322            Ok(answer) => answer,
323            Err(err) => return ToolResult::err(json!(err.to_string())),
324        };
325
326        let selection = match &answer {
327            PlanModePromptResponse::Single { selection, .. } => selection.as_str(),
328        };
329        if selection == "Keep planning" {
330            return ToolResult::ok(json!({
331                "approved": false,
332                "plan_path": report.display_path,
333                "answer": answer,
334            }));
335        }
336
337        let note = match &answer {
338            PlanModePromptResponse::Single { note, .. } => note.clone(),
339        };
340
341        if let Err(err) = set_plan_mode_enabled_state(&self.state, false) {
342            return ToolResult::err(json!(err.to_string()));
343        }
344
345        if selection == "Start in fresh context" {
346            return ToolResult::ok(json!({
347                "approved": true,
348                "plan_path": report.display_path,
349                "execution_mode": "fresh_context",
350            }));
351        }
352
353        ToolResult::ok(json!({
354            "approved": true,
355            "answer": answer,
356            "confirmation_display": plan_exit_confirmation_display(selection, note.as_deref()),
357            "plan_path": report.display_path,
358            "execution_mode": "current_session",
359            "next_turn_input": plan_exit_next_turn_input(&report.display_path, note.as_deref()),
360        }))
361    }
362}
363
364fn plan_mode_provider(
365    state: Arc<Mutex<PlanModeState>>,
366    prompt: Option<Arc<dyn PlanModePrompt>>,
367) -> StaticToolProvider<PlanModeTools> {
368    StaticToolProvider::new(
369        vec![plan_exit_tool_definition()],
370        PlanModeTools { state, prompt },
371    )
372}
373
374#[async_trait::async_trait]
375impl StaticToolExecute for PlanModeTools {
376    async fn execute(&self, call: ToolCall<'_>) -> ToolResult {
377        match call.name {
378            "plan_exit" => self.execute_plan_exit(call.context).await,
379            other => ToolResult::err_fmt(format_args!("Unknown tool: {other}")),
380        }
381    }
382}
383
384fn plan_exit_tool_definition() -> ToolDefinition {
385    ToolDefinition::raw(
386        "tool:plan_exit",
387        "plan_exit",
388        "Ask whether to exit plan mode.",
389        plan_exit_input_schema(),
390        plan_exit_output_schema(),
391    )
392    .with_examples(vec!["await plan.exit({})?".into()])
393    .with_lashlang_binding(LashlangToolBinding::new(["plan"], "exit"))
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}