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 {