import { Plugin, PluginKey, Selection } from '@milkdown/prose/state' import { $prose, $ctx, $markSchema } from '@milkdown/kit/utils' import type { Ctx } from '@milkdown/kit/core' import type { EditorView } from '@milkdown/prose/view' const COPILOT_PLUGIN_KEY = new PluginKey('milkdown-copilot') const DEBOUNCE_MS = 500 let enabled = true interface CopilotState { from: number to: number suggestion: string } interface CopilotConfig { fetchSuggestion: (prefix: string, suffix: string, signal?: AbortSignal) => Promise debounceMs?: number } const initialState: CopilotState = { from: 0, to: 0, suggestion: '' } export const copilotConfigCtx = $ctx({ fetchSuggestion: async () => '', debounceMs: DEBOUNCE_MS }, 'copilotConfig') export const copilotGhostMark = $markSchema('copilot_ghost', () => ({ excludes: '_', inclusive: true, parseDOM: [{ tag: 'span[data-copilot-ghost]' }], toDOM: () => ['span', { 'data-copilot-ghost': '', class: 'copilot-ghost-text' }, 0] })) let debounceTimer: ReturnType | null = null let abortController: AbortController | null = null let currentCtx: Ctx | null = null function clearGhostText(view: EditorView) { const state = COPILOT_PLUGIN_KEY.getState(view.state) if (state && state.suggestion && state.from < state.to) { const tr = view.state.tr .delete(state.from, state.to) .setMeta(COPILOT_PLUGIN_KEY, { ...initialState }) view.dispatch(tr) } } function insertGhostText(view: EditorView, suggestion: string, from: number) { if (!currentCtx || !suggestion) return const schema = view.state.schema const markType = schema.marks.copilot_ghost if (!markType) { console.error('[Copilot] copilot_ghost mark not found in schema') return } const tr = view.state.tr tr.insertText(suggestion, from) const endPos = from + suggestion.length tr.addMark(from, endPos, markType.create()) tr.setMeta(COPILOT_PLUGIN_KEY, { from, to: endPos, suggestion }) view.dispatch(tr) } function doFetchSuggestion(view: EditorView, pos: number, prefix: string, suffix: string) { if (!currentCtx) return const config = currentCtx.get(copilotConfigCtx.key) if (abortController) { abortController.abort() abortController = null } abortController = new AbortController() config.fetchSuggestion(prefix, suffix, abortController.signal) .then(suggestion => { if (view.state.selection.from !== pos) return if (suggestion) { insertGhostText(view, suggestion, pos) } }) .catch(e => { if (e.name !== 'AbortError') { console.error('[Copilot] Error:', e) } }) .finally(() => { abortController = null }) } function scheduleFetch(view: EditorView, pos: number, prefix: string, suffix: string) { if (!enabled) return if (debounceTimer) { clearTimeout(debounceTimer) debounceTimer = null } debounceTimer = setTimeout(() => { debounceTimer = null doFetchSuggestion(view, pos, prefix, suffix) }, DEBOUNCE_MS) } function acceptSuggestion(view: EditorView) { const state = COPILOT_PLUGIN_KEY.getState(view.state) if (!state?.suggestion || state.from >= state.to) return false const tr = view.state.tr.removeMark(state.from, state.to, view.state.schema.marks.copilot_ghost) const endPos = Math.min(state.to, tr.doc.content.size) tr.setSelection(Selection.near(tr.doc.resolve(endPos))) tr.setMeta(COPILOT_PLUGIN_KEY, { ...initialState }) view.dispatch(tr) return true } function rejectSuggestion(view: EditorView) { const state = COPILOT_PLUGIN_KEY.getState(view.state) if (!state?.suggestion) return false clearGhostText(view) return true } export const copilotPlugin = $prose((ctx) => { currentCtx = ctx return new Plugin({ key: COPILOT_PLUGIN_KEY, state: { init: () => ({ ...initialState }), apply: (tr, value) => { const meta = tr.getMeta(COPILOT_PLUGIN_KEY) if (meta !== undefined) { return meta } if (tr.docChanged && value.suggestion) { return { ...initialState } } return value } }, props: { handleKeyDown: (view, event) => { const state = COPILOT_PLUGIN_KEY.getState(view.state) if (event.key === 'Tab' && state?.suggestion) { event.preventDefault() return acceptSuggestion(view) } if (event.key === 'Escape' && state?.suggestion) { event.preventDefault() return rejectSuggestion(view) } if (state?.suggestion && event.key !== 'Shift' && event.key !== 'Control' && event.key !== 'Alt' && event.key !== 'Meta') { clearGhostText(view) } return false }, handleClick: (view, pos) => { const state = COPILOT_PLUGIN_KEY.getState(view.state) if (!state?.suggestion) return false if (pos >= state.from && pos < state.to) { return acceptSuggestion(view) } clearGhostText(view) return false } }, view: () => ({ update: (view, prevState) => { if (view.state.doc.eq(prevState.doc) && view.state.selection.eq(prevState.selection)) { return } const state = COPILOT_PLUGIN_KEY.getState(view.state) if (state?.suggestion) { return } if (!view.state.doc.eq(prevState.doc)) { const { from, to } = view.state.selection if (from !== to) return const doc = view.state.doc const prefix = doc.textBetween(0, from) const suffix = doc.textBetween(to, doc.content.size) scheduleFetch(view, from, prefix, suffix) } } }) }) }) export { COPILOT_PLUGIN_KEY } export function isCopilotEnabled(): boolean { return enabled } export function setCopilotEnabled(value: boolean): void { enabled = value if (!value) { if (debounceTimer) { clearTimeout(debounceTimer) debounceTimer = null } if (abortController) { abortController.abort() abortController = null } } }