1
import type { Editor, Range } from "@tiptap/core";
2
import { streamGenerate, AiStreamError } from "@/client/lib/ai/api";
3
import { extractDocumentTitle, extractGenerateContext } from "@/client/lib/ai/context";
4
import { parseAiBlocksFromText, isSingleInlineParagraph, getInlineTextFromParagraph } from "@/client/lib/ai/blocks";
5
import { toast } from "@/client/components/toast-store";
6
import {
7
  appendAiGenerateChunk,
8
  beginAiGenerate,
9
  endAiGenerate,
10
  getAiGenerateSession,
11
  registerAiGenerateAbort,
12
  unregisterAiGenerateAbort,
13
} from "../../extensions/ai-generate-indicator";
14
import type { EditorRuntimeSnapshot } from "../../editor-runtime-context";
15
import type { AiGenerateIntent } from "@/shared/types";
16
 
17
const activeGenerateControllers = new WeakMap<Editor, AbortController>();
18
 
19
export async function startGenerateAtRange(opts: {
20
  editor: Editor;
21
  range: Range;
22
  intent: AiGenerateIntent;
23
  runtime: EditorRuntimeSnapshot;
24
}): Promise<void> {
25
  const { editor, range, intent, runtime } = opts;
26
  if (!runtime.workspaceId || !runtime.pageId) return;
27
 
28
  const previous = activeGenerateControllers.get(editor);
29
  if (previous) previous.abort();
30
 
31
  editor.chain().focus(null, { scrollIntoView: false }).deleteRange(range).run();
32
 
33
  const startPos = editor.state.selection.from;
34
  const context = extractGenerateContext(editor.state, startPos);
35
  const pageTitle = extractDocumentTitle(editor.state.doc);
36
  const indicatorLabel = INDICATOR_LABELS[intent];
37
 
38
  const sessionId = beginAiGenerate(editor.view, startPos, indicatorLabel);
39
 
40
  const controller = new AbortController();
41
  activeGenerateControllers.set(editor, controller);
42
  registerAiGenerateAbort(sessionId, () => controller.abort());
43
 
44
  const iter = streamGenerate(
45
    runtime.workspaceId,
46
    runtime.pageId,
47
    {
48
      intent,
49
      beforeBlock: context.beforeBlock,
50
      afterBlock: context.afterBlock,
51
      pageTitle,
52
    },
53
    controller.signal,
54
  );
55
 
56
  let received = "";
57
  try {
58
    for await (const chunk of iter) {
59
      if (controller.signal.aborted) break;
60
      if (!chunk.text) continue;
61
      received += chunk.text;
62
      appendAiGenerateChunk(editor.view, sessionId, chunk.text);
63
    }
64
    if (controller.signal.aborted) {
65
      rollbackIfClean(editor, sessionId);
66
      return;
67
    }
68
    if (received.trim().length === 0) {
69
      rollbackIfClean(editor, sessionId);
70
      toast.error("The model returned no content. Try again or switch models.");
71
      return;
72
    }
73
    finalizeIfClean(editor, sessionId, received);
74
  } catch (err) {
75
    rollbackIfClean(editor, sessionId);
76
    if (controller.signal.aborted) return;
77
    const message = err instanceof AiStreamError ? err.message : "AI generation failed";
78
    toast.error(message);
79
  } finally {
80
    unregisterAiGenerateAbort(sessionId);
81
    endAiGenerate(editor.view, sessionId);
82
    if (activeGenerateControllers.get(editor) === controller) {
83
      activeGenerateControllers.delete(editor);
84
    }
85
  }
86
}
87
 
88
// Concurrent edits to the generated range mark the session dirty (range
89
// length diverged from the streamed accumulator, or the range collapsed). In
90
// that case never delete or replace — the doc now holds user/peer content
91
// and we must not overwrite it. Clean rollback uses the plugin's mapped
92
// `from`/`to`, not the original ints.
93
function rollbackIfClean(editor: Editor, sessionId: string): void {
94
  const session = getAiGenerateSession(editor.state);
95
  if (!session || session.sessionId !== sessionId) return;
96
  if (session.dirty) return;
97
  if (session.to <= session.from) return;
98
  const tr = editor.state.tr.delete(session.from, session.to);
99
  editor.view.dispatch(tr);
100
}
101
 
102
function finalizeIfClean(editor: Editor, sessionId: string, raw: string): void {
103
  const session = getAiGenerateSession(editor.state);
104
  if (!session || session.sessionId !== sessionId) return;
105
  if (session.dirty) {
106
    toast.info("Generation kept as-is — external edits detected.");
107
    return;
108
  }
109
  const { from, to } = session;
110
  if (to <= from) return;
111
 
112
  const blocks = parseAiBlocksFromText(raw);
113
  if (blocks.length === 0) return;
114
  if (isSingleInlineParagraph(blocks)) {
115
    const text = getInlineTextFromParagraph(blocks[0]);
116
    if (text === raw) return;
117
    const tr = editor.state.tr.insertText(text, from, to);
118
    editor.view.dispatch(tr);
119
    return;
120
  }
121
 
122
  const nodes = blocks
123
    .map((block) => {
124
      try {
125
        return editor.schema.nodeFromJSON(block);
126
      } catch {
127
        return null;
128
      }
129
    })
130
    .filter((node): node is NonNullable<typeof node> => node !== null);
131
  if (nodes.length === 0) return;
132
 
133
  // When the host block is a paragraph whose entire inline content is exactly
134
  // the streamed range, consume the host rather than splitting it, so the new
135
  // paragraphs don't leave an empty shell around them.
136
  let replaceFrom = from;
137
  let replaceTo = to;
138
  const $from = editor.state.doc.resolve(from);
139
  const depth = $from.depth;
140
  if (depth > 0) {
141
    const parent = $from.parent;
142
    if (parent.type.name === "paragraph") {
143
      const parentStart = $from.before(depth);
144
      const parentEnd = $from.after(depth);
145
      if (parentStart + 1 === from && parentEnd - 1 === to) {
146
        replaceFrom = parentStart;
147
        replaceTo = parentEnd;
148
      }
149
    }
150
  }
151
 
152
  const tr = editor.state.tr.replaceWith(replaceFrom, replaceTo, nodes);
153
  editor.view.dispatch(tr);
154
}
155
 
156
const INDICATOR_LABELS: Record<AiGenerateIntent, string> = {
157
  continue: "Continuing…",
158
  explain: "Explaining…",
159
  brainstorm: "Brainstorming…",
160
};