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 "../styles/ai-suggestion.css";
5
 
6
interface ActiveSession {
7
  sessionId: string;
8
  from: number;
9
  to: number;
10
  expectedLength: number;
11
  dirty: boolean;
12
  label: string;
13
}
14
 
15
interface IndicatorState {
16
  active: ActiveSession | null;
17
  inflight: { sessionId: string } | null;
18
}
19
 
20
type IndicatorMeta =
21
  | { type: "begin"; sessionId: string; pos: number; label: string }
22
  | { type: "chunk"; sessionId: string; text: string }
23
  | { type: "hide-indicator"; sessionId: string }
24
  | { type: "end"; sessionId: string };
25
 
26
export const aiGenerateIndicatorKey = new PluginKey<IndicatorState>("aiGenerateIndicator");
27
 
28
const abortHandlers = new Map<string, () => void>();
29
 
30
export function registerAiGenerateAbort(sessionId: string, handler: () => void): void {
31
  abortHandlers.set(sessionId, handler);
32
}
33
 
34
export function unregisterAiGenerateAbort(sessionId: string): void {
35
  abortHandlers.delete(sessionId);
36
}
37
 
38
function triggerAiGenerateAbort(sessionId: string): void {
39
  const handler = abortHandlers.get(sessionId);
40
  if (handler) handler();
41
}
42
 
43
export function beginAiGenerate(view: EditorView, pos: number, label = "Generating…"): string {
44
  const sessionId = createSessionId();
45
  view.dispatch(
46
    view.state.tr.setMeta(aiGenerateIndicatorKey, { type: "begin", sessionId, pos, label } satisfies IndicatorMeta),
47
  );
48
  return sessionId;
49
}
50
 
51
// Inserts the streaming chunk at the plugin's currently-tracked range end and
52
// stamps the transaction with chunk meta so the plugin's apply() advances `to`
53
// by exactly text.length without flagging the session dirty.
54
//
55
// Once the session is dirty (an external edit landed inside the generated
56
// range, or the range was deleted/collapsed), bail out: never recreate AI
57
// content where a collaborator/user has taken over. The text accumulator in
58
// the controller still advances so the stream can be tracked to completion,
59
// but no further mutations land in the document.
60
export function appendAiGenerateChunk(view: EditorView, sessionId: string, text: string): void {
61
  if (!text) return;
62
  const session = getAiGenerateSession(view.state);
63
  if (!session || session.sessionId !== sessionId) return;
64
  if (session.dirty) return;
65
  const tr = view.state.tr.insertText(text, session.to);
66
  tr.setMeta(aiGenerateIndicatorKey, { type: "chunk", sessionId, text } satisfies IndicatorMeta);
67
  view.dispatch(tr);
68
}
69
 
70
export function hideAiGenerateIndicator(view: EditorView, sessionId: string): void {
71
  view.dispatch(
72
    view.state.tr.setMeta(aiGenerateIndicatorKey, { type: "hide-indicator", sessionId } satisfies IndicatorMeta),
73
  );
74
}
75
 
76
export function endAiGenerate(view: EditorView, sessionId: string): void {
77
  view.dispatch(view.state.tr.setMeta(aiGenerateIndicatorKey, { type: "end", sessionId } satisfies IndicatorMeta));
78
}
79
 
80
export function isAiGenerateInflight(state: EditorState): boolean {
81
  const plugin = aiGenerateIndicatorKey.getState(state);
82
  return Boolean(plugin?.inflight);
83
}
84
 
85
export interface AiGenerateSessionSnapshot {
86
  sessionId: string;
87
  from: number;
88
  to: number;
89
  expectedLength: number;
90
  dirty: boolean;
91
}
92
 
93
export function getAiGenerateSession(state: EditorState): AiGenerateSessionSnapshot | null {
94
  const plugin = aiGenerateIndicatorKey.getState(state);
95
  if (!plugin?.active) return null;
96
  const { sessionId, from, to, expectedLength, dirty } = plugin.active;
97
  return { sessionId, from, to, expectedLength, dirty };
98
}
99
 
