How AI Writes Inside a Notion-Style Editor: Streamed Tool Calls, Partial JSON and Suggestion Marks

Architecture reconstructed from public sources. See Sources at the end.

A chat UI has it easy. The model streams text, you append it to a string, React re-renders a <p>. Now put the same model inside a document editor. The user selects three paragraphs, types "make this more concise", and expects the AI to edit those exact paragraphs in place, keep the headings and lists intact, show what changed, and let them accept or throw it all away. Appending a string is no longer enough. The model has to produce edits to a structured document, and the editor has to apply them while they are still streaming.

Notion does not publish how its editor AI works. BlockNote, an open source Notion-style block editor built on ProseMirror and Tiptap, does: its AI extension (@blocknote/xl-ai) is public code built on the Vercel AI SDK. So this case study uses Notion for the product behaviour and BlockNote's source (version 0.57.0, published 30 October 2026) for the architecture.

This pairs with Monday's post, The Vercel AI SDK for Frontend Engineers: streamText, useChat and Providers. That post covers the SDK primitives; this one shows a real product built on them.

TL;DR

An in-editor AI feature serializes the relevant blocks (as HTML, with block IDs and cursor position) into the prompt, and forces the model to answer with a single tool call whose input is a list of document operations: add, update, delete. The client reads that tool call while its JSON is still streaming, repairs the partial JSON, validates each operation, converts it into ProseMirror steps, and replays those steps as if a human were typing. The design decision that makes it work is that every AI edit is written as a suggestion (insertion and deletion marks) instead of a real change, so the document is always in a valid state, the user can review a diff, and accept or reject is a single transaction.

What the user sees

From the browser, the flow looks like this in Notion and in BlockNote alike:

1. You select text and choose "Edit with AI" (Notion's label), or open the AI menu.
2. You type a prompt, or pick a preset like "Fix spelling" or "Make shorter".
3. A short "thinking" state appears. Then text starts changing in place. In BlockNote you see an AI cursor (labelled "AI" by default) move through the document, select text, and type the replacement character by character.
4. Removed text stays visible with a strike-through, and new text is highlighted. Nothing is final yet.
5. When the AI finishes, you get Accept and Reject (Notion's help guide describes accept, discard and try again).
6. After accepting, one Ctrl+Z takes you back to the document as it was before the AI touched it, not through every character the AI typed.

In DevTools (BlockNote with a backend), you see one streamed POST per request. The interesting part is not text deltas but a growing tool call input, like {"operations":[{"type":"update","id":"abc$","block":"<p>..., arriving a few characters at a time.

Requirements

Here is how an interviewer would frame "design AI editing inside a block-based editor".

Functional

• Users can ask the AI to write new content at the cursor, or to transform a selection (rewrite, translate, shorten, format).
• The AI can add, update and delete blocks, including formatting (bold, links, headings, lists).
• Changes appear progressively while the model is generating.
• The user reviews the result as a diff and can accept, reject, retry or abort.
• Works in a collaborative document without leaking half-written AI edits to other users.

Non-functional

• Perceived latency: something should move on screen soon after the request starts. Time to first visible edit matters more than total time.
• Correctness: a half-streamed <ul><li> must never corrupt the document.
• Cost: only send the blocks the model needs, not the whole workspace, because input tokens are billed per request.
• Reliability: invented block IDs, invalid JSON or a dead stream must never destroy the user's original text.
• Undo integrity: AI edits must not pollute the undo stack with hundreds of tiny steps.

High-level architecture

Diagram
100%
flowchart LR
    subgraph Browser
        UI["AI menu<br/>(prompt, status, accept/reject)"]
        EXT["AI extension<br/>(Chat instance + state machine)"]
        DSB["Document state builder<br/>(blocks to HTML, IDs, cursor)"]
        EXEC["Stream tool executor<br/>(partial JSON to operations)"]
        PM["ProseMirror editor<br/>(suggestion marks, agent cursor)"]
    end
    subgraph Server
        API["POST /api/chat<br/>streamText + toolChoice required"]
    end
    LLM["LLM provider"]

    UI --> EXT
    EXT --> DSB
    DSB --> EXT
    EXT -- "messages + document state<br/>+ tool definitions" --> API
    API --> LLM
    LLM -- "streamed tool call input" --> API
    API -- "UI message stream (SSE)" --> EXT
    EXT --> EXEC
    EXEC -- "agent steps" --> PM
    PM --> UI
Zoom or open fullscreen for a closer look
visualized by

Walkthrough:

1. AI menu and extension. A thin UI over a small state machine with states closed, user-input, thinking, ai-writing, user-reviewing, plus an error state where retry() is available.
2. Document state builder. Before each request, the client serializes the document. The default format is HTML per block (aiDocumentFormats.html). JSON and Markdown formats exist in the code but are exported with an _experimental_ prefix. The builder includes block IDs and the cursor position, and, if text is selected, the selected blocks separately.
3. Server route. The backend is intentionally thin. BlockNote's docs show a Next.js route that calls streamText with a format-specific system prompt, injects the document state into the user messages, converts the client's serialized tool definitions into AI SDK tools, and sets toolChoice: "required" so the model cannot reply with plain prose. Tool calls are not executed on the server. They are forwarded to the client.
4. Stream tool executor. On the client, the tool call's input streams in. The executor repairs the partial JSON, figures out which operations are new or have grown, validates them, and hands them to per-operation executors (add, update, delete).
5. ProseMirror editor. Operations become ProseMirror transactions that insert insertion and deletion marks, plus metadata that drives the AI cursor.

The key boundary: the server never touches the document. It only adds the API key and system prompt; validation and edit logic live in the browser, next to the editor state. BlockNote even ships a ClientSideTransport that calls the provider directly from the browser, and a fetchViaProxy helper for routing provider requests through your server.

Deep dives

1. One tool, many operations, applied while streaming

The model gets exactly one tool, applyDocumentOperations, whose input is an object with an operations array. Each item has a type (add, update, delete) and the fields that operation needs. The system prompt tells the model to copy block IDs exactly, "including the trailing $".

That trailing $ is a small, clever detail. Block IDs are suffixed with $ before they go into the prompt (suffixIDs in the source), and the update validator rejects any ID that does not end in $. Think about what streaming JSON looks like: at some moment the model has emitted "id":"8f2a. A JSON repair function will happily close that string and give you id: "8f2a", which looks like a real value. The $ suffix turns "is this ID finished?" into a one-character check. The source does not state this reason, so treat it as my reading of the code.

Partial JSON handling uses the AI SDK's parsePartialJson, which returns a state: successful-parse, repaired-parse, failed-parse or undefined-input. A repaired-parse means "I had to close some brackets and strings for you", so the last operation is possibly partial. The executor tracks this with three flags on each operation:

• isPossiblyPartial: the last item in the array may still grow.
• isUpdateToPreviousOperation: this is the same operation as last time, with more content.
• the operation itself, already validated against its tool's schema.

If an operation fails validation mid-stream (no type yet, ID not finished), the code simply skips it with the comment "probably a partial operation that's not valid yet". If the stream ends while the last operation is still partial, the executor raises an error instead of applying a truncated edit.

The plumbing is plain Web Streams: one TransformStream turns JSON strings into operations, a second executes them.

2. Turning operations into "human-like" typing

Applying a partial update naively would be ugly. Imagine replacing a paragraph every time a new character arrives: the whole block would flicker. BlockNote does three things to avoid that.

Throttle by size. For a partial update, the executor only applies it once the serialized block has grown by at least 50 characters since the last application (STEP_SIZE = 50 in createUpdateBlockTool). Final chunks always apply.

Complete the HTML. The streamed block field is HTML, often cut off mid-tag like <p>hello <spa. A helper called getPartialHTML strips an unfinished tag or entity at the end, then runs the rest through DOMParser so the browser closes the open tags for you. Same trick as parsePartialJson, applied one level down.

Diff, then replay as an agent. The completed HTML is converted to blocks, diffed against the current block to get minimal ProseMirror ReplaceSteps, and each step is split into "agent steps": select the text being replaced, replace it with the first character, then insert the rest one character at a time. With delays on (the default, withDelays: true), select waits about 100ms, replace about 200ms and each inserted character about 10ms, each with a random jitter between 0.85x and 1.15x. That is why the AI cursor looks like it is typing instead of teleporting.

In React terms: the network delivers bursty updates, and the executor turns them into a smooth animation queue.

One more guard: if a partial update yields a single replace step, the executor waits for more content, so a translation is not rewritten sentence by sentence.

3. Suggestions, rebasing, undo and collaboration

Every AI edit is written as a suggestion, using marks from @handlewithcare/prosemirror-suggest-changes, a Google Docs-style suggestions library for ProseMirror. That library defines insertion, deletion and modification marks, and two commands: applySuggestions (keep insertions, remove deleted content) and revertSuggestions (remove insertions, keep deleted content). BlockNote's agent code adds these marks itself while replaying steps, so deleted text stays in the document with a strike-through until you decide.

The catch: the editor now contains deletion-marked text the model never saw, so positions no longer line up. BlockNote solves this with a rebaseTool: compute the "projection" (the document with suggestions applied), apply the model's operation to the projection, then map each step back onto the real document using the inverted mapping. If you have used git rebase, it is the same idea: replay changes made against one base onto a different base. The same tool is used for formats like Markdown that cannot represent every block feature.

Undo. Each agent step is dispatched with addToHistory: false, so hundreds of typing steps never reach the undo stack. On accept, the code captures the marked-up document, reverts suggestions to recover the original, then replaces the document with the accepted result as one history entry. The comment in acceptChanges spells out the goal: undo should go straight from the final state back to the original. On reject, suggestions are reverted, again without touching history.

Collaboration. When the editor is connected to a Yjs document, invokeAI first calls fork() on BlockNote's fork extension so AI modifications do not affect the shared document. Accept calls merge({ keepChanges: true }); reject calls merge({ keepChanges: false }). Your teammates never see the AI typing into their copy, only the accepted result.

Scroll. Auto-scroll follows the AI cursor and stops once a scroll event starts before the previous scrollend, which means a human is scrolling.

Diagram
100%
stateDiagram-v2
    [*] --> closed
    closed --> user_input: open AI menu
    user_input --> thinking: invokeAI(prompt)
    thinking --> ai_writing: first operation applied
    ai_writing --> user_reviewing: stream finished
    ai_writing --> error: invalid op or network error
    thinking --> error
    error --> thinking: retry()
    user_reviewing --> closed: acceptChanges() (fork merged)
    user_reviewing --> closed: rejectChanges() (fork discarded)
    ai_writing --> closed: abort()
Zoom or open fullscreen for a closer look
visualized by

Data flow: one "make this concise" request

Diagram
100%
sequenceDiagram
    participant U as User
    participant E as AI extension
    participant S as /api/chat
    participant M as LLM
    participant X as Stream tool executor
    participant P as ProseMirror

    U->>E: select 3 blocks, prompt "make this concise"
    E->>E: status thinking, fork Yjs doc
    E->>E: build document state (HTML, IDs with $, cursor)
    E->>S: POST messages + tool definitions
    S->>M: streamText(system prompt, tools, toolChoice required)
    M-->>S: tool-input deltas for applyDocumentOperations
    S-->>E: UI message stream (tool call parts)
    loop every chunk of tool input
        E->>X: accumulated JSON string
        X->>X: parsePartialJson, validate, new or updated ops
        X->>P: rebase, diff, agent steps (addToHistory false)
        P-->>U: AI cursor types, insertion and deletion marks
    end
    E->>U: status user-reviewing
    U->>E: Accept
    E->>P: apply suggestions as one history entry
    E->>E: merge fork into shared doc
Zoom or open fullscreen for a closer look
visualized by

Step by step:

1. The user submits a prompt. Status becomes thinking; in collaborative mode the Yjs document is forked.
2. The selected blocks and context are serialized to HTML, IDs get a $ suffix, and the cursor block is marked.
3. The client sends messages plus tool definitions; injectDocumentStateMessages adds the document state on the server.
4. The server calls streamText with toolChoice: "required" and returns a UI message stream.
5. The model streams one applyDocumentOperations input. The executor repairs the JSON, drops operations that are not valid yet, and emits new or grown ones.
6. Each operation is rebased, diffed into ProseMirror steps and replayed as agent steps with insertion and deletion marks.
7. When the stream ends, tool results go back into the chat (for follow-ups) and status becomes user-reviewing. Accept applies suggestions in one undoable transaction and merges the fork; reject reverts both.

Trade-offs

Confirmed vs inferred

Confirmed by sources (BlockNote docs, @blocknote/xl-ai 0.57.0 source, Notion and Tiptap docs):

• BlockNote AI is built on the Vercel AI SDK, is model-agnostic, and is labelled early preview.
• The backend uses streamText with toolChoice: "required", injects document state into user messages, and forwards tool calls to the client.
• HTML is the default format; JSON and Markdown are exported as experimental.
• IDs are suffixed with $ and the update validator rejects IDs without it.
• parsePartialJson parsing, skipping of invalid partial operations, and the error on a stream ending mid-operation.
• The 50-character throttle, DOMParser HTML completion, and agent delays (about 100ms, 200ms, 10ms, 0.85x to 1.15x jitter, withDelays on by default).
• Suggestion marks via prosemirror-suggest-changes, history-free agent steps, the accept and reject logic, and the Yjs fork and merge.
• Notion: "Edit with AI" on a selection or an agent that writes on the page; accept, discard or try again.
• Tiptap's AI Toolkit streams HTML with streamHtml and offers tracked changes, preview and review modes.

Inferred or reconstructed:

• Anything about Notion's internal architecture, including whether it uses tool calls, HTML, suggestions or a fork model. Notion has not published this.
• The reason for the $ suffix (detecting that a streamed ID is complete).
• That artificial delays exist mainly for readability rather than to hide network burstiness.
• That a production product would trim document context for cost. I did not measure token usage.

Build a mini version in JS/TS

The smallest working version, using AI SDK 7 (ai 7.x, which needs Node.js 22 or newer) with plain text blocks instead of HTML, and React state instead of ProseMirror. It keeps three ideas: one forced tool call, partial JSON applied while streaming, and drafts held until accepted.

Shared types and the tool schema:

// lib/ops.ts
import { z } from 'zod';

export const operationSchema = z.object({
  type: z.enum(['update', 'add']),
  id: z
    .string()
    .describe('Existing block id, copied exactly, including the trailing $'),
  position: z.enum(['before', 'after']).nullable().describe('Only for add'),
  text: z.string(),
});

export const editSchema = z.object({ operations: z.array(operationSchema) });

export type Operation = z.infer<typeof operationSchema>;
export type Block = {
  id: string;
  text: string;
  draft?: string;
  status?: 'inserted' | 'changed';
};

The route forces the tool call and forwards only the tool input deltas as plain text:

// app/api/ai-edit/route.ts
import {
  streamText,
  tool,
  createTextStreamResponse,
  type TextStreamPart,
  type ToolSet,
} from 'ai';
import { openai } from '@ai-sdk/openai';
import { editSchema, type Block } from '@/lib/ops';

export async function POST(req: Request) {
  const { prompt, blocks, cursorId } = (await req.json()) as {
    prompt: string;
    blocks: Block[];
    cursorId: string;
  };
  const doc = blocks.map((b) => `${b.id}$: ${b.text}`).join('\n');

  const result = streamText({
    model: openai('gpt-4.1'),
    instructions:
      'You edit a document made of blocks. Always answer with applyDocumentOperations. Copy ids exactly, including the trailing $.',
    prompt: `Document:\n${doc}\n\nCursor block: ${cursorId}$\n\nRequest: ${prompt}`,
    tools: {
      applyDocumentOperations: tool({
        description: 'Add or update blocks',
        inputSchema: editSchema,
      }),
    },
    toolChoice: 'required',
    abortSignal: req.signal,
  });

  const deltas = result.stream.pipeThrough(
    new TransformStream<TextStreamPart<ToolSet>, string>({
      transform(part, controller) {
        if (part.type === 'tool-input-delta') controller.enqueue(part.delta);
      },
    }),
  );
  return createTextStreamResponse({ stream: deltas });
}

On the client, repair the growing JSON and re-emit the last operation every time it grows:

// lib/stream-operations.ts
import { parsePartialJson } from 'ai';
import type { Operation } from './ops';

export type OpEvent = { op: Partial<Operation>; index: number };

export async function* streamOperations(
  res: Response,
): AsyncGenerator<OpEvent> {
  const reader = res.body!.pipeThrough(new TextDecoderStream()).getReader();
  let json = '';
  let seen = 0;
  while (true) {
    const { done, value } = await reader.read();
    if (done) return;
    json += value;
    const { value: parsed, state } = await parsePartialJson(json);
    if (state === 'failed-parse' || state === 'undefined-input') continue;
    const ops =
      (parsed as { operations?: Partial<Operation>[] })?.operations ?? [];
    // The previous last op may have grown, so start one step back.
    for (let i = Math.max(0, seen - 1); i < ops.length; i++) {
      yield { op: ops[i], index: i };
    }
    seen = ops.length;
  }
}

Apply operations as drafts with a pure function (safe under React Strict Mode, which can call state updaters twice):

// lib/apply-op.ts
import type { Block, Operation } from './ops';

export function applyOp(
  doc: Block[],
  op: Partial<Operation>,
  index: number,
  runId: string,
): Block[] {
  if (!op.id?.endsWith('$') || op.text === undefined) return doc; // id still streaming
  const targetId = op.id.slice(0, -1);
  if (op.type === 'update') {
    return doc.map((b) =>
      b.id === targetId ? { ...b, draft: op.text, status: 'changed' } : b,
    );
  }
  if (
    op.type !== 'add' ||
    (op.position !== 'before' && op.position !== 'after')
  )
    return doc;

  const newId = `${runId}-${index}`;
  if (doc.some((b) => b.id === newId)) {
    return doc.map((b) => (b.id === newId ? { ...b, draft: op.text } : b));
  }
  const at = doc.findIndex((b) => b.id === targetId);
  if (at < 0) return doc;
  const insertAt = op.position === 'after' ? at + 1 : at;
  const block: Block = {
    id: newId,
    text: '',
    draft: op.text,
    status: 'inserted',
  };
  return [...doc.slice(0, insertAt), block, ...doc.slice(insertAt)];
}

export const accept = (doc: Block[]): Block[] =>
  doc.map((b) => ({ id: b.id, text: b.draft ?? b.text }));

export const reject = (doc: Block[]): Block[] =>
  doc
    .filter((b) => b.status !== 'inserted')
    .map((b) => ({ id: b.id, text: b.text }));

And a minimal editor that renders the diff:

// components/ai-editor.tsx
'use client';
import { useState } from 'react';
import { applyOp, accept, reject } from '@/lib/apply-op';
import { streamOperations } from '@/lib/stream-operations';
import type { Block } from '@/lib/ops';

export function AiEditor({ initial }: { initial: Block[] }) {
  const [doc, setDoc] = useState(initial);
  const [phase, setPhase] = useState<'idle' | 'writing' | 'reviewing'>('idle');

  async function run(prompt: string, cursorId: string) {
    setPhase('writing');
    const runId = crypto.randomUUID();
    const res = await fetch('/api/ai-edit', {
      method: 'POST',
      body: JSON.stringify({ prompt, blocks: doc, cursorId }),
    });
    for await (const { op, index } of streamOperations(res)) {
      setDoc((d) => applyOp(d, op, index, runId));
    }
    setPhase('reviewing');
  }

  return (
    <div>
      {doc.map((b) => (
        <p key={b.id}>
          {b.draft !== undefined && b.text && <del>{b.text}</del>}{' '}
          {b.draft !== undefined ? <ins>{b.draft}</ins> : b.text}
        </p>
      ))}
      <button
        disabled={phase !== 'idle'}
        onClick={() => run('Make this more concise', doc[0].id)}>
        Make concise
      </button>
      {phase === 'reviewing' && (
        <>
          <button
            onClick={() => {
              setDoc(accept);
              setPhase('idle');
            }}>
            Accept
          </button>
          <button
            onClick={() => {
              setDoc(reject);
              setPhase('idle');
            }}>
            Reject
          </button>
        </>
      )}
    </div>
  );
}

Next steps: a real editor model with marks instead of draft fields, an AbortController behind a Stop button, and a word-level diff so <del> and <ins> cover only what changed.

Interview angle

If the prompt is "Design Notion AI's 'edit with AI' feature" in a frontend system design round, here is a 6-step answer:

1. Clarify scope. Generate at cursor vs transform a selection, single user vs collaborative, review required or auto-apply. Name the non-functional goals: time to first visible edit, valid document at every frame, clean undo.
2. Choose the output contract. Explain why free text is not enough and propose one forced tool call carrying a list of add, update and delete operations keyed by block ID. Mention the serialization format (HTML or Markdown) and its token cost.
3. Draw the architecture. Thin server proxy (streamText, API key, system prompt), client-side executor next to the editor state. Justify why tool calls execute on the client.
4. Deep dive the streaming client. Partial JSON repair, detecting complete operations, throttling partial updates, completing partial HTML, and replaying diffs as small steps so the UI does not flicker.
5. Handle review and safety. Suggestion marks for a visible diff, a rebase step so model positions map onto a marked-up document, history-free intermediate steps, one-step accept, and a forked collaborative document.
6. Cover failure modes. Hallucinated IDs (validate and skip), stream ending mid-operation (error, keep the original), user edits during writing (lock the editor or rebase), abort and retry.

Sources

• BlockNote AI overview, BlockNote docs, accessed 2 November 2026
• BlockNote AI: Backend Integration, BlockNote docs, accessed 2 November 2026
• BlockNote AI Reference, BlockNote docs, accessed 2 November 2026
• BlockNote source repository (packages/xl-ai), read from the published @blocknote/xl-ai 0.57.0 package, released 30 October 2026
• @blocknote/xl-ai on npm, version 0.57.0
• BlockNote PR #1903: custom AI executors, June to August 2025 (history of the backend execution model)
• prosemirror-suggest-changes, Handle with Care Collective, accessed 2 November 2026
• Use Notion AI to write better, more efficient notes and docs, Notion Help Center, accessed 2 November 2026
• Tiptap AI Toolkit: Insert content workflow, Tiptap docs, accessed 2 November 2026
• AI SDK 7 migration guide, Vercel, accessed 2 November 2026
• AI SDK 6 migration guide, Vercel, accessed 2 November 2026
• ai package on npm, version 7.x (7.0.0 released 25 June 2026)