Add rewrite selection editor command

src/editor-rewrite/command.ts

@@ -0,0 +1,56 @@
+import { MarkdownView, Notice, type Editor, type Plugin } from "obsidian";
+
+import { RewriteSelectionModal } from "./modal";
+import type { RewriteSession } from "./model";
+
+export interface RewriteSelectionCommandHost extends Plugin {
+  settings: {
+    codexPath: string;
+  };
+  vaultPath: string;
+}
+
+export function registerRewriteSelectionCommand(plugin: RewriteSelectionCommandHost): void {
+  plugin.addCommand({
+    id: "rewrite-selection",
+    name: "Rewrite selection",
+    editorCallback: (editor, view) => {
+      if (!(view instanceof MarkdownView) || !view.file) {
+        new Notice("Rewrite selection requires an active markdown note.");
+        return;
+      }
+
+      const originalText = editor.getSelection();
+      if (!originalText.trim()) {
+        new Notice("Select text to rewrite first.");
+        return;
+      }
+
+      const session: RewriteSession = {
+        filePath: view.file.path,
+        targetRange: {
+          from: clonePosition(editor.getCursor("from")),
+          to: clonePosition(editor.getCursor("to")),
+        },
+        originalText,
+        noteText: editor.getValue(),
+        contextMode: "note",
+        instruction: "",
+        status: "editing-prompt",
+        streamText: "",
+        replacementText: null,
+      };
+
+      new RewriteSelectionModal(plugin.app, {
+        codexPath: plugin.settings.codexPath,
+        cwd: plugin.vaultPath,
+        editor,
+        session,
+      }).open();
+    },
+  });
+}
+
+function clonePosition(position: ReturnType<Editor["getCursor"]>): ReturnType<Editor["getCursor"]> {
+  return { line: position.line, ch: position.ch };
+}

src/editor-rewrite/diff.ts

@@ -0,0 +1,66 @@
+export function buildSelectionUnifiedDiff(filePath: string, originalText: string, replacementText: string): string {
+  const originalLines = textLines(originalText);
+  const replacementLines = textLines(replacementText);
+  const changes = lineChanges(originalLines, replacementLines);
+  const oldCount = Math.max(originalLines.length, 1);
+  const newCount = Math.max(replacementLines.length, 1);
+
+  return [
+    `diff --git a/${filePath} b/${filePath}`,
+    `--- a/${filePath}`,
+    `+++ b/${filePath}`,
+    `@@ -1,${oldCount} +1,${newCount} @@`,
+    ...changes.map((change) => `${change.prefix}${change.text}`),
+  ].join("\n");
+}
+
+interface DiffLine {
+  prefix: " " | "+" | "-";
+  text: string;
+}
+
+function textLines(text: string): string[] {
+  if (text.length === 0) return [];
+  const lines = text.split("\n");
+  if (lines.at(-1) === "") lines.pop();
+  return lines;
+}
+
+function lineChanges(originalLines: string[], replacementLines: string[]): DiffLine[] {
+  const lengths = lcsLengths(originalLines, replacementLines);
+  const changes: DiffLine[] = [];
+  let oldIndex = 0;
+  let newIndex = 0;
+
+  while (oldIndex < originalLines.length || newIndex < replacementLines.length) {
+    if (oldIndex < originalLines.length && newIndex < replacementLines.length && originalLines[oldIndex] === replacementLines[newIndex]) {
+      changes.push({ prefix: " ", text: originalLines[oldIndex] ?? "" });
+      oldIndex += 1;
+      newIndex += 1;
+    } else if (
+      newIndex < replacementLines.length &&
+      (oldIndex === originalLines.length || lengths[oldIndex]?.[newIndex + 1] >= lengths[oldIndex + 1]?.[newIndex])
+    ) {
+      changes.push({ prefix: "+", text: replacementLines[newIndex] ?? "" });
+      newIndex += 1;
+    } else if (oldIndex < originalLines.length) {
+      changes.push({ prefix: "-", text: originalLines[oldIndex] ?? "" });
+      oldIndex += 1;
+    }
+  }
+
+  return changes.length > 0 ? changes : [{ prefix: " ", text: "" }];
+}
+
+function lcsLengths(left: string[], right: string[]): number[][] {
+  const rows = Array.from({ length: left.length + 1 }, () => Array.from({ length: right.length + 1 }, () => 0));
+  for (let leftIndex = left.length - 1; leftIndex >= 0; leftIndex -= 1) {
+    for (let rightIndex = right.length - 1; rightIndex >= 0; rightIndex -= 1) {
+      rows[leftIndex][rightIndex] =
+        left[leftIndex] === right[rightIndex]
+          ? (rows[leftIndex + 1]?.[rightIndex + 1] ?? 0) + 1
+          : Math.max(rows[leftIndex + 1]?.[rightIndex] ?? 0, rows[leftIndex]?.[rightIndex + 1] ?? 0);
+    }
+  }
+  return rows;
+}