100
export function createAiGenerateIndicatorPlugin(): Plugin<IndicatorState> {
101
  return new Plugin<IndicatorState>({
102
    key: aiGenerateIndicatorKey,
103
    state: {
104
      init(): IndicatorState {
105
        return { active: null, inflight: null };
106
      },
107
      apply(tr: Transaction, value: IndicatorState): IndicatorState {
108
        const meta = tr.getMeta(aiGenerateIndicatorKey) as IndicatorMeta | undefined;
109
        if (meta?.type === "begin") {
110
          return {
111
            active: {
112
              sessionId: meta.sessionId,
113
              from: meta.pos,
114
              to: meta.pos,
115
              expectedLength: 0,
116
              dirty: false,
117
              label: meta.label,
118
            },
119
            inflight: { sessionId: meta.sessionId },
120
          };
121
        }
122
        if (meta?.type === "chunk") {
123
          if (!value.active || value.active.sessionId !== meta.sessionId) return value;
124
          // Defense in depth: even if a chunk meta tr reaches apply() while
125
          // the session is dirty (helper guard bypassed, or some other
126
          // dispatch path), do not advance `to`/`expectedLength`. The
127
          // invariant requires no AI text be appended once ownership is
128
          // unprovable.
129
          if (value.active.dirty) return value;
130
          const expectedLength = value.active.expectedLength + meta.text.length;
131
          const to = value.active.to + meta.text.length;
132
          return {
133
            ...value,
134
            active: { ...value.active, to, expectedLength },
135
          };
136
        }
137
        if (meta?.type === "hide-indicator") {
138
          if (!value.active || value.active.sessionId !== meta.sessionId) return value;
139
          return { ...value, active: null };
140
        }
141
        if (meta?.type === "end") {
142
          const active = value.active && value.active.sessionId === meta.sessionId ? null : value.active;
143
          const inflight = value.inflight && value.inflight.sessionId === meta.sessionId ? null : value.inflight;
144
          if (active === value.active && inflight === value.inflight) return value;
145
          return { active, inflight };
146
        }
147
        if (!value.active) return value;
148
        if (tr.docChanged) {
149
          // Foreign (non-chunk) doc-changing transaction. Walk each step and
150
          // check whether its replaced range overlaps the active range in
151
          // pre-step coordinates; advance the active range through the step's
152
          // map; repeat. Length-divergence alone misses same-length
153
          // replacements (select 1 char, type 1 char).
154
          //
155
          // Mapping bias is "exclusive" at both ends: from uses +1 so a
156
          // foreign insertion at exactly `from` keeps the new content
157
          // OUTSIDE the range (the start advances past it). to uses -1 so a
158
          // foreign insertion at exactly `to` likewise stays outside (the
159
          // end stops before it). This matches the strict overlap test
160
          // (`oldEnd > from && oldStart < to`) which excludes zero-width
161
          // boundary steps from flagging dirty: range and dirty agree that
162
          // boundary content is not ours. Chunk transactions bypass this
163
          // branch via chunk meta, so they still extend the range from `to`.
164
          //
165
          // Collapsed mapped ranges count as dirty — never resurrect
166
          // deleted territory.
167
          let from = value.active.from;
168
          let to = value.active.to;
169
          let dirty = value.active.dirty;
170
          for (let i = 0; i < tr.steps.length; i++) {
171
            const stepMap = tr.steps[i].getMap();
172
            stepMap.forEach((oldStart, oldEnd) => {
173
              if (oldEnd > from && oldStart < to) dirty = true;
174
            });
175
            from = stepMap.map(from, 1);
176
            to = stepMap.map(to, -1);
177
          }
178
          const collapsed = to <= from;
179
          if (collapsed) dirty = true;
180
          return {
181
            ...value,
182
            active: { ...value.active, from, to: collapsed ? from : to, dirty },
183
          };
184
        }
185
        return value;
186
      },
187
    },
188
    props: {
189
      decorations(state) {
190
        const plugin = aiGenerateIndicatorKey.getState(state);
191
        if (!plugin?.active) return null;
192
        const indicatorPos = plugin.active.to;
193
        return DecorationSet.create(state.doc, [
194
          Decoration.widget(indicatorPos, () => renderIndicator(plugin.active!.label), {
195
            side: 1,
196
            ignoreSelection: true,
197
            key: `ai-generate-indicator:${indicatorPos}:${plugin.active.label}`,
198
          }),
199
        ]);
200
      },
201
      handleDOMEvents: {
202
        mousedown(view, event) {
203
          const target = event.target as HTMLElement | null;
204
          if (!target) return false;
205
          const button = target.closest<HTMLButtonElement>('[data-ai-generate-action="cancel"]');
206
          if (!button) return false;
207
          event.preventDefault();
208
          const plugin = aiGenerateIndicatorKey.getState(view.state);
209
          const sessionId = plugin?.inflight?.sessionId;
210
          if (sessionId) triggerAiGenerateAbort(sessionId);
211
          return true;
212
        },
213
      },
214
      handleKeyDown(view, event) {
215
        if (event.key !== "Escape") return false;
216
        const plugin = aiGenerateIndicatorKey.getState(view.state);
217
        const sessionId = plugin?.inflight?.sessionId;
218
        if (!sessionId) return false;
219
        triggerAiGenerateAbort(sessionId);
220
        return true;
221
      },
222
    },
223
  });
224
}
225
 
226
export const AiGenerateIndicator = Extension.create({
227
  name: "aiGenerateIndicator",
228
 
229
  addProseMirrorPlugins() {
230
    return [createAiGenerateIndicatorPlugin()];
231
  },
232
});
233
 
234
function renderIndicator(label: string): HTMLElement {
235
  const el = document.createElement("span");
236
  el.className = "tiptap-ai-generate-indicator";
237
  el.setAttribute("contenteditable", "false");
238
  el.setAttribute("aria-live", "polite");
239
 
240
  const dot = document.createElement("span");
241
  dot.className = "tiptap-ai-generate-indicator-dot";
242
  el.appendChild(dot);
243
 
244
  const text = document.createElement("span");
245
  text.className = "tiptap-ai-generate-indicator-label";
246
  text.textContent = label;
247
  el.appendChild(text);
248
 
249
  const cancel = document.createElement("button");
250
  cancel.type = "button";
251
  cancel.className = "tiptap-ai-generate-cancel";
252
  cancel.dataset.aiGenerateAction = "cancel";
253
  cancel.setAttribute("aria-label", "Cancel generation (Escape)");
254
 
255
  const kbd = document.createElement("kbd");
256
  kbd.className = "tiptap-ai-kbd";
257
  kbd.textContent = "esc";
258
  cancel.appendChild(kbd);
259
 
260
  const verb = document.createElement("span");
261
  verb.className = "tiptap-ai-generate-cancel-verb";
262
  verb.textContent = "cancel";
263
  cancel.appendChild(verb);
264
 
265
  el.appendChild(cancel);
266
 
267
  return el;
268
}
269
 
270
function createSessionId(): string {
271
  if (typeof crypto !== "undefined" && typeof crypto.randomUUID === "function") {
272
    return crypto.randomUUID();
273
  }
274
  return `ai-gen-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 10)}`;
275
}