Skip to content

Commit 58d8bed

Browse files
bloveclaude
blove
and
claude
authored
test(examples): citations e2e coverage + fix ag-ui citation delivery (#774)
* fix(ag-ui): deliver citations to the client The ag-ui example's attach_citations node only set additional_kwargs.citations, which the ag-ui protocol drops when streaming message text — so citations never reached the UI. Emit them as state.citations keyed by the AI message id, and re-apply bridgeCitationsState in the MESSAGES_SNAPSHOT reducer handler so the final message (whose id differs from the streamed chunk id) keeps its citations when the snapshot swaps the streamed messages. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * test(examples): e2e coverage for the citations & sources flow Aimock fixtures + Playwright specs for chat and ag-ui: assert inline pill markers resolve, the Sources panel is collapsed-by-default and expands to the cited detail cards, and focusing a marker opens the portaled preview card. The real search_documents tool's no-match fallback yields a deterministic 3-source set. Twin specs keep chat + ag-ui in sync. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1 parent ecd2947 commit 58d8bed

7 files changed

Lines changed: 192 additions & 2 deletions

File tree

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
// SPDX-License-Identifier: MIT
2+
import { test, expect } from '@playwright/test';
3+
import { sendPromptAndWait } from './test-helpers';
4+
5+
// Twin of examples/chat's citations spec. AG-UI delivers citations differently
6+
// (the backend surfaces them as state.citations[messageId], which the ag-ui
7+
// adapter mirrors and bridgeCitationsState() maps onto Message.citations) — but
8+
// the rendered surface is the shared @threadplane/chat markers/preview/panel, so
9+
// the assertions match. The graph runs the REAL search_documents tool; its
10+
// no-match fallback returns the first 3 corpus docs, so the set is deterministic:
11+
// ng-signals-overview, ng-signals-rxjs, ng-control-flow.
12+
const PROMPT = 'cite your sources on angular signals';
13+
14+
test('inline citation markers render as resolved pills', async ({ page }) => {
15+
const bubble = await sendPromptAndWait(page, PROMPT);
16+
const markers = bubble.locator('.chat-citation-marker');
17+
await expect(markers.first()).toBeVisible();
18+
expect(await markers.count()).toBeGreaterThanOrEqual(3);
19+
await expect(bubble.locator('a.chat-citation-marker').first()).toHaveAttribute('href', /angular\.dev/);
20+
await expect(bubble.locator('.chat-citation-marker--unresolved')).toHaveCount(0);
21+
});
22+
23+
test('sources panel is collapsed by default and expands to the cited sources', async ({ page }) => {
24+
const bubble = await sendPromptAndWait(page, PROMPT);
25+
26+
await expect(bubble.locator('.chat-citations')).toBeVisible();
27+
await expect(bubble.locator('.chat-citations__count')).toHaveText('3');
28+
29+
await expect(bubble.locator('.chat-citations__header')).toBeVisible();
30+
await expect(bubble.locator('.chat-citations__list')).toHaveCount(0);
31+
32+
await bubble.locator('.chat-citations__header').click();
33+
const cards = bubble.locator('.chat-citations-card');
34+
await expect(cards).toHaveCount(3);
35+
await expect(cards.nth(0)).toContainText('Signals');
36+
await expect(cards.nth(1)).toContainText('RxJS interop');
37+
await expect(cards.nth(2)).toContainText('control flow');
38+
await expect(cards.nth(0)).toHaveAttribute('href', /angular\.dev\/guide\/signals/);
39+
});
40+
41+
test('focusing a marker opens the portaled provenance preview card', async ({ page }) => {
42+
const bubble = await sendPromptAndWait(page, PROMPT);
43+
44+
await bubble.locator('a.chat-citation-marker').first().focus();
45+
46+
const preview = page.locator('.chat-citation-preview');
47+
await expect(preview).toBeVisible();
48+
await expect(preview.locator('.chat-citation-preview__domain')).toContainText('angular.dev');
49+
await expect(preview.locator('.chat-citation-preview__open')).toBeVisible();
50+
});

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
{
2+
"fixtures": [
3+
{
4+
"match": { "userMessage": "cite your sources on angular signals", "hasToolResult": true },
5+
"response": {
6+
"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]."
7+
}
8+
},
9+
{
10+
"match": { "userMessage": "cite your sources on angular signals" },
11+
"response": {
12+
"toolCalls": [
13+
{
14+
"name": "search_documents",
15+
"arguments": { "query": "authoritative overview of angular signals reactivity and control flow" }
16+
}
17+
]
18+
}
19+
}
20+
]
21+
}

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -534,6 +534,13 @@ class State(TypedDict):
534534
# channel must exist here so the graph retains the client catalog across
535535
# the generate → should_continue → attach_citations path.
536536
tools: Optional[list]
537+
# Per-message citations, keyed by AI message id. Unlike the LangGraph
538+
# transport (which reads AIMessage.additional_kwargs.citations directly),
539+
# the ag-ui protocol streams message TEXT without additional_kwargs, so
540+
# citations must travel as STATE. The ag-ui-langgraph adapter mirrors this
541+
# channel to the client, where bridgeCitationsState() reads
542+
# state.citations[messageId] onto Message.citations for rendering.
543+
citations: Optional[dict]
537544

538545

539546
async def generate(state: State, config: RunnableConfig) -> dict:
@@ -830,7 +837,12 @@ async def attach_citations(state: State) -> dict:
830837
tool_calls=getattr(last, "tool_calls", []) or [],
831838
response_metadata=getattr(last, "response_metadata", {}) or {},
832839
),
833-
]
840+
],
841+
# Also surface citations as STATE, keyed by the AI message id. The
842+
# ag-ui protocol drops additional_kwargs when streaming message text,
843+
# so this STATE channel is the only path by which the ag-ui client
844+
# (via bridgeCitationsState) can render them.
845+
"citations": {last.id: citations},
834846
}
835847