src/editor-rewrite/modal.ts

@@ -0,0 +1,170 @@
+import { Modal, Notice, type App, type Editor } from "obsidian";
+
+import { renderUnifiedDiff } from "../ui/turn-diff";
+import { buildSelectionUnifiedDiff } from "./diff";
+import { canApplyRewrite, type RewriteContextMode, type RewriteSession } from "./model";
+import { buildRewritePrompt } from "./prompt";
+import { runRewriteSelection } from "./runner";
+
+export interface RewriteSelectionModalOptions {
+  codexPath: string;
+  cwd: string;
+  editor: Editor;
+  session: RewriteSession;
+}
+
+export class RewriteSelectionModal extends Modal {
+  private instructionEl: HTMLTextAreaElement | null = null;
+  private contextEl: HTMLSelectElement | null = null;
+  private generateButton: HTMLButtonElement | null = null;
+  private applyButton: HTMLButtonElement | null = null;
+  private statusEl: HTMLElement | null = null;
+  private previewEl: HTMLElement | null = null;
+  private diffEl: HTMLElement | null = null;
+
+  constructor(
+    app: App,
+    private readonly options: RewriteSelectionModalOptions,
+  ) {
+    super(app);
+  }
+
+  onOpen(): void {
+    this.setTitle("Rewrite selection");
+    const { contentEl } = this;
+    contentEl.empty();
+    contentEl.addClass("codex-panel-rewrite");
+
+    this.instructionEl = contentEl.createEl("textarea", {
+      cls: "codex-panel-rewrite__instruction",
+      attr: { placeholder: "How should Codex rewrite the selected text?" },
+    });
+    this.instructionEl.value = this.options.session.instruction;
+    this.instructionEl.oninput = () => this.syncControls();
+
+    const controls = contentEl.createDiv({ cls: "codex-panel-rewrite__controls" });
+    const contextLabel = controls.createEl("label", { cls: "codex-panel-rewrite__context" });
+    contextLabel.createSpan({ text: "Context" });
+    this.contextEl = contextLabel.createEl("select");
+    this.addContextOption("note", "Selection + note context");
+    this.addContextOption("selection", "Selection only");
+    this.contextEl.value = this.options.session.contextMode === "selection" ? "selection" : "note";
+    this.contextEl.onchange = () => {
+      this.options.session.contextMode = this.contextEl?.value === "selection" ? "selection" : "note";
+    };
+
+    const actions = controls.createDiv({ cls: "codex-panel-rewrite__actions" });
+    this.generateButton = actions.createEl("button", { text: "Generate", attr: { type: "button" } });
+    this.generateButton.onclick = () => void this.generate();
+    const cancelButton = actions.createEl("button", { text: "Cancel", attr: { type: "button" } });
+    cancelButton.onclick = () => {
+      this.options.session.status = "cancelled";
+      this.close();
+    };
+
+    this.statusEl = contentEl.createDiv({ cls: "codex-panel-rewrite__status" });
+    this.previewEl = contentEl.createEl("pre", { cls: "codex-panel-rewrite__preview" });
+    this.diffEl = contentEl.createDiv({ cls: "codex-panel-rewrite__diff" });
+
+    const footer = contentEl.createDiv({ cls: "codex-panel-rewrite__footer" });
+    this.applyButton = footer.createEl("button", {
+      text: "Apply",
+      cls: "mod-cta",
+      attr: { type: "button" },
+    });
+    this.applyButton.onclick = () => this.apply();
+
+    this.setStatus("Enter an instruction, then generate a patch.");
+    this.syncControls();
+    this.instructionEl.focus();
+  }
+
+  onClose(): void {
+    this.contentEl.empty();
+  }
+
+  private addContextOption(value: Exclude<RewriteContextMode, "note-and-thread">, label: string): void {
+    this.contextEl?.createEl("option", { text: label, attr: { value } });
+  }
+
+  private async generate(): Promise<void> {
+    const instruction = this.instructionEl?.value.trim() ?? "";
+    if (!instruction) {
+      new Notice("Enter a rewrite instruction first.");
+      this.instructionEl?.focus();
+      return;
+    }
+
+    this.options.session.instruction = instruction;
+    this.options.session.contextMode = this.contextEl?.value === "selection" ? "selection" : "note";
+    this.options.session.status = "generating";
+    this.options.session.streamText = "";
+    this.options.session.replacementText = null;
+    this.previewEl?.setText("");
+    this.diffEl?.empty();
+    this.setStatus("Generating patch...");
+    this.syncControls();
+
+    try {
+      const output = await runRewriteSelection({
+        codexPath: this.options.codexPath,
+        cwd: this.options.cwd,
+        prompt: buildRewritePrompt(this.options.session),
+        onPreview: (text) => this.updatePreview(text),
+      });
+      this.options.session.replacementText = output.replacementText;
+      this.options.session.status = "preview";
+      this.renderDiff();
+      this.setStatus("Review the patch before applying it.");
+    } catch (error) {
+      this.options.session.status = "failed";
+      this.setStatus(error instanceof Error ? error.message : String(error));
+    } finally {
+      this.syncControls();
+    }
+  }
+
+  private updatePreview(text: string): void {
+    this.options.session.streamText = text;
+    this.previewEl?.setText(text);
+  }
+
+  private renderDiff(): void {
+    const replacement = this.options.session.replacementText;
+    if (replacement === null || !this.diffEl) return;
+    this.diffEl.empty();
+    renderUnifiedDiff(
+      this.diffEl,
+      buildSelectionUnifiedDiff(this.options.session.filePath, this.options.session.originalText, replacement),
+    );
+  }
+
+  private apply(): void {
+    const replacement = this.options.session.replacementText;
+    if (replacement === null) return;
+
+    const { editor, session } = this.options;
+    const currentText = editor.getRange(session.targetRange.from, session.targetRange.to);
+    if (!canApplyRewrite(currentText, session.originalText)) {
+      new Notice("Selection changed. Generate the rewrite again before applying.");
+      return;
+    }
+
+    editor.replaceRange(replacement, session.targetRange.from, session.targetRange.to, "codex-panel-rewrite");
+    session.status = "applied";
+    this.close();
+  }
+
+  private setStatus(text: string): void {
+    this.statusEl?.setText(text);
+  }
+
+  private syncControls(): void {
+    const generating = this.options.session.status === "generating";
+    const hasInstruction = Boolean(this.instructionEl?.value.trim());
+    if (this.instructionEl) this.instructionEl.disabled = generating;
+    if (this.contextEl) this.contextEl.disabled = generating;
+    if (this.generateButton) this.generateButton.disabled = generating || !hasInstruction;
+    if (this.applyButton) this.applyButton.disabled = generating || this.options.session.replacementText === null;
+  }
+}

