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

[blove]

Summary

Follow-up to #769 (citations & sources UI). Adds end-to-end coverage for the citations flow in both example apps — and, in doing so, uncovered and fixed a real bug: ag-ui citations never reached the UI.

Bug fix — ag-ui citation delivery

The ag-ui example's attach_citations node only set AIMessage.additional_kwargs.citations. The ag-ui protocol drops additional_kwargs when streaming message text, so the frontend (which reads state.citations[messageId] via bridgeCitationsState) never received them — citations silently didn't render in ag-ui.

• Backend (examples/ag-ui/python/src/graph.py): add a citations channel to State and emit state.citations keyed by the AI message id.
• Frontend (libs/ag-ui/src/lib/reducer.ts): re-apply bridgeCitationsState in the MESSAGES_SNAPSHOT handler. 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 the citations were dropped on the message swap. Covered by a new reducer unit test reproducing that exact ordering.

E2E coverage (chat + ag-ui twins)

Aimock fixtures + Playwright specs asserting the rendered system:

• Inline pill markers resolve (anchors, none unresolved)
• Sources panel is collapsed by default, expands to the 3 cited detail cards with correct titles + angular.dev links
• Focusing a marker opens the portaled preview card with the right domain + open action

The real search_documents tool's no-match fallback returns the first 3 corpus docs, so the citation set is deterministic (ng-signals-overview, ng-signals-rxjs, ng-control-flow) with no live LLM.

Test Plan

• examples/chat citations e2e — 3/3 pass locally
• examples/ag-ui citations e2e — 3/3 pass locally (was 0/3 before the fix)
• npx nx test ag-ui green (incl. new reducer test); 0 lint errors
• CI: both examples' e2e jobs (chat gating, ag-ui advisory)

🤖 Generated with Claude Code

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
threadplane	Ready	Preview, Comment	Jul 7, 2026 6:52pm

[github-actions]

Automated approval: this PR received an intelligent (AI) code review. See the review comments on this PR.

[github-actions]

Claude finished @blove's task in 2m 35s —— View job

---

Code Review

• Gather context (PR diff, changed files)
• Review backend fix (graph.py)
• Review frontend reducer fix (reducer.ts, reducer.spec.ts)
• Review e2e specs and fixtures
• Post findings

---

The bug fix is correct and the reasoning is well-documented. A few things worth noting:

Bug fix — reducer (libs/ag-ui/src/lib/reducer.ts)

The MESSAGES_SNAPSHOT re-bridge at line 304 is the right call. store.state() already holds the latest state from any prior STATE_SNAPSHOT / STATE_DELTA, so bridgeCitationsState will find the citations keyed by the final message id. The unit test in reducer.spec.ts directly reproduces the ordering and is a solid regression guard. No issues here.

Backend (examples/ag-ui/python/src/graph.py)

citations: Optional[dict] in State is correct. The {last.id: citations} emission at line 845 is safe — LangGraph always assigns an ID to AI messages before attach_citations runs. No issues.

E2e specs — two findings

1. expect(await markers.count()) lacks auto-retry — potential flakiness

In both new specs (examples/ag-ui/angular/e2e/citations.spec.ts:18 and examples/chat/angular/e2e/citations.spec.ts:15):

// current — no retry
expect(await markers.count()).toBeGreaterThanOrEqual(3);

This is a raw Jest-style assertion on a resolved Promise<number>. It does not retry. The preceding await expect(markers.first()).toBeVisible() ensures at least one marker is visible, but the remaining markers may still be mid-render when the count is sampled. Every other count assertion in these specs uses the Playwright await expect(...) form. In practice citations arrive as a batch so this probably doesn't flake, but it's inconsistent with the surrounding code.

Suggested fix in both files: Fix this →

// safer — auto-retries until the locator matches or timeout
await expect(markers).toHaveCount(4); // 4 inline refs across 3 sources

2. Missing attachBrowserHygiene — inconsistency with existing spec pattern

Seven of the existing e2e specs call attachBrowserHygiene to capture and assert on console errors. The two new citations specs don't, so unhandled Angular errors or network failures would go unnoticed. This is low-severity but an easy gap to close — if the citation pipeline throws in the client, the test could still pass.

Identical fixture files

Both examples/ag-ui/angular/e2e/fixtures/citations.json and examples/chat/angular/e2e/fixtures/citations.json are identical (same SHA). This is intentional — the aimock fixture intercepts at the AI API layer, not the transport layer, so the same LLM round-trip applies to both. The ag-ui spec's header comment explains this clearly.

Overall

The core bug fix is correct, well-tested, and safe to merge. The two e2e observations above are low-risk improvements — the missing retry is the one I'd address before merging to avoid flake on CI.

[github-actions]

Raw expect(await markers.count()) doesn't auto-retry. If markers render fractionally after the first one is visible this can flake. The fixture has 4 inline references across 3 sources — prefer the retrying form:

Suggested change

	expect(await markers.count()).toBeGreaterThanOrEqual(3);
	await expect(markers).toHaveCount(4);

[github-actions]

Same issue as the ag-ui twin — raw expect(await …count()) has no retry. Prefer:

Suggested change

	expect(await markers.count()).toBeGreaterThanOrEqual(3);
	await expect(markers).toHaveCount(4);