fix: maybeSaveOutputToState overwrites OutputKey with empty string on function-call events

[dannovikov]

Closes issue #577

Description:

Any LLMAgent with OutputKey set that also uses tools has its output key overwritten to an empty string on every function-call and function-response event. This happens because maybeSaveOutputToState uses !event.Partial as its guard condition, which only filters out streaming chunks — function-call events pass through, their parts contain no .Text, so strings.Builder produces "", which overwrites StateDelta[OutputKey].

Go (agent/llmagent/llmagent.go:392):

if a.OutputKey != "" && !event.Partial && event.Content != nil && len(event.Content.Parts) > 0 {

Python (agents/llm_agent.py:819-821):

if self.output_key and event.is_final_response() and event.content and event.content.parts:

The proposed change replaces !event.Partial with event.IsFinalResponse() — a method that already exists in session.Event and mirrors the Python behavior exactly — so that OutputKey is only written on the agent's actual final text response.

Testing

Added two subtests to the existing TestLlmAgent_MaybeSaveOutputToState in llmagent_saveoutput_test.go:

• "skips for function call events" — FunctionCall part, Partial=false → asserts StateDelta stays empty
• "skips for function response events" — FunctionResponse part, Partial=false → asserts StateDelta stays empty

Both fail before the fix (StateDelta gets {"result": ""}), pass after.

[gemini-code-assist]

Summary of Changes

Hello @dannovikov, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a bug in the LLM agent's state management where the OutputKey was inadvertently cleared during intermediate function call and response events. The change ensures that the agent's output state is only updated with the final, meaningful text response, preventing data loss and aligning the Go implementation with its Python counterpart.

Highlights

• Fixes Agent Output Overwrite: Corrected an issue where the OutputKey in LLMAgent was incorrectly overwritten with an empty string during function-call and function-response events.
• Updated Output Saving Logic: Modified the maybeSaveOutputToState function in the Go llmagent to use !event.IsFinalResponse() instead of !event.Partial, ensuring output is saved only for final text responses and aligning with Python behavior.
• Enhanced Test Coverage: Added two new subtests to TestLlmAgent_MaybeSaveOutputToState to specifically verify that function-call and function-response events do not incorrectly modify the agent's state delta.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• agent/llmagent/llmagent.go
  • Modified the conditional logic within maybeSaveOutputToState to correctly identify when to save the agent's output.
• agent/llmagent/llmagent_saveoutput_test.go
  • Introduced new test cases to validate that maybeSaveOutputToState properly skips saving output for function call and function response events.

Activity

• No human activity has been recorded on this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request aims to fix a bug where maybeSaveOutputToState incorrectly overwrites an output key with an empty string for function-call events. The approach of using event.IsFinalResponse() is correct and aligns the Go implementation with the Python version. The added test cases are excellent and correctly target the bug. However, there appears to be a small but critical logic error in the implementation of the fix. The condition should likely be event.IsFinalResponse() instead of !event.IsFinalResponse() to ensure the output is saved only for final text responses.