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.

Additionally state changes can also be caused by _append_new_message_to_session however these never trigger plug-in methods as this happens separately to after_tool_callback (the only place state changes are currently captured).

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
5. Runner._handle_new_message — After handling a new message, 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)
• 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
• Added test_state_delta_in_run_async_triggers_on_state_change_callback — verifies caller-supplied state_delta in run_async triggers on_state_change_callback
• Added test_no_state_delta_does_not_trigger_on_state_change_callback — verifies callback is not called when no state_delta is provided
• All tests pass locally (51 BQ plugin tests, 12 plugin manager tests, 37 runner 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):

_append_new_message_to_session testing

Below is a scrip that instantiates an agent with an initial state and then runs the agent provoking an additional state change during run time, I have including a screenshot of how the BigQueryAgentAnalyticsPlugin plug-in has written these to BigQuery. You can see that STATE_DELTA is only logged once for each state change as expected.

import asyncio

from dotenv import load_dotenv
from google.adk.runners import InMemoryRunner
from google.genai import types

load_dotenv(override=True)

from test_agent.agent import app


async def main():
    runner = InMemoryRunner(app=app)
    session = await runner.session_service.create_session(
        app_name="test_agent", user_id="test_user"
    )

    # First message includes caller-supplied state_delta to test
    # on_state_change_callback via _handle_new_message path.
    initial_state = {
        "user_preferences": {
            "language": "en",
            "timezone": "Europe/London",
        },
        "session_source": "cli_test",
    }

    prompts = [
        "add this contact: name: mike, email: mike@mail.com, number: 01234 98765",
    ]

    for i, prompt in enumerate(prompts):
        print(f"\nUser > {prompt}")
        kwargs = {
            "user_id": "test_user",
            "session_id": session.id,
            "new_message": types.Content(
                role="user", parts=[types.Part(text=prompt)]
            ),
        }
        # Pass state_delta only on the first message
        if i == 0:
            kwargs["state_delta"] = initial_state
            print(f"  [caller state_delta: {initial_state}]")

        async for event in runner.run_async(**kwargs):
            if event.content and event.content.parts:
                for part in event.content.parts:
                    if part.text:
                        print(f"{event.author} > {part.text}")
            if event.actions and event.actions.state_delta:
                print(f"  [state_delta: {event.actions.state_delta}]")

    final_session = await runner.session_service.get_session(
        app_name="test_agent", user_id="test_user", session_id=session.id
    )
    print(f"\nFinal session state: {final_session.state}")


if __name__ == "__main__":
    asyncio.run(main())

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.