diff --git a/frontend/src/i18n/locales/en/common.json b/frontend/src/i18n/locales/en/common.json index 8010fc9..e3cd2f1 100644 --- a/frontend/src/i18n/locales/en/common.json +++ b/frontend/src/i18n/locales/en/common.json @@ -167,7 +167,12 @@ "insertWithCitationsEmpty": "Insert with citations", "insertTextOnly": "Insert text only", "generateError": "Could not generate text.", - "threadTitle": "Drafting assistant" + "threadTitle": "Drafting assistant", + "stop": "Stop", + "stopAria": "Stop generation", + "targetStale": "The original selection is no longer in the document. Reselect the text.", + "pendingTargetHint": "Will insert at the highlighted selection", + "submitHint": "Ctrl+Enter to generate" }, "aiDraftDisclaimer": "AI-generated text — review before using.", "aiDraftAttribution": "— AI-generated. Review before using." diff --git a/frontend/src/i18n/locales/es/common.json b/frontend/src/i18n/locales/es/common.json index ad39f8e..143c7ed 100644 --- a/frontend/src/i18n/locales/es/common.json +++ b/frontend/src/i18n/locales/es/common.json @@ -167,7 +167,12 @@ "insertWithCitationsEmpty": "Insertar con citas", "insertTextOnly": "Insertar solo texto", "generateError": "No se pudo generar el texto.", - "threadTitle": "Asistente de redacción" + "threadTitle": "Asistente de redacción", + "stop": "Detener", + "stopAria": "Detener generación", + "targetStale": "La selección original ya no está en el documento. Vuelve a seleccionar el texto.", + "pendingTargetHint": "Insertará en la selección marcada", + "submitHint": "Ctrl+Enter para generar" }, "aiDraftDisclaimer": "Texto generado por IA — revísalo antes de usarlo.", "aiDraftAttribution": "— Generado por IA. Revisa antes de usar." diff --git a/frontend/src/lib/ai-draft-store.test.ts b/frontend/src/lib/ai-draft-store.test.ts new file mode 100644 index 0000000..c1c0242 --- /dev/null +++ b/frontend/src/lib/ai-draft-store.test.ts @@ -0,0 +1,84 @@ +import { describe, expect, it, vi } from 'vitest'; +import { useAiDraftStore } from './ai-draft-store'; +import type { ChatMessage } from '@/lib/types'; + +function resetStore(): void { + useAiDraftStore.setState({ docs: {} }); +} + +describe('useAiDraftStore', () => { + it('isolates draft state per document', () => { + resetStore(); + const controllerA = useAiDraftStore.getState().startGenerate('doc-a', null); + useAiDraftStore.getState().applyChunk('doc-a', { type: 'text', delta: 'parcial A' }, controllerA); + + const controllerB = useAiDraftStore.getState().startGenerate('doc-b', null); + useAiDraftStore.getState().applyChunk('doc-b', { type: 'text', delta: 'parcial B' }, controllerB); + + const stateA = useAiDraftStore.getState().getDocState('doc-a'); + const stateB = useAiDraftStore.getState().getDocState('doc-b'); + expect(stateA.stream?.role).toBe('assistant'); + expect(stateB.stream?.role).toBe('assistant'); + if (stateA.stream?.role === 'assistant') { + expect(stateA.stream.content).toEqual(['parcial A']); + } + if (stateB.stream?.role === 'assistant') { + expect(stateB.stream.content).toEqual(['parcial B']); + } + }); + + it('stopGenerate aborts the controller and keeps the partial stream', () => { + resetStore(); + const controller = useAiDraftStore.getState().startGenerate('doc-1', { + from: 1, + to: 4, + text: 'abc', + }); + useAiDraftStore.getState().applyChunk('doc-1', { type: 'text', delta: 'parcial' }, controller); + const aborted = vi.fn(); + controller.signal.addEventListener('abort', aborted); + + useAiDraftStore.getState().stopGenerate('doc-1'); + + expect(aborted).toHaveBeenCalled(); + const state = useAiDraftStore.getState().getDocState('doc-1'); + expect(state.busy).toBe(false); + expect(state.abortController).toBeNull(); + expect(state.insertTarget).toEqual({ from: 1, to: 4, text: 'abc' }); + expect(state.stream?.role).toBe('assistant'); + if (state.stream?.role === 'assistant') { + expect(state.stream.content).toEqual(['parcial']); + } + }); + + it('keeps threadId across clearDraft so follow-ups reuse the same thread', () => { + resetStore(); + useAiDraftStore.getState().setThreadId('doc-1', 'thread-x'); + useAiDraftStore.getState().setPrompt('doc-1', 'Redacta'); + const controller = useAiDraftStore.getState().startGenerate('doc-1', null); + useAiDraftStore.getState().applyChunk('doc-1', { type: 'text', delta: 'hola' }, controller); + useAiDraftStore.getState().finishGenerate('doc-1', controller); + + useAiDraftStore.getState().clearDraft('doc-1'); + + const state = useAiDraftStore.getState().getDocState('doc-1'); + expect(state.threadId).toBe('thread-x'); + expect(state.stream).toBeNull(); + expect(state.prompt).toBe(''); + expect(state.busy).toBe(false); + }); + + it('ignores chunks from a superseded generate', () => { + resetStore(); + const first = useAiDraftStore.getState().startGenerate('doc-1', null); + const second = useAiDraftStore.getState().startGenerate('doc-1', null); + useAiDraftStore.getState().applyChunk('doc-1', { type: 'text', delta: 'viejo' }, first); + useAiDraftStore.getState().applyChunk('doc-1', { type: 'text', delta: 'nuevo' }, second); + + const stream = useAiDraftStore.getState().getDocState('doc-1').stream as ChatMessage | null; + expect(stream?.role).toBe('assistant'); + if (stream?.role === 'assistant') { + expect(stream.content).toEqual(['nuevo']); + } + }); +}); diff --git a/frontend/src/lib/ai-draft-store.ts b/frontend/src/lib/ai-draft-store.ts new file mode 100644 index 0000000..9be0662 --- /dev/null +++ b/frontend/src/lib/ai-draft-store.ts @@ -0,0 +1,150 @@ +/** + * Per-document ephemeral AI-draft state for the editor assistant (#73). + * + * Session-scoped (no persist): closing the panel unmounts the UI but the last + * stream, prompt, and thread id stay here so a reopen does not lose work or + * mint a duplicate chat thread. Mirrors `useChatStream` abort plumbing. + * + * Invariants: + * - State is keyed by editor `docId`; documents never share a slice. + * - `threadId` survives `clearDraft` so follow-up prompts stay in one thread. + * - `stopGenerate` keeps accumulated `stream` (partial draft after Detener). + */ +import { create } from 'zustand'; +import { applyChunk as mergeChatChunk } from '@/lib/api.mock'; +import type { ChatChunk, ChatMessage } from '@/lib/types'; + +/** Captured editor range a preset-targeted draft should replace. */ +export interface AiInsertTarget { + from: number; + to: number; + text: string; +} + +export interface DocAiDraftState { + prompt: string; + stream: ChatMessage | null; + busy: boolean; + error: string | null; + threadId: string | null; + insertTarget: AiInsertTarget | null; + abortController: AbortController | null; +} + +export const EMPTY_DOC_AI_DRAFT: DocAiDraftState = { + prompt: '', + stream: null, + busy: false, + error: null, + threadId: null, + insertTarget: null, + abortController: null, +}; + +interface AiDraftStore { + docs: Record; + getDocState(docId: string): DocAiDraftState; + setPrompt(docId: string, prompt: string): void; + setError(docId: string, error: string | null): void; + setThreadId(docId: string, threadId: string): void; + startGenerate(docId: string, insertTarget: AiInsertTarget | null): AbortController; + applyChunk(docId: string, chunk: ChatChunk, controller: AbortController): void; + stopGenerate(docId: string): void; + finishGenerate(docId: string, controller: AbortController): void; + clearDraft(docId: string): void; +} + +function withDoc( + docs: Record, + docId: string, + updater: (state: DocAiDraftState) => DocAiDraftState, +): Record { + const current = docs[docId] ?? EMPTY_DOC_AI_DRAFT; + return { ...docs, [docId]: updater(current) }; +} + +function isCurrentController(doc: DocAiDraftState, controller: AbortController): boolean { + return doc.abortController === controller; +} + +export const useAiDraftStore = create((set, get) => ({ + docs: {}, + getDocState(docId) { + return get().docs[docId] ?? EMPTY_DOC_AI_DRAFT; + }, + setPrompt(docId, prompt) { + set((state) => ({ + docs: withDoc(state.docs, docId, (doc) => ({ ...doc, prompt })), + })); + }, + setError(docId, error) { + set((state) => ({ + docs: withDoc(state.docs, docId, (doc) => ({ ...doc, error })), + })); + }, + setThreadId(docId, threadId) { + set((state) => ({ + docs: withDoc(state.docs, docId, (doc) => ({ ...doc, threadId })), + })); + }, + startGenerate(docId, insertTarget) { + const abortController = new AbortController(); + set((state) => ({ + docs: withDoc(state.docs, docId, (doc) => ({ + ...doc, + stream: null, + busy: true, + error: null, + insertTarget, + abortController, + })), + })); + return abortController; + }, + applyChunk(docId, chunk, controller) { + set((state) => { + const current = state.docs[docId] ?? EMPTY_DOC_AI_DRAFT; + if (!isCurrentController(current, controller)) return state; + return { + docs: withDoc(state.docs, docId, (doc) => ({ + ...doc, + stream: mergeChatChunk(doc.stream, chunk), + })), + }; + }); + }, + stopGenerate(docId) { + const controller = get().docs[docId]?.abortController; + controller?.abort(); + set((state) => ({ + docs: withDoc(state.docs, docId, (doc) => ({ + ...doc, + busy: false, + abortController: null, + })), + })); + }, + finishGenerate(docId, controller) { + set((state) => { + const current = state.docs[docId] ?? EMPTY_DOC_AI_DRAFT; + if (current.abortController !== null && !isCurrentController(current, controller)) { + return state; + } + return { + docs: withDoc(state.docs, docId, (doc) => ({ + ...doc, + busy: false, + abortController: null, + })), + }; + }); + }, + clearDraft(docId) { + set((state) => ({ + docs: withDoc(state.docs, docId, (doc) => ({ + ...EMPTY_DOC_AI_DRAFT, + threadId: doc.threadId, + })), + })); + }, +})); diff --git a/frontend/src/pages/EditorPage.test.tsx b/frontend/src/pages/EditorPage.test.tsx index 1332cdf..8f13bbd 100644 --- a/frontend/src/pages/EditorPage.test.tsx +++ b/frontend/src/pages/EditorPage.test.tsx @@ -77,6 +77,7 @@ vi.mock('@/pages/editor/ExportMenu', () => ({ ExportMenu: () => null })); vi.mock('@/pages/editor/extensions/LegalCitation', () => ({ LegalCitation: {} })); vi.mock('@/pages/editor/extensions/CommentMark', () => ({ CommentMark: {} })); vi.mock('@/pages/editor/extensions/AiGeneratedMark', () => ({ AiGeneratedMark: {} })); +vi.mock('@/pages/editor/extensions/PendingInsertHighlight', () => ({ PendingInsertHighlight: {} })); /** Navigates when `docId` prop changes so doc-switch effects run inside one router. */ function DocIdNavigator({ docId }: { docId: string }) { diff --git a/frontend/src/pages/EditorPage.tsx b/frontend/src/pages/EditorPage.tsx index 01eb438..002c5b7 100644 --- a/frontend/src/pages/EditorPage.tsx +++ b/frontend/src/pages/EditorPage.tsx @@ -48,6 +48,7 @@ import { ExportMenu } from '@/pages/editor/ExportMenu'; import { LegalCitation } from '@/pages/editor/extensions/LegalCitation'; import { CommentMark } from '@/pages/editor/extensions/CommentMark'; import { AiGeneratedMark } from '@/pages/editor/extensions/AiGeneratedMark'; +import { PendingInsertHighlight } from '@/pages/editor/extensions/PendingInsertHighlight'; import { useCommentStore } from '@/lib/comment-store'; import { cn } from '@/lib/utils'; import { DocumentList, DocumentPicker } from '@/pages/editor/DocumentList'; @@ -187,6 +188,7 @@ function EditorDocumentSurface({ docId }: { docId: string }) { // Inline comment anchors (#602). The note text lives in comment-store. CommentMark, AiGeneratedMark, + PendingInsertHighlight, ], content: initialDoc.content, editable: !isReadOnly, @@ -403,7 +405,12 @@ function EditorDocumentSurface({ docId }: { docId: string }) { {editor && aiPanelOpen && createPortal( - setAiPanelOpen(false)} />, + setAiPanelOpen(false)} + />, document.body, )} @@ -460,6 +467,9 @@ function EditorDocumentSurface({ docId }: { docId: string }) { '[&_.ProseMirror_.lex-comment--resolved]:decoration-amber-400/70', '[&_.ProseMirror_.lex-ai-generated]:rounded-sm [&_.ProseMirror_.lex-ai-generated]:bg-[hsl(var(--indigo-500)/0.18)]', '[&_.ProseMirror_.lex-ai-generated]:box-decoration-clone [&_.ProseMirror_.lex-ai-generated]:px-0.5', + '[&_.ProseMirror_.lex-pending-insert]:rounded-sm [&_.ProseMirror_.lex-pending-insert]:box-decoration-clone', + '[&_.ProseMirror_.lex-pending-insert]:outline [&_.ProseMirror_.lex-pending-insert]:outline-dashed', + '[&_.ProseMirror_.lex-pending-insert]:outline-1 [&_.ProseMirror_.lex-pending-insert]:outline-indigo-400/80', isReadOnly && 'cursor-default', )} > diff --git a/frontend/src/pages/editor/AiDraftPanel.test.tsx b/frontend/src/pages/editor/AiDraftPanel.test.tsx index 1d7ffc6..ecf3fb2 100644 --- a/frontend/src/pages/editor/AiDraftPanel.test.tsx +++ b/frontend/src/pages/editor/AiDraftPanel.test.tsx @@ -1,11 +1,15 @@ -import { render, screen, waitFor } from '@testing-library/react'; +import { fireEvent, render, screen, waitFor } from '@testing-library/react'; import userEvent from '@testing-library/user-event'; +import { getSchema } from '@tiptap/core'; +import StarterKit from '@tiptap/starter-kit'; import type { Editor } from '@tiptap/react'; import { MemoryRouter } from 'react-router-dom'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { AiDraftPanel } from './AiDraftPanel'; +import { AiGeneratedMark } from './extensions/AiGeneratedMark'; import { api } from '@/lib/api'; +import { useAiDraftStore } from '@/lib/ai-draft-store'; import { useUi } from '@/lib/store'; import type { ChatSource } from '@/lib/types'; @@ -25,44 +29,97 @@ vi.mock('@/lib/queries', () => ({ useModels: () => useModelsMock(), })); +const mockSelection = vi.hoisted(() => ({ + current: '', +})); + +vi.mock('@tiptap/react', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + useEditorState: () => mockSelection.current, + }; +}); + const insertContentMock = vi.fn().mockReturnThis(); +const insertContentAtMock = vi.fn().mockReturnThis(); const insertLegalCitationMock = vi.fn().mockReturnThis(); +const setPendingInsertRangeMock = vi.fn(); +const clearPendingInsertRangeMock = vi.fn(); const chainMock = { focus: vi.fn().mockReturnThis(), insertContent: insertContentMock, + insertContentAt: insertContentAtMock, insertLegalCitation: insertLegalCitationMock, run: vi.fn(), }; +const schema = getSchema([StarterKit, AiGeneratedMark]); + +const docTextBetween = vi.fn((_from: number, _to: number) => ''); + +const editorState = { + selection: { from: 0, to: 0, empty: true }, + doc: { + content: { size: 100 }, + textBetween: (from: number, to: number) => docTextBetween(from, to), + }, +}; + const editor = { chain: () => chainMock, - state: { selection: { from: 0, to: 0, empty: true }, doc: { textBetween: () => '' } }, + schema, + commands: { + setPendingInsertRange: setPendingInsertRangeMock, + clearPendingInsertRange: clearPendingInsertRangeMock, + }, + state: editorState, } as unknown as Editor; -vi.mock('@tiptap/react', () => ({ - useEditorState: () => '', -})); - -function renderPanel() { +function renderPanel(props?: { docId?: string; docTitle?: string }) { const onClose = vi.fn(); - render( + const view = render( - + , ); - return { onClose }; + return { onClose, ...view }; } -describe('AiDraftPanel trust surfaces', () => { +function resetEditorSelection(from = 0, to = 0, text = ''): void { + editorState.selection = { from, to, empty: from === to }; + mockSelection.current = text; + docTextBetween.mockImplementation((selFrom: number, selTo: number) => + selFrom === from && selTo === to ? text : '', + ); +} + +describe('AiDraftPanel', () => { beforeEach(() => { + vi.clearAllMocks(); navigateMock.mockReset(); insertContentMock.mockClear(); + insertContentAtMock.mockClear(); + insertLegalCitationMock.mockClear(); chainMock.run.mockClear(); + setPendingInsertRangeMock.mockClear(); + clearPendingInsertRangeMock.mockClear(); + useAiDraftStore.setState({ docs: {} }); + resetEditorSelection(); useUi.setState({ defaultModel: 'ollama:qwen2.5:7b' }); useModelsMock.mockReturnValue({ data: [{ id: 'ollama:qwen2.5:7b', available: true, label: 'qwen2.5:7b', vendor: 'ollama', kind: 'local' }], }); - vi.spyOn(api.chat, 'create').mockResolvedValue({ id: 'thread-1', title: 'Asistente', updatedAt: new Date().toISOString() }); + vi.spyOn(api.chat, 'create').mockResolvedValue({ + id: 'thread-1', + title: 'Asistente', + updatedAt: new Date().toISOString(), + }); }); it('shows the persistent AI draft disclaimer', () => { @@ -86,12 +143,10 @@ describe('AiDraftPanel trust surfaces', () => { await userEvent.click(screen.getByRole('button', { name: /insertar solo texto/i })); - expect(insertContentMock).toHaveBeenCalledWith([ - { - type: 'paragraph', - content: [{ type: 'text', text: 'Borrador IA', marks: [{ type: 'aiGenerated' }] }], - }, - ]); + expect(insertContentMock).toHaveBeenCalled(); + const inserted = insertContentMock.mock.calls[0]?.[0] as unknown[]; + expect(JSON.stringify(inserted)).toContain('Borrador IA'); + expect(JSON.stringify(inserted)).toContain('aiGenerated'); }); it('navigates when a draft source citation is clicked', async () => { @@ -120,4 +175,240 @@ describe('AiDraftPanel trust surfaces', () => { await userEvent.click(screen.getByRole('button', { name: /LOPDGDD/i })); expect(navigateMock).toHaveBeenCalledWith('/laws/BOE-A-2018-16673#art-28'); }); + + it('inserts a preset draft at the captured range, not the live cursor', async () => { + vi.spyOn(api.chat, 'send').mockImplementation(async function* () { + yield { type: 'text', delta: 'Cláusula mejorada' }; + yield { type: 'done' }; + }); + resetEditorSelection(5, 10, 'original'); + + renderPanel(); + await userEvent.click(screen.getByRole('button', { name: /^mejorar$/i })); + + await waitFor(() => { + expect(screen.getByRole('button', { name: /insertar solo texto/i })).toBeEnabled(); + }); + + editorState.selection = { from: 40, to: 40, empty: true }; + mockSelection.current = ''; + + await userEvent.click(screen.getByRole('button', { name: /insertar solo texto/i })); + + expect(insertContentAtMock).toHaveBeenCalledWith({ from: 5, to: 10 }, expect.any(Array)); + expect(insertContentMock).not.toHaveBeenCalled(); + }); + + it('shows a stale-target error and does not insert when the range no longer matches', async () => { + vi.spyOn(api.chat, 'send').mockImplementation(async function* () { + yield { type: 'text', delta: 'Cláusula mejorada' }; + yield { type: 'done' }; + }); + resetEditorSelection(5, 10, 'original'); + + renderPanel(); + await userEvent.click(screen.getByRole('button', { name: /^mejorar$/i })); + + await waitFor(() => { + expect(screen.getByRole('button', { name: /insertar solo texto/i })).toBeEnabled(); + }); + + docTextBetween.mockReturnValue('texto distinto'); + + await userEvent.click(screen.getByRole('button', { name: /insertar solo texto/i })); + + expect(insertContentAtMock).not.toHaveBeenCalled(); + expect(insertContentMock).not.toHaveBeenCalled(); + expect(screen.getByText(/la selección original ya no está/i)).toBeInTheDocument(); + }); + + it('keeps the completed draft after the panel is remounted', async () => { + vi.spyOn(api.chat, 'send').mockImplementation(async function* () { + yield { type: 'text', delta: 'Borrador persistente' }; + yield { type: 'done' }; + }); + + const { unmount } = renderPanel(); + await userEvent.type(screen.getByLabelText(/instrucción para redactar/i), 'Redacta una cláusula'); + await userEvent.click(screen.getByRole('button', { name: /^generar$/i })); + + await waitFor(() => { + expect(screen.getByText('Borrador persistente')).toBeInTheDocument(); + }); + + unmount(); + renderPanel(); + expect(screen.getByText('Borrador persistente')).toBeInTheDocument(); + }); + + it('stops an in-flight stream and keeps the partial draft', async () => { + vi.spyOn(api.chat, 'send').mockImplementation(async function* (_threadId, _content, opts) { + yield { type: 'text', delta: 'parcial' }; + await new Promise((_resolve, reject) => { + opts?.signal?.addEventListener('abort', () => { + reject(new DOMException('The operation was aborted.', 'AbortError')); + }); + }); + }); + + renderPanel(); + await userEvent.type(screen.getByLabelText(/instrucción para redactar/i), 'Redacta'); + await userEvent.click(screen.getByRole('button', { name: /^generar$/i })); + + await waitFor(() => { + expect(screen.getByText('parcial')).toBeInTheDocument(); + }); + + await userEvent.click(screen.getByRole('button', { name: /detener/i })); + + await waitFor(() => { + expect(useAiDraftStore.getState().getDocState('doc-1').busy).toBe(false); + }); + expect(screen.getByText('parcial')).toBeInTheDocument(); + }); + + it('aborts generation when the panel unmounts mid-stream', async () => { + vi.spyOn(api.chat, 'send').mockImplementation(async function* (_threadId, _content, opts) { + yield { type: 'text', delta: 'parcial' }; + await new Promise((_resolve, reject) => { + opts?.signal?.addEventListener('abort', () => { + reject(new DOMException('The operation was aborted.', 'AbortError')); + }); + }); + }); + + const { unmount } = renderPanel(); + await userEvent.type(screen.getByLabelText(/instrucción para redactar/i), 'Redacta'); + await userEvent.click(screen.getByRole('button', { name: /^generar$/i })); + + await waitFor(() => { + expect(screen.getByText('parcial')).toBeInTheDocument(); + }); + + unmount(); + + await waitFor(() => { + const state = useAiDraftStore.getState().getDocState('doc-1'); + expect(state.busy).toBe(false); + expect(state.abortController).toBeNull(); + }); + const stream = useAiDraftStore.getState().getDocState('doc-1').stream; + expect(stream?.role).toBe('assistant'); + if (stream?.role === 'assistant') { + expect(stream.content).toEqual(['parcial']); + } + }); + + it('creates one chat thread per document across remounts', async () => { + const createSpy = vi.spyOn(api.chat, 'create').mockResolvedValue({ + id: 'thread-1', + title: 'Contrato de arrendamiento', + updatedAt: new Date().toISOString(), + }); + vi.spyOn(api.chat, 'send').mockImplementation(async function* () { + yield { type: 'text', delta: 'uno' }; + yield { type: 'done' }; + }); + + const { unmount } = renderPanel(); + await userEvent.type(screen.getByLabelText(/instrucción para redactar/i), 'Primera'); + await userEvent.click(screen.getByRole('button', { name: /^generar$/i })); + await waitFor(() => { + expect(screen.getByText('uno')).toBeInTheDocument(); + }); + unmount(); + + renderPanel(); + await userEvent.click(screen.getByRole('button', { name: /^generar$/i })); + await waitFor(() => { + expect(createSpy).toHaveBeenCalledTimes(1); + }); + expect(createSpy).toHaveBeenCalledWith({ + title: 'Contrato de arrendamiento', + model: 'ollama:qwen2.5:7b', + }); + }); + + it('creates a separate thread for a different document', async () => { + const createSpy = vi.spyOn(api.chat, 'create'); + createSpy.mockResolvedValueOnce({ + id: 'thread-a', + title: 'Doc A', + updatedAt: new Date().toISOString(), + }); + createSpy.mockResolvedValueOnce({ + id: 'thread-b', + title: 'Doc B', + updatedAt: new Date().toISOString(), + }); + vi.spyOn(api.chat, 'send').mockImplementation(async function* () { + yield { type: 'text', delta: 'ok' }; + yield { type: 'done' }; + }); + + const first = renderPanel({ docId: 'doc-a', docTitle: 'Doc A' }); + await userEvent.type(screen.getByLabelText(/instrucción para redactar/i), 'A'); + await userEvent.click(screen.getByRole('button', { name: /^generar$/i })); + await waitFor(() => expect(screen.getByText('ok')).toBeInTheDocument()); + first.unmount(); + + renderPanel({ docId: 'doc-b', docTitle: 'Doc B' }); + await userEvent.type(screen.getByLabelText(/instrucción para redactar/i), 'B'); + await userEvent.click(screen.getByRole('button', { name: /^generar$/i })); + await waitFor(() => expect(createSpy).toHaveBeenCalledTimes(2)); + expect(createSpy).toHaveBeenNthCalledWith(1, { title: 'Doc A', model: 'ollama:qwen2.5:7b' }); + expect(createSpy).toHaveBeenNthCalledWith(2, { title: 'Doc B', model: 'ollama:qwen2.5:7b' }); + }); + + it('submits the prompt with Ctrl+Enter', async () => { + const sendSpy = vi.spyOn(api.chat, 'send').mockImplementation(async function* () { + yield { type: 'text', delta: 'ok' }; + yield { type: 'done' }; + }); + + renderPanel(); + const textarea = screen.getByLabelText(/instrucción para redactar/i); + await userEvent.type(textarea, 'Redacta'); + fireEvent.keyDown(textarea, { key: 'Enter', ctrlKey: true }); + + await waitFor(() => { + expect(sendSpy).toHaveBeenCalled(); + }); + }); + + it('does not submit on plain Enter', async () => { + const sendSpy = vi.spyOn(api.chat, 'send').mockImplementation(async function* () { + yield { type: 'done' }; + }); + + renderPanel(); + const textarea = screen.getByLabelText(/instrucción para redactar/i); + await userEvent.type(textarea, 'Redacta'); + fireEvent.keyDown(textarea, { key: 'Enter' }); + + expect(sendSpy).not.toHaveBeenCalled(); + }); + + it('closes on Escape and aborts when busy', async () => { + vi.spyOn(api.chat, 'send').mockImplementation(async function* (_threadId, _content, opts) { + yield { type: 'text', delta: 'parcial' }; + await new Promise((_resolve, reject) => { + opts?.signal?.addEventListener('abort', () => { + reject(new DOMException('The operation was aborted.', 'AbortError')); + }); + }); + }); + + const { onClose } = renderPanel(); + await userEvent.type(screen.getByLabelText(/instrucción para redactar/i), 'Redacta'); + await userEvent.click(screen.getByRole('button', { name: /^generar$/i })); + await waitFor(() => expect(screen.getByText('parcial')).toBeInTheDocument()); + + fireEvent.keyDown(window, { key: 'Escape' }); + + expect(onClose).toHaveBeenCalled(); + await waitFor(() => { + expect(useAiDraftStore.getState().getDocState('doc-1').busy).toBe(false); + }); + }); }); diff --git a/frontend/src/pages/editor/AiDraftPanel.tsx b/frontend/src/pages/editor/AiDraftPanel.tsx index dac597e..43b1138 100644 --- a/frontend/src/pages/editor/AiDraftPanel.tsx +++ b/frontend/src/pages/editor/AiDraftPanel.tsx @@ -1,83 +1,109 @@ /** - * AiDraftPanel — AI-assisted drafting docked in the editor (#601). + * AiDraftPanel — AI-assisted drafting docked in the editor (#601 / #73). * - * Reuses the existing agentic chat stack (sprints 1-4): it streams a grounded - * answer via `api.chat.send` and accumulates it with `applyChunk`, exactly like - * ChatPage. The streamed text can be inserted into the document, and the - * `source` events the RAG loop surfaces are turned into #599 typed-citation - * nodes — so AI-drafted text lands with clickable, corpus-resolved references. + * Reuses the existing agentic chat stack: streams a grounded answer via + * `api.chat.send` and accumulates it with the per-document draft store, + * exactly like ChatPage + `useChatStream`. Streamed text can be inserted + * into the document; RAG `source` events become #599 typed-citation nodes. * * Rendered as a right-docked drawer (no backdrop) so the editor stays - * interactive: the user can place the cursor, then insert. + * interactive. Draft state lives in `useAiDraftStore` so close/reopen + * does not discard a generation or mint duplicate chat threads. * * --- WHERE TO CHANGE IF X CHANGES --- - * - Streaming/contract → `api.chat.send` + `applyChunk` (`lib/api`). + * - Streaming/contract → `api.chat.send` + store `applyChunk`. * - Source → citation mapping → `citationFromSource` in `./citation-utils`. + * - Markdown insertion → `markdownDraftToBlocks` in `./draft-insert-utils`. * - Model selection → mirrors ChatPage (`useModels` + `useUi.defaultModel`). */ -import { useMemo, useRef, useState } from 'react'; -import type { Editor, JSONContent } from '@tiptap/react'; +import { useCallback, useEffect, useMemo } from 'react'; +import type { Editor } from '@tiptap/react'; import { useEditorState } from '@tiptap/react'; import { useNavigate } from 'react-router-dom'; import { useTranslation } from 'react-i18next'; -import { Sparkles, X, Send, Loader2 } from 'lucide-react'; +import { Sparkles, X, Send, Square } from 'lucide-react'; import { Button } from '@/components/ui'; import { CitationCard } from '@/components/domain/CitationCard'; import { api } from '@/lib/api'; -import { applyChunk } from '@/lib/api.mock'; import { chatSourceHref } from '@/lib/chat-sources'; +import { isSentinelDocumentTitle } from '@/lib/editor-store'; +import { EMPTY_DOC_AI_DRAFT, useAiDraftStore, type AiInsertTarget } from '@/lib/ai-draft-store'; import { useModels } from '@/lib/queries'; import { useUi } from '@/lib/store'; -import type { ChatMessage as ChatMessageT, ChatSource } from '@/lib/types'; +import type { ChatSource } from '@/lib/types'; import { citationFromSource, type CitationAttrs } from './citation-utils'; +import { isInsertTargetValid, markdownDraftToBlocks } from './draft-insert-utils'; interface AiDraftPanelProps { editor: Editor; + docId: string; + docTitle: string; onClose: () => void; } -/** Split assistant text into paragraph nodes on blank lines (lazy markdown). */ -function draftToBlocks(text: string): JSONContent[] { - return text - .split(/\n{2,}/) - .map((para) => para.trim()) - .filter(Boolean) - .map((para) => ({ - type: 'paragraph', - content: [{ type: 'text', text: para, marks: [{ type: 'aiGenerated' }] }], - })); -} - interface Preset { id: string; labelKey: 'editor.aiDraft.presetImprove' | 'editor.aiDraft.presetSummary' | 'editor.aiDraft.presetExplain'; build: (selection: string) => string; } -export function AiDraftPanel({ editor, onClose }: AiDraftPanelProps) { - const { t } = useTranslation(); +const PRESETS: Preset[] = [ + { + id: 'improve', + labelKey: 'editor.aiDraft.presetImprove', + build: (s) => `Mejora la redacción de este texto legal, conservando su sentido:\n\n${s}`, + }, + { + id: 'summary', + labelKey: 'editor.aiDraft.presetSummary', + build: (s) => `Resume de forma concisa este texto:\n\n${s}`, + }, + { + id: 'explain', + labelKey: 'editor.aiDraft.presetExplain', + build: (s) => `Explica en lenguaje claro este texto:\n\n${s}`, + }, +]; - /** Selection-based quick actions. All operate on the current editor selection. */ - const presets: Preset[] = useMemo( - () => [ - { id: 'improve', labelKey: 'editor.aiDraft.presetImprove', build: (s) => `Mejora la redacción de este texto legal, conservando su sentido:\n\n${s}` }, - { id: 'summary', labelKey: 'editor.aiDraft.presetSummary', build: (s) => `Resume de forma concisa este texto:\n\n${s}` }, - { id: 'explain', labelKey: 'editor.aiDraft.presetExplain', build: (s) => `Explica en lenguaje claro este texto:\n\n${s}` }, - ], - [], - ); +function isAbortError(exc: unknown): boolean { + return exc instanceof DOMException && exc.name === 'AbortError'; +} + +function draftingThreadTitle(docTitle: string, fallback: string): string { + return isSentinelDocumentTitle(docTitle) ? fallback : docTitle; +} + +function captureSelectionTarget(editor: Editor): AiInsertTarget | null { + const { from, to, empty } = editor.state.selection; + if (empty) return null; + const text = editor.state.doc.textBetween(from, to, '\n'); + if (!text) return null; + return { from, to, text }; +} + +export function AiDraftPanel({ editor, docId, docTitle, onClose }: AiDraftPanelProps) { + const { t } = useTranslation(); const navigate = useNavigate(); const { data: models = [] } = useModels(); const defaultModel = useUi((s) => s.defaultModel); - // Same resolution as ChatPage: prefer the configured default if it's - // available, else the first available provider; empty string = none. const model = useMemo(() => { const current = models.find((m) => m.id === defaultModel); if (current?.available) return defaultModel; return models.find((m) => m.available)?.id ?? ''; }, [models, defaultModel]); - // Reactive selected-document text (drives the selection presets + indicator). + const draft = useAiDraftStore((s) => s.docs[docId] ?? EMPTY_DOC_AI_DRAFT); + const setPrompt = useAiDraftStore((s) => s.setPrompt); + const setError = useAiDraftStore((s) => s.setError); + const setThreadId = useAiDraftStore((s) => s.setThreadId); + const startGenerate = useAiDraftStore((s) => s.startGenerate); + const applyChunk = useAiDraftStore((s) => s.applyChunk); + const stopGenerate = useAiDraftStore((s) => s.stopGenerate); + const finishGenerate = useAiDraftStore((s) => s.finishGenerate); + const clearDraft = useAiDraftStore((s) => s.clearDraft); + + const { prompt, stream, busy, error, insertTarget } = draft; + const selectionText = useEditorState({ editor, selector: (snap) => { @@ -86,70 +112,123 @@ export function AiDraftPanel({ editor, onClose }: AiDraftPanelProps) { }, }); - const [prompt, setPrompt] = useState(''); - const [stream, setStream] = useState(null); - const [busy, setBusy] = useState(false); - const [error, setError] = useState(null); - // One thread reused for the whole drafting session (lazy-created). - const threadIdRef = useRef(null); - const draftText = stream?.role === 'assistant' ? stream.content.join('\n\n').trim() : ''; const sources: ChatSource[] = stream?.role === 'assistant' ? stream.sources : []; const streamError = stream?.role === 'assistant' ? stream.error?.detail : null; const displayError = error ?? streamError; - const generate = async (text: string) => { - const content = text.trim(); - if (busy || !content) return; - if (!model) { - setError(t('editor.aiDraft.noModel')); - return; - } - setBusy(true); - setStream(null); - setError(null); - try { - let threadId = threadIdRef.current; - if (!threadId) { - const created = await api.chat.create({ title: t('editor.aiDraft.threadTitle'), model }); - threadId = created.id; - threadIdRef.current = threadId; + const generate = useCallback( + async (text: string, target: AiInsertTarget | null) => { + const content = text.trim(); + if (useAiDraftStore.getState().getDocState(docId).busy || !content) return; + if (!model) { + setError(docId, t('editor.aiDraft.noModel')); + return; } - let current: ChatMessageT | null = null; - for await (const chunk of api.chat.send(threadId, content, { model })) { - current = applyChunk(current, chunk); - setStream(current); + const controller = startGenerate(docId, target); + try { + let threadId = useAiDraftStore.getState().getDocState(docId).threadId; + if (!threadId) { + const created = await api.chat.create({ + title: draftingThreadTitle(docTitle, t('editor.aiDraft.threadTitle')), + model, + }); + if (useAiDraftStore.getState().getDocState(docId).abortController !== controller) { + return; + } + threadId = created.id; + setThreadId(docId, threadId); + } + for await (const chunk of api.chat.send(threadId, content, { model, signal: controller.signal })) { + applyChunk(docId, chunk, controller); + } + } catch (exc) { + if (isAbortError(exc)) return; + setError(docId, exc instanceof Error ? exc.message : t('editor.aiDraft.generateError')); + } finally { + finishGenerate(docId, controller); } - } catch (exc) { - setError(exc instanceof Error ? exc.message : t('editor.aiDraft.generateError')); - } finally { - setBusy(false); - } - }; + }, + [applyChunk, docId, docTitle, finishGenerate, model, setError, setThreadId, startGenerate, t], + ); const runPreset = (preset: Preset) => { - if (!selectionText) return; - void generate(preset.build(selectionText)); + const target = captureSelectionTarget(editor); + if (!target) return; + void generate(preset.build(target.text), target); }; + const requestClose = useCallback(() => { + if (useAiDraftStore.getState().getDocState(docId).busy) { + stopGenerate(docId); + } + onClose(); + }, [docId, onClose, stopGenerate]); + + useEffect(() => { + return () => { + if (useAiDraftStore.getState().getDocState(docId).busy) { + useAiDraftStore.getState().stopGenerate(docId); + } + }; + }, [docId]); + + useEffect(() => { + const onKey = (e: KeyboardEvent) => { + if (e.key !== 'Escape') return; + e.preventDefault(); + requestClose(); + }; + window.addEventListener('keydown', onKey); + return () => window.removeEventListener('keydown', onKey); + }, [requestClose]); + + useEffect(() => { + if (busy && insertTarget) { + editor.commands.setPendingInsertRange({ from: insertTarget.from, to: insertTarget.to }); + return () => { + editor.commands.clearPendingInsertRange(); + }; + } + editor.commands.clearPendingInsertRange(); + return undefined; + }, [busy, editor, insertTarget]); + const insertDraft = (withCitations: boolean) => { if (!draftText) return; - let chain = editor.chain().focus().insertContent(draftToBlocks(draftText)); + if (insertTarget && !isInsertTargetValid(editor, insertTarget)) { + setError(docId, t('editor.aiDraft.targetStale')); + return; + } + const blocks = markdownDraftToBlocks(draftText, editor.schema); + let chain = editor.chain().focus(); + chain = insertTarget + ? chain.insertContentAt({ from: insertTarget.from, to: insertTarget.to }, blocks) + : chain.insertContent(blocks); if (withCitations) { - const citations = sources - .map(citationFromSource) - .filter((c): c is CitationAttrs => c !== null); + const citations = sources.map(citationFromSource).filter((c): c is CitationAttrs => c !== null); if (citations.length > 0) { - chain = chain.insertContent({ type: 'paragraph', content: [{ type: 'text', text: `${t('editor.aiDraft.sourcesPrefix')} ` }] }); + chain = chain.insertContent({ + type: 'paragraph', + content: [{ type: 'text', text: `${t('editor.aiDraft.sourcesPrefix')} ` }], + }); for (const citation of citations) { chain = chain.insertLegalCitation(citation); } } } chain.run(); + clearDraft(docId); onClose(); }; + const onPromptKeyDown = (e: React.KeyboardEvent) => { + if (e.nativeEvent.isComposing) return; + if (e.key !== 'Enter' || (!e.ctrlKey && !e.metaKey)) return; + e.preventDefault(); + if (!busy) void generate(prompt, null); + }; + const citableCount = sources.filter((s) => s.target?.lawId).length; return ( @@ -161,7 +240,13 @@ export function AiDraftPanel({ editor, onClose }: AiDraftPanelProps) {
{t('editor.aiDraft.title')} -
@@ -175,7 +260,6 @@ export function AiDraftPanel({ editor, onClose }: AiDraftPanelProps) {

{t('editor.aiDraftDisclaimer')}

- {/* Selection-based quick actions. */}
{t('editor.aiDraft.selectionSection')}
{selectionText ? ( @@ -185,8 +269,11 @@ export function AiDraftPanel({ editor, onClose }: AiDraftPanelProps) { ) : (
{t('editor.aiDraft.selectTextHint')}
)} + {insertTarget && ( +
{t('editor.aiDraft.pendingTargetHint')}
+ )}
- {presets.map((preset) => ( + {PRESETS.map((preset) => (
- {/* Free-text drafting prompt. */}
{t('editor.aiDraft.promptSection')}