Description
Late message-linked citations can produce an empty duplicate assistant message through public AG-UI streaming aggregation.
After m1/First and m2/Second have completed, an annotations event targeting m1 produces three final messages: m1/First without its citation, m2/Second, and an empty m1 carrying the citation. This affects default AGUIChatClient.get_response(stream=True) and Agent(client=client).run(stream=True). Both non-streaming controls return the expected two messages.
Code Sample
A mocked AG-UI endpoint emits this ordinary SSE event sequence between RUN_STARTED and RUN_FINISHED:
[
{"type":"TEXT_MESSAGE_START","messageId":"m1","role":"assistant"},
{"type":"TEXT_MESSAGE_CONTENT","messageId":"m1","delta":"First"},
{"type":"TEXT_MESSAGE_END","messageId":"m1"},
{"type":"TEXT_MESSAGE_START","messageId":"m2","role":"assistant"},
{"type":"TEXT_MESSAGE_CONTENT","messageId":"m2","delta":"Second"},
{"type":"TEXT_MESSAGE_END","messageId":"m2"},
{"type":"CUSTOM","name":"annotations","value":{"messageId":"m1","annotations":[{"type":"citation","url":"https://example.test/source"}]}}
]After consuming either public response stream and awaiting get_final_response(), m1 should retain its citation without another m1 message.
Error Messages / Stack Traces
No exception is raised. The error is the final message list:
- Expected: m1/First with its citation, m2/Second.
- Actual: m1/First without its citation, m2/Second, m1/empty with the citation.
Package Versions
Verified from source at b9d24c8:
- agent-framework-core: 1.20.0
- agent-framework-ag-ui: 1.5.0
Python Version
Python 3.12.
Additional Context
This follows #8948 and #7460. The AG-UI inner finalizer handles late annotations by message ID, but the default function-loop stream and mapped Agent stream each perform separate generic aggregation outside that finalizer.
I have a focused fix prepared with regression coverage for both public paths, completed tool rounds, structured output, buffered middleware, and compatibility with existing client and Agent finalizers.
The proposed private hooks retain generic aggregation by default and let AG-UI attach its message-linked annotations at the whole-loop and mapped-agent boundaries. Provider per-model-turn finalizers remain unchanged, and mapped AgentResponseUpdate objects remain authoritative. Please review whether these extension points fit the core aggregation contracts.
Description
Late message-linked citations can produce an empty duplicate assistant message through public AG-UI streaming aggregation.
After m1/First and m2/Second have completed, an annotations event targeting m1 produces three final messages: m1/First without its citation, m2/Second, and an empty m1 carrying the citation. This affects default AGUIChatClient.get_response(stream=True) and Agent(client=client).run(stream=True). Both non-streaming controls return the expected two messages.
Code Sample
A mocked AG-UI endpoint emits this ordinary SSE event sequence between RUN_STARTED and RUN_FINISHED:
[ {"type":"TEXT_MESSAGE_START","messageId":"m1","role":"assistant"}, {"type":"TEXT_MESSAGE_CONTENT","messageId":"m1","delta":"First"}, {"type":"TEXT_MESSAGE_END","messageId":"m1"}, {"type":"TEXT_MESSAGE_START","messageId":"m2","role":"assistant"}, {"type":"TEXT_MESSAGE_CONTENT","messageId":"m2","delta":"Second"}, {"type":"TEXT_MESSAGE_END","messageId":"m2"}, {"type":"CUSTOM","name":"annotations","value":{"messageId":"m1","annotations":[{"type":"citation","url":"https://example.test/source"}]}} ]After consuming either public response stream and awaiting get_final_response(), m1 should retain its citation without another m1 message.
Error Messages / Stack Traces
No exception is raised. The error is the final message list:
Package Versions
Verified from source at b9d24c8:
Python Version
Python 3.12.
Additional Context
This follows #8948 and #7460. The AG-UI inner finalizer handles late annotations by message ID, but the default function-loop stream and mapped Agent stream each perform separate generic aggregation outside that finalizer.
I have a focused fix prepared with regression coverage for both public paths, completed tool rounds, structured output, buffered middleware, and compatibility with existing client and Agent finalizers.
The proposed private hooks retain generic aggregation by default and let AG-UI attach its message-linked annotations at the whole-loop and mapped-agent boundaries. Provider per-model-turn finalizers remain unchanged, and mapped AgentResponseUpdate objects remain authoritative. Please review whether these extension points fit the core aggregation contracts.