diff --git a/components/AIMode/DocumentPane.tsx b/components/AIMode/DocumentPane.tsx
index 1847f8aa0..f99dd8a7b 100644
--- a/components/AIMode/DocumentPane.tsx
+++ b/components/AIMode/DocumentPane.tsx
@@ -18,6 +18,7 @@ import { noteDiffPersistableDoc } from '@/components/Notebook/NoteReview/noteDif
import { useNoteAgentReview } from '@/components/Notebook/NoteReview/useNoteAgentReview';
import { Button } from '@/components/ui/Button';
import { Loader } from '@/components/ui/Loader';
+import { DraftBlockPreview } from './DraftBlockPreview';
import { DocumentPaneSkeleton } from '@/components/skeletons/AIModeSkeleton';
import { useUpdateNote } from '@/hooks/useNote';
import type { AgentChat } from '@/types/agentChat';
@@ -60,7 +61,7 @@ export function DocumentPane({
readOnly = false,
className,
}: DocumentPaneProps) {
- const { note, content, loading, error, status, draftText, phaseLabel } = document;
+ const { note, content, loading, error, status, draftText, draftBlocks, phaseLabel } = document;
const noteId = note?.id ?? null;
const writing = status === 'drafting' || status === 'working';
@@ -196,19 +197,29 @@ export function DocumentPane({
{/* Mounted once per note: the editor's content prop is only read on
creation, and later versions arrive through the review. */}
-
+
+
+
- {status === 'drafting' && draftText && }
+ {status === 'drafting' && draftText && (
+
+ )}
{status === 'working' && document.hasWrittenVersion && (
@@ -268,26 +279,33 @@ function EmptyDocument() {
}
/** The section being written, appended below the settled content. */
-function DraftSection({ text }: { readonly text: string }) {
- const paragraphs = text.split(/\n{2,}/).filter((paragraph) => paragraph.trim().length > 0);
+function DraftSection({
+ text,
+ blocks,
+ editor,
+ hasSavedContent,
+}: {
+ readonly text: string;
+ readonly blocks: AIModeDocument['draftBlocks'];
+ readonly editor: Editor | null;
+ readonly hasSavedContent: boolean;
+}) {
return (
-
-
- Writing
+
+
+
+ Drafting
+
+ ·
+
+ Preview updates live
+
- {paragraphs.map((paragraph, index) => (
-
- {paragraph}
- {index === paragraphs.length - 1 && (
-
- )}
-
- ))}
+
);
}
diff --git a/components/AIMode/DraftBlockPreview.tsx b/components/AIMode/DraftBlockPreview.tsx
new file mode 100644
index 000000000..c2d9d6265
--- /dev/null
+++ b/components/AIMode/DraftBlockPreview.tsx
@@ -0,0 +1,50 @@
+'use client';
+
+import { useEffect, useState } from 'react';
+import { generateHTML, type Editor, type JSONContent } from '@tiptap/core';
+import sanitizeHtml from 'sanitize-html';
+import { ASSISTANT_HTML_SANITIZE_OPTIONS } from '@/components/AgentChat/MarkdownMessage';
+
+/** A presentation-only snapshot: never apply partial blocks to the live editor. */
+export function DraftBlockPreview({
+ blocks,
+ editor,
+ fallbackText,
+}: {
+ readonly blocks: JSONContent[] | null;
+ readonly editor: Editor | null;
+ readonly fallbackText: string;
+}) {
+ const [html, setHtml] = useState
(null);
+
+ useEffect(() => {
+ if (!editor || editor.isDestroyed || blocks == null) return;
+ try {
+ const rendered = generateHTML(
+ { type: 'doc', content: blocks },
+ editor.extensionManager.extensions
+ );
+ setHtml(sanitizeHtml(rendered, ASSISTANT_HTML_SANITIZE_OPTIONS));
+ } catch {
+ // A node type, mark, or text value may still be incomplete. Keep the
+ // last renderable snapshot until the next fragment supplies it.
+ }
+ }, [blocks, editor]);
+
+ if (html == null) {
+ return (
+
+ {fallbackText}
+
+ );
+ }
+
+ return (
+
+ );
+}
diff --git a/components/AIMode/useAIModeDocument.ts b/components/AIMode/useAIModeDocument.ts
index 71ecaa4c4..f2f8a7647 100644
--- a/components/AIMode/useAIModeDocument.ts
+++ b/components/AIMode/useAIModeDocument.ts
@@ -1,11 +1,13 @@
'use client';
import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
+import type { JSONContent } from '@tiptap/core';
import { NoteService } from '@/services/note.service';
import type { NoteWithContent } from '@/types/note';
import {
isActiveExecutionStatus,
type ChatExecution,
+ type ChatStreamItem,
type ChatNoteRef,
type AgentChat,
} from '@/types/agentChat';
@@ -40,6 +42,8 @@ export interface AIModeDocument {
readonly hasWrittenVersion: boolean;
/** Prose of the `edit_note` call being composed, paragraphs split by blank lines. */
readonly draftText: string | null;
+ readonly draftBlocks: JSONContent[] | null;
+ readonly draftKey: string;
/** What the assistant is doing, for the in-progress row when there is no draft. */
readonly phaseLabel: string | null;
/** Deep link to the note in the notebook, once its organization is known. */
@@ -64,13 +68,15 @@ function chatHasEditedNote(chat: AgentChat | null): boolean {
}
/** The `edit_note` draft the active turn is composing, if any. */
-function currentEditDraft(execution: ChatExecution | null): string | null {
+function currentEditDraft(
+ execution: ChatExecution | null
+): Extract | null {
if (execution == null || !isActiveExecutionStatus(execution.status)) return null;
const items = execution.stream?.items ?? [];
for (let index = items.length - 1; index >= 0; index -= 1) {
const item = items[index];
if (item.type === 'tool_draft' && item.tool === 'edit_note') {
- return item.text.length > 0 ? item.text : null;
+ return item;
}
}
return null;
@@ -119,7 +125,10 @@ export function useAIModeDocument({
if (noteId != null) fetchNote();
}, [noteId, fetchNote]);
- const draftText = currentEditDraft(latestExecution);
+ const draft = currentEditDraft(latestExecution);
+ const draftText = draft?.text || null;
+ const draftBlocks = Array.isArray(draft?.blocks) ? draft.blocks : null;
+ const draftKey = `${latestExecution?.stream?.id}:${draft?.id}`;
const turnActive = latestExecution != null && isActiveExecutionStatus(latestExecution.status);
const phaseLabel = turnActive ? (latestExecution?.phase?.label ?? null) : null;
@@ -146,6 +155,8 @@ export function useAIModeDocument({
status,
hasWrittenVersion,
draftText,
+ draftBlocks,
+ draftKey,
phaseLabel,
notebookHref,
reload: fetchNote,
diff --git a/components/AgentChat/ActivityFeed.tsx b/components/AgentChat/ActivityFeed.tsx
index 8dbd695ea..3529fbd79 100644
--- a/components/AgentChat/ActivityFeed.tsx
+++ b/components/AgentChat/ActivityFeed.tsx
@@ -215,6 +215,7 @@ function StreamedTextRow({
label,
text,
streaming,
+ autoExpand = true,
className,
bodyClassName,
icon: Icon,
@@ -222,13 +223,14 @@ function StreamedTextRow({
readonly label: string;
readonly text: string;
readonly streaming: boolean;
+ readonly autoExpand?: boolean;
readonly className: string;
readonly bodyClassName?: string;
readonly icon?: ComponentType<{ className?: string }>;
}) {
const [userExpanded, setUserExpanded] = useState(null);
const hasText = text.length > 0;
- const expanded = hasText && (userExpanded ?? streaming);
+ const expanded = hasText && (userExpanded ?? (streaming && autoExpand));
// Settled rows re-render on every stream delta; strip once per text value.
const preview = useMemo(() => stripMarkdown(text), [text]);
@@ -400,6 +402,7 @@ function ActivityItemBody({
label={humanizeLabel(item.label)}
text={item.text}
streaming={streaming}
+ autoExpand={false}
icon={TOOL_ICONS[item.tool] ?? Wrench}
className="text-gray-800 hover:text-gray-600 [--shine:theme(colors.gray.800)]"
bodyClassName="text-sm text-gray-500"
diff --git a/components/AgentChat/MarkdownMessage.tsx b/components/AgentChat/MarkdownMessage.tsx
index 9c041dec5..e7020e4ad 100644
--- a/components/AgentChat/MarkdownMessage.tsx
+++ b/components/AgentChat/MarkdownMessage.tsx
@@ -14,7 +14,7 @@ const md = new MarkdownIt({
breaks: true,
});
-const SANITIZE_OPTIONS: sanitizeHtml.IOptions = {
+export const ASSISTANT_HTML_SANITIZE_OPTIONS: sanitizeHtml.IOptions = {
allowedTags: [
'p',
'br',
@@ -48,6 +48,7 @@ const SANITIZE_OPTIONS: sanitizeHtml.IOptions = {
allowedAttributes: {
a: ['href', 'title', 'target', 'rel'],
code: ['class'],
+ ol: ['start'],
th: ['align'],
td: ['align'],
},
@@ -111,7 +112,10 @@ export function MarkdownMessage({
revealCarryTo,
}: MarkdownMessageProps) {
const shown = useTextReveal(content, revealKey, revealCarryTo);
- const html = useMemo(() => sanitizeHtml(md.render(shown), SANITIZE_OPTIONS), [shown]);
+ const html = useMemo(
+ () => sanitizeHtml(md.render(shown), ASSISTANT_HTML_SANITIZE_OPTIONS),
+ [shown]
+ );
return (