[quifox (quifox)]

Motivation & Context

Late message-linked AG-UI citations can create an empty duplicate assistant message through public streaming aggregation. With text for m1 and m2, followed by annotations for m1, the final response should contain the original two messages with the citation attached to m1.

Description & Review Guide

• What are the major changes? Add separate private finalization seams for the complete function-invocation loop and mapped Agent updates. Their defaults retain generic aggregation; AG-UI opts into its message-linked annotation finalizer at both boundaries. Reuse a shared Chat/Agent helper, preserving response-level metadata through generic aggregation.
• What is the impact of these changes? Public Client and Agent streaming responses retain late citations without duplicate messages. Run-scoped identity handling prevents citations from attaching to an earlier tool round that reused the same protocol message ID. Buffered text transformations, structured output and completed tool history remain intact.
• What do you want reviewers to focus on? The separation from per-provider-turn finalization, unchanged generic message grouping and custom Agent override precedence, and the public-path regression assertions for stream consumption, final-response-only, tool rounds and buffered transforms.

Related Issue

Fixes #9093. This continues the existing repair discussed in the maintainer's scope comment; it does not introduce a public extension contract.

AI Assistance

• No material AI assistance was used.
• This is an AI-assisted contribution. I reviewed, understood, and verified all submitted content and accept responsibility for it.

AI assistance details: Assistance was used for issue analysis, implementation, tests, and contribution text.

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR links to an agreed issue with no competing open PR, or the Related Issue section documents a trivial-change or repository-automation exception.
• This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.

[Copilot]

Copilot review overview

🟢 Approval recommended

The scoped aggregation changes preserve existing fallback behavior and are covered across the affected streaming paths.

Review effort: Balanced
Findings: None

What changed in this PR

Preserves late AG-UI citations across client, tool-loop, and agent streaming aggregation without creating duplicate messages.

Changes:

• Adds private whole-loop and agent-level finalization hooks.
• Applies annotation-aware aggregation at AG-UI streaming boundaries.
• Adds regression coverage for metadata, tools, structured output, and middleware.

File	Description
python/​packages/​core/​agent_framework/​_tools.py	Adds whole-loop finalization hook.
python/​packages/​core/​agent_framework/​_clients.py	Adds mapped-agent finalization hook.
python/​packages/​core/​agent_framework/​_agents.py	Delegates default agent aggregation to compatible clients.
python/​packages/​ag-ui/​agent_framework_ag_ui/​_event_converters.py	Generalizes annotation-aware aggregation for chat and agent responses.
python/​packages/​ag-ui/​agent_framework_ag_ui/​_client.py	Implements AG-UI-specific finalization hooks.
python/​packages/​core/​tests/​core/​test_function_invocation_logic.py	Tests whole-loop metadata and structured output.
python/​packages/​core/​tests/​core/​test_agents.py	Tests client and agent extension-point precedence.
python/​packages/​ag-ui/​tests/​ag_ui/​test_event_converters.py	Tests annotation-only metadata aggregation.
python/​packages/​ag-ui/​tests/​ag_ui/​test_ag_ui_client.py	Adds public-path citation regressions.

---

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

[quifox (quifox)]

Eduard van Valkenburg (@eavanvalkenburg) Following your scope guidance, updated the existing #9094.

Rechecked the repair against pinned main cf734bc4692ccfc944496a5b3182b3d65483c44d through the public Client/Agent HTTP SSE paths, including explicit stream consumption and final-response-only. With unmodified production code, those late-annotation cases reproduced the extra empty m1 (6 failed, 4 controls passed); the repaired public-path regressions pass.

The per-provider-response finalizer does not cover either outer boundary: the function loop aggregates all provider turns and local results, and Agent aggregates the mapped updates separately. Reusing the per-turn finalizer for those scopes would also rerun provider-specific processing over multiple turns. This revision therefore uses two separate private seams with generic defaults, with AG-UI supplying message-linked citation handling at both. Existing custom Agent finalizers keep precedence; protocol-only clients keep the generic fallback. Generic message_id grouping and Content.id semantics are unchanged.

Annotation-only updates keep their response metadata through identity-free aggregation copies. The helper scopes protocol IDs by run, restoring the public IDs on return. A local occurrence-identified tool result provides the boundary when middleware removes empty run markers; server-side tool results do not. Tests also cover buffered uppercase/redacted/replacement content, repeated citation batches, structured output, usage/continuation metadata, and exactly-once tool execution plus next-request history.

Offline unit validation on the pinned snapshot: core 9,338 passed (22 skipped, 2 expected failures; 92.07% coverage), AG-UI 1,848 passed (91.92%), OpenAI 775 passed. Source typing, Ruff, syntax and targeted post-adjustment regressions passed. Test typing passes with mypy, pyrefly, ty and pyright for both core and AG-UI. Zuban reports the same 447 core and 28 AG-UI diagnostics on unmodified and modified pinned snapshots; a normalized diagnostic comparison finds no additions. No live deployed-server or model-service validation is claimed.

[Eduard van Valkenburg (eavanvalkenburg)]

quifox (@quifox) A citation metadata update can erase the previous message — When an empty rebuilt message carries citation metadata and replace mode, this code treats it as metadata-only but keeps the replacement behavior. Aggregation can then replace the previous message with the empty one, so earlier text disappears. Please keep replacement updates out of the metadata-only path and add a regression with citations on the final redacted message.