Skip to content

feat(plugins): wire on_state_change_callback into plugin framework#4395

Open
mportdata wants to merge 2 commits intogoogle:mainfrom
mportdata:feat/on-state-change-callback
Open

feat(plugins): wire on_state_change_callback into plugin framework#4395
mportdata wants to merge 2 commits intogoogle:mainfrom
mportdata:feat/on-state-change-callback

Conversation

@mportdata
Copy link

Description

Fixes #4393

Problem

The plugin framework provides hooks for various lifecycle events but has no mechanism to notify plugins when session state changes occur via event.actions.state_delta. The BigQueryAgentAnalyticsPlugin already implements on_state_change_callback(), demonstrating real demand, but the framework never invokes it because plumbing is missing at three layers.

Solution

Wire on_state_change_callback through the full plugin stack:

  1. BasePlugin — Add default no-op on_state_change_callback method (consistent with all other callbacks)
  2. PluginManager — Add "on_state_change_callback" to PluginCallbackName and add run_on_state_change_callback dispatcher
  3. Runner._exec_with_plugin — After yielding the final event, detect non-empty state_delta and invoke the callback

Design decisions:

  • After yield: event is delivered without delay; the notification is purely observational
  • dict() copy: prevents plugins from mutating the event's state_delta
  • Conditional import: defensive against circular imports (pattern used elsewhere in codebase)
  • event.actions is never None: guaranteed by default_factory=EventActions in event.py

Test plan

  • Added on_state_change_callback override to FullOverridePlugin in test_base_plugin.py
  • Added assertion to test_base_plugin_default_callbacks_return_none verifying default returns None
  • Added assertion to test_base_plugin_all_callbacks_can_be_overridden verifying override works
  • Added on_state_change_callback handler to TestPlugin in test_plugin_manager.py
  • Updated test_all_callbacks_are_supported to include new callback
  • Added test_run_on_state_change_callback — basic invocation returns None, callback logged
  • Added test_run_on_state_change_callback_calls_all_plugins — both plugins' call_logs contain the callback
  • Added test_run_on_state_change_callback_wraps_exceptions — exception wrapped in RuntimeError with chained cause
  • All 15 plugin tests pass locally

Checklist

  • Read and followed the Contributing guidelines
  • Ran autoformat.sh for formatting
  • Changes are focused and minimal (~127 lines across 5 files)

Add plumbing so that plugins are notified when an event carries
session state changes (non-empty state_delta). This closes a gap
where BasePlugin had no default method, PluginManager had no
dispatcher, and the runner never triggered the callback.

Fixes google#4393
@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @mportdata, I'm Gemini Code Assist1! 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 significantly enhances the plugin framework by introducing a new callback mechanism, on_state_change_callback, which allows plugins to react to session state changes. Previously, there was no direct way for plugins to be notified when event.actions.state_delta occurred. This change addresses that gap by plumbing the new callback through the BasePlugin, PluginManager, and Runner components, enabling observational state change notifications for all registered plugins.

Highlights

  • New Callback Introduction: Introduced a new on_state_change_callback to the BasePlugin class, providing a default no-op implementation for plugins to override when state changes occur.
  • Plugin Manager Integration: Integrated the on_state_change_callback into the PluginManager, including it in the PluginCallbackName enumeration and adding a dedicated dispatcher method run_on_state_change_callback.
  • Runner Invocation: Implemented the invocation of on_state_change_callback within the Runner._exec_with_plugin method, ensuring plugins are notified of state changes after an event with a non-empty state_delta is yielded.

🧠 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/plugins/base_plugin.py
    • Added async def on_state_change_callback method with callback_context and state_delta parameters, including comprehensive docstrings.
  • src/google/adk/plugins/plugin_manager.py
    • Extended PluginCallbackName to include "on_state_change_callback".
    • Implemented async def run_on_state_change_callback to dispatch the new callback to all registered plugins.
  • src/google/adk/runners.py
    • Modified _exec_with_plugin to check for final_event.actions.state_delta after yielding an event.
    • Conditionally imports CallbackContext and calls plugin_manager.run_on_state_change_callback with a copy of the state_delta if changes are detected.
  • tests/unittests/plugins/test_base_plugin.py
    • Updated FullOverridePlugin to include an override for on_state_change_callback.
    • Added new test assertions to test_base_plugin_default_callbacks_return_none and test_base_plugin_all_callbacks_can_be_overridden to cover the new callback.
  • tests/unittests/plugins/test_plugin_manager.py
    • Added on_state_change_callback to TestPlugin for testing purposes.
    • Modified test_all_callbacks_are_supported to include the new callback in its checks.
    • Introduced three new tests: test_run_on_state_change_callback, test_run_on_state_change_callback_calls_all_plugins, and test_run_on_state_change_callback_wraps_exceptions, to thoroughly validate the new callback's functionality and error handling.
Activity
  • Extensive unit tests were added across test_base_plugin.py and test_plugin_manager.py to cover the new on_state_change_callback functionality, including default behavior, overrides, dispatching to multiple plugins, and exception handling.
  • All 15 existing plugin tests were confirmed to pass locally after these changes.
  • The changes adhere to contributing guidelines and have been formatted using autoformat.sh.
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.

@adk-bot adk-bot added the core [Component] This issue is related to the core interface and implementation label Feb 6, 2026
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request effectively introduces the on_state_change_callback into the plugin framework, providing a much-needed mechanism for plugins to react to state changes. The implementation is well-structured across the BasePlugin, PluginManager, and Runner layers, and is accompanied by a comprehensive set of unit tests. My review includes a couple of suggestions to improve docstring accuracy and type hint consistency, which will enhance clarity for future developers using this new callback.

Comment on lines 382 to 384
This callback is invoked after an event with a non-empty
``state_delta`` is yielded from the runner. It is observational:
returning a value has no effect on execution flow.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The docstring states that "returning a value has no effect on execution flow." This is not entirely accurate. The PluginManager will short-circuit and stop calling this callback on subsequent plugins if any plugin returns a non-None value. While the runner ignores the returned value, this short-circuiting is an effect on the execution flow.

To avoid confusion for developers implementing this callback, I suggest clarifying this behavior.

Suggested change
This callback is invoked after an event with a non-empty
``state_delta`` is yielded from the runner. It is observational:
returning a value has no effect on execution flow.
This callback is invoked after an event with a non-empty
``state_delta`` is yielded from the runner. It is observational, but
returning a non-`None` value will short-circuit subsequent plugins.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 1375419. Updated the docstring to clarify the short-circuiting behaviour:

This callback is invoked after an event with a non-empty
``state_delta`` is yielded from the runner. It is observational, but
returning a non-`None` value will short-circuit subsequent plugins.

*,
callback_context: CallbackContext,
state_delta: dict[str, Any],
) -> Optional[None]:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The return type hint is Optional[None], but this method returns the result of _run_callbacks, which can return any value if a plugin decides to return something. This creates a type inconsistency, as _run_callbacks is typed to return Optional[Any]. To align with the actual behavior and the return type of the underlying _run_callbacks method, this should be Optional[Any].

Suggested change
) -> Optional[None]:
) -> Optional[Any]:

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 1375419. Changed the return type from Optional[None] to Optional[Any] to align with the return type of _run_callbacks.

- Clarify docstring: non-None return short-circuits subsequent plugins
- Fix return type: Optional[None] -> Optional[Any] to match _run_callbacks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core [Component] This issue is related to the core interface and implementation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Plugin framework: on_state_change_callback is never invoked

2 participants