src/editor-rewrite/model.ts

@@ -0,0 +1,24 @@
+import type { EditorPosition } from "obsidian";
+
+export type RewriteContextMode = "selection" | "note" | "note-and-thread";
+
+export type RewriteStatus = "editing-prompt" | "generating" | "preview" | "applied" | "cancelled" | "failed";
+
+export interface RewriteSession {
+  filePath: string;
+  targetRange: {
+    from: EditorPosition;
+    to: EditorPosition;
+  };
+  originalText: string;
+  noteText: string;
+  contextMode: RewriteContextMode;
+  instruction: string;
+  status: RewriteStatus;
+  streamText: string;
+  replacementText: string | null;
+}
+
+export function canApplyRewrite(currentText: string, originalText: string): boolean {
+  return currentText === originalText;
+}

src/editor-rewrite/output.ts

@@ -0,0 +1,32 @@
+import type { Turn } from "../generated/app-server/v2/Turn";
+
+export interface RewriteOutput {
+  replacementText: string;
+}
+
+export function parseRewriteOutput(text: string): RewriteOutput | null {
+  try {
+    const parsed = JSON.parse(text.trim()) as unknown;
+    if (!parsed || typeof parsed !== "object") return null;
+    const replacementText = (parsed as { replacementText?: unknown }).replacementText;
+    if (typeof replacementText !== "string") return null;
+    return { replacementText };
+  } catch {
+    return null;
+  }
+}
+
+export function rewriteOutputFromTurn(turn: Turn): RewriteOutput | null {
+  const text = lastAgentMessageText(turn);
+  return text ? parseRewriteOutput(text) : null;
+}
+
+function lastAgentMessageText(turn: Turn): string | null {
+  for (let index = turn.items.length - 1; index >= 0; index -= 1) {
+    const item = turn.items[index];
+    if (item?.type !== "agentMessage") continue;
+    const text = item.text.trim();
+    if (text) return text;
+  }
+  return null;
+}

