1
import { Extension } from "@tiptap/core";
2
import { Plugin, PluginKey, type EditorState, type Transaction } from "@tiptap/pm/state";
3
import { Decoration, DecorationSet, type EditorView } from "@tiptap/pm/view";
4
import { parseAiBlocksFromText, isSingleInlineParagraph, getInlineTextFromParagraph } from "@/client/lib/ai/blocks";
5
import "../styles/ai-suggestion.css";
6
 
7
interface AiSuggestionState {
8
  active: {
9
    sessionId: string;
10
    from: number;
11
    to: number;
12
    buffer: string;
13
    status: "streaming" | "ready" | "error";
14
    error?: string;
15
  } | null;
16
}
17
 
18
export const aiSuggestionKey = new PluginKey<AiSuggestionState>("aiSuggestion");
19
 
20
const rewriteAbortHandlers = new Map<string, () => void>();
21
 
22
export function registerAiRewriteAbort(sessionId: string, handler: () => void): void {
23
  rewriteAbortHandlers.set(sessionId, handler);
24
}
25
 
26
export function unregisterAiRewriteAbort(sessionId: string): void {
27
  rewriteAbortHandlers.delete(sessionId);
28
}
29
 
30
function triggerAiRewriteAbort(sessionId: string): void {
31
  const handler = rewriteAbortHandlers.get(sessionId);
32
  if (handler) handler();
33
}
34
 
35
type AiSuggestionMeta =
36
  | { type: "start"; sessionId: string; from: number; to: number }
37
  | { type: "append"; sessionId: string; text: string }
38
  | { type: "finish"; sessionId: string }
39
  | { type: "error"; sessionId: string; message: string }
40
  | { type: "cancel" }
41
  | { type: "commit" };
42
 
43
export function startAiSuggestion(view: EditorView, from: number, to: number): string {
44
  const sessionId = createSessionId();
45
  view.dispatch(
46
    view.state.tr.setMeta(aiSuggestionKey, { type: "start", sessionId, from, to } satisfies AiSuggestionMeta),
47
  );
48
  return sessionId;
49
}
50
 
51
export function appendAiSuggestion(view: EditorView, sessionId: string, text: string): void {
52
  view.dispatch(view.state.tr.setMeta(aiSuggestionKey, { type: "append", sessionId, text } satisfies AiSuggestionMeta));
53
}
54
 
55
export function finishAiSuggestion(view: EditorView, sessionId: string): void {
56
  view.dispatch(view.state.tr.setMeta(aiSuggestionKey, { type: "finish", sessionId } satisfies AiSuggestionMeta));
57
}
58
 
59
export function errorAiSuggestion(view: EditorView, sessionId: string, message: string): void {
60
  view.dispatch(
61
    view.state.tr.setMeta(aiSuggestionKey, { type: "error", sessionId, message } satisfies AiSuggestionMeta),
62
  );
63
}
64
 
65
export function cancelAiSuggestion(view: EditorView): void {
66
  view.dispatch(view.state.tr.setMeta(aiSuggestionKey, { type: "cancel" } satisfies AiSuggestionMeta));
67
}
68
 
69
export function commitAiSuggestion(view: EditorView): boolean {
70
  const current = aiSuggestionKey.getState(view.state);
71
  if (!current?.active) return false;
72
  const { from, to, buffer } = current.active;
73
 
74
  const tr = view.state.tr;
75
  const blocks = parseAiBlocksFromText(buffer);
76
  if (blocks.length === 0) {
77
    tr.delete(from, to);
78
  } else if (isSingleInlineParagraph(blocks)) {
79
    const text = getInlineTextFromParagraph(blocks[0]);
80
    if (text) {
81
      tr.insertText(text, from, to);
82
    } else {
83
      tr.delete(from, to);
84
    }
85
  } else {
86
    const nodes = blocks
87
      .map((block) => {
88
        try {
89
          return view.state.schema.nodeFromJSON(block);
90
        } catch {
91
          return null;
92
        }
93
      })
94
      .filter((node): node is NonNullable<typeof node> => node !== null);
95
    if (nodes.length === 0) {
96
      tr.insertText(buffer, from, to);
97
    } else {
98
      tr.replaceWith(from, to, nodes);
99
    }
100
  }
101
  tr.setMeta(aiSuggestionKey, { type: "cancel" } satisfies AiSuggestionMeta);
102
  view.dispatch(tr);
103
  return true;
104
}
105
 