836848

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
// SPDX-License-Identifier: MIT
2+
import { test, expect } from '@playwright/test';
3+
import { sendPromptAndWait } from './test-helpers';
4+
5+
// Matches fixtures/citations.json. The graph runs the REAL search_documents
6+
// tool; its no-match fallback returns the first 3 corpus docs, so the citation
7+
// set is deterministic: ng-signals-overview, ng-signals-rxjs, ng-control-flow.
8+
const PROMPT = 'cite your sources on angular signals';
9+
10+
test('inline citation markers render as resolved pills', async ({ page }) => {
11+
const bubble = await sendPromptAndWait(page, PROMPT);
12+
const markers = bubble.locator('.chat-citation-marker');
13+
await expect(markers.first()).toBeVisible();
14+
// 4 inline references across 3 distinct sources.
15+
expect(await markers.count()).toBeGreaterThanOrEqual(3);
16+
// Resolved markers with a URL render as anchors linking out; none unresolved.
17+
await expect(bubble.locator('a.chat-citation-marker').first()).toHaveAttribute('href', /angular\.dev/);
18+
await expect(bubble.locator('.chat-citation-marker--unresolved')).toHaveCount(0);
19+
});
20+
21+
test('sources panel is collapsed by default and expands to the cited sources', async ({ page }) => {
22+
const bubble = await sendPromptAndWait(page, PROMPT);
23+
24+
await expect(bubble.locator('.chat-citations')).toBeVisible();
25+
await expect(bubble.locator('.chat-citations__count')).toHaveText('3');
26+
27+
// Collapsed by default: header shows, list is absent.
28+
await expect(bubble.locator('.chat-citations__header')).toBeVisible();
29+
await expect(bubble.locator('.chat-citations__list')).toHaveCount(0);
30+
31+
// Expand → three detail cards in citation order.
32+
await bubble.locator('.chat-citations__header').click();
33+
const cards = bubble.locator('.chat-citations-card');
34+
await expect(cards).toHaveCount(3);
35+
await expect(cards.nth(0)).toContainText('Signals');
36+
await expect(cards.nth(1)).toContainText('RxJS interop');
37+
await expect(cards.nth(2)).toContainText('control flow');
38+
// Cards are links to the source (the card element itself is the <a>).
39+
await expect(cards.nth(0)).toHaveAttribute('href', /angular\.dev\/guide\/signals/);
40+
});
41+
42+
test('focusing a marker opens the portaled provenance preview card', async ({ page }) => {
43+
const bubble = await sendPromptAndWait(page, PROMPT);
44+
45+
// Focus (deterministic across pointer types) opens the preview.
46+
await bubble.locator('a.chat-citation-marker').first().focus();
47+
48+
// The preview is portaled to the body-level overlay container, not the bubble.
49+
const preview = page.locator('.chat-citation-preview');
50+
await expect(preview).toBeVisible();
51+
await expect(preview.locator('.chat-citation-preview__domain')).toContainText('angular.dev');
52+
await expect(preview.locator('.chat-citation-preview__open')).toBeVisible();
53+
});

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
{
2+
"fixtures": [
3+
{
4+
"match": { "userMessage": "cite your sources on angular signals", "hasToolResult": true },
5+
"response": {
6+
"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]."
7+
}
8+
},
9+
{
10+
"match": { "userMessage": "cite your sources on angular signals" },
11+
"response": {
12+
"toolCalls": [
13+
{
14+
"name": "search_documents",
15+
"arguments": { "query": "authoritative overview of angular signals reactivity and control flow" }
16+
}
17+
]
18+
}
19+
}
20+
]
21+
}

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,33 @@ describe('reduceEvent', () => {
6363
expect(store.messages()[0].content).toBe('hi there');
6464
});
6565

66+
it('MESSAGES_SNAPSHOT re-applies citations from prior STATE onto the final message', () => {
67+
// Reproduces the ag-ui citation-delivery ordering: STATE_SNAPSHOT carries
68+
// citations keyed by the final AI message id, but arrives BEFORE the
69+
// MESSAGES_SNAPSHOT that swaps the streamed chunk-id message for the final
70+
// one. Without re-bridging in the MESSAGES_SNAPSHOT handler the citations
71+
// are dropped on the swap.
72+
const store = makeStore();
73+
reduceEvent({
74+
type: 'STATE_SNAPSHOT',
75+
snapshot: {
76+
citations: {
77+
'resp-final': [
78+
{ id: 'ng-signals-overview', index: 1, title: 'Signals — Angular guide', url: 'https://angular.dev/guide/signals' },
79+
],
80+
},
81+
},
82+
} as any, store);
83+
reduceEvent({
84+
type: 'MESSAGES_SNAPSHOT',
85+
messages: [{ id: 'resp-final', role: 'assistant', content: 'Signals are reactive [^ng-signals-overview].' }],
86+
} as any, store);
87+
88+
const msg = store.messages().find((m) => m.id === 'resp-final');
89+
expect(msg?.citations?.length).toBe(1);
90+
expect(msg?.citations?.[0]).toMatchObject({ id: 'ng-signals-overview', title: 'Signals — Angular guide' });
91+
});
92+
6693
it('TOOL_CALL_START appends a running tool call', () => {
6794
const store = makeStore();
6895
reduceEvent({ type: 'TOOL_CALL_START', toolCallId: 't1', toolCallName: 'search' } as any, store);

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

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

0 commit comments

Comments
 (0)