fix(chat): bridge defers surfaceUpdate dispatch until components have ids

libs/chat/src/lib/a2ui/partial-args-bridge.spec.ts

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: MIT
 import { describe, it, expect, beforeEach } from 'vitest';
 import { TestBed } from '@angular/core/testing';
+import type { A2uiMessage } from '@ngaf/a2ui';
 import { createPartialArgsBridge } from './partial-args-bridge';
 import { createA2uiSurfaceStore, type A2uiSurfaceStore } from './surface-store';
 
@@ -61,9 +62,12 @@ describe('createPartialArgsBridge', () => {
     expect(store.surfaces().get('s')?.dataModel).toEqual({ msg: 'hi' });
   });
 
-  it('marks tool_call_id as live in the store', () => {
+  it('marks tool_call_id as live in the store once the initial surface pair dispatches', () => {
     const bridge = createPartialArgsBridge(store);
-    bridge.push('tc-5', '{"envelopes":[{"surfaceUpdate":{"surfaceId":"s","components":[]}}]}');
+    bridge.push(
+      'tc-5',
+      '{"envelopes":[{"surfaceUpdate":{"surfaceId":"s","components":[{"id":"root","type":"text","props":{}}]}}]}',
+    );
     expect(store.isPartialLive('tc-5')).toBe(true);
   });
 
@@ -91,4 +95,65 @@ describe('createPartialArgsBridge', () => {
     bridge.push('tc-8', '{"envelopes":[{"surfaceUpdate":{"surfaceId":"s","components":[{"id":"only","type":"text","props":{}}]}}]}');
     expect(store.surfaces().get('s')?.components.has('only')).toBe(true);
   });
+
+  it('incremental push waits for a component id before mounting the surface', () => {
+    const bridge = createPartialArgsBridge(store);
+    // 1: object started, no id on first component yet.
+    bridge.push('tc-9', '{"envelopes":[{"surfaceUpdate":{"surfaceId":"s","components":[{');
+    expect(store.surfaces().has('s')).toBe(false);
+    // 2: started the "id" key but no value yet.
+    bridge.push('tc-9', '{"envelopes":[{"surfaceUpdate":{"surfaceId":"s","components":[{"');
+    expect(store.surfaces().has('s')).toBe(false);
+    // 3: id, type, props all present and component object is closed.
+    bridge.push(
+      'tc-9',
+      '{"envelopes":[{"surfaceUpdate":{"surfaceId":"s","components":[{"id":"root","type":"text","props":{}}]}}',
+    );
+    expect(store.surfaces().get('s')?.components.has('root')).toBe(true);
+  });
+
+  it('mounts the surface on first complete push and applies a dataModelUpdate on a later push', () => {
+    const bridge = createPartialArgsBridge(store);
+    const surfaceUpdate = JSON.stringify({
+      surfaceUpdate: { surfaceId: 's', components: [{ id: 'root', type: 'text', props: {} }] },
+    });
+    const dataModelUpdate = JSON.stringify({
+      dataModelUpdate: { surfaceId: 's', contents: [{ key: 'greeting', valueString: 'hello' }] },
+    });
+    bridge.push('tc-10', '{"envelopes":[' + surfaceUpdate + ']}');
+    expect(store.surfaces().get('s')?.components.has('root')).toBe(true);
+    expect(store.surfaces().get('s')?.dataModel).toEqual({});
+    bridge.push('tc-10', '{"envelopes":[' + surfaceUpdate + ',' + dataModelUpdate + ']}');
+    expect(store.surfaces().get('s')?.dataModel).toEqual({ greeting: 'hello' });
+  });
+
+  it('synthetic beginRendering targets the first component when multiple have ids and none is "root"', () => {
+    // Spy on applyPartialArgs to inspect the synthesised beginRendering envelope.
+    const captured: A2uiMessage[][] = [];
+    const orig = store.applyPartialArgs.bind(store);
+    (store as { applyPartialArgs: typeof store.applyPartialArgs }).applyPartialArgs = (
+      toolCallId: string,
+      envs: readonly A2uiMessage[],
+    ) => {
+      captured.push(envs.slice() as A2uiMessage[]);
+      orig(toolCallId, envs);
+    };
+    const bridge = createPartialArgsBridge(store);
+    bridge.push(
+      'tc-11',
+      '{"envelopes":[{"surfaceUpdate":{"surfaceId":"s","components":[{"id":"alpha","type":"text","props":{}},{"id":"beta","type":"text","props":{}}]}}]}',
+    );
+    const surface = store.surfaces().get('s');
+    expect(surface).toBeTruthy();
+    expect(surface!.components.has('alpha')).toBe(true);
+    expect(surface!.components.has('beta')).toBe(true);
+    // First dispatch should be [surfaceUpdate, synthesised beginRendering with root="alpha"].
+    expect(captured.length).toBeGreaterThan(0);
+    const firstBatch = captured[0];
+    const beginEnv = firstBatch.find((e) => 'beginRendering' in e) as
+      | { beginRendering: { surfaceId: string; root: string } }
+      | undefined;
+    expect(beginEnv).toBeTruthy();
+    expect(beginEnv!.beginRendering.root).toBe('alpha');
+  });
 });

libs/chat/src/lib/a2ui/partial-args-bridge.ts

