1
import { describe, expect, it } from "vitest";
2
import { getSchema } from "@tiptap/core";
3
import { StarterKit } from "@tiptap/starter-kit";
4
import { EditorState } from "@tiptap/pm/state";
5
import type { EditorView } from "@tiptap/pm/view";
6
import {
7
  aiGenerateIndicatorKey,
8
  appendAiGenerateChunk,
9
  createAiGenerateIndicatorPlugin,
10
  type AiGenerateSessionSnapshot,
11
} from "@/client/components/editor/extensions/ai-generate-indicator";
12
 
13
const schema = getSchema([StarterKit.configure({ undoRedo: false })]);
14
 
15
function makeState(): EditorState {
16
  const doc = schema.nodeFromJSON({
17
    type: "doc",
18
    content: [{ type: "paragraph", content: [{ type: "text", text: "Hello world." }] }],
19
  });
20
  return EditorState.create({ schema, doc, plugins: [createAiGenerateIndicatorPlugin()] });
21
}
22
 
23
function getSession(state: EditorState): AiGenerateSessionSnapshot | null {
24
  const plugin = aiGenerateIndicatorKey.getState(state);
25
  if (!plugin?.active) return null;
26
  const { sessionId, from, to, expectedLength, dirty } = plugin.active;
27
  return { sessionId, from, to, expectedLength, dirty };
28
}
29
 
30
describe("AiGenerateIndicator session range", () => {
31
  it("tracks range advancement on chunk transactions", () => {
32
    let state = makeState();
33
    const begin = state.tr.setMeta(aiGenerateIndicatorKey, {
34
      type: "begin",
35
      sessionId: "s1",
36
      pos: 6,
37
      label: "Generating…",
38
    });
39
    state = state.apply(begin);
40
 
41
    const chunk1 = state.tr
42
      .insertText("AAA", 6)
43
      .setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" });
44
    state = state.apply(chunk1);
45
 
46
    let session = getSession(state);
47
    expect(session).toMatchObject({ from: 6, to: 9, expectedLength: 3, dirty: false });
48
 
49
    const chunk2 = state.tr
50
      .insertText("BB", 9)
51
      .setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "BB" });
52
    state = state.apply(chunk2);
53
 
54
    session = getSession(state);
55
    expect(session).toMatchObject({ from: 6, to: 11, expectedLength: 5, dirty: false });
56
  });
57
 
58
  it("maps the range past external edits before the range without flagging dirty", () => {
59
    let state = makeState();
60
    state = state.apply(
61
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
62
    );
63
    state = state.apply(
64
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
65
    );
66
 
67
    // External insertion at pos 1 (before the range) — no plugin meta.
68
    state = state.apply(state.tr.insertText("X", 1));
69
 
70
    const session = getSession(state);
71
    expect(session).toMatchObject({ from: 7, to: 10, expectedLength: 3, dirty: false });
72
  });
73
 
74
  it("does not absorb a foreign insertion at exactly the start boundary", () => {
75
    let state = makeState();
76
    state = state.apply(
77
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
78
    );
79
    state = state.apply(
80
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
81
    );
82
 
83
    // Foreign insertion at exactly `from`. The user's character must stay
84
    // outside the tracked range so rollback/finalize cannot delete or replace
85
    // it. Range advances to [7, 10), still length 3, dirty stays false.
86
    state = state.apply(state.tr.insertText("Y", 6));
87
 
88
    const session = getSession(state);
89
    expect(session?.from).toBe(7);
90
    expect(session?.to).toBe(10);
91
    expect(session?.dirty).toBe(false);
92
    expect(state.doc.textBetween(session!.from, session!.to)).toBe("AAA");
93
  });
94
 
95
  it("does not absorb a foreign insertion at exactly the end boundary", () => {
96
    let state = makeState();
97
    state = state.apply(
98
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
99
    );
100
    state = state.apply(
101
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
102
    );
103
 
104
    // Foreign insertion at exactly `to`. The boundary character must stay
105
    // outside the tracked range. Range stays [6, 9), dirty stays false.
106
    state = state.apply(state.tr.insertText("Z", 9));
107
 
108
    const session = getSession(state);
109
    expect(session?.from).toBe(6);
110
    expect(session?.to).toBe(9);
111
    expect(session?.dirty).toBe(false);
112
    expect(state.doc.textBetween(session!.from, session!.to)).toBe("AAA");
113
  });
114
 
115
  it("does not flag dirty for external edits after the range", () => {
116
    let state = makeState();
117
    state = state.apply(
118
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
119
    );
120
    state = state.apply(
121
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
122
    );
123
 
124
    // External insertion at the very end (after the range).
125
    state = state.apply(state.tr.insertText("Z", state.doc.content.size - 1));
126
 
127
    const session = getSession(state);
128
    expect(session?.from).toBe(6);
129
    expect(session?.to).toBe(9);
130
    expect(session?.dirty).toBe(false);
131
  });
