feat(plugins): wire on_state_change_callback into plugin framework

[mportdata]

Description

Fixes #4393

Problem

Changes to state don't trigger plug-in methods the same way other events do.

Pattern for events triggering plug-in methods

Each plug-in in Google ADK inherits the BasePlugin class (BigQueryAgentAnalyticsPlugin, ContextFilterPlugin, etc.).

BasePlugin has several no-op methods such as before_run_callback, on_event_callback, etc. that are inherited by all Google ADK plug-ins. Each of these methods is associated with an event that can occur in Google ADK. Since they functionally do nothing, they are overwritten by plug-ins that inherit them when an action needs to be taken before, on or after an event.

For example, on_user_message_callback is a no-op method defined on BasePlugin. This is overwritten when defining BigQueryAgentAnalyticsPlugin so that the event is logged (using _log_event, another BigQueryAgentAnalyticsPlugin method).

This pattern is adhered to when having plug-in methods triggered by events.

The PluginManager class is used to call a specific method on all registered plug-ins. This works as methods associated with events are uniformly named across plug-ins.

The Runner class is what is used to call plug-in methods (via PluginManager) upon each event.

How state_delta is currently managed

This pattern is currently not followed when capturing changes to state in BigQueryAgentAnalyticsPlugin.

Currently, to capture state changes in the BigQueryAgentAnalyticsPlugin, when after_tool_callback is triggered we check to see if tool_context.actions.state_delta exists. If it does then a STATE_DELTA event is logged.

Solution

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 each event, detect non-empty state_delta and invoke the callback
4. BigQueryAgentAnalyticsPlugin — Remove inline STATE_DELTA logging from after_tool_callback and route it through on_state_change_callback instead

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)
• Single path — previously STATE_DELTA was logged inline within after_tool_callback; this has been removed in favour of routing all STATE_DELTA logging through on_state_change_callback, consistent with how the framework handles other event types

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
• Added test_after_tool_callback_no_inline_state_delta — verifies after_tool_callback no longer logs STATE_DELTA inline
• Added test_on_state_change_callback_logs_correctly — verifies STATE_DELTA is logged via on_state_change_callback
• All plugin tests pass locally (51 BQ plugin tests, 12 plugin manager tests)

ADK Web and BigQuery testing

When tested locally with ADK web I have verified STATE_DELTA events are still written to BigQuery the expected number of times with the fields and values as they were before (inline with what is described in the documentation also):

Checklist

• Read and followed the Contributing guidelines
• Ran autoformat.sh for formatting
• Changes are focused and minimal

[gemini-code-assist]

Summary of Changes

Hello @mportdata, 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 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. ↩

[gemini-code-assist]

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.

[gemini-code-assist]

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.

[mportdata]

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.

[gemini-code-assist]

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]:

[mportdata]

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