diff --git a/src/features/selection-rewrite/model.ts b/src/features/selection-rewrite/model.ts index 003b4e8c..e712e7f4 100644 --- a/src/features/selection-rewrite/model.ts +++ b/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(["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; +} diff --git a/src/features/selection-rewrite/popover.tsx b/src/features/selection-rewrite/popover.tsx index 59906f29..6687ad09 100644 --- a/src/features/selection-rewrite/popover.tsx +++ b/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(); diff --git a/src/features/selection-rewrite/prompt.ts b/src/features/selection-rewrite/prompt.ts index 69869021..466a1224 100644 --- a/src/features/selection-rewrite/prompt.ts +++ b/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 : ""}`; } diff --git a/tests/features/selection-rewrite/selection-rewrite.test.ts b/tests/features/selection-rewrite/selection-rewrite.test.ts index ee9902f0..a29a6b50 100644 --- a/tests/features/selection-rewrite/selection-rewrite.test.ts +++ b/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[0]["editor"]; getRange: ReturnType; + getValue: ReturnType; replaceRange: ReturnType; } { 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[0]["editor"], getRange, + getValue, replaceRange, }; }