Follow-up to #416 (merged in #481), which landed AgentReply.Attachments for the user-message producing path only. The attach_image tool, the ReplyAttachmentBuffer, and the drain-on-final logic are all wired solely into UserMessageHandler. Before extending attachments to the other producing paths (scheduled tasks, subagents, A2A results), three things need addressing — the first two are correctness, the third is structure.
1. Re-key the buffer per-turn / per-reply correlation (not per-session)
ReplyAttachmentBuffer is keyed by sessionId. That's safe today because the agent serializes producers within a session (IAgentWorkSerializer.AcquireForUserAsync + ISessionTracker.BeginSession cancelling the prior loop), so only one stage+drain cycle runs per session at a time.
It stops being safe the moment a second producer can be in flight for the same primary session concurrently — which is exactly what the deferred paths do:
- A2A results and subagent completions fold back to the primary
sessionId and publish their own replies asynchronously, potentially while the user's own loop is running.
- If
attach_image is registered on those paths, they all stage under the one sessionId key. Whichever path publishes a final reply first calls Drain(sessionId) and scoops up every staged attachment, including ones another concurrent producer staged for a different reply → attachments land on the wrong bubble.
Fix: key the buffer by something narrower than session (turn id / reply correlation id) so each producer drains only its own staged attachments.
2. Clear/expire staged attachments on cancel and ClearContext
Staging happens mid-loop; draining happens at the final reply. If a turn calls attach_image and is then cancelled before publishing (user sends a new message mid-loop), the loop throws OperationCanceledException, never publishes, and never drains — so the staged attachment lingers and the next final reply for that session drains it onto the wrong answer.
The buffer is also never cleared on session-cancel or ClearContext, unlike SessionClientCapabilityStore (which ClearContextHandler clears).
Fix: clear staged attachments on session-cancel and in ClearContextHandler; consider a TTL so orphaned stages can't linger indefinitely. (Per-turn keying from #1 largely subsumes the cancelled-turn leak, but the explicit clear is still worth it.)
3. Factor the drain step before duplicating it across handlers
The "build the final AgentReply + drain attachments onto it" step currently lives inline in UserMessageHandler.PublishReplyAsync. Each other producing path has its own reply-publishing method. Extending naively means copy-pasting the drain into ScheduledTaskHandler, SubagentRunner, and the A2A handlers (4+ sites).
Fix: extract the drain step into a small shared helper that every final-reply publish site calls, rather than duplicating it.
Scope note
None of this is a bug in the shipped slice — the user-message path is correct as-is. This is purely the prerequisite hardening for the deferred extension documented in design/client-rendering-capabilities.md.
Follow-up to #416 (merged in #481), which landed
AgentReply.Attachmentsfor the user-message producing path only. Theattach_imagetool, theReplyAttachmentBuffer, and the drain-on-final logic are all wired solely intoUserMessageHandler. Before extending attachments to the other producing paths (scheduled tasks, subagents, A2A results), three things need addressing — the first two are correctness, the third is structure.1. Re-key the buffer per-turn / per-reply correlation (not per-session)
ReplyAttachmentBufferis keyed bysessionId. That's safe today because the agent serializes producers within a session (IAgentWorkSerializer.AcquireForUserAsync+ISessionTracker.BeginSessioncancelling the prior loop), so only one stage+drain cycle runs per session at a time.It stops being safe the moment a second producer can be in flight for the same primary session concurrently — which is exactly what the deferred paths do:
sessionIdand publish their own replies asynchronously, potentially while the user's own loop is running.attach_imageis registered on those paths, they all stage under the onesessionIdkey. Whichever path publishes a final reply first callsDrain(sessionId)and scoops up every staged attachment, including ones another concurrent producer staged for a different reply → attachments land on the wrong bubble.Fix: key the buffer by something narrower than session (turn id / reply correlation id) so each producer drains only its own staged attachments.
2. Clear/expire staged attachments on cancel and ClearContext
Staging happens mid-loop; draining happens at the final reply. If a turn calls
attach_imageand is then cancelled before publishing (user sends a new message mid-loop), the loop throwsOperationCanceledException, never publishes, and never drains — so the staged attachment lingers and the next final reply for that session drains it onto the wrong answer.The buffer is also never cleared on session-cancel or
ClearContext, unlikeSessionClientCapabilityStore(whichClearContextHandlerclears).Fix: clear staged attachments on session-cancel and in
ClearContextHandler; consider a TTL so orphaned stages can't linger indefinitely. (Per-turn keying from #1 largely subsumes the cancelled-turn leak, but the explicit clear is still worth it.)3. Factor the drain step before duplicating it across handlers
The "build the final
AgentReply+ drain attachments onto it" step currently lives inline inUserMessageHandler.PublishReplyAsync. Each other producing path has its own reply-publishing method. Extending naively means copy-pasting the drain intoScheduledTaskHandler,SubagentRunner, and the A2A handlers (4+ sites).Fix: extract the drain step into a small shared helper that every final-reply publish site calls, rather than duplicating it.
Scope note
None of this is a bug in the shipped slice — the user-message path is correct as-is. This is purely the prerequisite hardening for the deferred extension documented in
design/client-rendering-capabilities.md.