1
import type { Editor } from "@tiptap/core";
2
import { streamRewrite, AiStreamError } from "@/client/lib/ai/api";
3
import { extractDocumentTitle, extractRewriteContext } from "@/client/lib/ai/context";
4
import {
5
  appendAiSuggestion,
6
  cancelAiSuggestion,
7
  errorAiSuggestion,
8
  finishAiSuggestion,
9
  registerAiRewriteAbort,
10
  startAiSuggestion,
11
  unregisterAiRewriteAbort,
12
} from "../extensions/ai-suggestion";
13
import type { AiRewriteAction } from "@/shared/types";
14
import type { EditorRuntimeSnapshot } from "../editor-runtime-context";
15
 
16
const activeRewriteControllers = new WeakMap<Editor, AbortController>();
17
 
18
export async function runRewrite(opts: {
19
  editor: Editor;
20
  action: AiRewriteAction;
21
  runtime: EditorRuntimeSnapshot;
22
}): Promise<void> {
23
  const { editor, action, runtime } = opts;
24
  if (!runtime.workspaceId || !runtime.pageId) return;
25
 
26
  const selection = editor.state.selection;
27
  if (selection.empty) return;
28
  const { from, to } = selection;
29
 
30
  const context = extractRewriteContext(editor.state);
31
  if (!context.selectedText) return;
32
 
33
  const previous = activeRewriteControllers.get(editor);
34
  if (previous) previous.abort();
35
 
36
  const controller = new AbortController();
37
  activeRewriteControllers.set(editor, controller);
38
 
39
  const pageTitle = extractDocumentTitle(editor.state.doc);
40
  const sessionId = startAiSuggestion(editor.view, from, to);
41
  registerAiRewriteAbort(sessionId, () => controller.abort());
42
 
43
  const iter = streamRewrite(
44
    runtime.workspaceId,
45
    runtime.pageId,
46
    {
47
      action,
48
      selectedText: context.selectedText,
49
      parentBlock: context.parentBlock,
50
      beforeBlock: context.beforeBlock,
51
      afterBlock: context.afterBlock,
52
      pageTitle,
53
    },
54
    controller.signal,
55
  );
56
 
57
  let received = "";
58
  let pending = "";
59
  let rafId: number | null = null;
60
  const flush = () => {
61
    rafId = null;
62
    if (pending.length === 0) return;
63
    if (controller.signal.aborted) {
64
      pending = "";
65
      return;
66
    }
67
    appendAiSuggestion(editor.view, sessionId, pending);
68
    pending = "";
69
  };
70
 
71
  try {
72
    try {
73
      for await (const chunk of iter) {
74
        if (controller.signal.aborted) break;
75
        if (!chunk.text) continue;
76
        received += chunk.text;
77
        pending += chunk.text;
78
        if (rafId === null) rafId = requestAnimationFrame(flush);
79
      }
80
    } finally {
81
      if (rafId !== null) cancelAnimationFrame(rafId);
82
      flush();
83
    }
84
    if (controller.signal.aborted) return;
85
    if (received.trim().length === 0) {
86
      errorAiSuggestion(editor.view, sessionId, "The model returned no content. Try again or switch models.");
87
    } else {
88
      finishAiSuggestion(editor.view, sessionId);
89
    }
90
  } catch (err) {
91
    if (controller.signal.aborted) return;
92
    if (err instanceof AiStreamError) {
93
      errorAiSuggestion(editor.view, sessionId, err.message);
94
    } else {
95
      errorAiSuggestion(editor.view, sessionId, "AI rewrite failed");
96
    }
97
  } finally {
98
    unregisterAiRewriteAbort(sessionId);
99
    if (activeRewriteControllers.get(editor) === controller) {
100
      activeRewriteControllers.delete(editor);
101
    }
102
  }
103
}
104
 
105
export function dismissRewrite(editor: Editor): void {
106
  const controller = activeRewriteControllers.get(editor);
107
  if (controller) {
108
    controller.abort();
109
    activeRewriteControllers.delete(editor);
110
  }
111
  cancelAiSuggestion(editor.view);
112
}