Show formatted notes as the assistant writes

components/AIMode/DocumentPane.tsx

@@ -18,6 +18,7 @@
 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';
@@ -51,7 +52,7 @@
  * below the editor, and a turn with no draft shows an in-progress row so the
  * page never sits frozen.
  */
 export function DocumentPane({
   document,
   chat,
   tab,
@@ -60,7 +61,7 @@
   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 @@
 
             {/* Mounted once per note: the editor's content prop is only read on
                 creation, and later versions arrive through the review. */}
-            <BlockEditorClientWrapper
-              key={noteId ?? 'none'}
-              content={content.content}
-              contentJson={content.contentJson}
-              editable={!readOnly}
-              locked={locked}
-              requireTitle={false}
-              autofocus={false}
-              onUpdate={readOnly ? undefined : handleEditorUpdate}
-              setEditor={setEditor}
-            />
+            <div className={writing && !document.hasWrittenVersion ? 'hidden' : undefined}>
+              <BlockEditorClientWrapper
+                key={noteId ?? 'none'}
+                content={content.content}
+                contentJson={content.contentJson}
+                editable={!readOnly}
+                locked={locked}
+                requireTitle={false}
+                autofocus={false}
+                onUpdate={readOnly ? undefined : handleEditorUpdate}
+                setEditor={setEditor}
+              />
+            </div>
 
-            {status === 'drafting' && draftText && <DraftSection text={draftText} />}
+            {status === 'drafting' && draftText && (
+              <DraftSection
+                text={draftText}
+                blocks={draftBlocks}
+                editor={editor}
+                key={document.draftKey}
+                hasSavedContent={document.hasWrittenVersion}
+              />
+            )}
 
             {status === 'working' && document.hasWrittenVersion && (
               <InProgressRow label={phaseLabel ?? 'Working'} />
@@ -268,26 +279,33 @@
 }
 
 /** 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 (
     <section
-      aria-live="polite"
       aria-label="Section being written"
-      className="prose prose-sm prose-neutral mt-6 max-w-none border-t border-dashed border-primary-200 pt-5"
+      className={cn(hasSavedContent && 'mt-8 border-t border-gray-100 pt-6')}
     >
-      <div className="mb-3 flex items-center gap-2 text-[11px] font-semibold uppercase tracking-wider text-primary-600">
-        <Loader size="sm" className="!h-2.5 !w-2.5" />
-        Writing
+      <div role="status" className="mb-6 flex items-center gap-2 text-xs text-gray-500">
+        <Loader size="sm" className="!h-3 !w-3 text-primary-500" />
+        <span>
+          Drafting
+          <span className="mx-2 text-gray-300" aria-hidden="true">
+            ·
+          </span>
+          Preview updates live
+        </span>
       </div>
-      {paragraphs.map((paragraph, index) => (
-        <p key={index} className="whitespace-pre-wrap">
-          {paragraph}
-          {index === paragraphs.length - 1 && (
-            <span className="ml-0.5 inline-block h-[1em] w-[2px] translate-y-[2px] animate-pulse bg-primary-500" />
-          )}
-        </p>
-      ))}
+      <DraftBlockPreview blocks={blocks} editor={editor} fallbackText={text} />
     </section>
   );
 }

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<string | null>(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 (
+      <div className="whitespace-pre-wrap text-base leading-[1.8] text-gray-800">
+        {fallbackText}
+      </div>
+    );
+  }
+
+  return (
+    <div
+      className="text-base leading-[1.8] text-gray-800 break-words [&_p]:my-4 [&_h1]:text-2xl [&_h1]:font-semibold [&_h1]:mb-6 [&_h2]:text-xl [&_h2]:font-semibold [&_h2]:mt-8 [&_h2]:mb-3 [&_h3]:font-semibold [&_h3]:mt-6 [&_h3]:mb-2 [&_ul]:my-4 [&_ul]:list-disc [&_ul]:pl-5 [&_ol]:my-4 [&_ol]:list-decimal [&_ol]:pl-5 [&_li]:my-2 [&_blockquote]:border-l-2 [&_blockquote]:pl-4 [&_a]:text-primary-600 [&_a]:underline [&_pre]:overflow-x-auto [&_pre]:whitespace-pre [&_table]:block [&_table]:overflow-x-auto [&_td]:border [&_td]:p-2 [&_th]:border [&_th]:p-2"
+      // Generated by Tiptap and sanitized with the same policy as assistant messages.
+      // eslint-disable-next-line react/no-danger
+      dangerouslySetInnerHTML={{ __html: html }}
+    />
+  );
+}

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<ChatStreamItem, { type: 'tool_draft' }> | 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,

components/AgentChat/ActivityFeed.tsx

@@ -215,20 +215,22 @@ function StreamedTextRow({
   label,
   text,
   streaming,
+  autoExpand = true,
   className,
   bodyClassName,
   icon: Icon,
 }: {
   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<boolean | null>(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"

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 (
     <div

hooks/useAgentChat.ts

@@ -126,7 +126,13 @@ function phaseForDelta(delta: ChatStreamDelta | undefined): ExecutionPhase {
 function newStreamItem(delta: ChatStreamDelta, maximum: number): ChatStreamItem {
   const base = { id: delta.id, text: delta.delta.slice(0, maximum), at: delta.at };
   return delta.type === 'tool_draft'
-    ? { ...base, type: 'tool_draft', tool: delta.tool, label: delta.label }
+    ? {
+        ...base,
+        type: 'tool_draft',
+        tool: delta.tool,
+        label: delta.label,
+        blocks: Array.isArray(delta.blocks) ? delta.blocks : undefined,
+      }
     : { ...base, type: delta.type };
 }
 
@@ -143,6 +149,13 @@ function appendStreamDeltas(
       items.push(newStreamItem(delta, maximum));
     } else if (existing.type === delta.type) {
       existing.text = `${existing.text}${delta.delta}`.slice(0, maximum);
+      if (
+        existing.type === 'tool_draft' &&
+        delta.type === 'tool_draft' &&
+        Array.isArray(delta.blocks)
+      ) {
+        existing.blocks = delta.blocks;
+      }
     } else {
       // An id changing type indicates an incompatible/corrupt frame. Recover
       // from the server checkpoint instead of combining unlike content.

types/agentChat.ts

@@ -9,6 +9,8 @@
  * against the backend contract.
  */
 
+import type { JSONContent } from '@tiptap/core';
+
 import type { EffortLevel } from './agentModels';
 
 export type ExecutionStatus =
@@ -79,6 +81,8 @@ export interface ChatToolCallActivity {
  */
 export interface ChatToolDraftActivity {
   type: 'tool_draft';
+  /** Complete formatted preview snapshot; replaces itself on each frame. */
+  blocks?: JSONContent[];
   /**
    * The prose extracted from the arguments so far. Empty for tools whose
    * arguments aren't prose — a search query is written in an instant, so only
@@ -253,6 +257,8 @@ export type ChatStreamDelta =
   | (ChatStreamDeltaBase & { type: 'narration' | 'thinking' })
   | (ChatStreamDeltaBase & {
       type: 'tool_draft';
+      /** Complete formatted preview snapshot, not an appended delta. */
+      blocks?: JSONContent[];
       /** Machine name; empty when the provider skipped the block-start event. */
       tool: string;
       /** Human copy supplied by the backend; always rendered verbatim. */