test(examples): citations e2e coverage + fix ag-ui citation delivery

examples/ag-ui/angular/e2e/citations.spec.ts

@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: MIT
+import { test, expect } from '@playwright/test';
+import { sendPromptAndWait } from './test-helpers';
+
+// Twin of examples/chat's citations spec. AG-UI delivers citations differently
+// (the backend surfaces them as state.citations[messageId], which the ag-ui
+// adapter mirrors and bridgeCitationsState() maps onto Message.citations) — but
+// the rendered surface is the shared @threadplane/chat markers/preview/panel, so
+// the assertions match. The graph runs the REAL search_documents tool; its
+// no-match fallback returns the first 3 corpus docs, so the set is deterministic:
+// ng-signals-overview, ng-signals-rxjs, ng-control-flow.
+const PROMPT = 'cite your sources on angular signals';
+
+test('inline citation markers render as resolved pills', async ({ page }) => {
+  const bubble = await sendPromptAndWait(page, PROMPT);
+  const markers = bubble.locator('.chat-citation-marker');
+  await expect(markers.first()).toBeVisible();
+  expect(await markers.count()).toBeGreaterThanOrEqual(3);
+  await expect(bubble.locator('a.chat-citation-marker').first()).toHaveAttribute('href', /angular\.dev/);
+  await expect(bubble.locator('.chat-citation-marker--unresolved')).toHaveCount(0);
+});
+
+test('sources panel is collapsed by default and expands to the cited sources', async ({ page }) => {
+  const bubble = await sendPromptAndWait(page, PROMPT);
+
+  await expect(bubble.locator('.chat-citations')).toBeVisible();
+  await expect(bubble.locator('.chat-citations__count')).toHaveText('3');
+
+  await expect(bubble.locator('.chat-citations__header')).toBeVisible();
+  await expect(bubble.locator('.chat-citations__list')).toHaveCount(0);
+
+  await bubble.locator('.chat-citations__header').click();
+  const cards = bubble.locator('.chat-citations-card');
+  await expect(cards).toHaveCount(3);
+  await expect(cards.nth(0)).toContainText('Signals');
+  await expect(cards.nth(1)).toContainText('RxJS interop');
+  await expect(cards.nth(2)).toContainText('control flow');
+  await expect(cards.nth(0)).toHaveAttribute('href', /angular\.dev\/guide\/signals/);
+});
+
+test('focusing a marker opens the portaled provenance preview card', async ({ page }) => {
+  const bubble = await sendPromptAndWait(page, PROMPT);
+
+  await bubble.locator('a.chat-citation-marker').first().focus();
+
+  const preview = page.locator('.chat-citation-preview');
+  await expect(preview).toBeVisible();
+  await expect(preview.locator('.chat-citation-preview__domain')).toContainText('angular.dev');
+  await expect(preview.locator('.chat-citation-preview__open')).toBeVisible();
+});

examples/ag-ui/angular/e2e/fixtures/citations.json

@@ -0,0 +1,21 @@
+{
+  "fixtures": [
+    {
+      "match": { "userMessage": "cite your sources on angular signals", "hasToolResult": true },
+      "response": {
+        "content": "Angular signals are a reactivity primitive that tracks reads and notifies consumers on change [^ng-signals-overview]. They interoperate with RxJS through `toSignal()` and `toObservable()` [^ng-signals-rxjs], and modern templates express reactive UI with built-in control flow like `@if` and `@for` [^ng-control-flow]. Signals pair naturally with that control flow for local state [^ng-signals-overview]."
+      }
+    },
+    {
+      "match": { "userMessage": "cite your sources on angular signals" },
+      "response": {
+        "toolCalls": [
+          {
+            "name": "search_documents",
+            "arguments": { "query": "authoritative overview of angular signals reactivity and control flow" }
+          }
+        ]
+      }
+    }
+  ]
+}

examples/ag-ui/python/src/graph.py

