feat(word-diff): implement word-level diffing for granular text changes (#2817)

Author: andrii-harbour

packages/ai/src/ai-actions/__tests__/editor/word-diff.test.ts

@@ -0,0 +1,122 @@
+import { describe, it, expect } from 'vitest';
+import { tokenizeWords, computeWordDiff, getWordChanges } from '../../editor/word-diff';
+
+describe('tokenizeWords', () => {
+  it('should tokenize a basic sentence', () => {
+    const tokens = tokenizeWords('The quick fox');
+    expect(tokens).toEqual([
+      { text: 'The', offset: 0 },
+      { text: ' ', offset: 3 },
+      { text: 'quick', offset: 4 },
+      { text: ' ', offset: 9 },
+      { text: 'fox', offset: 10 },
+    ]);
+  });
+
+  it('should handle multiple spaces between words', () => {
+    const tokens = tokenizeWords('hello  world');
+    expect(tokens).toEqual([
+      { text: 'hello', offset: 0 },
+      { text: '  ', offset: 5 },
+      { text: 'world', offset: 7 },
+    ]);
+  });
+
+  it('should handle leading and trailing whitespace', () => {
+    const tokens = tokenizeWords('  hello ');
+    expect(tokens).toEqual([
+      { text: '  ', offset: 0 },
+      { text: 'hello', offset: 2 },
+      { text: ' ', offset: 7 },
+    ]);
+  });
+
+  it('should return empty array for empty string', () => {
+    expect(tokenizeWords('')).toEqual([]);
+  });
+
+  it('should handle a single word', () => {
+    expect(tokenizeWords('hello')).toEqual([{ text: 'hello', offset: 0 }]);
+  });
+
+  it('should handle punctuation attached to words', () => {
+    const tokens = tokenizeWords('Hello, world!');
+    expect(tokens).toEqual([
+      { text: 'Hello,', offset: 0 },
+      { text: ' ', offset: 6 },
+      { text: 'world!', offset: 7 },
+    ]);
+  });
+});
+
+describe('computeWordDiff', () => {
+  it('should return empty array for identical strings', () => {
+    expect(computeWordDiff('hello world', 'hello world')).toEqual([]);
+  });
+
+  it('should detect a single word replacement', () => {
+    const changes = getWordChanges('The quick fox', 'The fast fox');
+    expect(changes).toEqual([{ type: 'replace', oldFrom: 4, oldTo: 9, newText: 'fast' }]);
+  });
+
+  it('should detect multiple word replacements', () => {
+    const changes = getWordChanges(
+      'The quick brown fox jumps over the lazy dog',
+      'The fast brown fox leaps over the lazy cat',
+    );
+    expect(changes).toEqual([
+      { type: 'replace', oldFrom: 4, oldTo: 9, newText: 'fast' },
+      { type: 'replace', oldFrom: 20, oldTo: 25, newText: 'leaps' },
+      { type: 'replace', oldFrom: 40, oldTo: 43, newText: 'cat' },
+    ]);
+  });
+
+  it('should detect word insertion', () => {
+    const changes = getWordChanges('The fox', 'The quick fox');
+    expect(changes).toHaveLength(1);
+    expect(changes[0].type).toBe('insert');
+    // "quick " is inserted (word + trailing space before "fox")
+    expect(changes[0]).toHaveProperty('newText', 'quick ');
+  });
+
+  it('should detect word deletion', () => {
+    const changes = getWordChanges('The quick fox', 'The fox');
+    expect(changes).toHaveLength(1);
+    expect(changes[0].type).toBe('delete');
+    // "quick " (word + space) is removed as a contiguous block
+    expect(changes[0]).toHaveProperty('oldFrom', 4);
+    expect(changes[0]).toHaveProperty('oldTo', 10);
+  });
+
+  it('should handle complete rewrite', () => {
+    const changes = getWordChanges('hello world', 'goodbye earth');
+    // Each word is replaced separately since the space is a shared separator
+    expect(changes.length).toBeGreaterThanOrEqual(1);
+    expect(changes.every((c) => c.type === 'replace')).toBe(true);
+  });
+
+  it('should handle empty old text', () => {
+    const diff = computeWordDiff('', 'hello');
+    expect(diff).toEqual([{ type: 'insert', insertAt: 0, newText: 'hello' }]);
+  });
+
+  it('should handle empty new text', () => {
+    const diff = computeWordDiff('hello', '');
+    expect(diff).toEqual([{ type: 'delete', oldFrom: 0, oldTo: 5 }]);
+  });
+
+  it('should handle both empty', () => {
+    expect(computeWordDiff('', '')).toEqual([]);
+  });
+
+  it('should preserve whitespace tokens as equal', () => {
+    const diff = computeWordDiff('a b c', 'a x c');
+    const changes = diff.filter((op) => op.type !== 'equal');
+    expect(changes).toEqual([{ type: 'replace', oldFrom: 2, oldTo: 3, newText: 'x' }]);
+  });
+
+  it('should handle sentence with punctuation changes', () => {
+    const changes = getWordChanges('The company shall provide services.', 'The company must provide services.');
+    expect(changes).toEqual([{ type: 'replace', oldFrom: 12, oldTo: 17, newText: 'must' }]);
+  });
+});

packages/ai/src/ai-actions/editor/editor-adapter.ts

@@ -1,6 +1,7 @@
 import type { Editor, FoundMatch, MarkType } from '../../shared';
 import type { Node as ProseMirrorNode, Mark } from 'prosemirror-model';
 import { generateId, stripListPrefix } from '../../shared';
+import { getWordChanges } from './word-diff';
 
 /**
  * Default highlight color for text selections.
@@ -630,6 +631,17 @@ export class EditorAdapter {
     const replacementEnd = suggestedText.length - suffix;
     const replacementText = suggestedText.slice(prefix, replacementEnd);
 
+    // Try word-level diff for more granular tracked changes
+    const wordChanges = getWordChanges(originalText.slice(prefix, originalTextLength - suffix), replacementText);
+
+    if (wordChanges.length > 1) {
+      // Multiple word-level changes: apply each separately in reverse order
+      // so that earlier positions remain valid while we modify later ones.
+      this.applyGranularChanges(changeFrom, changeTo, wordChanges);
+      return;
+    }
+
+    // 0 or 1 word change: use existing single-replacement logic (better mark handling)
     const segments = this.collectTextSegments(changeFrom, changeTo);
     const nodes = this.buildTextNodes(changeFrom, changeTo, replacementText, segments);
     const tr = state.tr.delete(changeFrom, changeTo);
@@ -642,6 +654,95 @@ export class EditorAdapter {
     this.editor.dispatch(tr);
   }
 
+  /**
+   * Applies multiple word-level changes in a single transaction.
+   * Changes are applied in reverse document order to preserve positions.
+   *
+   * @param rangeFrom - Start of the overall change range in document positions
+   * @param rangeTo - End of the overall change range in document positions
+   * @param changes - Word diff operations with character offsets relative to the range text
+   * @private
+   */
+  private applyGranularChanges(
+    rangeFrom: number,
+    rangeTo: number,
+    changes: Array<
+      | { type: 'replace'; oldFrom: number; oldTo: number; newText: string }
+      | { type: 'delete'; oldFrom: number; oldTo: number }
+      | { type: 'insert'; insertAt: number; newText: string }
+    >,
+  ): void {
+    const { state } = this.editor;
+    if (!state) {
+      return;
+    }
+
+    // Pre-compute all document positions from the current (unmodified) state.
+    // Character offsets in changes are relative to the range text (rangeFrom..rangeTo).
+    const mappedChanges = changes.map((change) => {
+      if (change.type === 'insert') {
+        return {
+          ...change,
+          docPos: this.mapCharOffsetToPosition(rangeFrom, rangeTo, change.insertAt),
+        };
+      }
+      return {
+        ...change,
+        docFrom: this.mapCharOffsetToPosition(rangeFrom, rangeTo, change.oldFrom),
+        docTo: this.mapCharOffsetToPosition(rangeFrom, rangeTo, change.oldTo),
+      };
+    });
+
+    // Apply changes in forward order, remapping pre-computed positions through
+    // steps added during this loop so that length changes from earlier
+    // replacements are accounted for.
+    const tr = state.tr;
+    const baseSteps = tr.steps.length;
+    for (let i = 0; i < mappedChanges.length; i++) {
+      const change = mappedChanges[i];
+
+      // Remap pre-computed positions through steps added in this loop
+      const remap = (pos: number) => {
+        for (let s = baseSteps; s < tr.steps.length; s++) {
+          const step = tr.steps[s] as { getMap?: () => { map: (p: number) => number } } | undefined;
+          if (step && typeof step.getMap === 'function') {
+            pos = step.getMap().map(pos);
+          }
+        }
+        return pos;
+      };
+
+      if (change.type === 'delete') {
+        tr.delete(remap(change.docFrom), remap(change.docTo));
+      } else if (change.type === 'insert') {
+        const marks = this.getMarksAtPosition(change.docPos);
+        const node = state.schema.text(change.newText, marks);
+        tr.insert(remap(change.docPos), node);
+      } else {
+        // replace: use replaceWith for a single atomic step when available,
+        // fall back to delete+insert for test mocks that lack replaceWith.
+        const from = remap(change.docFrom);
+        const to = remap(change.docTo);
+        const segments = this.collectTextSegments(from, to);
+        const nodes = this.buildTextNodes(from, to, change.newText, segments);
+        if (typeof (tr as Record<string, unknown>).replaceWith === 'function') {
+          (
+            tr as unknown as { replaceWith: (from: number, to: number, content: ProseMirrorNode[]) => void }
+          ).replaceWith(from, to, nodes);
+        } else {
+          tr.delete(from, to);
+          let insertPos = from;
+          for (const node of nodes) {
+            tr.insert(insertPos, node);
+            insertPos += node.nodeSize;
+          }
+        }
+      }
+    }
+
+    this.editor.dispatch(tr);
+  }
+
   /**
    * Replaces text in the document while intelligently preserving ProseMirror marks.
    * Uses a diffing algorithm to minimize document changes by only replacing changed portions.