132
 
133
  it("flags dirty when an external edit changes the range length (insert inside)", () => {
134
    let state = makeState();
135
    state = state.apply(
136
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
137
    );
138
    state = state.apply(
139
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
140
    );
141
 
142
    // External insertion at pos 7 (inside the AAA range).
143
    state = state.apply(state.tr.insertText("Y", 7));
144
 
145
    const session = getSession(state);
146
    expect(session?.dirty).toBe(true);
147
    // Range length now 4, expectedLength still 3.
148
    expect(session && session.to - session.from).toBe(4);
149
  });
150
 
151
  it("flags dirty when the range collapses to zero", () => {
152
    let state = makeState();
153
    state = state.apply(
154
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
155
    );
156
    state = state.apply(
157
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
158
    );
159
 
160
    // External delete that wipes out the entire generated range.
161
    state = state.apply(state.tr.delete(6, 9));
162
 
163
    const session = getSession(state);
164
    expect(session?.dirty).toBe(true);
165
    expect(session && session.to <= session.from).toBe(true);
166
  });
167
 
168
  it("matrix cell A: left-adjacent non-zero deletion ending exactly at `from` is clean", () => {
169
    let state = makeState();
170
    state = state.apply(
171
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
172
    );
173
    state = state.apply(
174
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
175
    );
176
 
177
    // delete([4, 6)) is left-adjacent: oldEnd === from. Strict overlap test
178
    // excludes it (6 > 6 is false). Range maps to [4, 7) and stays clean.
179
    state = state.apply(state.tr.delete(4, 6));
180
 
181
    const session = getSession(state);
182
    expect(session?.from).toBe(4);
183
    expect(session?.to).toBe(7);
184
    expect(session?.dirty).toBe(false);
185
    expect(state.doc.textBetween(session!.from, session!.to)).toBe("AAA");
186
  });
187
 
188
  it("matrix cell C: deletion crossing the start boundary is dirty", () => {
189
    let state = makeState();
190
    state = state.apply(
191
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
192
    );
193
    state = state.apply(
194
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
195
    );
196
 
197
    // delete([5, 7)) crosses the start boundary: oldStart < from < oldEnd.
198
    // Overlap test fires (7 > 6 && 5 < 9), session goes dirty, range shrinks.
199
    state = state.apply(state.tr.delete(5, 7));
200
 
201
    const session = getSession(state);
202
    expect(session?.dirty).toBe(true);
203
    expect(session?.from).toBe(5);
204
    expect(session?.to).toBe(7);
205
  });
206
 
207
  it("matrix cell G: deletion crossing the end boundary is dirty", () => {
208
    let state = makeState();
209
    state = state.apply(
210
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
211
    );
212
    state = state.apply(
213
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
214
    );
215
 
216
    // delete([8, 10)) crosses the end boundary: oldStart < to < oldEnd.
217
    // Overlap test fires (10 > 6 && 8 < 9), session goes dirty, range shrinks.
218
    state = state.apply(state.tr.delete(8, 10));
219
 
220
    const session = getSession(state);
221
    expect(session?.dirty).toBe(true);
222
    expect(session?.from).toBe(6);
223
    expect(session?.to).toBe(8);
224
  });
225
 
226
  it("matrix cell I: right-adjacent non-zero deletion starting exactly at `to` is clean", () => {
227
    let state = makeState();
228
    state = state.apply(
229
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
230
    );
231
    state = state.apply(
232
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
233
    );
234
 
235
    // delete([9, 11)) is right-adjacent: oldStart === to. Strict overlap test
236
    // excludes it (9 < 9 is false). Range stays at [6, 9) and clean.
237
    state = state.apply(state.tr.delete(9, 11));
238
 
239
    const session = getSession(state);
240
    expect(session?.from).toBe(6);
241
    expect(session?.to).toBe(9);
242
    expect(session?.dirty).toBe(false);
243
    expect(state.doc.textBetween(session!.from, session!.to)).toBe("AAA");
244
  });
245
 
246
  it("chunk meta on a dirty session does not advance `to` or `expectedLength`", () => {
247
    let state = makeState();
248
    state = state.apply(
249
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
250
    );
251
    state = state.apply(
252
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
253
    );
254
 
255
    // Foreign edit inside the range flips dirty.
256
    state = state.apply(state.tr.delete(6, 7));
257
    const beforeSession = getSession(state);
258
    expect(beforeSession?.dirty).toBe(true);
259
 
260
    // Even if a chunk meta tr reaches apply() while dirty (helper bypassed,
261
    // alternate dispatch path, etc.), `to` and `expectedLength` must not
262
    // advance. The invariant disallows new AI text once ownership is
263
    // unprovable.
264
    state = state.apply(state.tr.setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "XYZ" }));