@@ -534,6 +534,13 @@ class State(TypedDict):
     # channel must exist here so the graph retains the client catalog across
     # the generate → should_continue → attach_citations path.
     tools: Optional[list]
+    # Per-message citations, keyed by AI message id. Unlike the LangGraph
+    # transport (which reads AIMessage.additional_kwargs.citations directly),
+    # the ag-ui protocol streams message TEXT without additional_kwargs, so
+    # citations must travel as STATE. The ag-ui-langgraph adapter mirrors this
+    # channel to the client, where bridgeCitationsState() reads
+    # state.citations[messageId] onto Message.citations for rendering.
+    citations: Optional[dict]
 
 
 async def generate(state: State, config: RunnableConfig) -> dict:
@@ -830,7 +837,12 @@ async def attach_citations(state: State) -> dict:
                 tool_calls=getattr(last, "tool_calls", []) or [],
                 response_metadata=getattr(last, "response_metadata", {}) or {},
             ),
-        ]
+        ],
+        # Also surface citations as STATE, keyed by the AI message id. The
+        # ag-ui protocol drops additional_kwargs when streaming message text,
+        # so this STATE channel is the only path by which the ag-ui client
+        # (via bridgeCitationsState) can render them.
+        "citations": {last.id: citations},
     }
 
 

examples/chat/angular/e2e/citations.spec.ts

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: MIT
+import { test, expect } from '@playwright/test';
+import { sendPromptAndWait } from './test-helpers';
+
+// Matches fixtures/citations.json. The graph runs the REAL search_documents
+// tool; its no-match fallback returns the first 3 corpus docs, so the citation
+// set is deterministic: ng-signals-overview, ng-signals-rxjs, ng-control-flow.
+const PROMPT = 'cite your sources on angular signals';
+
+test('inline citation markers render as resolved pills', async ({ page }) => {
+  const bubble = await sendPromptAndWait(page, PROMPT);
+  const markers = bubble.locator('.chat-citation-marker');
+  await expect(markers.first()).toBeVisible();
+  // 4 inline references across 3 distinct sources.
+  expect(await markers.count()).toBeGreaterThanOrEqual(3);
+  // Resolved markers with a URL render as anchors linking out; none unresolved.
+  await expect(bubble.locator('a.chat-citation-marker').first()).toHaveAttribute('href', /angular\.dev/);
+  await expect(bubble.locator('.chat-citation-marker--unresolved')).toHaveCount(0);
+});
+
+test('sources panel is collapsed by default and expands to the cited sources', async ({ page }) => {
+  const bubble = await sendPromptAndWait(page, PROMPT);
+
+  await expect(bubble.locator('.chat-citations')).toBeVisible();
+  await expect(bubble.locator('.chat-citations__count')).toHaveText('3');
+
+  // Collapsed by default: header shows, list is absent.
+  await expect(bubble.locator('.chat-citations__header')).toBeVisible();
+  await expect(bubble.locator('.chat-citations__list')).toHaveCount(0);
+
+  // Expand → three detail cards in citation order.
+  await bubble.locator('.chat-citations__header').click();
+  const cards = bubble.locator('.chat-citations-card');
+  await expect(cards).toHaveCount(3);
+  await expect(cards.nth(0)).toContainText('Signals');
+  await expect(cards.nth(1)).toContainText('RxJS interop');
+  await expect(cards.nth(2)).toContainText('control flow');
+  // Cards are links to the source (the card element itself is the <a>).
+  await expect(cards.nth(0)).toHaveAttribute('href', /angular\.dev\/guide\/signals/);
+});
+
+test('focusing a marker opens the portaled provenance preview card', async ({ page }) => {
+  const bubble = await sendPromptAndWait(page, PROMPT);
+
+  // Focus (deterministic across pointer types) opens the preview.
+  await bubble.locator('a.chat-citation-marker').first().focus();
+
+  // The preview is portaled to the body-level overlay container, not the bubble.
+  const preview = page.locator('.chat-citation-preview');
+  await expect(preview).toBeVisible();
+  await expect(preview.locator('.chat-citation-preview__domain')).toContainText('angular.dev');
+  await expect(preview.locator('.chat-citation-preview__open')).toBeVisible();
+});