src/editor-rewrite/prompt.ts

@@ -0,0 +1,67 @@
+import type { RewriteContextMode, RewriteSession } from "./model";
+
+const MAX_NOTE_CONTEXT_CHARS = 20_000;
+
+export const REWRITE_SERVICE_NAME = "codex-panel-rewrite-selection";
+
+export const REWRITE_DEVELOPER_INSTRUCTIONS = [
+  "You rewrite selected Obsidian note text.",
+  "Return only JSON matching the requested schema.",
+  "Do not include Markdown fences, explanations, alternatives, or comments outside JSON.",
+  "The only editable target is the selected text. Use any provided note context only for consistency.",
+  "Preserve Obsidian-specific syntax such as wikilinks, block ids, callouts, frontmatter-like text, and Dataview blocks unless the user's instruction explicitly asks to change them.",
+].join("\n");
+
+export function buildRewritePrompt(session: RewriteSession): string {
+  return [
+    "Rewrite the selected text according to the user's instruction.",
+    "",
+    "Rules:",
+    "- Return a JSON object with exactly one field: replacementText.",
+    "- replacementText must be the full replacement for the selected text.",
+    "- Do not edit outside the selected text.",
+    "- Preserve the language and style of the surrounding note unless instructed otherwise.",
+    "",
+    "Target:",
+    `- File: ${session.filePath}`,
+    `- Selection: ${positionLabel(session.targetRange.from)} to ${positionLabel(session.targetRange.to)}`,
+    `- Context mode: ${contextModeLabel(session.contextMode)}`,
+    "",
+    "User instruction:",
+    session.instruction,
+    "",
+    "Selected text:",
+    fenced(session.originalText),
+    ...noteContextLines(session),
+  ].join("\n");
+}
+
+function noteContextLines(session: RewriteSession): string[] {
+  if (session.contextMode === "selection") return [];
+  return [
+    "",
+    "Current note context:",
+    fenced(truncateNoteContext(session.noteText)),
+    "",
+    "Reminder: use the note context only to make the selected-text replacement coherent.",
+  ];
+}
+
+function contextModeLabel(mode: RewriteContextMode): string {
+  if (mode === "selection") return "Selection only";
+  if (mode === "note") return "Selection + note context";
+  return "Selection + note context + active thread";
+}
+
+function positionLabel(position: { line: number; ch: number }): string {
+  return `L${position.line + 1}:C${position.ch + 1}`;
+}
+
+function fenced(text: string): string {
+  return ["```text", text, "```"].join("\n");
+}
+
+function truncateNoteContext(text: string): string {
+  if (text.length <= MAX_NOTE_CONTEXT_CHARS) return text;
+  return `${text.slice(0, MAX_NOTE_CONTEXT_CHARS - 1)}…`;
+}

