Model selection rewrite lifecycle as state machine

src/features/selection-rewrite/model.ts

@@ -2,7 +2,6 @@ import type { EditorPosition } from "obsidian";
 import type { ReasoningEffort } from "../../domain/catalog/metadata";
 
 type SelectionRewriteStatus = SelectionRewriteState["status"];
-const TERMINAL_SELECTION_REWRITE_STATUSES = new Set<SelectionRewriteStatus>(["applied", "cancelled"]);
 const APPLY_CONTEXT_RADIUS = 1_000;
 
 interface SelectionRewriteTextRange {
@@ -74,6 +73,13 @@ export type SelectionRewriteLifecycleEvent =
   | { type: "cancelled" }
   | { type: "applied" };
 
+type SelectionRewriteLifecycleEventType = SelectionRewriteLifecycleEvent["type"];
+type SelectionRewriteLifecycleTransition = (state: SelectionRewriteState, event: SelectionRewriteLifecycleEvent) => SelectionRewriteState;
+type SelectionRewriteLifecycleTransitionTable = Record<
+  SelectionRewriteStatus,
+  Record<SelectionRewriteLifecycleEventType, SelectionRewriteLifecycleTransition>
+>;
+
 export interface SelectionRewriteApplyContext {
   currentText: string;
   currentNoteText: string;
@@ -115,47 +121,121 @@ export function transitionSelectionRewriteState(
   state: SelectionRewriteState,
   event: SelectionRewriteLifecycleEvent,
 ): SelectionRewriteState {
-  switch (event.type) {
-    case "generation-started":
-      if (TERMINAL_SELECTION_REWRITE_STATUSES.has(state.status)) return state;
-      return {
-        ...state,
-        instruction: event.instruction,
-        status: "generating",
-        streamText: "",
-        replacementText: null,
-        debugText: null,
-      };
-    case "preview-updated":
-      if (state.status !== "generating") return state;
-      return { ...state, streamText: event.text };
-    case "generation-succeeded":
-      if (state.status !== "generating") return state;
-      return {
-        ...state,
-        status: "preview",
-        streamText: "",
-        replacementText: event.replacementText,
-      };
-    case "generation-failed":
-      if (state.status !== "generating") return state;
-      return {
-        ...state,
-        status: "failed",
-        streamText: "",
-        debugText: event.debugText,
-      };
-    case "cancelled":
-      return {
-        ...state,
-        status: "cancelled",
-      };
-    case "applied":
-      return {
-        ...state,
-        status: "applied",
-      };
-  }
+  return selectionRewriteLifecycleTransitions[state.status][event.type](state, event);
+}
+
+const keepSelectionRewriteState: SelectionRewriteLifecycleTransition = (state) => state;
+
+const generationStartedTransition: SelectionRewriteLifecycleTransition = (state, event) => ({
+  ...selectionRewriteBaseState(state),
+  instruction: requireGenerationInstruction(event),
+  status: "generating",
+  streamText: "",
+  replacementText: null,
+  debugText: null,
+});
+
+const previewUpdatedTransition: SelectionRewriteLifecycleTransition = (state, event) => ({
+  ...selectionRewriteBaseState(state),
+  status: "generating",
+  streamText: requirePreviewText(event),
+  replacementText: null,
+  debugText: null,
+});
+
+const generationSucceededTransition: SelectionRewriteLifecycleTransition = (state, event) => ({
+  ...selectionRewriteBaseState(state),
+  status: "preview",
+  streamText: "",
+  replacementText: requireReplacementText(event),
+  debugText: null,
+});
+
+const generationFailedTransition: SelectionRewriteLifecycleTransition = (state, event) => ({
+  ...selectionRewriteBaseState(state),
+  status: "failed",
+  streamText: "",
+  replacementText: null,
+  debugText: optionalDebugText(event),
+});
+
+const cancelledTransition: SelectionRewriteLifecycleTransition = (state) => ({
+  ...selectionRewriteBaseState(state),
+  status: "cancelled",
+  streamText: state.streamText,
+  replacementText: state.replacementText,
+  debugText: state.debugText,
+});
+
+const appliedTransition: SelectionRewriteLifecycleTransition = (state) => ({
+  ...selectionRewriteBaseState(state),
+  status: "applied",
+  streamText: state.streamText,
+  replacementText: state.replacementText,
+  debugText: state.debugText,
+});
+
+const editableSelectionRewriteTransitions = {
+  "generation-started": generationStartedTransition,
+  "preview-updated": keepSelectionRewriteState,
+  "generation-succeeded": keepSelectionRewriteState,
+  "generation-failed": keepSelectionRewriteState,
+  cancelled: cancelledTransition,
+  applied: appliedTransition,
+} satisfies Record<SelectionRewriteLifecycleEventType, SelectionRewriteLifecycleTransition>;
+
+const terminalSelectionRewriteTransitions = {
+  "generation-started": keepSelectionRewriteState,
+  "preview-updated": keepSelectionRewriteState,
+  "generation-succeeded": keepSelectionRewriteState,
+  "generation-failed": keepSelectionRewriteState,
+  cancelled: cancelledTransition,
+  applied: appliedTransition,
+} satisfies Record<SelectionRewriteLifecycleEventType, SelectionRewriteLifecycleTransition>;
+
+const selectionRewriteLifecycleTransitions: SelectionRewriteLifecycleTransitionTable = {
+  "editing-prompt": editableSelectionRewriteTransitions,
+  generating: {
+    "generation-started": generationStartedTransition,
+    "preview-updated": previewUpdatedTransition,
+    "generation-succeeded": generationSucceededTransition,
+    "generation-failed": generationFailedTransition,
+    cancelled: cancelledTransition,
+    applied: appliedTransition,
+  },
+  preview: editableSelectionRewriteTransitions,
+  failed: editableSelectionRewriteTransitions,
+  cancelled: terminalSelectionRewriteTransitions,
+  applied: terminalSelectionRewriteTransitions,
+};
+
+function requireGenerationInstruction(event: SelectionRewriteLifecycleEvent): string {
+  if ("instruction" in event) return event.instruction;
+  throw new Error(`Selection rewrite lifecycle event ${event.type} does not include an instruction.`);
+}
+
+function requirePreviewText(event: SelectionRewriteLifecycleEvent): string {
+  if ("text" in event) return event.text;
+  throw new Error(`Selection rewrite lifecycle event ${event.type} does not include preview text.`);
+}
+
+function requireReplacementText(event: SelectionRewriteLifecycleEvent): string {
+  if ("replacementText" in event) return event.replacementText;
+  throw new Error(`Selection rewrite lifecycle event ${event.type} does not include replacement text.`);
+}
+
+function optionalDebugText(event: SelectionRewriteLifecycleEvent): string | null {
+  return "debugText" in event ? event.debugText : null;
+}
+
+function selectionRewriteBaseState(state: SelectionRewriteState): SelectionRewriteBaseState {
+  return {
+    filePath: state.filePath,
+    targetRange: state.targetRange,
+    originalText: state.originalText,
+    noteText: state.noteText,
+    instruction: state.instruction,
+  };
 }
 
 function selectionRewriteRangeContextFingerprint(text: string, offsets: SelectionRewriteTextRangeOffsets): string {

tests/features/selection-rewrite/selection-rewrite.test.ts

@@ -7,6 +7,7 @@ import { buildSelectionUnifiedDiff } from "../../../src/features/selection-rewri
 import {
   canApplySelectionRewrite,
   transitionSelectionRewriteState,
+  type SelectionRewriteLifecycleEvent,
   type SelectionRewriteState,
 } from "../../../src/features/selection-rewrite/model";
 import { selectionRewriteOutputParseResultFromText } from "../../../src/features/selection-rewrite/output";
@@ -221,6 +222,50 @@ describe("selection rewrite lifecycle", () => {
     expect(transitionSelectionRewriteState(rewriteState(), { type: "cancelled" }).status).toBe("cancelled");
     expect(transitionSelectionRewriteState(rewriteState({ replacementText: "New text." }), { type: "applied" }).status).toBe("applied");
   });
+
+  it.each([{ status: "editing-prompt" }, { status: "preview" }, { status: "failed" }] satisfies {
+    status: SelectionRewriteState["status"];
+  }[])("allows regeneration from $status state", ({ status }) => {
+    const next = transitionSelectionRewriteState(selectionRewriteStateWithStatus(status), {
+      type: "generation-started",
+      instruction: "Try again.",
+    });
+
+    expect(next).toMatchObject({
+      status: "generating",
+      instruction: "Try again.",
+      streamText: "",
+      replacementText: null,
+      debugText: null,
+    });
+  });
+
+  it.each([
+    { status: "editing-prompt", event: previewUpdatedEvent() },
+    { status: "editing-prompt", event: generationSucceededEvent() },
+    { status: "editing-prompt", event: generationFailedEvent() },
+    { status: "preview", event: previewUpdatedEvent() },
+    { status: "preview", event: generationSucceededEvent() },
+    { status: "preview", event: generationFailedEvent() },
+    { status: "failed", event: previewUpdatedEvent() },
+    { status: "failed", event: generationSucceededEvent() },
+    { status: "failed", event: generationFailedEvent() },
+    { status: "cancelled", event: { type: "generation-started", instruction: "late" } },
+    { status: "cancelled", event: previewUpdatedEvent() },
+    { status: "cancelled", event: generationSucceededEvent() },
+    { status: "cancelled", event: generationFailedEvent() },
+    { status: "applied", event: { type: "generation-started", instruction: "late" } },
+    { status: "applied", event: previewUpdatedEvent() },
+    { status: "applied", event: generationSucceededEvent() },
+    { status: "applied", event: generationFailedEvent() },
+  ] satisfies { status: SelectionRewriteState["status"]; event: SelectionRewriteLifecycleEvent }[])(
+    "preserves state identity for ignored transition from $status via $event.type",
+    ({ status, event }) => {
+      const state = selectionRewriteStateWithStatus(status);
+
+      expect(transitionSelectionRewriteState(state, event)).toBe(state);
+    },
+  );
 });
 
 describe("selection rewrite runner lifecycle", () => {
@@ -604,6 +649,35 @@ function rewriteState(overrides: Partial<SelectionRewriteState> = {}): Selection
   } as SelectionRewriteState;
 }
 
+function selectionRewriteStateWithStatus(status: SelectionRewriteState["status"]): SelectionRewriteState {
+  switch (status) {
+    case "editing-prompt":
+      return rewriteState({ status: "editing-prompt", streamText: "", replacementText: null, debugText: null });
+    case "generating":
+      return rewriteState({ status: "generating", streamText: "draft", replacementText: null, debugText: null });
+    case "preview":
+      return rewriteState({ status: "preview", streamText: "", replacementText: "Preview text.", debugText: null });
+    case "failed":
+      return rewriteState({ status: "failed", streamText: "", replacementText: null, debugText: "debug" });
+    case "cancelled":
+      return rewriteState({ status: "cancelled", streamText: "partial", replacementText: null, debugText: "debug" });
+    case "applied":
+      return rewriteState({ status: "applied", streamText: "", replacementText: "Applied text.", debugText: null });
+  }
+}
+
+function previewUpdatedEvent(): SelectionRewriteLifecycleEvent {
+  return { type: "preview-updated", text: "late preview" };
+}
+
+function generationSucceededEvent(): SelectionRewriteLifecycleEvent {
+  return { type: "generation-succeeded", replacementText: "late replacement" };
+}
+
+function generationFailedEvent(): SelectionRewriteLifecycleEvent {
+  return { type: "generation-failed", debugText: "late debug" };
+}
+
 function openPopover(popover: SelectionRewritePopover): void {
   void act(() => {
     popover.open();