examples/chat/angular/e2e/fixtures/citations.json

@@ -0,0 +1,21 @@
+{
+  "fixtures": [
+    {
+      "match": { "userMessage": "cite your sources on angular signals", "hasToolResult": true },
+      "response": {
+        "content": "Angular signals are a reactivity primitive that tracks reads and notifies consumers on change [^ng-signals-overview]. They interoperate with RxJS through `toSignal()` and `toObservable()` [^ng-signals-rxjs], and modern templates express reactive UI with built-in control flow like `@if` and `@for` [^ng-control-flow]. Signals pair naturally with that control flow for local state [^ng-signals-overview]."
+      }
+    },
+    {
+      "match": { "userMessage": "cite your sources on angular signals" },
+      "response": {
+        "toolCalls": [
+          {
+            "name": "search_documents",
+            "arguments": { "query": "authoritative overview of angular signals reactivity and control flow" }
+          }
+        ]
+      }
+    }
+  ]
+}

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

@@ -63,6 +63,33 @@ describe('reduceEvent', () => {
     expect(store.messages()[0].content).toBe('hi there');
   });
 
+  it('MESSAGES_SNAPSHOT re-applies citations from prior STATE onto the final message', () => {
+    // Reproduces the ag-ui citation-delivery ordering: STATE_SNAPSHOT carries
+    // citations keyed by the final AI message id, but arrives BEFORE the
+    // MESSAGES_SNAPSHOT that swaps the streamed chunk-id message for the final
+    // one. Without re-bridging in the MESSAGES_SNAPSHOT handler the citations
+    // are dropped on the swap.
+    const store = makeStore();
+    reduceEvent({
+      type: 'STATE_SNAPSHOT',
+      snapshot: {
+        citations: {
+          'resp-final': [
+            { id: 'ng-signals-overview', index: 1, title: 'Signals — Angular guide', url: 'https://angular.dev/guide/signals' },
+          ],
+        },
+      },
+    } as any, store);
+    reduceEvent({
+      type: 'MESSAGES_SNAPSHOT',
+      messages: [{ id: 'resp-final', role: 'assistant', content: 'Signals are reactive [^ng-signals-overview].' }],
+    } as any, store);
+
+    const msg = store.messages().find((m) => m.id === 'resp-final');
+    expect(msg?.citations?.length).toBe(1);
+    expect(msg?.citations?.[0]).toMatchObject({ id: 'ng-signals-overview', title: 'Signals — Angular guide' });
+  });
+
   it('TOOL_CALL_START appends a running tool call', () => {
     const store = makeStore();
     reduceEvent({ type: 'TOOL_CALL_START', toolCallId: 't1', toolCallName: 'search' } as any, store);

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

@@ -295,7 +295,13 @@ export function reduceEvent(event: BaseEvent, store: ReducerStore): void {
         const { toolCalls: _dropped, ...rest } = m;
         return { ...rest, toolCallIds: ids } as unknown as Message;
       });
-      store.messages.set(messages);
+      // Re-apply per-message citations from the already-received STATE. A
+      // MESSAGES_SNAPSHOT replaces the streamed messages wholesale — and the
+      // final snapshot message id (str(AIMessage.id), e.g. "resp-…") differs
+      // from the streaming chunk id the earlier STATE_SNAPSHOT bridged against,
+      // so without re-bridging here the citations (keyed by the final id) would
+      // be dropped on the message swap.
+      store.messages.set(bridgeCitationsState({ state: store.state() }, messages));
       if (snapshotToolCalls.length > 0) {
         store.toolCalls.update((prev) => {
           // Merge: keep existing entries (they may carry richer state from