@@ -20,8 +20,13 @@ interface BridgeState {
   parser: ReturnType<typeof createPartialJsonParser>;
   /** Number of envelopes already dispatched to the store. */
   dispatchedCount: number;
-  /** Have we synthesised a beginRendering for this turn yet? */
-  synthesised: boolean;
+  /**
+   * Have we dispatched the initial surfaceUpdate + synthesised beginRendering
+   * pair for this turn yet? Until true, dispatch is deferred — we wait for
+   * the first surfaceUpdate to have at least one component with an `id` so
+   * `pickRoot` can target a real root and the surface actually mounts.
+   */
+  surfacePairDispatched: boolean;
   /** surfaceId the synthesised beginRendering targets (to avoid double-mounting). */
   synthesisedSurfaceId: string | null;
   /** Once true, all subsequent pushes are ignored. */
@@ -183,7 +188,7 @@ export function createPartialArgsBridge(store: A2uiSurfaceStore): PartialArgsBri
       s = {
         parser: createPartialJsonParser(),
         dispatchedCount: 0,
-        synthesised: false,
+        surfacePairDispatched: false,
         synthesisedSurfaceId: null,
         poisoned: false,
       };
@@ -223,39 +228,54 @@ export function createPartialArgsBridge(store: A2uiSurfaceStore): PartialArgsBri
     const envelopes = normalizeEnvelopeArgs(materialised);
     if (!envelopes) return;
 
-    // Dispatch envelopes that haven't been dispatched yet. partial-json
-    // returns partially-built objects too; skip incomplete ones (those
-    // missing the discriminator top-level key).
-    const ready: A2uiMessage[] = [];
-    let realDispatched = 0;
-    for (let i = 0; i < envelopes.length; i++) {
+    // Phase 1: defer initial dispatch until the first envelope is a complete
+    // surfaceUpdate whose components include at least one entry with an `id`
+    // — otherwise pickRoot returns null and synthesis silently no-ops, leaving
+    // the surface unmounted forever. Once we have a pickable root, dispatch
+    // the surfaceUpdate AND a synthesised beginRendering as an atomic pair.
+    if (!state.surfacePairDispatched) {
+      const firstEnv = envelopes[0] as A2uiMessage | undefined;
+      if (!firstEnv || !('surfaceUpdate' in firstEnv)) return;
+      if (!isStructurallyComplete(firstEnv)) return;
+      const upd = (firstEnv as { surfaceUpdate: A2uiSurfaceUpdate }).surfaceUpdate;
+      if (upd.components.length === 0) return;
+      const root = pickRoot(upd.components);
+      if (!root) return;
+      state.surfacePairDispatched = true;
+      state.dispatchedCount = 1; // index 0 = the surfaceUpdate we just sent
+      state.synthesisedSurfaceId = upd.surfaceId;
+      store.applyPartialArgs(toolCallId, [
+        firstEnv,
+        { beginRendering: { surfaceId: upd.surfaceId, root } },
+      ]);
+    }
+
+    // Phase 2: dispatch any newly-complete envelopes beyond the initial pair.
+    const newEnvelopes: A2uiMessage[] = [];
+    for (let i = state.dispatchedCount; i < envelopes.length; i++) {
       const env = envelopes[i] as A2uiMessage;
-      if (!isStructurallyComplete(env)) continue;
-      // Only dispatch envelopes BEYOND those already sent.
-      if (i < state.dispatchedCount) {
-        // already dispatched; do not re-send.
+      if (!isStructurallyComplete(env)) {
+        // Stop at the first not-yet-complete envelope; later siblings can't
+        // exist before earlier ones complete (envelopes are an ordered list).
+        break;
+      }
+      // Skip a "real" beginRendering for the synthesised surface — the store
+      // already treats it as idempotent (just re-applies styles), but we
+      // advance past it so dispatchedCount stays in sync with the index.
+      if (
+        'beginRendering' in env &&
+        (env as { beginRendering: { surfaceId?: string } }).beginRendering.surfaceId ===
+          state.synthesisedSurfaceId
+      ) {
+        state.dispatchedCount = i + 1;
         continue;
       }
-      ready.push(env);
-      realDispatched++;
-      // Inject synthesised beginRendering immediately after first surfaceUpdate.
-      const synth = synthesisedBegin(env, state);
-      if (synth) ready.push(synth);
+      newEnvelopes.push(env);
+      state.dispatchedCount = i + 1;
+    }
+    if (newEnvelopes.length > 0) {
+      store.applyPartialArgs(toolCallId, newEnvelopes);
     }
-    if (ready.length === 0) return;
-
-    state.dispatchedCount += realDispatched;
-    store.applyPartialArgs(toolCallId, ready);
-  }
-
-  function synthesisedBegin(env: A2uiMessage, state: BridgeState): A2uiMessage | null {
-    if (state.synthesised || !('surfaceUpdate' in env)) return null;
-    const upd = (env as { surfaceUpdate: A2uiSurfaceUpdate }).surfaceUpdate;
-    const root = pickRoot(upd.components);
-    if (!root) return null;
-    state.synthesised = true;
-    state.synthesisedSurfaceId = upd.surfaceId;
-    return { beginRendering: { surfaceId: upd.surfaceId, root } };
   }
 
   function isPoisoned(toolCallId: string): boolean {