fix(chat): reclassify replaced content during stage playback

Author: blove

apps/website/e2e/home-stage.spec.ts

@@ -43,6 +43,23 @@ const progress = (page: Page) =>
   );
 
 test.describe('homepage stage', () => {
+  test('gradual playback replaces the render introduction with an editable surface', async ({ page }) => {
+    test.skip(process.env['STAGE_LIVE_FRAME'] !== 'true', 'requires the matching stage replay deployment');
+    await page.setViewportSize({ width: 1440, height: 900 });
+    await page.goto('/');
+    const act = page.locator('[data-stage-act]');
+    await expect(act).toHaveAttribute('data-state', 'ready');
+    const frame = page.frameLocator('.stage-frame-iframe');
+    for (const p of [0.87, 0.89, 0.91, 0.93, 0.95, 1]) {
+      await scrollAct(page, p);
+      await expect(act).toHaveAttribute('data-interactive', '', { timeout: 30_000 });
+    }
+    const notes = frame.getByRole('textbox', { name: 'Follow-up notes' });
+    await notes.fill('Verify the next retention review.');
+    await expect(notes).toHaveValue('Verify the next retention review.');
+    await expect(frame.locator('a2ui-surface')).toHaveCount(1);
+  });
+
   test('settled replay allows State inspection and form edits; page motion resumes playback', async ({ page }) => {
     test.skip(process.env['STAGE_LIVE_FRAME'] !== 'true', 'requires the matching stage replay deployment');
     await page.setViewportSize({ width: 1440, height: 900 });

libs/chat/src/lib/streaming/content-classifier.spec.ts

@@ -125,6 +125,31 @@ describe('ContentClassifier', () => {
   });
 
   describe('type transitions', () => {
+    it('reclassifies a longer A2UI replacement after a streamed introduction', () => {
+      const c = setup();
+      c.update("I'll render a compact cleanup report UI now.");
+      expect(c.type()).toBe('markdown');
+
+      c.update('---a2ui_JSON---\n' + JSON.stringify({
+        version: 'v0.9',
+        createSurface: { surfaceId: 'cleanup', catalogId: 'https://a2ui.org/specification/v0_9/catalogs/basic/catalog.json' },
+      }) + '\n' + JSON.stringify({
+        version: 'v0.9',
+        updateComponents: { surfaceId: 'cleanup', components: [{ id: 'root', component: 'Text', text: 'Cleanup report' }] },
+      }) + '\n');
+
+      expect(c.type()).toBe('a2ui');
+      expect(c.a2uiSurfaces().has('cleanup')).toBe(true);
+      expect(c.markdown()).toBe('');
+    });
+
+    it('reparses an equal-length JSON replacement instead of keeping stale elements', () => {
+      const c = setup();
+      c.update('{"root":"first","elements":{}}');
+      c.update('{"root":"other","elements":{}}');
+      expect(c.spec()?.root).toBe('other');
+    });
+
     it('never downgrades from markdown', () => {
       const c = setup();
       c.update('Hello');

libs/chat/src/lib/streaming/content-classifier.ts

@@ -46,6 +46,7 @@ export function createContentClassifier(): ContentClassifier {
   const errorsSignal = signal<string[]>([]);
 
   let processedLength = 0;
+  let previousContent = '';
   let store: ParseTreeStore | null = null;
   let jsonStartIndex = 0;
 
@@ -148,15 +149,13 @@ export function createContentClassifier(): ContentClassifier {
     // NG0600 forbids writing signals during change detection; untracked()
     // opts out of the reactive graph for this imperative push-based update.
     untracked(() => {
-      // If content shrunk vs. last seen length, the underlying message was
-      // replaced (e.g. via langgraph RemoveMessage / id-match content
-      // replacement followed by force-refresh-from-server). Reset state so
-      // the new content is classified fresh — otherwise the classifier
-      // keeps the streamed (pre-mutation) markdown/json type and the UI
-      // never updates.
-      if (content.length < processedLength) {
+      // A same-id replacement can be longer than its streamed introduction
+      // (e.g. the render tool replaces prose with an A2UI payload). Only
+      // append-only content can reuse the existing classification and parser.
+      if (!content.startsWith(previousContent)) {
         resetState();
       }
+      previousContent = content;
       const currentType = typeSignal();
 
       if (currentType === 'pending') {