import { createApp, h, reactive } from 'vue' import { parserCtx, serializerCtx } from '@milkdown/kit/core' import { $ctx, $node, $prose, $remark, $view } from '@milkdown/kit/utils' import { Plugin, PluginKey, Selection } from '@milkdown/prose/state' import { type Node as ProseNode, Slice, type Schema } from '@milkdown/prose/model' import { Decoration, DecorationSet, type EditorView, type NodeView } from '@milkdown/prose/view' import ProBlockCrepe from '../components/ProBlockCrepe.vue' import { extractDocBlockContextFromMarkdown } from '../utils/docBlock.js' import { extractTextFromOCR, getOcrCache } from '../utils/ocrCache' import { PRO_BLOCK_NODE_TYPE, PRO_DISPLAY_LABEL, PRO_TRIGGER_TEXT, parseProBlockSyntax, serializeProBlockSyntax } from '../utils/proBlock.js' const PRO_BLOCK_INPUT_PLUGIN_KEY = new PluginKey('milkdown-pro-block-input') const PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY = new PluginKey('milkdown-pro-block-highlight') const PRO_BLOCK_INPUT_META = 'pro-block-input-meta' const PRO_CONTEXT_LIMIT = 32 * 1024 const FALLBACK_BLOCK_SEPARATOR = '\n\n' const FALLBACK_LEAF_TEXT = '\n' const IMAGE_NODE_TYPES = new Set(['image', 'image-block', 'imageBlock']) interface ProBlockConfig { fetchSuggestionStream: (payload: { prefix: string suffix: string languageId: string signal?: AbortSignal model?: string temperature?: number onChunk?: (chunk: string) => void }) => Promise getProModel: () => string showError: (message: string) => void isThinkingActive: () => boolean setThinkingActive: (value: boolean) => void } interface ProRequestPayload { prefix: string suffix: string languageId: string blocked: boolean } function serializeRangeToMarkdown( doc: ProseNode, from: number, to: number, schema: Schema, serializer: (content: ProseNode) => string ): string { if (from >= to) return '' const fallback = doc.textBetween(from, to, FALLBACK_BLOCK_SEPARATOR, FALLBACK_LEAF_TEXT) if (typeof serializer !== 'function') return fallback const slice = doc.slice(from, to) if (slice.content.size <= 0) return '' try { const sliceDoc = schema.topNodeType.createAndFill(undefined, slice.content) return sliceDoc ? serializer(sliceDoc) : fallback } catch { return fallback } } function transformProBlockChildren(node: any) { if (!node || !Array.isArray(node.children)) return node.children = node.children.map((child: any) => { const replacement = parseProParagraphNode(child) if (replacement) return replacement transformProBlockChildren(child) return child }) } function parseProParagraphNode(node: any) { if (!node || node.type !== 'paragraph' || !Array.isArray(node.children)) return null if (node.children.length !== 1) return null const textNode = node.children[0] if (!textNode || textNode.type !== 'text' || typeof textNode.value !== 'string') return null const parsed = parseProBlockSyntax(textNode.value) if (!parsed) return null return { type: PRO_BLOCK_NODE_TYPE, content: parsed.content, autoStart: false, } } function findProBlockReplacements(doc: ProseNode) { const replacements: Array<{ from: number; to: number; content: string; autoStart: boolean }> = [] doc.descendants((node, pos) => { if (node.type.name !== 'paragraph' || node.childCount !== 1) return true const firstChild = node.firstChild if (!firstChild?.isText) return true const parsed = parseProBlockSyntax(node.textContent) if (!parsed) return true replacements.push({ from: pos, to: pos + node.nodeSize, content: parsed.content, autoStart: parsed.autoStart, }) return false }) return replacements } function buildOcrContext(doc: ProseNode) { const lines: string[] = [] doc.descendants((node) => { if (!IMAGE_NODE_TYPES.has(node.type.name)) return true const src = typeof node.attrs?.src === 'string' ? node.attrs.src : '' if (!src) return true const ocrText = getOcrCache(src) if (!ocrText) return true const preview = extractTextFromOCR(ocrText, 120) if (!preview) return true const label = typeof node.attrs?.alt === 'string' && node.attrs.alt.trim() ? node.attrs.alt.trim() : 'image' lines.push(`![${label}](${src}) `) return true }) return lines.join('\n') } function createHighlightDecorations(doc: ProseNode, from: number, to: number) { if (to <= from) return DecorationSet.empty const decorations = [ Decoration.inline(from, to, { class: 'pro-block-accepted-highlight' }), ] doc.nodesBetween(from, to, (node, pos) => { if (!node.isBlock) return true const nodeFrom = pos const nodeTo = pos + node.nodeSize if (nodeFrom >= from && nodeTo <= to) { decorations.push(Decoration.node(nodeFrom, nodeTo, { class: 'pro-block-accepted-block' })) } return true }) return DecorationSet.create(doc, decorations) } function replaceWithParsedMarkdownSlice(tr: any, from: number, to: number, parsedDoc: ProseNode) { if (!parsedDoc || parsedDoc.content.size <= 0) return null const parsedSlice = Slice.maxOpen(parsedDoc.content) if (!parsedSlice || parsedSlice.size <= 0) return null tr.replaceRange(from, to, parsedSlice) const startPos = tr.mapping.map(from, -1) const endPos = tr.mapping.map(to, 1) if (endPos <= startPos) return null return { from: startPos, to: endPos } } function insertProBlockNode(view: EditorView, autoStart: boolean) { const proBlockType = view.state.schema.nodes[PRO_BLOCK_NODE_TYPE] if (!proBlockType) return false const { from, to } = view.state.selection const blockNode = proBlockType.create({ content: '', autoStart }) const tr = view.state.tr.replaceRangeWith(from, to, blockNode) const nextPos = Math.min(from + blockNode.nodeSize, tr.doc.content.size) tr.setSelection(Selection.near(tr.doc.resolve(nextPos), 1)) view.dispatch(tr.scrollIntoView()) view.focus() return true } function replaceParagraphWithProBlock( view: EditorView, from: number, to: number, content: string, autoStart: boolean ) { const proBlockType = view.state.schema.nodes[PRO_BLOCK_NODE_TYPE] if (!proBlockType) return false const blockNode = proBlockType.create({ content, autoStart: autoStart && content === '', }) const tr = view.state.tr.replaceWith(from, to, blockNode) const nextPos = Math.min(from + blockNode.nodeSize, tr.doc.content.size) tr.setMeta(PRO_BLOCK_INPUT_META, true) tr.setSelection(Selection.near(tr.doc.resolve(nextPos), 1)) view.dispatch(tr.scrollIntoView()) view.focus() return true } function tryHandleProTriggerTextInput(view: EditorView, insertedText: string) { if (!insertedText || !insertedText.includes(']')) return false const { state } = view const { $from, from, to } = state.selection const paragraph = $from.parent if (!paragraph || paragraph.type.name !== 'paragraph') return false if (!$from.sameParent(state.selection.$to)) return false const paragraphDepth = $from.depth const paragraphStart = $from.start(paragraphDepth) const startOffset = from - paragraphStart const endOffset = to - paragraphStart const nextText = `${paragraph.textContent.slice(0, startOffset)}${insertedText}${paragraph.textContent.slice(endOffset)}` const parsed = parseProBlockSyntax(nextText) if (!parsed?.autoStart || parsed.content) return false const blockFrom = $from.before(paragraphDepth) const blockTo = blockFrom + paragraph.nodeSize return replaceParagraphWithProBlock(view, blockFrom, blockTo, parsed.content, parsed.autoStart) } function isAbortError(error: unknown) { return Boolean(error && typeof error === 'object' && 'name' in error && (error as { name?: string }).name === 'AbortError') } export const proBlockConfigCtx = $ctx({ fetchSuggestionStream: async () => '', getProModel: () => '', showError: () => {}, isThinkingActive: () => false, setThinkingActive: () => {}, }, 'proBlockConfig') class ProBlockNodeView implements NodeView { node: ProseNode view: EditorView getPos: (() => number) | boolean dom: HTMLElement app: ReturnType | null = null props: Record parser: (markdown: string) => Promise serializer: (content: ProseNode) => string config: ProBlockConfig abortController: AbortController | null = null requestSeq = 0 redoCount = 0 destroyed = false highlightTimer: ReturnType | null = null constructor( node: ProseNode, view: EditorView, getPos: (() => number) | boolean, parser: (markdown: string) => Promise, serializer: (content: ProseNode) => string, config: ProBlockConfig ) { this.node = node this.view = view this.getPos = getPos this.parser = parser this.serializer = serializer this.config = config this.dom = document.createElement('div') this.dom.className = 'pro-block-node-view' this.props = reactive({ stage: node.attrs.content ? 'done' : 'idle', previewContent: '', savedContent: node.attrs.content || '', isBusy: false, title: PRO_DISPLAY_LABEL, activateAction: () => { void this.startThinking(false) }, discardAction: () => { this.discardResult() }, redoAction: () => { void this.startThinking(true) }, acceptAction: () => { void this.acceptResult() }, }) this.mount() queueMicrotask(() => { if (!this.destroyed) { this.consumeAutoStart() } }) } mount() { this.app = createApp({ render: () => h(ProBlockCrepe, { ...this.props }), }) this.app.mount(this.dom) } getPosValue() { if (typeof this.getPos === 'function') { try { const pos = this.getPos() if (typeof pos === 'number') return pos } catch { // Ignore and fall back to DOM lookup. } } try { const pos = this.view.posAtDOM(this.dom, 0) return typeof pos === 'number' ? pos : undefined } catch { return undefined } } updateAttrs(patch: Record) { const pos = this.getPosValue() if (pos === undefined) return const nextAttrs = { ...this.node.attrs, ...patch } this.view.dispatch(this.view.state.tr.setNodeMarkup(pos, undefined, nextAttrs)) } setStage(stage: 'idle' | 'thinking' | 'streaming' | 'done', previewContent = '') { this.props.stage = stage this.props.previewContent = previewContent this.props.isBusy = stage === 'thinking' || stage === 'streaming' } buildRequestPayload(): ProRequestPayload { const pos = this.getPosValue() if (pos === undefined) { return { prefix: '', suffix: '', languageId: 'markdown', blocked: true, } } const doc = this.view.state.doc const schema = this.view.state.schema const prefixMarkdown = serializeRangeToMarkdown(doc, 0, pos, schema, this.serializer) || doc.textBetween(0, pos, FALLBACK_BLOCK_SEPARATOR, FALLBACK_LEAF_TEXT) const suffixStart = Math.min(pos + this.node.nodeSize, doc.content.size) const suffixMarkdown = serializeRangeToMarkdown(doc, suffixStart, doc.content.size, schema, this.serializer) || doc.textBetween(suffixStart, doc.content.size, FALLBACK_BLOCK_SEPARATOR, FALLBACK_LEAF_TEXT) const ocrContext = buildOcrContext(doc) const docContext = extractDocBlockContextFromMarkdown(`${prefixMarkdown}\n\n${suffixMarkdown}`, 1600) const fullPrefix = [ocrContext, docContext, prefixMarkdown].filter(Boolean).join('\n\n') return { prefix: fullPrefix, suffix: suffixMarkdown, languageId: 'markdown', blocked: fullPrefix.length + suffixMarkdown.length > PRO_CONTEXT_LIMIT, } } consumeAutoStart() { if (!this.node.attrs.autoStart) return this.updateAttrs({ autoStart: false }) void this.startThinking(false) } stopThinkingLock() { this.config.setThinkingActive(false) } abortActive(reason = 'abort') { if (!this.abortController) return this.abortController.abort(reason) this.abortController = null this.stopThinkingLock() } discardResult() { this.abortActive('discard') this.redoCount = 0 this.props.savedContent = '' this.setStage('idle', '') this.updateAttrs({ content: '', autoStart: false }) } async startThinking(isRedo: boolean) { if (this.props.isBusy) return const previousContent = this.node.attrs.content || '' let requestSeq = this.requestSeq let requestStarted = false try { if (this.config.isThinkingActive()) { this.config.showError('当前已有一个 PRO 思考正在进行,请先等待它完成。') return } const payload = this.buildRequestPayload() if (payload.blocked) { this.config.showError('PRO 模式上下文超过 32KB,无法继续思考。') return } requestSeq = this.requestSeq + 1 this.requestSeq = requestSeq this.abortController = new AbortController() this.config.setThinkingActive(true) requestStarted = true this.setStage('thinking', '') const temperature = Math.min(1.2, 0.7 + (isRedo ? (this.redoCount + 1) * 0.2 : 0)) const result = await this.config.fetchSuggestionStream({ ...payload, signal: this.abortController.signal, model: this.config.getProModel(), temperature, onChunk: (chunk) => { if (this.destroyed || this.requestSeq !== requestSeq) return const nextContent = `${this.props.previewContent || ''}${chunk}` this.setStage('streaming', nextContent) }, }) if (this.destroyed || this.requestSeq !== requestSeq) return const content = String(result || this.props.previewContent || '').trim() this.redoCount = isRedo ? this.redoCount + 1 : 0 this.props.savedContent = content this.setStage(content ? 'done' : 'idle', '') this.updateAttrs({ content, autoStart: false }) } catch (error) { if (!isAbortError(error)) { this.config.showError(error instanceof Error ? error.message : 'PRO 模式思考失败') } this.props.savedContent = previousContent this.setStage(previousContent ? 'done' : 'idle', '') } finally { if (requestStarted && this.requestSeq === requestSeq) { this.abortController = null this.stopThinkingLock() } } } async acceptResult() { if (this.props.isBusy) return const source = String(this.props.savedContent || this.node.attrs.content || '').trim() if (!source) return const pos = this.getPosValue() if (pos === undefined) return const from = pos const to = pos + this.node.nodeSize const tr = this.view.state.tr let insertedRange: { from: number; to: number } | null = null try { const parsedDoc = await this.parser(source) insertedRange = replaceWithParsedMarkdownSlice(tr, from, to, parsedDoc) } catch { insertedRange = null } if (!insertedRange) { tr.insertText(source, from, to) insertedRange = { from, to: from + source.length } } const endPos = Math.min(insertedRange.to, tr.doc.content.size) tr.setSelection(Selection.near(tr.doc.resolve(endPos), 1)) tr.setMeta(PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY, { set: insertedRange }) this.view.dispatch(tr.scrollIntoView()) this.view.focus() if (this.highlightTimer) { clearTimeout(this.highlightTimer) } this.highlightTimer = setTimeout(() => { if (this.destroyed) return this.view.dispatch(this.view.state.tr.setMeta(PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY, { clear: true })) }, 1000) } update(node: ProseNode) { if (node.type !== this.node.type) return false this.node = node if (!this.props.isBusy) { this.props.savedContent = node.attrs.content || '' this.props.previewContent = '' this.props.stage = node.attrs.content ? 'done' : 'idle' } return true } stopEvent(event: Event) { const target = event.target as Node | null return Boolean(target && this.dom.contains(target)) } ignoreMutation() { return true } destroy() { this.destroyed = true this.abortActive('destroy') if (this.highlightTimer) { clearTimeout(this.highlightTimer) this.highlightTimer = null } this.app?.unmount() this.app = null } } export const proBlockRemark = $remark('proBlockRemark', () => () => { return (tree: any) => { transformProBlockChildren(tree) } }) export const proBlockNode = $node(PRO_BLOCK_NODE_TYPE, () => ({ group: 'block', atom: true, isolating: true, selectable: true, draggable: false, marks: '', attrs: { content: { default: '' }, autoStart: { default: false }, }, parseDOM: [ { tag: 'div[data-pro-block="true"]', getAttrs: (dom) => ({ content: (dom as HTMLElement).getAttribute('data-pro-content') || '', autoStart: ((dom as HTMLElement).getAttribute('data-pro-autostart') || '') === 'true', }), }, ], toDOM: (node) => [ 'div', { 'data-pro-block': 'true', 'data-pro-content': node.attrs.content || '', 'data-pro-autostart': String(Boolean(node.attrs.autoStart)), }, ], parseMarkdown: { match: (node) => node.type === PRO_BLOCK_NODE_TYPE, runner: (state, node, type) => { state.addNode(type, { content: String(node.content || ''), autoStart: false, }) }, }, toMarkdown: { match: (node) => node.type.name === PRO_BLOCK_NODE_TYPE, runner: (state, node) => { state.addNode('paragraph', [{ type: 'text', value: serializeProBlockSyntax(node.attrs.content), }]) }, }, leafText: (node) => serializeProBlockSyntax(node.attrs.content), })) export const proBlockView = $view(proBlockNode, (ctx) => { const parser = ctx.get(parserCtx) const serializer = ctx.get(serializerCtx) const config = ctx.get(proBlockConfigCtx.key) return (node, view, getPos) => new ProBlockNodeView(node, view, getPos, parser, serializer, config) }) export const proBlockInputPlugin = $prose(() => { return new Plugin({ key: PRO_BLOCK_INPUT_PLUGIN_KEY, props: { handleTextInput: (view, _from, _to, text) => { return tryHandleProTriggerTextInput(view, text) }, handleKeyDown: (view, event) => { const pressed = (event.ctrlKey || event.metaKey) && event.shiftKey && event.key.toLowerCase() === 'p' if (!pressed) return false event.preventDefault() return insertProBlockNode(view, false) }, }, appendTransaction: (transactions, _oldState, newState) => { if (!transactions.some((transaction) => transaction.docChanged)) return null if (transactions.some((transaction) => transaction.getMeta(PRO_BLOCK_INPUT_META))) return null const proBlockType = newState.schema.nodes[PRO_BLOCK_NODE_TYPE] if (!proBlockType) return null const replacements = findProBlockReplacements(newState.doc) if (replacements.length === 0) return null let tr = newState.tr for (let index = replacements.length - 1; index >= 0; index -= 1) { const replacement = replacements[index] tr = tr.replaceWith( replacement.from, replacement.to, proBlockType.create({ content: replacement.content, autoStart: replacement.autoStart && replacement.content === '', }) ) } if (!tr.docChanged) return null tr.setMeta(PRO_BLOCK_INPUT_META, true) return tr }, }) }) export const proBlockHighlightPlugin = $prose(() => { return new Plugin({ key: PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY, state: { init: () => DecorationSet.empty, apply: (tr, value) => { const meta = tr.getMeta(PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY) if (meta?.clear) { return DecorationSet.empty } if (meta?.set) { return createHighlightDecorations(tr.doc, meta.set.from, meta.set.to) } return value.map(tr.mapping, tr.doc) }, }, props: { decorations: (state) => PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY.getState(state), }, }) }) export function insertProBlockAtSelection(view: EditorView, autoStart = false) { return insertProBlockNode(view, autoStart) } export { PRO_TRIGGER_TEXT }