265
 
266
    const afterSession = getSession(state);
267
    expect(afterSession?.dirty).toBe(true);
268
    expect(afterSession?.to).toBe(beforeSession?.to);
269
    expect(afterSession?.expectedLength).toBe(beforeSession?.expectedLength);
270
  });
271
 
272
  it("flags dirty when a deletion partially overlaps the range start", () => {
273
    let state = makeState();
274
    state = state.apply(
275
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
276
    );
277
    state = state.apply(
278
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
279
    );
280
 
281
    // Delete one character of the range (the first A). Step source [6,7)
282
    // overlaps [6,9), so dirty must flip on. Range shrinks to [6,8).
283
    state = state.apply(state.tr.delete(6, 7));
284
 
285
    const session = getSession(state);
286
    expect(session?.dirty).toBe(true);
287
    expect(session?.from).toBe(6);
288
    expect(session?.to).toBe(8);
289
  });
290
 
291
  it("flags dirty when a deletion partially overlaps the range end", () => {
292
    let state = makeState();
293
    state = state.apply(
294
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
295
    );
296
    state = state.apply(
297
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
298
    );
299
 
300
    // Delete the last character of the range. Step source [8,9) overlaps
301
    // [6,9). Range shrinks to [6,8).
302
    state = state.apply(state.tr.delete(8, 9));
303
 
304
    const session = getSession(state);
305
    expect(session?.dirty).toBe(true);
306
    expect(session?.from).toBe(6);
307
    expect(session?.to).toBe(8);
308
  });
309
 
310
  it("flags dirty when a same-length replacement lands inside the range", () => {
311
    let state = makeState();
312
    state = state.apply(
313
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
314
    );
315
    state = state.apply(
316
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
317
    );
318
 
319
    // Replace the middle "A" with "Y": same total length but the user's edit
320
    // is inside the generated range. Length-divergence would miss this.
321
    state = state.apply(state.tr.replaceWith(7, 8, state.schema.text("Y")));
322
 
323
    const session = getSession(state);
324
    expect(session?.dirty).toBe(true);
325
    // Mapped length unchanged (still 3) but the session is dirty.
326
    expect(session && session.to - session.from).toBe(3);
327
  });
328
 
329
  it("appendAiGenerateChunk skips dispatch once the session is dirty (collapsed range)", () => {
330
    let state = makeState();
331
    state = state.apply(
332
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
333
    );
334
    state = state.apply(
335
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
336
    );
337
    // External delete collapses the generated range.
338
    state = state.apply(state.tr.delete(6, 9));
339
 
340
    const session = getSession(state);
341
    expect(session?.dirty).toBe(true);
342
    expect(session && session.to <= session.from).toBe(true);
343
 
344
    let dispatched = 0;
345
    const fakeView = {
346
      state,
347
      dispatch: () => {
348
        dispatched += 1;
349
      },
350
    } as unknown as EditorView;
351
 
352
    appendAiGenerateChunk(fakeView, "s1", "should-not-land");
353
    expect(dispatched).toBe(0);
354
  });
355
 
356
  it("appendAiGenerateChunk skips dispatch once the session is dirty (length unchanged)", () => {
357
    let state = makeState();
358
    state = state.apply(
359
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
360
    );
361
    state = state.apply(
362
      state.tr.insertText("AAA", 6).setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId: "s1", text: "AAA" }),
363
    );
364
    // Same-length replacement inside the range still flags dirty under the
365
    // step-walk detector. Subsequent chunks must not land.
366
    state = state.apply(state.tr.replaceWith(7, 8, state.schema.text("Y")));
367
 
368
    const session = getSession(state);
369
    expect(session?.dirty).toBe(true);
370
 
371
    let dispatched = 0;
372
    const fakeView = {
373
      state,
374
      dispatch: () => {
375
        dispatched += 1;
376
      },
377
    } as unknown as EditorView;
378
 
379
    appendAiGenerateChunk(fakeView, "s1", "extra");
380
    expect(dispatched).toBe(0);
381
  });
382
 
383
  it("clears active session on end", () => {
384
    let state = makeState();
385
    state = state.apply(
386
      state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId: "s1", pos: 6, label: "x" }),
387
    );
388
    state = state.apply(state.tr.setMeta(aiGenerateIndicatorKey, { type: "end", sessionId: "s1" }));
389
 
390
    expect(getSession(state)).toBeNull();
391
  });
392
});