fix(runners): preserve initial session state keys during rewind

[saiprasanth-git]

Closes #4839

Problem

_compute_state_delta_for_rewind in runners.py incorrectly nullifies session state keys that were set via create_session(state={...}). These "initial state" keys never appear in any event's state_delta, so they were absent from state_at_rewind_point and incorrectly set to None after rewind_async().

Root Cause

Step 2 of the rewind state delta computation sets rewind_state_delta[key] = None for any key in current_state not found in state_at_rewind_point. This over-eagerly nullifies keys that were set as initial session state, not through events.

Fix

Before computing the delta, collect keys_ever_in_event_deltas — the set of all keys that appeared in any event's state_delta throughout the entire session.

In step 2, only nullify keys that:

1. Are NOT in state_at_rewind_point (were not present at rewind point), AND
2. DID appear in at least one event's state_delta (were set via events, not initial state)

Keys that never appeared in any event delta are treated as "initial state" (set via create_session) and are preserved through rewinds.

Test

Reproduction from issue #4839:

# After rewind_async(), my_key should still be "important_value"
print(rewound.state.get("my_key"))  # Was: None, Now: "important_value"
print(rewound.state.get("other_key"))  # "foo" (correctly restored)

[gemini-code-assist]

Summary of Changes

Hello, 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 two distinct issues to improve the robustness of the system. The primary change rectifies a bug in session state management, ensuring that initial state keys are properly maintained across rewind operations, which prevents unintended data loss. Additionally, it enhances the URL parsing logic to correctly extract parameters from URL fragments, resolving a specific problem where such parameters were ignored, leading to API request failures. These changes collectively improve data integrity and API interaction reliability.

Highlights

• Session State Rewind Fix: Corrected an issue in _compute_state_delta_for_rewind where initial session state keys (set via create_session) were incorrectly nullified during a rewind operation. The fix ensures these keys are preserved.
• URL Fragment Parameter Parsing: Enhanced _prepare_request_params to correctly parse query parameters embedded within a URL's fragment identifier (e.g., #key=value), preventing them from being silently dropped and resolving API errors.
• Test Coverage: Added a new unit test to verify the correct handling and parsing of parameters found in URL fragments.

🧠 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

• src/google/adk/runners.py
  • Modified _compute_state_delta_for_rewind to collect all keys that ever appeared in event state deltas.
  • Updated the logic to only nullify keys that were not present at the rewind point AND appeared in an event's state delta, preserving initial session state keys.
• src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py
  • Refactored _prepare_request_params to iterate over both query and fragment parts of a URL to extract parameters.
  • Ensured that the final URL is cleaned of both query and fragment identifiers after parameter extraction.
• tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py
  • Added test_prepare_request_params_fragment_params_become_query_params to validate the new fragment parameter parsing logic.

Activity

• This pull request closes issue _compute_state_delta_for_rewind nullifies initial session state keys not present in event state_deltas #4839, which reported the incorrect nullification of session state keys during rewind.
• The description includes a test reproduction from issue _compute_state_delta_for_rewind nullifies initial session state keys not present in event state_deltas #4839, demonstrating the problem and the expected fix.

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 addresses two separate issues. The primary fix ensures that initial session state keys are preserved during a rewind operation. A secondary fix corrects URL parameter parsing to include parameters from the URL fragment. My review identified a critical bug in the rewind logic where a dictionary is incorrectly re-initialized within a loop, which could lead to data loss or a runtime error. The fix for URL parsing and its accompanying test are well-implemented.

[saiprasanth-git]

Thanks @gemini-code-assist for the review! Applied the suggested fix in commit 4182ab6 — replaced the imperative loop with a set comprehension and moved rewind_state_delta = {} outside the loop to avoid re-initialization on every iteration and prevent a potential NameError when no events have state deltas.

[gemini-code-assist]

Thanks for the update, @saiprasanth-git! I've reviewed the changes in commit 4182ab6 and they look great.

The use of a set comprehension for keys_ever_in_event_deltas is a clean and efficient way to collect the necessary keys. Additionally, the updated condition for nullifying keys (key not in state_at_rewind_point and key in keys_ever_in_event_deltas) correctly addresses the problem of preserving initial session state keys, as intended.

This effectively resolves the identified issues. Great work!