src/editor-rewrite/runner.ts

@@ -0,0 +1,103 @@
+import { AppServerClient } from "../app-server/client";
+import type { ServerNotification } from "../generated/app-server/ServerNotification";
+import type { JsonValue } from "../generated/app-server/serde_json/JsonValue";
+import type { ThreadItem } from "../generated/app-server/v2/ThreadItem";
+import type { Turn } from "../generated/app-server/v2/Turn";
+import { REWRITE_DEVELOPER_INSTRUCTIONS, REWRITE_SERVICE_NAME } from "./prompt";
+import { rewriteOutputFromTurn, type RewriteOutput } from "./output";
+
+const REWRITE_TIMEOUT_MS = 120_000;
+
+const REWRITE_OUTPUT_SCHEMA: JsonValue = {
+  type: "object",
+  properties: {
+    replacementText: {
+      type: "string",
+    },
+  },
+  required: ["replacementText"],
+  additionalProperties: false,
+};
+
+export interface RunRewriteSelectionOptions {
+  codexPath: string;
+  cwd: string;
+  prompt: string;
+  onPreview?: (text: string) => void;
+}
+
+export async function runRewriteSelection(options: RunRewriteSelectionOptions): Promise<RewriteOutput> {
+  let threadId: string | null = null;
+  let expectedTurnId: string | null = null;
+  let preview = "";
+  let completed = false;
+  let timeout: ReturnType<Window["setTimeout"]> | null = null;
+  let rejectCompletedTurn: ((error: Error) => void) | null = null;
+  let handleNotification: (notification: ServerNotification) => void = () => undefined;
+  const completedItems: ThreadItem[] = [];
+
+  const completedTurn = new Promise<Turn>((resolve, reject) => {
+    rejectCompletedTurn = reject;
+    timeout = window.setTimeout(() => {
+      if (completed) return;
+      completed = true;
+      reject(new Error("Timed out while rewriting the selection."));
+    }, REWRITE_TIMEOUT_MS);
+
+    handleNotification = (notification): void => {
+      if (completed) return;
+      if (notification.method === "item/agentMessage/delta") {
+        if (!threadId || notification.params.threadId !== threadId) return;
+        if (expectedTurnId && notification.params.turnId !== expectedTurnId) return;
+        preview += notification.params.delta;
+        options.onPreview?.(preview);
+        return;
+      }
+      if (notification.method === "item/completed") {
+        if (!threadId || notification.params.threadId !== threadId) return;
+        if (expectedTurnId && notification.params.turnId !== expectedTurnId) return;
+        completedItems.push(notification.params.item);
+        return;
+      }
+      if (notification.method === "turn/completed") {
+        if (!threadId || notification.params.threadId !== threadId) return;
+        if (expectedTurnId && notification.params.turn.id !== expectedTurnId) return;
+        completed = true;
+        resolve(turnWithCollectedItems(notification.params.turn, completedItems));
+      }
+    };
+  });
+
+  let client!: AppServerClient;
+  client = new AppServerClient(options.codexPath, options.cwd, {
+    onNotification: (notification) => handleNotification(notification),
+    onServerRequest: (request) => client.rejectServerRequest(request.id, -32601, "Selection rewrite does not handle server requests."),
+    onLog: () => undefined,
+    onExit: () => {
+      if (completed) return;
+      completed = true;
+      rejectCompletedTurn?.(new Error("Codex rewrite app-server exited."));
+    },
+  });
+
+  try {
+    await client.connect();
+    const threadResponse = await client.startEphemeralThread(options.cwd, REWRITE_SERVICE_NAME, REWRITE_DEVELOPER_INSTRUCTIONS);
+    threadId = threadResponse.thread.id;
+    const turnResponse = await client.startStructuredTurn(threadId, options.cwd, options.prompt, REWRITE_OUTPUT_SCHEMA);
+    expectedTurnId = turnResponse.turn.id;
+    const turn = turnResponse.turn.status === "completed" ? turnWithCollectedItems(turnResponse.turn, completedItems) : await completedTurn;
+    const output = rewriteOutputFromTurn(turn);
+    if (!output) throw new Error("Codex did not return a valid rewrite patch.");
+    return output;
+  } finally {
+    completed = true;
+    if (timeout) window.clearTimeout(timeout);
+    client.disconnect();
+  }
+}
+
+function turnWithCollectedItems(turn: Turn, completedItems: ThreadItem[]): Turn {
+  if (turn.items.length > 0 || completedItems.length === 0) return turn;
+  return { ...turn, items: completedItems };
+}

