fix: disable plan mode when moving to next workflow step (#525)

2026-04-03 15:07:46 +01:00
parent 4c0c9a7fb0
commit 82eb59d91e
6 changed files with 510 additions and 69 deletions
@@ -21,6 +21,7 @@ import {
  usePendingPRFeedback,
 } from "@/hooks/domains/comments/use-pending-comments";
 import { buildContextItems } from "../chat-context-items";
+import { useAutoDisablePlanMode, usePlanLayoutHandlers } from "./use-plan-mode-helpers";
 import type { ContextItem } from "@/lib/types/context";
 import type { DiffComment } from "@/lib/diff/types";
 import type { PlanComment, PRFeedbackComment } from "@/lib/state/slices/comments";
@@ -60,18 +61,22 @@ type AutoApplyPlanLayoutOpts = {
  resolvedSessionId: string | null;
  taskId: string | null;
  sessionMetaPlanMode: boolean;
+  currentStepHasPlanMode: boolean;
  setActiveDocument: (sid: string, doc: ActiveDocument | null) => void;
  applyBuiltInPreset: (preset: BuiltInPreset) => void;
  setPlanMode: (sid: string, enabled: boolean) => void;
  addContextFile: (sid: string, file: { path: string; name: string }) => void;
 };

-/** Auto-apply plan layout when session metadata has plan_mode enabled. */
+/** Auto-apply plan layout when session metadata has plan_mode enabled AND the current
+ *  step actually configures enable_plan_mode. The step check guards against stale
+ *  sessionMetaPlanMode from deepMerge hydration preserving deleted metadata keys. */
 function useAutoApplyPlanLayout(opts: AutoApplyPlanLayoutOpts) {
  const {
    resolvedSessionId,
    taskId,
    sessionMetaPlanMode,
+    currentStepHasPlanMode,
    setActiveDocument,
    applyBuiltInPreset,
    setPlanMode,
@@ -80,8 +85,17 @@ function useAutoApplyPlanLayout(opts: AutoApplyPlanLayoutOpts) {
  const autoAppliedPlanSessionRef = useRef<string | null>(null);
  useEffect(() => {
    if (!resolvedSessionId || !taskId) return;
+    // Reset the guard when plan mode is disabled so future plan-mode steps
+    // in the same session can be auto-applied (e.g. after proceeding away and back).
+    if (!sessionMetaPlanMode && autoAppliedPlanSessionRef.current === resolvedSessionId) {
+      autoAppliedPlanSessionRef.current = null;
+      return;
+    }
    if (autoAppliedPlanSessionRef.current === resolvedSessionId) return;
-    if (sessionMetaPlanMode) {
+    // Only auto-apply if both the session metadata AND the current step agree on plan mode.
+    // sessionMetaPlanMode can be stale (deepMerge hydration preserves deleted keys),
+    // so we cross-check with the step's actual configuration.
+    if (sessionMetaPlanMode && currentStepHasPlanMode) {
      autoAppliedPlanSessionRef.current = resolvedSessionId;
      setActiveDocument(resolvedSessionId, { type: "plan", taskId });
      applyBuiltInPreset("plan");
@@ -92,6 +106,7 @@ function useAutoApplyPlanLayout(opts: AutoApplyPlanLayoutOpts) {
    resolvedSessionId,
    taskId,
    sessionMetaPlanMode,
+    currentStepHasPlanMode,
    setActiveDocument,
    applyBuiltInPreset,
    setPlanMode,
@@ -115,13 +130,19 @@ export function usePlanMode(resolvedSessionId: string | null, taskId: string | n
  const planModeFromStore = useAppStore((state) =>
    resolvedSessionId ? (state.chatInput.planModeBySessionId[resolvedSessionId] ?? false) : false,
  );
-
  const sessionMetaPlanMode = useAppStore((state) =>
    resolvedSessionId
      ? state.taskSessions.items[resolvedSessionId]?.metadata?.plan_mode === true
      : false,
  );
-
+  const currentStepHasPlanMode = useAppStore((s) => {
+    if (!taskId) return false;
+    const task = s.kanban.tasks.find((t) => t.id === taskId);
+    const stepId = task?.workflowStepId;
+    if (!stepId) return false;
+    const step = s.kanban.steps.find((st) => st.id === stepId);
+    return step?.events?.on_enter?.some((a) => a.type === "enable_plan_mode") ?? false;
+  });
  const planModeEnabled = planModeFromStore;
  const planLayoutVisible = activeDocument?.type === "plan";

@@ -138,68 +159,37 @@ export function usePlanMode(resolvedSessionId: string | null, taskId: string | n
    resolvedSessionId,
    taskId,
    sessionMetaPlanMode,
+    currentStepHasPlanMode,
    setActiveDocument,
    applyBuiltInPreset,
    setPlanMode,
    addContextFile,
  });

+  useAutoDisablePlanMode({
+    resolvedSessionId,
+    taskId,
+    sessionMetaPlanMode,
+    planModeFromStore,
+    applyBuiltInPreset,
+    closeDocument,
+    setActiveDocument,
+    setPlanMode,
+    removeContextFile,
+  });
+
  const refocusChatAfterLayout = useRefocusChatAfterLayout();
-
-  // Toggle only the plan layout (plan panel + preset) without changing chat input plan mode state
-  const togglePlanLayout = useCallback(
-    (show: boolean) => {
-      if (!resolvedSessionId || !taskId) return;
-      if (show) {
-        setActiveDocument(resolvedSessionId, { type: "plan", taskId });
-        applyBuiltInPreset("plan");
-      } else {
-        applyBuiltInPreset("default");
-        closeDocument(resolvedSessionId);
-        setActiveDocument(resolvedSessionId, null);
-      }
-      refocusChatAfterLayout();
-    },
-    [
-      resolvedSessionId,
-      taskId,
-      setActiveDocument,
-      applyBuiltInPreset,
-      closeDocument,
-      refocusChatAfterLayout,
-    ],
-  );
-
-  // Full plan mode toggle: layout + chat input state (context file, border, submit style)
-  const handlePlanModeChange = useCallback(
-    (enabled: boolean) => {
-      if (!resolvedSessionId || !taskId) return;
-      if (enabled) {
-        setActiveDocument(resolvedSessionId, { type: "plan", taskId });
-        applyBuiltInPreset("plan");
-        setPlanMode(resolvedSessionId, true);
-        addContextFile(resolvedSessionId, { path: PLAN_CONTEXT_PATH, name: "Plan" });
-      } else {
-        applyBuiltInPreset("default");
-        closeDocument(resolvedSessionId);
-        setActiveDocument(resolvedSessionId, null);
-        setPlanMode(resolvedSessionId, false);
-        removeContextFile(resolvedSessionId, PLAN_CONTEXT_PATH);
-      }
-      refocusChatAfterLayout();
-    },
-    [
-      resolvedSessionId,
-      taskId,
-      setActiveDocument,
-      applyBuiltInPreset,
-      closeDocument,
-      setPlanMode,
-      addContextFile,
-      removeContextFile,
-      refocusChatAfterLayout,
-    ],
-  );
+  const { togglePlanLayout, handlePlanModeChange } = usePlanLayoutHandlers({
+    resolvedSessionId,
+    taskId,
+    setActiveDocument,
+    applyBuiltInPreset,
+    closeDocument,
+    setPlanMode,
+    addContextFile,
+    removeContextFile,
+    refocusChatAfterLayout,
+  });

  return {
    planModeEnabled,
@@ -0,0 +1,155 @@
+"use client";
+
+import { useCallback, useEffect, useRef } from "react";
+import type { ActiveDocument } from "@/lib/state/slices/ui/types";
+import type { BuiltInPreset } from "@/lib/state/layout-manager/presets";
+
+const PLAN_CONTEXT_PATH = "plan:context";
+
+// --- Auto-disable plan mode ---
+
+type AutoDisablePlanOpts = {
+  resolvedSessionId: string | null;
+  taskId: string | null;
+  sessionMetaPlanMode: boolean;
+  planModeFromStore: boolean;
+  applyBuiltInPreset: (preset: BuiltInPreset) => void;
+  closeDocument: (sid: string) => void;
+  setActiveDocument: (sid: string, doc: ActiveDocument | null) => void;
+  setPlanMode: (sid: string, enabled: boolean) => void;
+  removeContextFile: (sid: string, path: string) => void;
+};
+
+/**
+ * Auto-disable plan mode when the backend clears plan_mode from session metadata.
+ *
+ * Triggers when `sessionMetaPlanMode` transitions from true to false — the backend explicitly
+ * cleared plan_mode (e.g., on_exit: disable_plan_mode, or entering a non-plan-mode step).
+ *
+ * User-initiated step moves (proceed button, stepper) disable plan mode directly in their
+ * click handlers rather than relying on this effect.
+ */
+export function useAutoDisablePlanMode(opts: AutoDisablePlanOpts) {
+  const {
+    resolvedSessionId,
+    taskId,
+    sessionMetaPlanMode,
+    planModeFromStore,
+    applyBuiltInPreset,
+    closeDocument,
+    setActiveDocument,
+    setPlanMode,
+    removeContextFile,
+  } = opts;
+
+  const prevSessionMetaPlanRef = useRef(sessionMetaPlanMode);
+
+  useEffect(() => {
+    const wasPlanMode = prevSessionMetaPlanRef.current;
+    prevSessionMetaPlanRef.current = sessionMetaPlanMode;
+
+    if (!resolvedSessionId || !taskId) return;
+
+    if (wasPlanMode && !sessionMetaPlanMode && planModeFromStore) {
+      applyBuiltInPreset("default");
+      closeDocument(resolvedSessionId);
+      setActiveDocument(resolvedSessionId, null);
+      setPlanMode(resolvedSessionId, false);
+      removeContextFile(resolvedSessionId, PLAN_CONTEXT_PATH);
+    }
+  }, [
+    resolvedSessionId,
+    taskId,
+    sessionMetaPlanMode,
+    planModeFromStore,
+    applyBuiltInPreset,
+    closeDocument,
+    setActiveDocument,
+    setPlanMode,
+    removeContextFile,
+  ]);
+}
+
+// --- Plan layout handlers ---
+
+type PlanLayoutHandlersOpts = {
+  resolvedSessionId: string | null;
+  taskId: string | null;
+  setActiveDocument: (sid: string, doc: ActiveDocument | null) => void;
+  applyBuiltInPreset: (preset: BuiltInPreset) => void;
+  closeDocument: (sid: string) => void;
+  setPlanMode: (sid: string, enabled: boolean) => void;
+  addContextFile: (sid: string, file: { path: string; name: string }) => void;
+  removeContextFile: (sid: string, path: string) => void;
+  refocusChatAfterLayout: () => void;
+};
+
+/** Returns togglePlanLayout and handlePlanModeChange callbacks. */
+export function usePlanLayoutHandlers(opts: PlanLayoutHandlersOpts) {
+  const {
+    resolvedSessionId,
+    taskId,
+    setActiveDocument,
+    applyBuiltInPreset,
+    closeDocument,
+    setPlanMode,
+    addContextFile,
+    removeContextFile,
+    refocusChatAfterLayout,
+  } = opts;
+
+  const togglePlanLayout = useCallback(
+    (show: boolean) => {
+      if (!resolvedSessionId || !taskId) return;
+      if (show) {
+        setActiveDocument(resolvedSessionId, { type: "plan", taskId });
+        applyBuiltInPreset("plan");
+      } else {
+        applyBuiltInPreset("default");
+        closeDocument(resolvedSessionId);
+        setActiveDocument(resolvedSessionId, null);
+      }
+      refocusChatAfterLayout();
+    },
+    [
+      resolvedSessionId,
+      taskId,
+      setActiveDocument,
+      applyBuiltInPreset,
+      closeDocument,
+      refocusChatAfterLayout,
+    ],
+  );
+
+  const handlePlanModeChange = useCallback(
+    (enabled: boolean) => {
+      if (!resolvedSessionId || !taskId) return;
+      if (enabled) {
+        setActiveDocument(resolvedSessionId, { type: "plan", taskId });
+        applyBuiltInPreset("plan");
+        setPlanMode(resolvedSessionId, true);
+        addContextFile(resolvedSessionId, { path: PLAN_CONTEXT_PATH, name: "Plan" });
+      } else {
+        applyBuiltInPreset("default");
+        closeDocument(resolvedSessionId);
+        setActiveDocument(resolvedSessionId, null);
+        setPlanMode(resolvedSessionId, false);
+        removeContextFile(resolvedSessionId, PLAN_CONTEXT_PATH);
+      }
+      refocusChatAfterLayout();
+    },
+    [
+      resolvedSessionId,
+      taskId,
+      setActiveDocument,
+      applyBuiltInPreset,
+      closeDocument,
+      setPlanMode,
+      addContextFile,
+      removeContextFile,
+      refocusChatAfterLayout,
+    ],
+  );
+
+  return { togglePlanLayout, handlePlanModeChange };
+}
@@ -7,6 +7,10 @@ import { Button } from "@kandev/ui/button";
 import { IconArrowRight } from "@tabler/icons-react";
 import { moveTask } from "@/lib/api";
 import { StepCapabilityIcons } from "@/components/step-capability-icons";
+import { useAppStore } from "@/components/state-provider";
+import { useContextFilesStore } from "@/lib/state/context-files-store";
+import { useLayoutStore } from "@/lib/state/layout-store";
+import { useDockviewStore } from "@/lib/state/dockview-store";
 import type { KanbanStepEvents } from "@/lib/state/slices/kanban/types";

 type Step = {
@@ -20,6 +24,38 @@ type Step = {
  is_start_step?: boolean;
 };

+const PLAN_CONTEXT_PATH = "plan:context";
+
+/** Returns a callback that disables plan mode for the active session of a task. */
+function useDisablePlanMode() {
+  const activeSessionId = useAppStore((s) => s.tasks.activeSessionId);
+  const planModeEnabled = useAppStore((s) =>
+    activeSessionId ? (s.chatInput.planModeBySessionId[activeSessionId] ?? false) : false,
+  );
+  const setPlanMode = useAppStore((s) => s.setPlanMode);
+  const setActiveDocument = useAppStore((s) => s.setActiveDocument);
+  const closeDocument = useLayoutStore((s) => s.closeDocument);
+  const removeContextFile = useContextFilesStore((s) => s.removeFile);
+  const applyBuiltInPreset = useDockviewStore((s) => s.applyBuiltInPreset);
+
+  return useCallback(() => {
+    if (!activeSessionId || !planModeEnabled) return;
+    applyBuiltInPreset("default");
+    closeDocument(activeSessionId);
+    setActiveDocument(activeSessionId, null);
+    setPlanMode(activeSessionId, false);
+    removeContextFile(activeSessionId, PLAN_CONTEXT_PATH);
+  }, [
+    activeSessionId,
+    planModeEnabled,
+    setPlanMode,
+    setActiveDocument,
+    closeDocument,
+    removeContextFile,
+    applyBuiltInPreset,
+  ]);
+}
+
 type WorkflowStepperProps = {
  steps: Step[];
  currentStepId: string | null;
@@ -36,6 +72,7 @@ const WorkflowStepper = memo(function WorkflowStepper({
  isArchived,
 }: WorkflowStepperProps) {
  const [movingToStepId, setMovingToStepId] = useState<string | null>(null);
+  const disablePlanMode = useDisablePlanMode();

  const sortedSteps = useMemo(() => [...steps].sort((a, b) => a.position - b.position), [steps]);

@@ -47,6 +84,7 @@ const WorkflowStepper = memo(function WorkflowStepper({
  const handleMove = useCallback(
    async (stepId: string) => {
      if (!taskId || !workflowId) return;
+      disablePlanMode();
      setMovingToStepId(stepId);
      try {
        await moveTask(taskId, {
@@ -60,7 +98,7 @@ const WorkflowStepper = memo(function WorkflowStepper({
        setMovingToStepId(null);
      }
    },
-    [taskId, workflowId],
+    [taskId, workflowId, disablePlanMode],
  );

  if (sortedSteps.length === 0) return null;
@@ -551,6 +551,11 @@ export class SessionPage {
    return tab.locator(".tabler-icon-star").first();
  }

+  /** "Move to next step" button in the chat status bar. */
+  proceedNextStepButton(): Locator {
+    return this.page.getByTestId("proceed-next-step");
+  }
+
  /** Click a task in the sidebar by title. */
  async clickTaskInSidebar(title: string): Promise<void> {
    const taskRow = this.sidebar.locator("[role='button']").filter({ hasText: title });
@@ -0,0 +1,206 @@
+import { test, expect } from "../../fixtures/test-base";
+import { KanbanPage } from "../../pages/kanban-page";
+import { SessionPage } from "../../pages/session-page";
+
+test.describe("Manual proceed to next workflow step", () => {
+  /**
+   * Regression test: moving a task out of a plan-mode step must disable plan mode
+   * and show the next step's auto-start prompt in chat.
+   *
+   * Workflow: Spec (auto_start_agent) -> Work (auto_start_agent) -> Done
+   *
+   * 1. Create task via API in Spec step — agent auto-starts and completes
+   * 2. Navigate to session page, manually enable plan mode
+   * 3. Click "proceed to next step" button (Work)
+   * 4. Assert: plan mode is disabled (plan panel hidden, default input placeholder)
+   * 5. Assert: stepper shows Work as current step
+   * 6. Assert: Work step's auto-start prompt response is visible in chat
+   */
+  test("proceed from plan-mode step disables plan mode and shows next step prompt", async ({
+    testPage,
+    apiClient,
+    seedData,
+  }) => {
+    const workflow = await apiClient.createWorkflow(
+      seedData.workspaceId,
+      "Proceed Plan Step Workflow",
+    );
+
+    const specStep = await apiClient.createWorkflowStep(workflow.id, "Spec", 0);
+    const workStep = await apiClient.createWorkflowStep(workflow.id, "Work", 1);
+    await apiClient.createWorkflowStep(workflow.id, "Done", 2);
+
+    // Spec: auto-start agent (completes quickly so we can proceed)
+    await apiClient.updateWorkflowStep(specStep.id, {
+      prompt: 'e2e:message("spec done")\n{{task_prompt}}',
+      events: {
+        on_enter: [{ type: "auto_start_agent" }],
+      },
+    });
+
+    // Work: auto-start agent with a response we can assert on
+    await apiClient.updateWorkflowStep(workStep.id, {
+      prompt: 'e2e:delay(2000)\ne2e:message("work step response")\n{{task_prompt}}',
+      events: {
+        on_enter: [{ type: "auto_start_agent" }],
+      },
+    });
+
+    await apiClient.saveUserSettings({
+      workspace_id: seedData.workspaceId,
+      workflow_filter_id: workflow.id,
+      enable_preview_on_click: false,
+    });
+
+    // Create task via API in Spec step — triggers auto_start_agent
+    await apiClient.createTask(seedData.workspaceId, "Plan Proceed Task", {
+      workflow_id: workflow.id,
+      workflow_step_id: specStep.id,
+      agent_profile_id: seedData.agentProfileId,
+      repository_ids: [seedData.repositoryId],
+    });
+
+    // Navigate to task session page
+    const kanban = new KanbanPage(testPage);
+    await kanban.goto();
+
+    const card = kanban.taskCardInColumn("Plan Proceed Task", specStep.id);
+    await expect(card).toBeVisible({ timeout: 15_000 });
+    await card.click();
+    await expect(testPage).toHaveURL(/\/t\//, { timeout: 15_000 });
+
+    const session = new SessionPage(testPage);
+    await session.waitForLoad();
+
+    // Wait for agent to complete its turn (idle input visible)
+    await expect(session.idleInput()).toBeVisible({ timeout: 30_000 });
+
+    // Stepper shows Spec as current step
+    await expect(session.stepperStep("Spec")).toHaveAttribute("aria-current", "step", {
+      timeout: 10_000,
+    });
+
+    // The "proceed to next step" button should be visible (showing "Work")
+    const proceedBtn = session.proceedNextStepButton();
+    await expect(proceedBtn).toBeVisible({ timeout: 10_000 });
+
+    // --- Enable plan mode manually (simulates being in a plan-mode workflow step) ---
+    await testPage.waitForTimeout(1_000);
+    await session.togglePlanMode();
+    await expect(session.planPanel).toBeVisible({ timeout: 15_000 });
+    await expect(session.planModeInput()).toBeVisible({ timeout: 10_000 });
+
+    // --- Click proceed to move to Work step ---
+    await proceedBtn.click();
+
+    // Plan mode should be disabled: plan panel hidden and input shows default placeholder
+    await expect(session.planPanel).not.toBeVisible({ timeout: 15_000 });
+    await expect(session.planModeInput()).not.toBeVisible({ timeout: 10_000 });
+
+    // Stepper shows Work as current step
+    await expect(session.stepperStep("Work")).toHaveAttribute("aria-current", "step", {
+      timeout: 15_000,
+    });
+
+    // Work step auto-start prompt should be visible in chat (user message bubble)
+    await expect(
+      session.chat.getByText("work step response", { exact: false }).first(),
+    ).toBeVisible({ timeout: 30_000 });
+
+    // Session returns to idle in default (non-plan) mode
+    await expect(session.idleInput()).toBeVisible({ timeout: 15_000 });
+  });
+
+  /**
+   * Regression test: the proceed button must re-enable after each step transition.
+   * Previously, isMoving stayed true after clicking proceed because the button
+   * reappeared (for the next step) before isMoving was reset.
+   *
+   * Workflow: Spec -> Work -> Review -> QA -> Done (all with auto_start_agent)
+   */
+  test("proceed button re-enables across multiple step transitions", async ({
+    testPage,
+    apiClient,
+    seedData,
+  }) => {
+    test.setTimeout(120_000);
+
+    const workflow = await apiClient.createWorkflow(
+      seedData.workspaceId,
+      "Multi Step Proceed Workflow",
+    );
+
+    const stepNames = ["Spec", "Work", "Review", "QA", "Done"];
+    const steps: { id: string; name: string }[] = [];
+    for (let i = 0; i < stepNames.length; i++) {
+      const step = await apiClient.createWorkflowStep(workflow.id, stepNames[i], i);
+      steps.push({ id: step.id, name: stepNames[i] });
+    }
+
+    // Configure all steps except Done with auto_start_agent and a fast prompt
+    for (const step of steps.slice(0, -1)) {
+      await apiClient.updateWorkflowStep(step.id, {
+        prompt: `e2e:message("${step.name} done")\n{{task_prompt}}`,
+        events: {
+          on_enter: [{ type: "auto_start_agent" }],
+        },
+      });
+    }
+
+    await apiClient.saveUserSettings({
+      workspace_id: seedData.workspaceId,
+      workflow_filter_id: workflow.id,
+      enable_preview_on_click: false,
+    });
+
+    // Create task in Spec step
+    await apiClient.createTask(seedData.workspaceId, "Multi Step Task", {
+      workflow_id: workflow.id,
+      workflow_step_id: steps[0].id,
+      agent_profile_id: seedData.agentProfileId,
+      repository_ids: [seedData.repositoryId],
+    });
+
+    const kanban = new KanbanPage(testPage);
+    await kanban.goto();
+
+    const card = kanban.taskCardInColumn("Multi Step Task", steps[0].id);
+    await expect(card).toBeVisible({ timeout: 15_000 });
+    await card.click();
+    await expect(testPage).toHaveURL(/\/t\//, { timeout: 15_000 });
+
+    const session = new SessionPage(testPage);
+    await session.waitForLoad();
+
+    // Walk through each step transition: Spec -> Work -> Review -> QA -> Done
+    for (let i = 0; i < steps.length - 1; i++) {
+      const currentName = steps[i].name;
+      const nextName = steps[i + 1].name;
+
+      // Wait for agent to complete in current step
+      await expect(session.idleInput()).toBeVisible({ timeout: 30_000 });
+
+      // Stepper shows current step
+      await expect(session.stepperStep(currentName)).toHaveAttribute("aria-current", "step", {
+        timeout: 10_000,
+      });
+
+      // Proceed button is visible AND enabled (not disabled by stale isMoving)
+      const proceedBtn = session.proceedNextStepButton();
+      await expect(proceedBtn).toBeVisible({ timeout: 10_000 });
+      await expect(proceedBtn).toBeEnabled({ timeout: 5_000 });
+
+      // Click proceed
+      await proceedBtn.click();
+
+      // Stepper updates to next step
+      await expect(session.stepperStep(nextName)).toHaveAttribute("aria-current", "step", {
+        timeout: 15_000,
+      });
+    }
+
+    // On Done step: proceed button should NOT be visible (final step)
+    await expect(session.idleInput()).toBeVisible({ timeout: 30_000 });
+    await expect(session.proceedNextStepButton()).not.toBeVisible({ timeout: 5_000 });
+  });
+});
@@ -4,8 +4,13 @@ import { useToast } from "@/components/toast-provider";
 import { getWebSocketClient } from "@/lib/ws/connection";
 import { setChatDraftContent } from "@/lib/local-storage";
 import { moveTask } from "@/lib/api/domains/kanban-api";
+import { useContextFilesStore } from "@/lib/state/context-files-store";
+import { useLayoutStore } from "@/lib/state/layout-store";
+import { useDockviewStore } from "@/lib/state/dockview-store";
 import type { ChatInputContainerHandle } from "@/components/task/chat/chat-input-container";

+const PLAN_CONTEXT_PATH = "plan:context";
+
 const AUTO_TRANSITION_ACTIONS = ["move_to_next", "move_to_previous", "move_to_step"];

 function useNextWorkflowStep(taskId: string | null) {
@@ -18,7 +23,11 @@ function useNextWorkflowStep(taskId: string | null) {
    return task?.workflowStepId ?? null;
  });

-  const [isMoving, setIsMoving] = useState(false);
+  // Track which step the user was on when they clicked proceed. isMoving is true
+  // while taskStepId still matches the step we're moving FROM. Once the WS event
+  // updates taskStepId, isMoving becomes false and the button is clickable again.
+  const [moveFromStepId, setMoveFromStepId] = useState<string | null>(null);
+  const isMoving = moveFromStepId != null && taskStepId === moveFromStepId;

  const sortedSteps = useMemo(() => [...steps].sort((a, b) => a.position - b.position), [steps]);

@@ -51,21 +60,19 @@ function useNextWorkflowStep(taskId: string | null) {

  const proceed = useCallback(async () => {
    if (!taskId || !workflowId || !nextStep) return;
-    setIsMoving(true);
+    setMoveFromStepId(taskStepId);
    try {
      await moveTask(taskId, {
        workflow_id: workflowId,
        workflow_step_id: nextStep.id,
        position: 0,
      });
-      // Keep isMoving=true until the WS event updates taskStepId and
-      // proceedStepName becomes null, hiding the button entirely.
    } catch (err) {
      console.error("Failed to proceed to next step:", err);
      toast({ description: "Failed to proceed to next step", variant: "error" });
-      setIsMoving(false);
+      setMoveFromStepId(null);
    }
-  }, [taskId, workflowId, nextStep, toast]);
+  }, [taskId, workflowId, nextStep, taskStepId, toast]);

  const proceedStepName = nextStep && !currentStepAutoTransitions ? nextStep.title : null;

@@ -115,6 +122,31 @@ After completing the implementation, provide a summary of what was done.
  }, [resolvedSessionId, taskId, handlePlanModeChange, chatInputRef]);
 }

+/** Directly disable plan mode state + layout, bypassing the MCP availability guard. */
+function useDirectDisablePlanMode(resolvedSessionId: string | null) {
+  const setPlanMode = useAppStore((s) => s.setPlanMode);
+  const setActiveDocument = useAppStore((s) => s.setActiveDocument);
+  const closeDocument = useLayoutStore((s) => s.closeDocument);
+  const removeContextFile = useContextFilesStore((s) => s.removeFile);
+  const applyBuiltInPreset = useDockviewStore((s) => s.applyBuiltInPreset);
+
+  return useCallback(() => {
+    if (!resolvedSessionId) return;
+    applyBuiltInPreset("default");
+    closeDocument(resolvedSessionId);
+    setActiveDocument(resolvedSessionId, null);
+    setPlanMode(resolvedSessionId, false);
+    removeContextFile(resolvedSessionId, PLAN_CONTEXT_PATH);
+  }, [
+    resolvedSessionId,
+    applyBuiltInPreset,
+    closeDocument,
+    setActiveDocument,
+    setPlanMode,
+    removeContextFile,
+  ]);
+}
+
 export function usePlanActions(opts: {
  resolvedSessionId: string | null;
  taskId: string | null;
@@ -128,9 +160,24 @@ export function usePlanActions(opts: {
    opts.handlePlanModeChange,
    opts.chatInputRef,
  );
-  const { proceedStepName, nextStepIsWorkStep, proceed, isMoving } = useNextWorkflowStep(
-    opts.taskId,
-  );
+  const {
+    proceedStepName,
+    nextStepIsWorkStep,
+    proceed: rawProceed,
+    isMoving,
+  } = useNextWorkflowStep(opts.taskId);
+
+  const disablePlanMode = useDirectDisablePlanMode(opts.resolvedSessionId);
+  const { planModeEnabled } = opts;
+  // Disable plan mode before proceeding so the layout switches back to default
+  // and the next step's auto-start prompt is visible in chat.
+  const proceed = useCallback(() => {
+    if (planModeEnabled) {
+      disablePlanMode();
+    }
+    rawProceed();
+  }, [planModeEnabled, disablePlanMode, rawProceed]);
+
  const implementPlanHandler =
    opts.planModeEnabled && !nextStepIsWorkStep ? handleImplementPlan : undefined;
  return { implementPlanHandler, proceedStepName, proceed, isMoving };