@dxos/plugin-assistant 0.10.0 → 0.11.1

package/src/components/ChatThread/ChatThread.tsx

@@ -4,8 +4,8 @@
 
 import React, { forwardRef, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 
+import { Identity } from '@dxos/halo';
 import { PublicKey } from '@dxos/keys';
-import { type Identity } from '@dxos/react-client/halo';
 import { type ThemedClassName, setRef } from '@dxos/react-ui';
 import { MarkdownStream, type MarkdownStreamController, type MarkdownStreamProps } from '@dxos/react-ui-markdown';
 import { type Message } from '@dxos/types';
@@ -14,7 +14,7 @@ import { keyToFallback } from '@dxos/util';
 import { type Assistant } from '../../types';
 import { type ChatEvent } from '../Chat';
 import { componentRegistry, createBlockRenderer } from './registry';
-import { MessageSyncer } from './sync';
+import { type MessageSpan, MessageSyncer } from './sync';
 
 const defaultOptions: MarkdownStreamProps['options'] = {
   autoScroll: true,
@@ -25,11 +25,13 @@ const defaultOptions: MarkdownStreamProps['options'] = {
 
 export type ChatThreadProps = ThemedClassName<
   {
-    identity?: Identity;
+    identity?: Identity.Info;
     messages?: Message.Message[];
     error?: Error;
     viewType?: Assistant.ChatView;
     onEvent?: (event: ChatEvent) => void;
+    /** Publishes the syncer's per-message document spans after each update (minimap, prompt nav). */
+    onSpans?: (spans: MessageSpan[]) => void;
   } & Pick<MarkdownStreamProps, 'options' | 'debug' | 'extensions' | 'footer'>
 >;
 
@@ -46,6 +48,7 @@ export const ChatThread = forwardRef<MarkdownStreamController | null, ChatThread
       extensions,
       viewType,
       onEvent,
+      onSpans,
     },
     forwardedRef,
   ) => {
@@ -59,7 +62,9 @@ export const ChatThread = forwardRef<MarkdownStreamController | null, ChatThread
     );
 
     const userHue = useMemo(
-      () => identity?.profile?.data?.hue || keyToFallback(identity?.identityKey ?? PublicKey.random()).hue,
+      () =>
+        identity?.data?.hue ||
+        keyToFallback(identity?.identityKey ? PublicKey.fromHex(identity.identityKey) : PublicKey.random()).hue,
       [identity],
     );
 
@@ -77,18 +82,41 @@ export const ChatThread = forwardRef<MarkdownStreamController | null, ChatThread
       initializedRef.current = true;
     }, [controller, error]);
 
+    // Widget callbacks are routed through the syncer's context, so they must be referentially stable
+    // — rebuilding the syncer replaces the document.
+    const onEventRef = useRef(onEvent);
+    useEffect(() => {
+      onEventRef.current = onEvent;
+    }, [onEvent]);
+    const handleRewind = useCallback((id: string) => onEventRef.current?.({ type: 'rewind', id }), []);
+
     // Update document.
     const renderer = useMemo(() => createBlockRenderer(viewType), [viewType]);