106
export function getAiSuggestionState(state: EditorState): AiSuggestionState | null {
107
  return aiSuggestionKey.getState(state) ?? null;
108
}
109
 
110
export function isAiRewriteInflight(state: EditorState): boolean {
111
  const plugin = aiSuggestionKey.getState(state);
112
  return plugin?.active?.status === "streaming";
113
}
114
 
115
export const AiSuggestion = Extension.create({
116
  name: "aiSuggestion",
117
 
118
  addProseMirrorPlugins() {
119
    return [
120
      new Plugin<AiSuggestionState>({
121
        key: aiSuggestionKey,
122
        state: {
123
          init(): AiSuggestionState {
124
            return { active: null };
125
          },
126
          apply(tr: Transaction, value: AiSuggestionState): AiSuggestionState {
127
            const meta = tr.getMeta(aiSuggestionKey) as AiSuggestionMeta | undefined;
128
 
129
            if (meta?.type === "start") {
130
              return {
131
                active: {
132
                  sessionId: meta.sessionId,
133
                  from: meta.from,
134
                  to: meta.to,
135
                  buffer: "",
136
                  status: "streaming",
137
                },
138
              };
139
            }
140
            if (meta?.type === "cancel" || meta?.type === "commit") {
141
              return { active: null };
142
            }
143
 
144
            if (!value.active) return value;
145
 
146
            if (meta?.type === "append") {
147
              if (meta.sessionId !== value.active.sessionId) return value;
148
              return {
149
                active: { ...value.active, buffer: value.active.buffer + meta.text },
150
              };
151
            }
152
            if (meta?.type === "finish") {
153
              if (meta.sessionId !== value.active.sessionId) return value;
154
              return { active: { ...value.active, status: "ready" } };
155
            }
156
            if (meta?.type === "error") {
157
              if (meta.sessionId !== value.active.sessionId) return value;
158
              return { active: { ...value.active, status: "error", error: meta.message } };
159
            }
160
 
161
            if (tr.docChanged) {
162
              const from = tr.mapping.map(value.active.from, -1);
163
              const to = tr.mapping.map(value.active.to, 1);
164
              if (to <= from) return { active: null };
165
              return { active: { ...value.active, from, to } };
166
            }
167
 
168
            return value;
169
          },
170
        },
171
        props: {
172
          decorations(state) {
173
            const plugin = aiSuggestionKey.getState(state);
174
            if (!plugin?.active) return null;
175
            const { from, to, buffer, status, error } = plugin.active;
176
            const decorations: Decoration[] = [
177
              Decoration.inline(from, to, { class: "tiptap-ai-target" }),
178
              Decoration.widget(to, () => renderPreview({ buffer, status, error }), {
179
                side: 1,
180
                ignoreSelection: true,
181
                key: `ai-suggestion-preview:${buffer.length}:${status}`,
182
              }),
183
            ];
184
            return DecorationSet.create(state.doc, decorations);
185
          },
186
          handleDOMEvents: {
187
            mousedown(view, event) {
188
              const target = event.target as HTMLElement | null;
189
              if (!target) return false;
190
              const button = target.closest<HTMLButtonElement>("[data-ai-suggestion-action]");
191
              if (!button) return false;
192
              event.preventDefault();
193
              const action = button.dataset.aiSuggestionAction;
194
              const plugin = aiSuggestionKey.getState(view.state);
195
              if (action === "accept") {
196
                commitAiSuggestion(view);
197
              } else if (action === "reject") {
198
                if (plugin?.active?.status === "streaming") {
199
                  triggerAiRewriteAbort(plugin.active.sessionId);
200
                }
201
                cancelAiSuggestion(view);
202
              }
203
              return true;
204
            },
205
          },
206
          handleKeyDown(view, event) {
207
            const plugin = aiSuggestionKey.getState(view.state);
208
            if (!plugin?.active) return false;
209
 
210
            if (event.key === "Escape") {
211
              if (plugin.active.status === "streaming") {
212
                triggerAiRewriteAbort(plugin.active.sessionId);
213
              }
214
              cancelAiSuggestion(view);
215
              return true;
216
            }
217
 
218
            if (event.key === "Enter" && plugin.active.status === "ready") {
219
              commitAiSuggestion(view);
220
              return true;
221
            }
222
 
223
            return false;
224
          },
225
        },
226
      }),
227
    ];
228
  },
229
});
230
 
