From 451beb5a457960a76e24059514c3592b0700a51d Mon Sep 17 00:00:00 2001 From: Mirone Date: Thu, 16 Apr 2026 01:10:48 +0900 Subject: [PATCH] feat: replace-selection mode for streaming plugin (#2335) * feat: replace-selection mode for streaming plugin * fix: apply empty-paragraph snap for collapsed selection in selection mode --- docs/api/crepe.md | 5 + docs/api/plugin-streaming.md | 16 +++ e2e/shim.d.ts | 2 +- e2e/src/crepe-streaming/main.ts | 2 +- e2e/tests/crepe/streaming.spec.ts | 126 +++++++++++++++++- packages/crepe/src/feature/ai/commands.ts | 12 +- .../src/__test__/streaming-plugin.spec.ts | 66 +++++++++ .../src/streaming-commands.ts | 71 ++++++---- .../plugins/plugin-streaming/src/types.ts | 11 +- storybook/stories/crepe/setup.ts | 12 +- 10 files changed, 283 insertions(+), 40 deletions(-) diff --git a/docs/api/crepe.md b/docs/api/crepe.md index 9004076bf0e..4cc778b01d3 100644 --- a/docs/api/crepe.md +++ b/docs/api/crepe.md @@ -606,6 +606,11 @@ workflow. Users supply a `provider` (an async generator that yields markdown tokens) and Crepe handles the rest: start streaming, push chunks, end streaming, and optionally hand off to diff review. +When the user has a text selection, `runAICmd` replaces the selected text +with the AI output. The provider receives the selected text in +`AIPromptContext.selection` for context-aware generation. When the +selection is empty, content is inserted at the cursor position. + ```typescript import { Crepe } from '@milkdown/crepe' import type { AIFeatureConfig } from '@milkdown/crepe/feature/ai' diff --git a/docs/api/plugin-streaming.md b/docs/api/plugin-streaming.md index f18d02210be..f93ddd90582 100644 --- a/docs/api/plugin-streaming.md +++ b/docs/api/plugin-streaming.md @@ -88,6 +88,22 @@ The insert strategy depends on where the cursor is when streaming starts: | Table cell | All content inserted as plain text, newlines collapsed to spaces | | Between blocks (depth 0) | Full markdown parse, inserted as block nodes | +## Replace Selection + +You can replace the current text selection with streamed content: + +```typescript +editor.action((ctx) => { + ctx.get(commandsCtx).call(startStreamingCmd.key, { insertAt: 'selection' }) +}) +``` + +When the selection is non-empty, the selected range is replaced by the streamed content as it arrives. When the selection is collapsed (empty), this behaves identically to `insertAt: 'cursor'`. + +The insert strategy is resolved based on the position at `selection.from`. For example, if the selection starts inside a paragraph, the `split-block` strategy is used; if it starts inside a code block, plain-text insertion is used. + +After streaming ends, aborting with `keep: false` restores the original document including the selected text. Diff review mode also works correctly — the diff shows the original selection being replaced. + ## Diff Review After Streaming When the diff plugin is also loaded (e.g. via `Crepe.Feature.AI` in Crepe, or by manually calling `editor.use(diff)` on a standalone editor), you can hand off to diff review mode after streaming ends: diff --git a/e2e/shim.d.ts b/e2e/shim.d.ts index 2e66572cbac..1ad47fe9f82 100644 --- a/e2e/shim.d.ts +++ b/e2e/shim.d.ts @@ -47,7 +47,7 @@ declare global { var __rejectChunk__: (index: number) => boolean var __startStreaming__: (options?: { - insertAt?: 'cursor' | number + insertAt?: 'cursor' | 'selection' | number }) => boolean var __pushChunk__: (token: string) => boolean var __endStreaming__: (options?: { diffReview?: boolean }) => boolean diff --git a/e2e/src/crepe-streaming/main.ts b/e2e/src/crepe-streaming/main.ts index 9d472627877..1e69ad4776e 100644 --- a/e2e/src/crepe-streaming/main.ts +++ b/e2e/src/crepe-streaming/main.ts @@ -18,7 +18,7 @@ setup(async () => { await crepe.create() globalThis.__startStreaming__ = (options?: { - insertAt?: 'cursor' | number + insertAt?: 'cursor' | 'selection' | number }) => crepe.editor.action(callCommand('StartStreaming', options)) globalThis.__pushChunk__ = (token: string) => crepe.editor.action(callCommand('PushChunk', token)) diff --git a/e2e/tests/crepe/streaming.spec.ts b/e2e/tests/crepe/streaming.spec.ts index c65919fa0fa..6793091ec73 100644 --- a/e2e/tests/crepe/streaming.spec.ts +++ b/e2e/tests/crepe/streaming.spec.ts @@ -11,7 +11,7 @@ test.beforeEach(async ({ page }) => { async function simulateStream( page: Page, tokens: string[], - options?: { delayMs?: number; insertAt?: 'cursor' | number } + options?: { delayMs?: number; insertAt?: 'cursor' | 'selection' | number } ) { const delayMs = options?.delayMs ?? 30 const insertAt = options?.insertAt @@ -552,3 +552,127 @@ test.describe('insert-at-cursor', () => { expect(markdown).toContain('After quote.') }) }) + +test.describe('replace-selection', () => { + test('replaces selected text within a paragraph', async ({ page }) => { + await setMarkdown(page, 'Hello world, this is a test.') + await waitNextFrame(page) + + // Select "world" by keyboard: Home → move right 6 → shift-select 5 + const editor = page.locator('.editor') + await editor.locator('p').first().click() + await page.keyboard.press('Home') + for (let i = 0; i < 6; i++) await page.keyboard.press('ArrowRight') + await page.keyboard.down('Shift') + for (let i = 0; i < 5; i++) await page.keyboard.press('ArrowRight') + await page.keyboard.up('Shift') + await waitNextFrame(page) + + await simulateStream(page, ['universe'], { insertAt: 'selection' }) + await waitForFlush(page, 'universe') + + await page.evaluate(() => window.__endStreaming__()) + await waitNextFrame(page) + + const markdown = await getMarkdown(page) + expect(markdown).toContain('Hello universe') + expect(markdown).not.toContain('world') + }) + + test('replaces selection with multi-block content', async ({ page }) => { + await setMarkdown(page, 'Start.\n\nMiddle paragraph.\n\nEnd.') + await waitNextFrame(page) + + // Select entire "Middle paragraph." text + const editor = page.locator('.editor') + await editor.locator('p').nth(1).click() + await page.keyboard.press('Home') + await page.keyboard.down('Shift') + await page.keyboard.press('End') + await page.keyboard.up('Shift') + await waitNextFrame(page) + + await simulateStream(page, ['# New Heading\n\nReplaced content.'], { + insertAt: 'selection', + }) + await waitForFlush(page, 'Replaced content.') + + await page.evaluate(() => window.__endStreaming__()) + await waitNextFrame(page) + + const markdown = await getMarkdown(page) + expect(markdown).toContain('Start.') + expect(markdown).toContain('New Heading') + expect(markdown).toContain('Replaced content.') + expect(markdown).toContain('End.') + expect(markdown).not.toContain('Middle paragraph.') + }) + + test('abort replace-selection restores original', async ({ page }) => { + await setMarkdown(page, 'Keep this text intact.') + await waitNextFrame(page) + + // Select "this text" + const editor = page.locator('.editor') + await editor.locator('p').first().click() + await page.keyboard.press('Home') + for (let i = 0; i < 5; i++) await page.keyboard.press('ArrowRight') + await page.keyboard.down('Shift') + for (let i = 0; i < 9; i++) await page.keyboard.press('ArrowRight') + await page.keyboard.up('Shift') + await waitNextFrame(page) + + await simulateStream(page, ['REPLACED'], { insertAt: 'selection' }) + await waitForFlush(page, 'REPLACED') + + await page.evaluate(() => window.__abortStreaming__({ keep: false })) + await waitNextFrame(page) + + const markdown = await getMarkdown(page) + expect(markdown).toContain('Keep this text intact.') + expect(markdown).not.toContain('REPLACED') + }) + + test('empty selection with insertAt selection behaves like cursor', async ({ + page, + }) => { + await setMarkdown(page, 'Before. After.') + await waitNextFrame(page) + + // Place cursor without selecting (collapsed selection) + const editor = page.locator('.editor') + await editor.locator('p').first().click() + await page.keyboard.press('Home') + for (let i = 0; i < 7; i++) await page.keyboard.press('ArrowRight') + await waitNextFrame(page) + + await simulateStream(page, [' Inserted.'], { insertAt: 'selection' }) + await waitForFlush(page, 'Inserted.') + + await page.evaluate(() => window.__endStreaming__()) + await waitNextFrame(page) + + const markdown = await getMarkdown(page) + expect(markdown).toContain('Before. Inserted.') + expect(markdown).toContain('After.') + }) + + test('collapsed selection in empty document snaps like cursor mode', async ({ + page, + }) => { + // Empty editor has a single empty paragraph — block content should + // replace it cleanly, same as insertAt: 'cursor'. + const editor = page.locator('.editor') + await editor.locator('p').first().click() + await waitNextFrame(page) + + await simulateStream(page, ['# Heading'], { insertAt: 'selection' }) + await waitForFlush(page, 'Heading') + + await page.evaluate(() => window.__endStreaming__()) + await waitNextFrame(page) + + const markdown = await getMarkdown(page) + expect(markdown.trim()).toBe('# Heading') + }) +}) diff --git a/packages/crepe/src/feature/ai/commands.ts b/packages/crepe/src/feature/ai/commands.ts index 9a35e3a902c..8b4215056ee 100644 --- a/packages/crepe/src/feature/ai/commands.ts +++ b/packages/crepe/src/feature/ai/commands.ts @@ -104,11 +104,9 @@ async function runProvider( /// provider asynchronously. The command returns synchronously; the /// provider runs in the background. /// -/// Note on selection: when the user has a text selection, the provider -/// receives the selected text in `AIPromptContext.selection`, but the -/// streaming output is *inserted at* the selection head — it does not -/// *replace* the selected text. A replace-selection mode requires -/// streaming plugin changes and is planned as a follow-up. +/// When the user has a text selection, the streamed output replaces the +/// selected text. The provider also receives the selected text in +/// `AIPromptContext.selection` for context-aware generation. export const runAICmd = $command('RunAI', (ctx) => { return (options?: RunAIOptions) => (state, dispatch) => { if (!options?.instruction) return false @@ -129,11 +127,11 @@ export const runAICmd = $command('RunAI', (ctx) => { // are side-effect-free so we can return true here. if (!dispatch) return true - // Start streaming at the cursor position. + // Start streaming — replaces the selection if non-empty. const commands = ctx.get(commandsCtx) const insertAt = state.selection.empty ? ('cursor' as const) - : state.selection.head + : ('selection' as const) if (!commands.call(startStreamingCmd.key, { insertAt })) return false // Everything after startStreamingCmd is wrapped in try/catch: if diff --git a/packages/plugins/plugin-streaming/src/__test__/streaming-plugin.spec.ts b/packages/plugins/plugin-streaming/src/__test__/streaming-plugin.spec.ts index 48ac68c557d..45cc8a3fa8b 100644 --- a/packages/plugins/plugin-streaming/src/__test__/streaming-plugin.spec.ts +++ b/packages/plugins/plugin-streaming/src/__test__/streaming-plugin.spec.ts @@ -194,6 +194,72 @@ describe('insert-at-cursor state transitions', () => { }) }) +describe('replace-selection state transitions', () => { + it('sets different insertPos and insertEndPos on start', () => { + const originalDoc = doc(p(text('hello world'))) + const state = applyStreamingAction(null, { + type: 'start', + originalDoc, + insertPos: 2, + insertEndPos: 8, + lastApplyTime: Date.now(), + }) + expect(state).not.toBeNull() + expect(state!.insertPos).toBe(2) + expect(state!.insertEndPos).toBe(8) + expect(state!.active).toBe(true) + }) + + it('updates insertEndPos on apply while preserving insertPos', () => { + let state: StreamingState | null = applyStreamingAction(null, { + type: 'start', + originalDoc: doc(p(text('hello world'))), + insertPos: 2, + insertEndPos: 8, + lastApplyTime: Date.now(), + }) + state = applyStreamingAction(state, { + type: 'push', + token: 'replacement', + }) + expect(state!.buffer).toBe('replacement') + expect(state!.insertPos).toBe(2) + expect(state!.insertEndPos).toBe(8) + + state = applyStreamingAction(state, { + type: 'apply', + lastApplyTime: Date.now(), + insertEndPos: 15, + }) + expect(state!.insertPos).toBe(2) + expect(state!.insertEndPos).toBe(15) + }) + + it('returns null on end', () => { + let state: StreamingState | null = applyStreamingAction(null, { + type: 'start', + originalDoc: doc(p(text('hello world'))), + insertPos: 2, + insertEndPos: 8, + lastApplyTime: Date.now(), + }) + state = applyStreamingAction(state, { type: 'end' }) + expect(state).toBeNull() + }) + + it('returns null on abort', () => { + let state: StreamingState | null = applyStreamingAction(null, { + type: 'start', + originalDoc: doc(p(text('hello world'))), + insertPos: 2, + insertEndPos: 8, + lastApplyTime: Date.now(), + }) + state = applyStreamingAction(state, { type: 'abort' }) + expect(state).toBeNull() + }) +}) + describe('buffer accumulation', () => { it('builds markdown incrementally', () => { let state: StreamingState | null = createStreamingState() diff --git a/packages/plugins/plugin-streaming/src/streaming-commands.ts b/packages/plugins/plugin-streaming/src/streaming-commands.ts index 3bfa4d7f905..3b83c1f23eb 100644 --- a/packages/plugins/plugin-streaming/src/streaming-commands.ts +++ b/packages/plugins/plugin-streaming/src/streaming-commands.ts @@ -27,31 +27,52 @@ export const startStreamingCmd = $command('StartStreaming', () => { let insertPos: number | undefined let insertEndPos: number | undefined if (options?.insertAt != null) { - const rawPos = - options.insertAt === 'cursor' - ? state.selection.head - : options.insertAt - if (!Number.isFinite(rawPos)) return false - insertPos = Math.max( - 0, - Math.min(Math.round(rawPos), state.doc.content.size) - ) - - // If cursor is inside a top-level empty textblock (e.g. the default - // empty paragraph in a new editor), snap the range to cover the whole - // block so that block-level content replaces it cleanly. - // Only depth === 1 — nested empty textblocks (inside list items, - // blockquotes, etc.) should not be snapped; the normal strategy - // handles them correctly at their original position. - const resolved = state.doc.resolve(insertPos) - if ( - resolved.parent.isTextblock && - !resolved.parent.type.spec.code && - resolved.parent.content.size === 0 && - resolved.depth === 1 - ) { - insertPos = resolved.before(resolved.depth) - insertEndPos = resolved.after(resolved.depth) + if (options.insertAt === 'selection') { + // Replace-selection mode: use the full selection range. + insertPos = state.selection.from + insertEndPos = state.selection.to + + // When the selection is collapsed, apply the same empty-paragraph + // snap as cursor mode so the behavior is truly identical. + if (state.selection.empty) { + const resolved = state.doc.resolve(insertPos) + if ( + resolved.parent.isTextblock && + !resolved.parent.type.spec.code && + resolved.parent.content.size === 0 && + resolved.depth === 1 + ) { + insertPos = resolved.before(resolved.depth) + insertEndPos = resolved.after(resolved.depth) + } + } + } else { + const rawPos = + options.insertAt === 'cursor' + ? state.selection.head + : options.insertAt + if (!Number.isFinite(rawPos)) return false + insertPos = Math.max( + 0, + Math.min(Math.round(rawPos), state.doc.content.size) + ) + + // If cursor is inside a top-level empty textblock (e.g. the default + // empty paragraph in a new editor), snap the range to cover the whole + // block so that block-level content replaces it cleanly. + // Only depth === 1 — nested empty textblocks (inside list items, + // blockquotes, etc.) should not be snapped; the normal strategy + // handles them correctly at their original position. + const resolved = state.doc.resolve(insertPos) + if ( + resolved.parent.isTextblock && + !resolved.parent.type.spec.code && + resolved.parent.content.size === 0 && + resolved.depth === 1 + ) { + insertPos = resolved.before(resolved.depth) + insertEndPos = resolved.after(resolved.depth) + } } } diff --git a/packages/plugins/plugin-streaming/src/types.ts b/packages/plugins/plugin-streaming/src/types.ts index 290cdf316a4..2dc49933afb 100644 --- a/packages/plugins/plugin-streaming/src/types.ts +++ b/packages/plugins/plugin-streaming/src/types.ts @@ -52,9 +52,14 @@ export interface StreamingConfig { /// Options for starting a streaming session. export interface StartStreamingOptions { - /// Insert at cursor position or a specific position instead of replacing - /// the whole document. 'cursor' resolves to current selection head. - insertAt?: 'cursor' | number + /// Insert at cursor position, replace the current selection, or insert + /// at a specific position instead of replacing the whole document. + /// - `'cursor'`: resolves to current `selection.head` (insert point). + /// - `'selection'`: resolves to `selection.from`/`selection.to`, + /// replacing the selected range. When the selection is collapsed + /// this behaves identically to `'cursor'`. + /// - `number`: absolute position in the document. + insertAt?: 'cursor' | 'selection' | number } /// Options for ending a streaming session. diff --git a/storybook/stories/crepe/setup.ts b/storybook/stories/crepe/setup.ts index 5cf6298a849..ce96e4e4569 100644 --- a/storybook/stories/crepe/setup.ts +++ b/storybook/stories/crepe/setup.ts @@ -7,7 +7,7 @@ import type { import { Crepe } from '@milkdown/crepe' import { abortAICmd, runAICmd } from '@milkdown/crepe/feature/ai' import all from '@milkdown/crepe/theme/common/style.css?inline' -import { commandsCtx } from '@milkdown/kit/core' +import { commandsCtx, editorViewCtx } from '@milkdown/kit/core' import { acceptAllDiffsCmd, clearDiffReviewCmd, @@ -407,7 +407,15 @@ export function setupStreamingDemo(config: setupConfig) { } startBtn.addEventListener('click', () => { - crepe.editor.action(callCommand(startStreamingCmd.key)) + // When the user has selected text, use replace-selection mode; + // otherwise fall back to replace-whole-doc mode. + crepe.editor.action((ctx) => { + const view = ctx.get(editorViewCtx) + const options = view.state.selection.empty + ? undefined + : { insertAt: 'selection' as const } + ctx.get(commandsCtx).call(startStreamingCmd.key, options) + }) setStreaming(true) const chars = Array.from(textarea.value)