-    const syncer = useMemo(() => controller && new MessageSyncer(controller, renderer), [controller, renderer]);
+    const syncer = useMemo(
+      () => controller && new MessageSyncer(controller, renderer, { onRewind: handleRewind }),
+      [controller, renderer, handleRewind],
+    );
     useEffect(() => {
       if (!syncer) {
         return;
       }
 
+      // Publish the context every pass: widget props read it from a CodeMirror state field, and
+      // `setContext` dispatches through `viewRef.current?` — a no-op before the view exists, so doing this
+      // once on mount silently loses it and every widget callback (e.g. rewind) dies on the optional call.
+      controller?.setContext(syncer.context);
+
+      // Publish the context every pass: widget props read it from a CodeMirror state field, and
+      // `setContext` dispatches through `viewRef.current?` — a no-op before the view exists, so doing this
+      // once on mount silently loses it and every widget callback (e.g. rewind) dies on the optional call.
+      controller?.setContext(syncer.context);
+
       if (syncer.update(messages)) {
         controller?.scrollToBottom('instant');
       }
-    }, [controller, syncer, messages]);
+      // Spans are valid synchronously after `update` (offsets are computed during the walk).
+      onSpans?.(syncer.getSpans());
+    }, [controller, syncer, messages, onSpans]);
 
     // Event adapter.
     const handleEvent = useCallback<NonNullable<MarkdownStreamProps['onEvent']>>(

package/src/components/ChatThread/Link.tsx

@@ -14,6 +14,7 @@ export type ObjectLinkProps = {
   dxn: URI.URI;
 };
 
+// TODO(burdon): Document.
 export const ObjectLink = ({ space, dxn }: ObjectLinkProps) => {
   const ref = useMemo(() => space.db.makeRef(dxn), [space, dxn]);
 

package/src/components/ChatThread/index.ts

@@ -2,4 +2,6 @@
 // Copyright 2025 DXOS.org
 //
 
+export type { MessageSpan } from './sync';
+
 export * from './ChatThread';

package/src/components/ChatThread/registry.tsx

@@ -9,9 +9,9 @@ import { ContentBlock, type Message } from '@dxos/types';
 import { AnchorWidget, type XmlWidgetRegistry, getXmlTextChild } from '@dxos/ui-editor';
 
 import { type Assistant } from '../../types';
-import { type BlockRenderer, type MessageThreadContext } from './sync';
-import { applyToolBlockToWidgetState } from './tool-widget-state';
+import { type BlockRenderer, type MessageThreadContext, applyToolBlockToWidgetState } from './sync';
 import {
+  BranchWidget,
   FallbackWidget,
   ReasoningWidget,
   ReferenceWidget,
@@ -88,8 +88,9 @@ export const componentRegistry: XmlWidgetRegistry = {
       return options?.length ? new SelectWidget(options) : null;
     },
   },
+  // Inline so a run of suggestions flows onto one wrapped line instead of one block per line.
   'suggestion': {
-    block: true,
+    block: false,
     factory: ({ children }) => {
       const text = getXmlTextChild(children);
       return text ? new SuggestionWidget(text) : null;
@@ -115,6 +116,10 @@ export const componentRegistry: XmlWidgetRegistry = {
   // React Widgets (portaled outside of the editor)
   //
 
+  'branch': {
+    block: true,
+    Component: BranchWidget,
+  },
   'summary': {
     block: true,
     Component: SummaryWidget,
@@ -168,8 +173,16 @@ export function createBlockRenderer(viewType: Assistant.ChatView | undefined): B
     if (!isBlockVisible(viewType, message, block)) {
       return;
     }
+
     let str = blockToMarkdownImpl(context, message, block);
     if (str && !block.pending) {
+      // Suggestions are inline widgets meant to flow: keeping them on one line lets a run of them wrap
+      // as chips instead of stacking one per line. A trailing space separates adjacent tags so the
+      // inline parser sees them as distinct.
+      if (block._tag === 'suggestion') {
+        return (str += ' ');
+      }
+
       // Use a blank line as the block separator so each rendered block parses as its own
       // markdown block. A single newline lets CommonMark absorb a following `<prompt>` (an
       // HTML type-7 tag, which can't interrupt an open paragraph) into the previous
@@ -208,7 +221,11 @@ const blockToMarkdownImpl = (context: MessageThreadContext, message: Message.Mes
         if (block.disposition === 'synthetic') {
           return renderXMLBlock('synthetic', { content: block.text, pending: block.pending });
         } else {
-          return `\n<prompt>${block.text}</prompt>`;
+          // Built by explicit concatenation, not an indented template: `trim` dedents by the minimum
+          // indent across all lines, and a multi-line prompt contributes lines at zero indent — so the
+          // source file's indentation survived into the document, where 4+ spaces is an indented code
+          // block and neither the prompt nor the toolbar parsed as an element any more.
+          return `\n<prompt>${block.text}</prompt>\n${renderBranchToolbar(message)}\n`;
         }
       } else {
         const text = block.text.trim();
@@ -291,6 +308,13 @@ const blockToMarkdownImpl = (context: MessageThreadContext, message: Message.Mes
   }
 };
 
+/**
+ * Mini toolbar below a user prompt: branch (soft fork) plus when the prompt was sent.
+ * `messageId` rather than `id` because a tag's `id` attribute would shadow the widget's own id.
+ */
+const renderBranchToolbar = (message: Message.Message) =>
+  `<branch messageId="${escapeXmlAttribute(message.id)}" created="${escapeXmlAttribute(message.created)}" />`;
+
 /**
  * Escape text embedded in generated XML so the mixed XML parser does not treat `&`, `<`, `>` as markup.
  */

package/src/components/ChatThread/sync/index.ts

@@ -0,0 +1,6 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+export * from './sync';
+export { applyToolBlockToWidgetState } from './tool-widget-state';

package/src/components/ChatThread/{sync.test.ts → sync/sync.test.ts}

@@ -5,7 +5,7 @@
 // @vitest-environment jsdom
 
 import { EditorView } from '@codemirror/view';
-import { describe, it } from '@effect/vitest';
+import { describe, expect, it } from '@effect/vitest';
 import * as Effect from 'effect/Effect';
 
 import { Obj } from '@dxos/echo';
@@ -14,8 +14,8 @@ import { type ContentBlock, type Message } from '@dxos/types';
 
 import { createMessage } from '#testing';
 
-import { createBlockRenderer } from './registry';
-import { type BlockRenderer, MessageSyncer, type MessageThreadContext, type TextModel } from './sync';
+import { createBlockRenderer } from '../registry';
+import { type BlockRenderer, MessageSyncer, type MessageThreadContext, type TextModel } from '../sync/sync';
 
 class TestDocument implements TextModel {
   private readonly _view = new EditorView({ extensions: [] });
@@ -56,13 +56,13 @@ describe('reducers', () => {
       ];
 
       syncer.update(messages);
-      expect(doc.content).toEqual('\n<prompt>Hello</prompt>\nHi there!\n');
+      expect(normalize(doc.content)).toEqual('\n<prompt>Hello</prompt>\n<branch />\n\nHi there!\n');
 
       Obj.update(messages[1], (obj) => {
         obj.blocks.push({ _tag: 'text', text: 'How can I help?' });
       });
       syncer.update(messages);
-      expect(doc.content).toEqual('\n<prompt>Hello</prompt>\nHi there!\nHow can I help?\n');
+      expect(normalize(doc.content)).toEqual('\n<prompt>Hello</prompt>\n<branch />\n\nHi there!\nHow can I help?\n');
     }),
   );
 
@@ -78,7 +78,7 @@ describe('reducers', () => {
       ];
 
       syncer.update(messages);
-      expect(doc.content).toEqual('\n<prompt>Hello</prompt>\nHi there!');
+      expect(normalize(doc.content)).toEqual('\n<prompt>Hello</prompt>\n<branch />\n\nHi there!');
 
       Obj.update(messages[1], (obj) => {
         const block = obj.blocks[0] as Mutable<ContentBlock.Text>;
@@ -91,7 +91,43 @@ describe('reducers', () => {
         obj.blocks.push({ _tag: 'text', text: 'How can I help?' });
       });
       syncer.update(messages);
-      expect(doc.content).toEqual('\n<prompt>Hello</prompt>\nHi there! How are you?\nHow can I help?\n');
+      expect(normalize(doc.content)).toEqual(
+        '\n<prompt>Hello</prompt>\n<branch />\n\nHi there! How are you?\nHow can I help?\n',
+      );
+    }),
+  );
+
+  it.effect(
+    'tracks per-message document ranges',
+    Effect.fn(function* ({ expect }) {
+      const doc = new TestDocument();
+      const syncer = new MessageSyncer(doc, createBlockRenderer('thinking'));
+
+      const messages = [
+        createMessage('user', [{ _tag: 'text', text: 'Hello' }]),
+        createMessage('assistant', [{ _tag: 'text', text: 'Hi there!' }]),
+      ];
+
+      syncer.update(messages);
+      const spans = syncer.getSpans();
+      expect(spans.map((range) => range.id)).toEqual([messages[0].id, messages[1].id]);
+      // Ranges tile the document contiguously and slice back to each message's rendered content.
+      expect(spans[0].from).toEqual(0);
+      expect(spans[1].from).toEqual(spans[0].to);
+      expect(spans[1].to).toEqual(doc.length);
+      expect(doc.content.slice(spans[0].from, spans[0].to)).toContain('<prompt>Hello</prompt>');
+      expect(doc.content.slice(spans[1].from, spans[1].to)).toContain('Hi there!');
+
+      // Appending to the last message extends its range through the append path.
+      Obj.update(messages[1], (obj) => {
+        obj.blocks.push({ _tag: 'text', text: 'How can I help?' });
+      });
+      syncer.update(messages);
+      const extended = syncer.getSpans();
+      expect(extended).toHaveLength(2);
+      expect(extended[0].to).toEqual(spans[0].to);
+      expect(extended[1].to).toEqual(doc.length);
+      expect(doc.content.slice(extended[1].from, extended[1].to)).toContain('How can I help?');
     }),
   );
 
@@ -255,3 +291,127 @@ describe('reducers', () => {
     }),
   );
 });
+
+describe('MessageSyncer tool widget rehydration', () => {
+  it.effect(
+    'restores tool widget state after reset replaces the document',
+    Effect.fnUntraced(function* (_) {
+      const updates: { id: string; value: any }[] = [];
+      class RecordingDocument extends TestDocument {
+        override updateWidget(id: string, value: any) {
+          updates.push({ id, value });
+        }
+      }
+
+      const document = new RecordingDocument();
+      const syncer = new MessageSyncer(document, createBlockRenderer('thinking'));
+      const messages = [
+        createMessage('assistant', [
+          { _tag: 'toolCall', toolCallId: 'abc', name: 'search', input: '{}', providerExecuted: false },
+        ]),
+        createMessage('user', [
+          { _tag: 'toolResult', toolCallId: 'abc', name: 'search', result: 'ok', providerExecuted: false },
+        ]),
+      ] as Message.Message[];
+
+      syncer.reset(messages);
+      // `reset` rehydrates in a promise continuation after `setContent`.
+      yield* Effect.promise(() => Promise.resolve());
+      yield* Effect.promise(() => Promise.resolve());
+
+      // Reduce the dispatches the way the widget state store does: each is either a value or a
+      // function of the previous state, applied in arrival order.
+      const forTool = updates.filter((update) => update.id === 'abc');
+      expect(forTool.length).toBeGreaterThan(0);
+      const state = forTool.reduce<{ blocks: ContentBlock.Any[] }>(
+        (previous, { value }) => (typeof value === 'function' ? value(previous) : value),
+        { blocks: [] },
+      );
+      expect(state.blocks.map((block) => block._tag)).toEqual(['toolCall', 'toolResult']);
+    }),
+  );
+
+  // Regression: the prompt was built with an indented `trim` template, which dedents by the minimum
+  // indent across all lines — so a multi-line prompt, contributing lines at zero indent, left the source
+  // file's indentation in the document. At 4+ spaces CommonMark reads an indented code block, so neither
+  // the prompt nor the toolbar parsed as an element and the turn rendered as raw text.
+  it.effect(
+    'a multi-line prompt emits unindented markup',
+    Effect.fn(function* ({ expect }) {
+      const doc = new TestDocument();
+      const syncer = new MessageSyncer(doc, createBlockRenderer('thinking'));
+
+      const prompt = createMessage('user', [{ _tag: 'text', text: 'this is a list:\n- foo\n- bar' }]);
+      syncer.update([prompt]);
+
+      const indented = doc.content.split('\n').filter((line) => /^\s+\S/.test(line));
+      expect(indented).toEqual([]);
+      expect(normalize(doc.content)).toContain('<prompt>this is a list:\n- foo\n- bar</prompt>\n<branch />');
+    }),
+  );
+
+  // Regression: `update` streams a suffix, so an append-only path leaves removed turns on screen. A
+  // rewind shrinks the thread, which has to fall back to a full reset.
+  it.effect(
+    'a shrinking message list replaces the document',
+    Effect.fn(function* ({ expect }) {
+      const doc = new TestDocument();
+      const syncer = new MessageSyncer(doc, createBlockRenderer('thinking'));
+
+      const prompt = createMessage('user', [{ _tag: 'text', text: 'Hello' }]);
+      const reply = createMessage('assistant', [{ _tag: 'text', text: 'Hi there!' }]);
+      syncer.update([prompt, reply]);
+      expect(doc.content).toContain('Hi there!');
+
+      // The rewind case: same first message, fewer turns.
+      const replaced = syncer.update([prompt]);
+      expect(replaced).toBe(true);
+      expect(doc.content).not.toContain('Hi there!');
+      expect(doc.content).toContain('Hello');
+    }),
+  );
+
+  it.effect(
+    'a message list diverging mid-thread replaces the document',
+    Effect.fn(function* ({ expect }) {
+      const doc = new TestDocument();
+      const syncer = new MessageSyncer(doc, createBlockRenderer('thinking'));
+
+      const prompt = createMessage('user', [{ _tag: 'text', text: 'Hello' }]);
+      const abandoned = createMessage('assistant', [{ _tag: 'text', text: 'Abandoned.' }]);
+      syncer.update([prompt, abandoned]);
+
+      // Same length, different second turn — the branch changed underneath us.
+      const retry = createMessage('assistant', [{ _tag: 'text', text: 'Retry.' }]);
+      expect(syncer.update([prompt, retry])).toBe(true);
+      expect(doc.content).not.toContain('Abandoned.');
+      expect(doc.content).toContain('Retry.');
+    }),
+  );
+
+  // Regression: widget callbacks reach the thread through the syncer's context, which the host has to
+  // publish to the editor via `setContext`. Nothing did, so a widget received `context: undefined` and
+  // the rewind button silently no-oped.
+  it.effect(
+    'exposes a context routing widget callbacks to the handlers',
+    Effect.fn(function* ({ expect }) {
+      const rewound: string[] = [];
+      const syncer = new MessageSyncer(new TestDocument(), createBlockRenderer('thinking'), {
+        onRewind: (id) => rewound.push(id),
+      });
+
+      expect(syncer.context).toBeDefined();
+      syncer.context.rewind('msg-1');
+      expect(rewound).toEqual(['msg-1']);
+    }),
+  );
+});
+
+/**
+ * Canonicalizes the rendered document for comparison:
+ * - the branch toolbar's attributes, since `messageId` is a fresh ULID and `created` a wall-clock
+ *   timestamp, so neither is stable across runs;
+ * - the run of leading newlines, which is presentational padding above the first block and is tuned
+ *   independently — these tests are about the block sequence and the incremental append path.
+ */
+const normalize = (content: string) => content.replace(/<branch\b[^>]*\/>/g, '<branch />').replace(/^\n+/, '\n');

package/src/components/ChatThread/{sync.ts → sync/sync.ts}

@@ -14,6 +14,17 @@ import { rehydrateToolWidgetsFromMessages } from './tool-widget-state';
  */
 export type TextModel = Pick<MarkdownStreamController, 'length' | 'setContent' | 'append' | 'updateWidget'>;
 
+/**
+ * Document offset range occupied by a message's rendered blocks. Positions are in the same
+ * space as the CodeMirror document (i.e. what {@link MarkdownStreamController.scrollTo} and
+ * {@link MarkdownStreamController.getVisibleRange} operate on).
+ */
+export type MessageSpan = {
+  id: string;
+  from: number;
+  to: number;
+};
+
 /**
  * Renders a block to markdown.
  *
@@ -29,11 +40,26 @@ export type BlockRenderer = (
   block: ContentBlock.Any,
 ) => string | undefined;
 
+/**
+ * Callbacks a widget can invoke through {@link MessageThreadContext}.
+ *
+ * React widgets are portaled by `MarkdownStream` without its `onEvent`, so interactive widgets reach
+ * the thread through the context they already receive in their props.
+ */
+export type MessageThreadHandlers = {
+  /** Soft-fork the thread from the given message. */
+  onRewind?: (messageId: string) => void;
+};
+
 /**
  * Thread context passed to renderer.
+ * This enables the renderer to "stream" content into the widget state.
  */
 export class MessageThreadContext implements Pick<MarkdownStreamController, 'updateWidget'> {
-  constructor(private readonly _widgetState?: XmlWidgetStateManager) {}
+  constructor(
+    private readonly _widgetState?: XmlWidgetStateManager,
+    private readonly _handlers: MessageThreadHandlers = {},
+  ) {}
 
   updateWidget<T>(id: string, value: StateDispatch<T>) {
     this._widgetState?.updateWidget(id, value);
@@ -43,6 +69,10 @@ export class MessageThreadContext implements Pick<MarkdownStreamController, 'upd
   getObjectLabel(_id: URI.URI) {
     return 'Object';
   }
+
+  rewind(messageId: string) {
+    this._handlers.onRewind?.(messageId);
+  }
 }
 
 /**
@@ -61,30 +91,60 @@ export class MessageThreadContext implements Pick<MarkdownStreamController, 'upd
  */
 export class MessageSyncer {
   private _threadId?: string;
+
+  /**
+   * Ids of the messages rendered so far, in document order. `update` streams a suffix, so it needs to
+   * know when the incoming list no longer extends what is on screen — a rewind drops turns, and an
+   * append-only path would leave them rendered.
+   */
+  private _renderedIds: string[] = [];
+
+  /** Cumulative block index (across all completed blocks in all messages). */
   private _completed = 0;
+
+  /** Chars of the in-flight block (at index `_completed`) already appended. */
   private _trailing = 0;
 
+  /** Document offset at the `_completed` boundary (length of all fully-appended blocks). */
+  private _completedOffset = 0;
+
+  /** Per-message document offset ranges, keyed by message id in document order. */
+  private readonly _spans = new Map<string, { from: number; to: number }>();
+
   private readonly _context: MessageThreadContext;
 
   constructor(
     private readonly _document: TextModel,
     private readonly _renderer: BlockRenderer,
+    handlers: MessageThreadHandlers = {},
   ) {
-    this._context = new MessageThreadContext(this._document);
+    this._context = new MessageThreadContext(this._document, handlers);
   }
 
   get context() {
     return this._context;
   }
 
+  /**
+   * Per-message document offset spans, in document order. Valid synchronously after
+   * {@link reset} or {@link update} (the offsets are derived from the same rendered buffer
+   * that is dispatched to the document).
+   */
+  getSpans(): MessageSpan[] {
+    return Array.from(this._spans, ([id, { from, to }]) => ({ id, from, to }));
+  }
+
   /**
    * Replace the document with the rendering of `messages`. Use on mount, on thread switch,
    * and from {@link update} when it detects an identity change in `messages[0]`.
    */
   reset(messages: Message.Message[] = []): void {
     this._threadId = messages[0]?.id;
+    this._renderedIds = messages.map((message) => message.id);
     this._completed = 0;
     this._trailing = 0;
+    this._completedOffset = 0;
+    this._spans.clear();
     const buffer = this._walk(messages);
     // Match the pre-rewrite behaviour: rendering from a steady state (initial mount with
     // non-empty messages, or thread switch) lands the entire content via `setContent` — which
@@ -97,16 +157,30 @@ export class MessageSyncer {
     });
   }
 
+  /**
+   * Whether `messages` still begins with everything already rendered, making this an append.
+   *
+   * A rewind removes turns from the middle or end of the thread, and the streaming path can only add
+   * text — so anything other than an extension has to go through {@link reset}.
+   */
+  #extendsRendered(messages: Message.Message[]): boolean {
+    if (messages.length < this._renderedIds.length) {
+      return false;
+    }
+    return this._renderedIds.every((id, index) => messages[index]?.id === id);
+  }
+
   /**
    * Stream the suffix of the rendered messages into the document.
    * Returns `true` if the document was replaced (initial mount or thread switch), `false`
    * if the call was a streaming append (or a no-op).
    */
   update(messages: Message.Message[]): boolean {
-    if (messages[0]?.id !== this._threadId) {
+    if (messages[0]?.id !== this._threadId || !this.#extendsRendered(messages)) {
       this.reset(messages);
       return true;
     }
+    this._renderedIds = messages.map((message) => message.id);
     const buffer = this._walk(messages);
     if (buffer.length > 0) {
       void this._document.append(buffer);
@@ -122,6 +196,9 @@ export class MessageSyncer {
   _walk(messages: Message.Message[]): string {
     let buffer = '';
     let index = 0;
+    // Absolute document offset at the `_completed` boundary; blocks before it were rendered on
+    // earlier calls and their offsets are already baked into `_completedOffset` (and `_ranges`).
+    let offset = this._completedOffset;
     outer: for (const message of messages) {
       for (const block of message.blocks) {
         if (index < this._completed) {
@@ -132,6 +209,11 @@ export class MessageSyncer {
         if (rendered.length > this._trailing) {
           buffer += rendered.slice(this._trailing);
         }
+        // The document never shrinks (see the `_trailing` guard below), so span/offset accounting
+        // must track the document, not a renderer output that regressed below what was appended.
+        const renderedLength = Math.max(rendered.length, this._trailing);
+        // The block occupies `[offset, offset + renderedLength)`; extend the message's range.
+        this._saveRange(message.id, offset, offset + renderedLength);
         if (block.pending) {
           // Stay on this block; record how far we've appended so the next call can resume.
           // `Math.max`-style guard against a non-monotonic renderer output without shrinking the doc.
@@ -142,9 +224,21 @@ export class MessageSyncer {
         }
         this._completed = index + 1;
         this._trailing = 0;
+        offset += renderedLength;
+        this._completedOffset = offset;
         index++;
       }
     }
     return buffer;
   }
+
+  /** Record (or extend) the offset range of a message. `from` is preserved; `to` never regresses. */
+  private _saveRange(id: string, from: number, to: number): void {
+    const existing = this._spans.get(id);
+    if (existing) {
+      existing.to = Math.max(existing.to, to);
+    } else {
+      this._spans.set(id, { from, to });
+    }
+  }
 }

package/src/components/ChatThread/{tool-widget-state.ts → sync/tool-widget-state.ts}

@@ -23,15 +23,13 @@ export const applyToolBlockToWidgetState = (context: ToolWidgetStateSink, block:
       });
       break;
     }
+
     case 'toolResult': {
       context.updateWidget<{ blocks: ContentBlock.Any[] }>(block.toolCallId, ({ blocks = [] } = { blocks: [] }) => ({
         blocks: [...blocks, block],
       }));
       break;
     }
-    default: {
-      break;
-    }
   }
 };