diff --git a/examples/ag-ui/angular/e2e/citations.spec.ts b/examples/ag-ui/angular/e2e/citations.spec.ts new file mode 100644 index 000000000..be2c51647 --- /dev/null +++ b/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(); +}); diff --git a/examples/ag-ui/angular/e2e/fixtures/citations.json b/examples/ag-ui/angular/e2e/fixtures/citations.json new file mode 100644 index 000000000..c2edd3f28 --- /dev/null +++ b/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" } + } + ] + } + } + ] +} diff --git a/examples/ag-ui/python/src/graph.py b/examples/ag-ui/python/src/graph.py index befc422a2..5b044763c 100644 --- a/examples/ag-ui/python/src/graph.py +++ b/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}, } diff --git a/examples/chat/angular/e2e/citations.spec.ts b/examples/chat/angular/e2e/citations.spec.ts new file mode 100644 index 000000000..2a42ec35b --- /dev/null +++ b/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 ). + 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(); +}); diff --git a/examples/chat/angular/e2e/fixtures/citations.json b/examples/chat/angular/e2e/fixtures/citations.json new file mode 100644 index 000000000..c2edd3f28 --- /dev/null +++ b/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" } + } + ] + } + } + ] +} diff --git a/libs/ag-ui/src/lib/reducer.spec.ts b/libs/ag-ui/src/lib/reducer.spec.ts index 6684d74e0..1645209ba 100644 --- a/libs/ag-ui/src/lib/reducer.spec.ts +++ b/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); diff --git a/libs/ag-ui/src/lib/reducer.ts b/libs/ag-ui/src/lib/reducer.ts index 38e0818f3..e453e44cb 100644 --- a/libs/ag-ui/src/lib/reducer.ts +++ b/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