1
import { describe, expect, it } from "vitest";
2
import { getSchema } from "@tiptap/core";
3
import { StarterKit } from "@tiptap/starter-kit";
4
import { EditorState, TextSelection } from "@tiptap/pm/state";
5
import { extractDocumentTitle, extractGenerateContext, extractRewriteContext } from "@/client/lib/ai/context";
6
 
7
const schema = getSchema([StarterKit.configure({ undoRedo: false })]);
8
 
9
function stateFromDoc(doc: Parameters<typeof schema.nodeFromJSON>[0], from?: number, to?: number): EditorState {
10
  const pmDoc = schema.nodeFromJSON(doc);
11
  const state = EditorState.create({ schema, doc: pmDoc });
12
  if (typeof from !== "number") return state;
13
  const resolvedFrom = state.doc.resolve(from);
14
  const resolvedTo = state.doc.resolve(to ?? from);
15
  return state.apply(state.tr.setSelection(new TextSelection(resolvedFrom, resolvedTo)));
16
}
17
 
18
describe("extractRewriteContext", () => {
19
  it("captures selected text, parent block, and both siblings for a middle paragraph", () => {
20
    const state = stateFromDoc(
21
      {
22
        type: "doc",
23
        content: [
24
          { type: "paragraph", content: [{ type: "text", text: "Intro paragraph." }] },
25
          { type: "paragraph", content: [{ type: "text", text: "Middle paragraph text." }] },
26
          { type: "paragraph", content: [{ type: "text", text: "Closing paragraph." }] },
27
        ],
28
      },
29
      19,
30
      26,
31
    );
32
 
33
    const ctx = extractRewriteContext(state);
34
    expect(ctx.selectedText).toBe("Middle ");
35
    expect(ctx.parentBlock).toBe("Middle paragraph text.");
36
    expect(ctx.beforeBlock).toBe("Intro paragraph.");
37
    expect(ctx.afterBlock).toBe("Closing paragraph.");
38
  });
39
 
40
  it("returns empty siblings at document edges", () => {
41
    const state = stateFromDoc(
42
      {
43
        type: "doc",
44
        content: [{ type: "paragraph", content: [{ type: "text", text: "Only paragraph." }] }],
45
      },
46
      1,
47
      5,
48
    );
49
 
50
    const ctx = extractRewriteContext(state);
51
    expect(ctx.selectedText).toBe("Only");
52
    expect(ctx.parentBlock).toBe("Only paragraph.");
53
    expect(ctx.beforeBlock).toBe("");
54
    expect(ctx.afterBlock).toBe("");
55
  });
56
});
57
 
58
describe("extractGenerateContext", () => {
59
  it("slices the current block at the cursor and joins with prev sibling", () => {
60
    const state = stateFromDoc(
61
      {
62
        type: "doc",
63
        content: [
64
          { type: "paragraph", content: [{ type: "text", text: "First thought." }] },
65
          { type: "paragraph", content: [{ type: "text", text: "Second thought." }] },
66
        ],
67
      },
68
      18,
69
    );
70
 
71
    const ctx = extractGenerateContext(state, 18);
72
    expect(ctx.beforeBlock).toBe("First thought.\n\nS");
73
    expect(ctx.afterBlock).toBe("econd thought.");
74
  });
75
 
76
  it("at the start of a paragraph, includes only prev sibling before and full current after", () => {
77
    const state = stateFromDoc(
78
      {
79
        type: "doc",
80
        content: [
81
          { type: "paragraph", content: [{ type: "text", text: "First thought." }] },
82
          { type: "paragraph", content: [{ type: "text", text: "Second thought." }] },
83
        ],
84
      },
85
      17,
86
    );
87
 
88
    const ctx = extractGenerateContext(state, 17);
89
    expect(ctx.beforeBlock).toBe("First thought.");
90
    expect(ctx.afterBlock).toBe("Second thought.");
91
  });
92
 
93
  it("at the end of a paragraph, includes the full current block before and next sibling after", () => {
94
    const state = stateFromDoc(
95
      {
96
        type: "doc",
97
        content: [
98
          { type: "paragraph", content: [{ type: "text", text: "First thought." }] },
99
          { type: "paragraph", content: [{ type: "text", text: "Second thought." }] },
100
        ],
101
      },
102
      15,
103
    );
104
 
105
    const ctx = extractGenerateContext(state, 15);
106
    expect(ctx.beforeBlock).toBe("First thought.");
107
    expect(ctx.afterBlock).toBe("Second thought.");
108
  });
109
 
110
  it("slices a nested list item's content at the cursor", () => {
111
    const state = stateFromDoc(
112
      {
113
        type: "doc",
114
        content: [
115
          {
116
            type: "bulletList",
117
            content: [
118
              {
119
                type: "listItem",
120
                content: [{ type: "paragraph", content: [{ type: "text", text: "Hello" }] }],
121
              },
122
            ],
123
          },
124
        ],
125
      },
126
      5,
127
    );
128
 
129
    const ctx = extractGenerateContext(state, 5);
130
    expect(ctx.beforeBlock).toBe("He");
131
    expect(ctx.afterBlock).toBe("llo");
132
  });
133
 
134
  it("uses prev/next blocks when the cursor lands in an empty paragraph", () => {
135
    const state = stateFromDoc(
136
      {
137
        type: "doc",
138
        content: [
139
          { type: "paragraph", content: [{ type: "text", text: "Head." }] },
140
          { type: "paragraph" },
141
          { type: "paragraph", content: [{ type: "text", text: "Tail." }] },
142
        ],
143
      },
144
      8,
145
    );
146
 
147
    const ctx = extractGenerateContext(state, 8);
148
    expect(ctx.beforeBlock).toBe("Head.");
149
    expect(ctx.afterBlock).toBe("Tail.");
150
  });
151
 
152
  it("walks past multiple empty blocks on either side", () => {
153
    const state = stateFromDoc(
154
      {
155
        type: "doc",
156
        content: [
157
          { type: "paragraph", content: [{ type: "text", text: "Anchor before." }] },
158
          { type: "paragraph" },
159
          { type: "paragraph" },
160
          { type: "paragraph" },
161
          { type: "paragraph" },
162
          { type: "paragraph" },
163
          { type: "paragraph", content: [{ type: "text", text: "Anchor after." }] },
164
        ],
165
      },
166
      19,
167
    );
168
 
169
    const ctx = extractGenerateContext(state, 19);
170
    expect(ctx.beforeBlock).toBe("Anchor before.");
171
    expect(ctx.afterBlock).toBe("Anchor after.");
172
  });
173
 
174
  it("returns empty context when surrounded only by empty blocks within walk limit", () => {
175
    const state = stateFromDoc(
176
      {
177
        type: "doc",
178
        content: Array.from({ length: 5 }, () => ({ type: "paragraph" })),
179
      },
180
      2,
181
    );
182
 
183
    const ctx = extractGenerateContext(state, 2);
184
    expect(ctx.beforeBlock).toBe("");
185
    expect(ctx.afterBlock).toBe("");
186
  });
187
});
188
 