src/main.ts

@@ -1,6 +1,7 @@
 import { Plugin, type WorkspaceLeaf } from "obsidian";
 
 import { VIEW_TYPE_CODEX_PANEL, VIEW_TYPE_CODEX_TURN_DIFF } from "./constants";
+import { registerRewriteSelectionCommand } from "./editor-rewrite/command";
 import { CodexPanelView } from "./panel/view";
 import { CodexTurnDiffView } from "./panel/turn-diff-view";
 import { DEFAULT_SETTINGS, getVaultPath, normalizeSettings, settingsMatchNormalizedData, type CodexPanelSettings } from "./settings/model";
@@ -43,6 +44,8 @@ export default class CodexPanelPlugin extends Plugin {
       },
     });
 
+    registerRewriteSelectionCommand(this);
+
     this.addSettingTab(new CodexPanelSettingTab(this.app, this));
   }
 

styles.css

@@ -1394,6 +1394,67 @@
   white-space: pre-wrap;
 }
 
+.codex-panel-rewrite {
+  display: flex;
+  flex-direction: column;
+  gap: var(--size-4-3, 12px);
+}
+
+.codex-panel-rewrite__instruction {
+  min-height: 96px;
+  resize: vertical;
+}
+
+.codex-panel-rewrite__controls,
+.codex-panel-rewrite__footer {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: var(--size-4-2, 8px);
+}
+
+.codex-panel-rewrite__context {
+  display: inline-flex;
+  align-items: center;
+  gap: var(--size-2-3, 6px);
+  color: var(--text-muted);
+  font-size: var(--font-ui-small);
+}
+
+.codex-panel-rewrite__actions {
+  display: inline-flex;
+  gap: var(--size-2-3, 6px);
+}
+
+.codex-panel-rewrite__status {
+  color: var(--text-muted);
+  font-size: var(--font-ui-small);
+}
+
+.codex-panel-rewrite__preview {
+  display: none;
+  max-height: 120px;
+  overflow: auto;
+  margin: 0;
+  padding: 8px;
+  border-radius: var(--radius-s);
+  background: var(--code-background, var(--background-secondary));
+  color: var(--text-muted);
+  font-family: var(--font-monospace);
+  font-size: var(--font-smaller);
+  white-space: pre-wrap;
+}
+
+.codex-panel-rewrite__preview:not(:empty) {
+  display: block;
+}
+
+.codex-panel-rewrite__diff .codex-panel__diff {
+  max-height: min(360px, 45vh);
+  margin: 0;
+  white-space: pre-wrap;
+}
+
 .codex-panel__user-input-question {
   display: grid;
   gap: var(--codex-panel-control-gap);

tests/editor-rewrite/editor-rewrite.test.ts

@@ -0,0 +1,108 @@
+import { describe, expect, it } from "vitest";
+
+import { buildSelectionUnifiedDiff } from "../../src/editor-rewrite/diff";
+import { canApplyRewrite, type RewriteSession } from "../../src/editor-rewrite/model";
+import { parseRewriteOutput, rewriteOutputFromTurn } from "../../src/editor-rewrite/output";
+import { buildRewritePrompt } from "../../src/editor-rewrite/prompt";
+import type { Turn } from "../../src/generated/app-server/v2/Turn";
+
+describe("editor rewrite output", () => {
+  it("parses valid rewrite JSON", () => {
+    expect(parseRewriteOutput('{"replacementText":"rewritten"}')).toEqual({ replacementText: "rewritten" });
+  });
+
+  it("rejects invalid rewrite JSON", () => {
+    expect(parseRewriteOutput("replacementText: rewritten")).toBeNull();
+    expect(parseRewriteOutput('{"replacementText":42}')).toBeNull();
+    expect(parseRewriteOutput('{"text":"rewritten"}')).toBeNull();
+  });
+
+  it("extracts the final rewrite JSON from a completed turn", () => {
+    expect(
+      rewriteOutputFromTurn(
+        turn([
+          { type: "agentMessage", id: "a1", text: '{"replacementText":"first"}', phase: "final_answer", memoryCitation: null },
+          { type: "agentMessage", id: "a2", text: '{"replacementText":"final"}', phase: "final_answer", memoryCitation: null },
+        ]),
+      ),
+    ).toEqual({ replacementText: "final" });
+  });
+});
+
+describe("editor rewrite prompt", () => {
+  it("omits note context in selection-only mode", () => {
+    const prompt = buildRewritePrompt(session({ contextMode: "selection" }));
+
+    expect(prompt).toContain("Selected text:");
+    expect(prompt).toContain("Rewrite this sentence.");
+    expect(prompt).not.toContain("Current note context:");
+  });
+
+  it("includes note context in note mode", () => {
+    const prompt = buildRewritePrompt(session({ contextMode: "note" }));
+
+    expect(prompt).toContain("Context mode: Selection + note context");
+    expect(prompt).toContain("Current note context:");
+    expect(prompt).toContain("# Heading");
+  });
+});
+
+describe("editor rewrite diff", () => {
+  it("renders replacements in a selection-scoped unified diff", () => {
+    const diff = buildSelectionUnifiedDiff("Note.md", "alpha\nbeta", "alpha\ngamma");
+
+    expect(diff).toContain("diff --git a/Note.md b/Note.md");
+    expect(diff).toContain("@@ -1,2 +1,2 @@");
+    expect(diff).toContain(" alpha");
+    expect(diff).toContain("-beta");
+    expect(diff).toContain("+gamma");
+  });
+
+  it("renders additions and deletions", () => {
+    expect(buildSelectionUnifiedDiff("Note.md", "", "added")).toContain("+added");
+    expect(buildSelectionUnifiedDiff("Note.md", "removed", "")).toContain("-removed");
+  });
+
+  it("renders unchanged text as context", () => {
+    expect(buildSelectionUnifiedDiff("Note.md", "same", "same")).toContain(" same");
+  });
+});
+
+describe("editor rewrite apply guard", () => {
+  it("allows apply only when the current range still matches the original text", () => {
+    expect(canApplyRewrite("original", "original")).toBe(true);
+    expect(canApplyRewrite("changed", "original")).toBe(false);
+  });
+});
+
+function session(overrides: Partial<RewriteSession> = {}): RewriteSession {
+  return {
+    filePath: "Note.md",
+    targetRange: {
+      from: { line: 1, ch: 0 },
+      to: { line: 1, ch: 22 },
+    },
+    originalText: "Rewrite this sentence.",
+    noteText: "# Heading\n\nRewrite this sentence.\n\nNext paragraph.",
+    contextMode: "note",
+    instruction: "Make it clearer.",
+    status: "editing-prompt",
+    streamText: "",
+    replacementText: null,
+    ...overrides,
+  };
+}
+
+function turn(items: Turn["items"], overrides: Partial<Turn> = {}): Turn {
+  return {
+    id: "turn",
+    items,
+    itemsView: "full",
+    status: "completed",
+    error: null,
+    startedAt: null,
+    completedAt: null,
+    durationMs: null,
+    ...overrides,
+  };
+}