From 3e75ad775d8b3a3d8a616f6944fe7394539bb3fb Mon Sep 17 00:00:00 2001 From: Brian Love Date: Tue, 12 May 2026 16:19:15 -0700 Subject: [PATCH] fix(chat): bridge defers surfaceUpdate dispatch until components have ids MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The partial-args bridge dispatched surfaceUpdate envelopes the moment they were structurally complete (components is an array) — typically very early in streaming when the first component object had not yet streamed its id field. pickRoot then returned null and synthesis was skipped. The bridge's dispatchedCount advanced past the surfaceUpdate envelope, so subsequent pushes never retried and the surface never mounted via the live-stream path. Now: defer dispatch until pickRoot(components) returns a real string, then dispatch the surfaceUpdate + synthesised beginRendering as an atomic pair. Subsequent envelopes flow through the existing index-based dispatch loop. A real beginRendering arriving later for the same surfaceId is treated as idempotent and skip-advanced past in the dispatch loop. --- .../src/lib/a2ui/partial-args-bridge.spec.ts | 69 ++++++++++++++- libs/chat/src/lib/a2ui/partial-args-bridge.ts | 84 ++++++++++++------- 2 files changed, 119 insertions(+), 34 deletions(-) diff --git a/libs/chat/src/lib/a2ui/partial-args-bridge.spec.ts b/libs/chat/src/lib/a2ui/partial-args-bridge.spec.ts index 7bdf891e..30643255 100644 --- a/libs/chat/src/lib/a2ui/partial-args-bridge.spec.ts +++ b/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'); + }); }); diff --git a/libs/chat/src/lib/a2ui/partial-args-bridge.ts b/libs/chat/src/lib/a2ui/partial-args-bridge.ts index 2ac7a6c3..b1420214 100644 --- a/libs/chat/src/lib/a2ui/partial-args-bridge.ts +++ b/libs/chat/src/lib/a2ui/partial-args-bridge.ts @@ -20,8 +20,13 @@ interface BridgeState { parser: ReturnType; /** 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 {