feat(ag-ui): bridge state.citations into Message.citations

Create bridgeCitationsState() normalizer and wire it into reduceEvent()
after STATE_SNAPSHOT and STATE_DELTA apply their patches, so that
state.citations[messageId] arrays flow through as Message.citations on
every affected assistant message.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: blove

libs/ag-ui/src/lib/bridge-citations-state.spec.ts

@@ -0,0 +1,49 @@
+// libs/ag-ui/src/lib/bridge-citations-state.spec.ts
+// SPDX-License-Identifier: MIT
+import { bridgeCitationsState } from './bridge-citations-state';
+import type { Message } from '@ngaf/chat';
+
+describe('bridgeCitationsState', () => {
+  const baseMsg = (id: string): Message => ({ id, role: 'assistant', content: 'x' });
+
+  it('returns messages unchanged when state has no citations', () => {
+    const msgs = [baseMsg('m1'), baseMsg('m2')];
+    const result = bridgeCitationsState({ state: {} }, msgs);
+    expect(result).toEqual(msgs);
+  });
+
+  it('merges citations into matching messages by id', () => {
+    const result = bridgeCitationsState(
+      { state: { citations: { m1: [{ id: 'a', title: 'A', url: 'https://a' }] } } },
+      [baseMsg('m1'), baseMsg('m2')],
+    );
+    expect(result[0].citations).toEqual([{ id: 'a', index: 1, title: 'A', url: 'https://a' }]);
+    expect(result[1].citations).toBeUndefined();
+  });
+
+  it('idempotent — same input produces same output', () => {
+    const state = { state: { citations: { m1: [{ id: 'a', title: 'A' }] } } };
+    const msgs = [baseMsg('m1')];
+    const a = bridgeCitationsState(state, msgs);
+    const b = bridgeCitationsState(state, a);
+    expect(b[0].citations).toEqual(a[0].citations);
+  });
+
+  it('coerces key spellings (href/source, name, excerpt)', () => {
+    const result = bridgeCitationsState(
+      { state: { citations: { m1: [{ name: 'N', href: 'https://h', excerpt: 'E' }] } } },
+      [baseMsg('m1')],
+    );
+    expect(result[0].citations).toEqual([
+      { id: 'c1', index: 1, title: 'N', url: 'https://h', snippet: 'E' },
+    ]);
+  });
+
+  it('handles string entries', () => {
+    const result = bridgeCitationsState(
+      { state: { citations: { m1: ['https://x'] } } },
+      [baseMsg('m1')],
+    );
+    expect(result[0].citations).toEqual([{ id: 'c1', index: 1, url: 'https://x' }]);
+  });
+});

libs/ag-ui/src/lib/bridge-citations-state.ts

@@ -0,0 +1,45 @@
+// libs/ag-ui/src/lib/bridge-citations-state.ts
+// SPDX-License-Identifier: MIT
+import type { Citation, Message } from '@ngaf/chat';
+
+interface ThreadStateLike {
+  state?: Record<string, unknown>;
+}
+
+export function bridgeCitationsState(thread: ThreadStateLike, messages: Message[]): Message[] {
+  const citationsByMsg = (thread.state as { citations?: unknown })?.citations;
+  if (!citationsByMsg || typeof citationsByMsg !== 'object') return messages;
+  const map = citationsByMsg as Record<string, unknown>;
+  return messages.map(msg => {
+    const raw = map[msg.id];
+    if (!Array.isArray(raw) || raw.length === 0) return msg;
+    return { ...msg, citations: raw.map((entry, i) => normalizeCitation(entry, i + 1)) };
+  });
+}
+
+function normalizeCitation(entry: unknown, fallbackIndex: number): Citation {
+  if (typeof entry === 'string') {
+    return { id: `c${fallbackIndex}`, index: fallbackIndex, url: entry };
+  }
+  const e = (entry ?? {}) as Record<string, unknown>;
+  const str = (key: string): string | undefined =>
+    typeof e[key] === 'string' ? (e[key] as string) : undefined;
+  const firstStr = (...keys: string[]): string | undefined => {
+    for (const k of keys) {
+      const v = str(k);
+      if (v !== undefined) return v;
+    }
+    return undefined;
+  };
+  return {
+    id: str('id') ?? str('refId') ?? `c${fallbackIndex}`,
+    index: typeof e['index'] === 'number' ? (e['index'] as number) : fallbackIndex,
+    title: firstStr('title', 'name'),
+    url: firstStr('url', 'href', 'source'),
+    snippet: firstStr('snippet', 'content', 'excerpt'),
+    extra:
+      typeof e['extra'] === 'object' && e['extra'] !== null
+        ? (e['extra'] as Record<string, unknown>)
+        : undefined,
+  };
+}

libs/ag-ui/src/lib/reducer.ts

@@ -10,6 +10,7 @@ import type {
 } from '@ngaf/chat';
 import type { BaseEvent } from '@ag-ui/client';
 import { applyPatch, type Operation } from 'fast-json-patch';
+import { bridgeCitationsState } from './bridge-citations-state';
 
 export interface ReducerStore {
   messages:  WritableSignal<Message[]>;
@@ -157,13 +158,16 @@ export function reduceEvent(event: BaseEvent, store: ReducerStore): void {
     }
     case 'STATE_SNAPSHOT': {
       const e = event as unknown as { snapshot: Record<string, unknown> };
-      store.state.set(e.snapshot ?? {});
+      const snapshot = e.snapshot ?? {};
+      store.state.set(snapshot);
+      store.messages.update(msgs => bridgeCitationsState({ state: snapshot }, msgs));
       return;
     }
     case 'STATE_DELTA': {
       const e = event as unknown as { delta: Operation[] };
       const next = applyPatch(deepClone(store.state()), e.delta).newDocument;
       store.state.set(next);
+      store.messages.update(msgs => bridgeCitationsState({ state: next }, msgs));
       return;
     }
     case 'MESSAGES_SNAPSHOT': {