Harden selection rewrite context checks

src/features/selection-rewrite/model.ts

@@ -3,6 +3,12 @@ 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 {
+  from: EditorPosition;
+  to: EditorPosition;
+}
 
 interface SelectionRewriteBaseState {
   filePath: string;
@@ -68,8 +74,41 @@ export type SelectionRewriteLifecycleEvent =
   | { type: "cancelled" }
   | { type: "applied" };
 
-export function canApplySelectionRewrite(currentText: string, originalText: string): boolean {
-  return currentText === originalText;
+export interface SelectionRewriteApplyContext {
+  currentText: string;
+  currentNoteText: string;
+}
+
+export interface SelectionRewriteTextRangeOffsets {
+  from: number;
+  to: number;
+}
+
+export function canApplySelectionRewrite(current: SelectionRewriteApplyContext, state: SelectionRewriteState): boolean {
+  if (current.currentText !== state.originalText) return false;
+
+  const originalOffsets = selectionRewriteTextRangeOffsets(state.noteText, state.targetRange, state.originalText);
+  const currentOffsets = selectionRewriteTextRangeOffsets(current.currentNoteText, state.targetRange, state.originalText);
+  if (!originalOffsets || !currentOffsets) return false;
+
+  return (
+    selectionRewriteRangeContextFingerprint(state.noteText, originalOffsets) ===
+    selectionRewriteRangeContextFingerprint(current.currentNoteText, currentOffsets)
+  );
+}
+
+export function selectionRewriteTextRangeOffsets(
+  text: string,
+  range: SelectionRewriteTextRange,
+  expectedText?: string,
+): SelectionRewriteTextRangeOffsets | null {
+  const from = editorPositionOffset(text, range.from);
+  const to = editorPositionOffset(text, range.to);
+  if (from !== null && to !== null && to >= from) return { from, to };
+  if (!expectedText) return null;
+
+  const fallbackFrom = text.indexOf(expectedText);
+  return fallbackFrom === -1 ? null : { from: fallbackFrom, to: fallbackFrom + expectedText.length };
 }
 
 export function transitionSelectionRewriteState(
@@ -118,3 +157,28 @@ export function transitionSelectionRewriteState(
       };
   }
 }
+
+function selectionRewriteRangeContextFingerprint(text: string, offsets: SelectionRewriteTextRangeOffsets): string {
+  const beforeStart = Math.max(0, offsets.from - APPLY_CONTEXT_RADIUS);
+  const afterEnd = Math.min(text.length, offsets.to + APPLY_CONTEXT_RADIUS);
+  return `${text.slice(beforeStart, offsets.from)}\0${text.slice(offsets.to, afterEnd)}`;
+}
+
+function editorPositionOffset(text: string, position: EditorPosition): number | null {
+  if (position.line < 0 || position.ch < 0) return null;
+  let line = 0;
+  let lineStart = 0;
+  while (lineStart <= text.length) {
+    if (line === position.line) {
+      const lineEnd = text.indexOf("\n", lineStart);
+      const end = lineEnd === -1 ? text.length : lineEnd;
+      if (lineStart + position.ch > end) return null;
+      return lineStart + position.ch;
+    }
+    const nextLine = text.indexOf("\n", lineStart);
+    if (nextLine === -1) return null;
+    line += 1;
+    lineStart = nextLine + 1;
+  }
+  return null;
+}

src/features/selection-rewrite/popover.tsx

@@ -141,7 +141,7 @@ export class SelectionRewritePopover {
     const { editor } = this.options;
     const state = this.session.state;
     const currentText = editor.getRange(state.targetRange.from, state.targetRange.to);
-    if (!canApplySelectionRewrite(currentText, state.originalText)) {
+    if (!canApplySelectionRewrite({ currentText, currentNoteText: editor.getValue() }, state)) {
       new Notice("Selection changed. Generate the rewrite again before applying.");
       this.session.setStatus("Selection changed. Generate the rewrite again before applying.");
       this.renderView();

src/features/selection-rewrite/prompt.ts

@@ -1,4 +1,4 @@
-import type { SelectionRewriteState } from "./model";
+import { selectionRewriteTextRangeOffsets, type SelectionRewriteState } from "./model";
 
 const MAX_NOTE_CONTEXT_CHARS = 20_000;
 
@@ -34,7 +34,7 @@ export function buildSelectionRewritePrompt(state: SelectionRewriteState): strin
     fenced(state.originalText),
     "",
     "Current note context:",
-    fenced(truncateNoteContext(state.noteText)),
+    fenced(selectionCenteredNoteContext(state)),
     "",
     "Reminder: use the note context only to make the selected-text replacement coherent.",
   ].join("\n");
@@ -54,7 +54,32 @@ function safeBacktickFence(text: string): string {
   return "`".repeat(longestRun + 1);
 }
 
-function truncateNoteContext(text: string): string {
+function selectionCenteredNoteContext(state: SelectionRewriteState): string {
+  const text = state.noteText;
   if (text.length <= MAX_NOTE_CONTEXT_CHARS) return text;
-  return `${text.slice(0, MAX_NOTE_CONTEXT_CHARS - 1)}…`;
+
+  const offsets = selectionRewriteTextRangeOffsets(text, state.targetRange, state.originalText);
+  if (!offsets) return `${text.slice(0, MAX_NOTE_CONTEXT_CHARS - 1)}…`;
+
+  const prefix = offsets.from > 0 ? "…\n" : "";
+  const suffix = offsets.to < text.length ? "\n…" : "";
+  const bodyBudget = MAX_NOTE_CONTEXT_CHARS - prefix.length - suffix.length;
+  const selectedLength = Math.max(0, offsets.to - offsets.from);
+  if (selectedLength >= bodyBudget) {
+    const end = Math.min(text.length, offsets.from + bodyBudget);
+    return `${offsets.from > 0 ? prefix : ""}${text.slice(offsets.from, end)}${end < text.length ? suffix : ""}`;
+  }
+
+  const surroundingBudget = Math.max(0, bodyBudget - selectedLength);
+  let start = Math.max(0, offsets.from - Math.floor(surroundingBudget / 2));
+  let end = Math.min(text.length, offsets.to + Math.ceil(surroundingBudget / 2));
+
+  const currentLength = end - start;
+  if (currentLength < bodyBudget) {
+    const missing = bodyBudget - currentLength;
+    start = Math.max(0, start - missing);
+    end = Math.min(text.length, end + (bodyBudget - (end - start)));
+  }
+
+  return `${start > 0 ? prefix : ""}${text.slice(start, end)}${end < text.length ? suffix : ""}`;
 }

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

@@ -86,6 +86,23 @@ describe("selection rewrite prompt", () => {
     expect(prompt).toContain("````text\n```ts");
     expect(prompt).toContain("`````text\n````markdown");
   });
+
+  it("centers long note context around the selection", () => {
+    const noteText = `START ONLY\n${"filler\n".repeat(4_000)}Near before\nRevise this sentence.\nNear after`;
+    const prompt = buildSelectionRewritePrompt(
+      rewriteState({
+        noteText,
+        targetRange: {
+          from: { line: 4_002, ch: 0 },
+          to: { line: 4_002, ch: 22 },
+        },
+      }),
+    );
+
+    expect(prompt).toContain("Near before");
+    expect(prompt).toContain("Near after");
+    expect(prompt).not.toContain("START ONLY");
+  });
 });
 
 describe("selection rewrite diff", () => {
@@ -129,9 +146,19 @@ describe("selection rewrite diff", () => {
 });
 
 describe("selection rewrite apply guard", () => {
-  it("allows apply only when the current range still matches the original text", () => {
-    expect(canApplySelectionRewrite("original", "original")).toBe(true);
-    expect(canApplySelectionRewrite("changed", "original")).toBe(false);
+  it("allows apply only when the current range and surrounding note context still match the original selection", () => {
+    const state = rewriteState();
+    expect(canApplySelectionRewrite({ currentText: "Revise this sentence.", currentNoteText: state.noteText }, state)).toBe(true);
+    expect(canApplySelectionRewrite({ currentText: "Changed sentence.", currentNoteText: state.noteText }, state)).toBe(false);
+    expect(
+      canApplySelectionRewrite(
+        {
+          currentText: "Revise this sentence.",
+          currentNoteText: "# Heading\n\nRevise this sentence.\n\nEdited next paragraph.",
+        },
+        state,
+      ),
+    ).toBe(false);
   });
 });
 
@@ -601,21 +628,25 @@ function popoverOptions(
   };
 }
 
-function editorFixture(options: { currentText?: string } = {}): {
+function editorFixture(options: { currentText?: string; currentNoteText?: string } = {}): {
   editor: ConstructorParameters<typeof SelectionRewritePopover>[0]["editor"];
   getRange: ReturnType<typeof vi.fn>;
+  getValue: ReturnType<typeof vi.fn>;
   replaceRange: ReturnType<typeof vi.fn>;
 } {
   const getRange = vi.fn(() => options.currentText ?? "Revise this sentence.");
+  const getValue = vi.fn(() => options.currentNoteText ?? rewriteState().noteText);
   const replaceRange = vi.fn();
   return {
     editor: {
       getCursor: () => ({ line: 1, ch: 22 }),
       getRange,
+      getValue,
       posToOffset: () => 0,
       replaceRange,
     } as unknown as ConstructorParameters<typeof SelectionRewritePopover>[0]["editor"],
     getRange,
+    getValue,
     replaceRange,
   };
 }