231
function renderPreview(args: { buffer: string; status: "streaming" | "ready" | "error"; error?: string }): HTMLElement {
232
  const wrapper = document.createElement("span");
233
  wrapper.className = `tiptap-ai-preview tiptap-ai-preview-${args.status}`;
234
  wrapper.setAttribute("contenteditable", "false");
235
  wrapper.setAttribute("data-ai-suggestion", "preview");
236
 
237
  const bubble = document.createElement("span");
238
  bubble.className = "tiptap-ai-preview-bubble";
239
  if (args.status === "error") {
240
    bubble.textContent = args.error ?? "AI request failed";
241
  } else if (args.status === "streaming" && args.buffer.length === 0) {
242
    bubble.className += " tiptap-ai-preview-bubble-pending";
243
    bubble.textContent = "Thinking…";
244
  } else {
245
    bubble.textContent = args.buffer;
246
  }
247
  wrapper.appendChild(bubble);
248
 
249
  const actions = document.createElement("span");
250
  actions.className = "tiptap-ai-preview-actions";
251
 
252
  const rejectVerb = args.status === "error" ? "dismiss" : args.status === "streaming" ? "cancel" : "reject";
253
  actions.appendChild(
254
    createPreviewAction({
255
      key: "esc",
256
      verb: rejectVerb,
257
      action: "reject",
258
      ariaLabel: `${capitalize(rejectVerb)} AI suggestion (Escape)`,
259
      extraClass: "tiptap-ai-preview-reject",
260
    }),
261
  );
262
 
263
  if (args.status === "ready") {
264
    const divider = document.createElement("span");
265
    divider.className = "tiptap-ai-preview-divider";
266
    divider.setAttribute("aria-hidden", "true");
267
    divider.textContent = "·";
268
    actions.appendChild(divider);
269
 
270
    actions.appendChild(
271
      createPreviewAction({
272
        key: "⏎",
273
        verb: "accept",
274
        action: "accept",
275
        ariaLabel: "Accept AI suggestion (Enter)",
276
        extraClass: "tiptap-ai-preview-accept",
277
      }),
278
    );
279
  }
280
 
281
  wrapper.appendChild(actions);
282
  return wrapper;
283
}
284
 
285
function createPreviewAction(opts: {
286
  key: string;
287
  verb: string;
288
  action: "accept" | "reject";
289
  ariaLabel: string;
290
  extraClass: string;
291
}): HTMLButtonElement {
292
  const btn = document.createElement("button");
293
  btn.type = "button";
294
  btn.className = `tiptap-ai-preview-action ${opts.extraClass}`;
295
  btn.dataset.aiSuggestionAction = opts.action;
296
  btn.setAttribute("aria-label", opts.ariaLabel);
297
 
298
  const kbd = document.createElement("kbd");
299
  kbd.className = "tiptap-ai-kbd";
300
  kbd.textContent = opts.key;
301
  btn.appendChild(kbd);
302
 
303
  const verbNode = document.createElement("span");
304
  verbNode.className = "tiptap-ai-preview-verb";
305
  verbNode.textContent = opts.verb;
306
  btn.appendChild(verbNode);
307
 
308
  return btn;
309
}
310
 
311
function capitalize(text: string): string {
312
  return text.length === 0 ? text : text[0].toUpperCase() + text.slice(1);
313
}
314
 
315
function createSessionId(): string {
316
  if (typeof crypto !== "undefined" && typeof crypto.randomUUID === "function") {
317
    return crypto.randomUUID();
318
  }
319
  return `ai-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 10)}`;
320
}