Skip to content

fix: accumulate text across tool calls for output_key in streaming mode#5599

Open
theonlychant wants to merge 4 commits intogoogle:mainfrom
theonlychant:fix/output-key-text-accumulation
Open

fix: accumulate text across tool calls for output_key in streaming mode#5599
theonlychant wants to merge 4 commits intogoogle:mainfrom
theonlychant:fix/output-key-text-accumulation

Conversation

@theonlychant
Copy link
Copy Markdown

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):


Problem

When using LlmAgent with output_key and StreamingMode.SSE, text streamed
before tool calls is lost. The __maybe_save_output_to_state() method only
processes is_final_response() events, which only fires after all tool
executions complete. This causes 60–70% of agent responses to be discarded in
scenarios where agents make tool calls mid-response.

Example: An agent that streams intro text → calls a tool → streams progress
→ calls a tool → streams conclusion will only save the conclusion to output_key.

Solution

Accumulate text parts across all streaming events (not just the final response)
and persist the full combined text to output_key state. This ensures intro
text, progress updates, and conclusions are all preserved regardless of tool
call interleaving.


Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Manual End-to-End (E2E) Tests:

Run the reproduction script from issue #5590. Before this fix, match ratio was
~34.9% (only conclusion saved). After this fix, all text parts across tool calls
are accumulated correctly.

Expected output after fix:

Expected: 662 chars (6 text parts)
Actual:   662 chars (from output_key)
Match:    100%
✅ PASS: Text accumulation working

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

@vuth-seb
Copy link
Copy Markdown

vuth-seb commented May 6, 2026

Hello, I tried your patch and it's not working yet. I see that it's accumulating the text but the last response in event.is_final_response() still replaced the state with the last stream part

@theonlychant
Copy link
Copy Markdown
Author

Hello, I tried your patch and it's not working yet. I see that it's accumulating the text but the last response in event.is_final_response() still replaced the state with the last stream part

My apologies friend I forgot to add that part in the patch you can check now, I've just pushed a fix to my branch fix/output-key-text-accumulation. The key change is accumulating text across all streaming events and then using the full accumulated text in the is_final_response() block rather than only the last chunk. Give it a try and let me know if it works for your case!

@vuth-seb
Copy link
Copy Markdown

vuth-seb commented May 6, 2026

this works perfectly! Thanks

@theonlychant
Copy link
Copy Markdown
Author

this works perfectly! Thanks

You're welcome glad to help

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Text accumulation issue for output_key

2 participants