189
describe("extractDocumentTitle", () => {
190
  it("returns the first h1 text", () => {
191
    const state = stateFromDoc({
192
      type: "doc",
193
      content: [
194
        { type: "heading", attrs: { level: 1 }, content: [{ type: "text", text: "Hello World" }] },
195
        { type: "paragraph", content: [{ type: "text", text: "body" }] },
196
      ],
197
    });
198
    expect(extractDocumentTitle(state.doc)).toBe("Hello World");
199
  });
200
 
201
  it("returns empty string when there is no h1", () => {
202
    const state = stateFromDoc({
203
      type: "doc",
204
      content: [
205
        { type: "heading", attrs: { level: 2 }, content: [{ type: "text", text: "Sub" }] },
206
        { type: "paragraph", content: [{ type: "text", text: "body" }] },
207
      ],
208
    });
209
    expect(extractDocumentTitle(state.doc)).toBe("");
210
  });
211
 
212
  it("returns empty string for an empty document", () => {
213
    const state = stateFromDoc({ type: "doc", content: [{ type: "paragraph" }] });
214
    expect(extractDocumentTitle(state.doc)).toBe("");
215
  });
216
 
217
  it("skips an empty h1 and uses the next one", () => {
218
    const state = stateFromDoc({
219
      type: "doc",
220
      content: [
221
        { type: "heading", attrs: { level: 1 } },
222
        { type: "heading", attrs: { level: 1 }, content: [{ type: "text", text: "Real Title" }] },
223
      ],
224
    });
225
    expect(extractDocumentTitle(state.doc)).toBe("Real Title");
226
  });
227
 
228
  it("returns the first h1 when multiple exist", () => {
229
    const state = stateFromDoc({
230
      type: "doc",
231
      content: [
232
        { type: "heading", attrs: { level: 1 }, content: [{ type: "text", text: "First" }] },
233
        { type: "heading", attrs: { level: 1 }, content: [{ type: "text", text: "Second" }] },
234
      ],
235
    });
236
    expect(extractDocumentTitle(state.doc)).toBe("First");
237
  });
238
});
239
 
240
describe("extractRewriteContext across empty siblings", () => {
241
  it("skips adjacent empty paragraphs to find before/after context", () => {
242
    const state = stateFromDoc(
243
      {
244
        type: "doc",
245
        content: [
246
          { type: "paragraph", content: [{ type: "text", text: "First." }] },
247
          { type: "paragraph" },
248
          { type: "paragraph", content: [{ type: "text", text: "Middle." }] },
249
          { type: "paragraph" },
250
          { type: "paragraph", content: [{ type: "text", text: "Last." }] },
251
        ],
252
      },
253
      12,
254
      19,
255
    );
256
 
257
    const ctx = extractRewriteContext(state);
258
    expect(ctx.parentBlock).toBe("Middle.");
259
    expect(ctx.beforeBlock).toBe("First.");
260
    expect(ctx.afterBlock).toBe("Last.");
261
  });
262
});