fix(agents): Ensure state deltas are visible in agent callbacks

[shitalbabasopatil]

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

• Closes: LlmAgent output_key state delta not visible to after_agent_callback or next agent's before_agent_callback in SequentialAgent #5566
• Related: LlmAgent output_key state delta not visible to after_agent_callback or next agent's before_agent_callback in SequentialAgent #5566

2. Or, if no issue exists, describe the change:

If applicable, please follow the issue templates to provide as much detail as
possible.

Problem:

1. The Core Issue
   During the execution lifecycle of LlmAgent, outputs generated by the agent are persisted to the session state using the __maybe_save_output_to_state() method. This method successfully stores the payload inside event.actions.state_delta on the model response event, and the Runner successfully persists this to the session using the SessionService's append_event method.

However, plugins or user-defined callbacks executed immediately afterwards (such as after_agent_callback and subsequent before_agent_callback invocations) are unable to read this newly committed state. Attempting to fetch state['result'] returns an empty value (or results in a KeyError), because the key genuinely does not exist in the session state dictionary at the time the callback is executed.

2. Root Cause Analysis
   The failure occurs due to how CallbackContext is instantiated within the execution loop in src/google/adk/agents/base_agent.py:

async def _handle_after_agent_callback(self, invocation_context: InvocationContext) -> Optional[Event]:
    callback_context = CallbackContext(invocation_context)
    ...

When CallbackContext is initialized without explicitly passing event_actions, it defaults to an empty EventActions() object:

class Context(ReadonlyContext):
  def __init__(...):
    self._event_actions = event_actions or EventActions()
    self._state = State(
        value=invocation_context.session.state,
        delta=self._event_actions.state_delta,
    )

As a result, the context relies entirely on invocation_context.session.state. While the Runner does call SessionService.append_event(), depending on the session service implementation (especially distributed or storage-backed ones like SpannerSessionService), the updated state might not immediately reflect inside the locally cached invocation_context.session.state dictionary at the precise moment the callback runs.

Even if it does update correctly in some environments, creating the CallbackContext without the pending state_delta severs the connection to the most recent changes that were just emitted by the _run_async_impl loop.

Solution:
To resolve this issue and maintain consistency with how context states are scoped throughout ADK, CallbackContext needs to proactively load the EventActions from the most recent event.

We updated src/google/adk/agents/base_agent.py to intercept the last emitted event and extract its actions object, passing it directly into the CallbackContext:

    last_event = None
    if ctx.session.events:
      last_event = ctx.session.events[-1]

    event_actions = None
    if last_event and last_event.actions:
      event_actions = last_event.actions

    callback_context = CallbackContext(ctx, event_actions=event_actions)

Testing Plan

To ensure this issue is permanently resolved and to prevent future regressions, two new asynchronous regression tests were added to tests/unittests/agents/test_base_agent.py. These tests simulate the entire InMemoryRunner execution flow, mimicking exactly how LlmAgent records and broadcasts state changes.

Scenario 1: test_run_async_after_agent_callback_state_visibility
Objective: Verify that an agent's own after_agent_callback correctly observes state changes generated within its own reason-act loop.
Implementation Details:

• A custom subclass _StateTestingAgent was created to override _run_async_impl. This mock agent yields a single Event configured with EventActions(state_delta={"test_key": "test_val"}).
• The after_agent_callback is wired to assert that ctx.state['test_key'] == 'test_val'.
• The test invokes InMemoryRunner.run_async(), enforcing that the Runner successfully applies the delta to the session in real time, and the callback resolves the "test_key" dynamically via CallbackContext.

Scenario 2: test_run_async_before_agent_callback_state_visibility
Objective: Verify that in a multi-agent orchestrated flow, the succeeding agent's before_agent_callback correctly observes the state changes left behind by the preceding agent.
Implementation Details:

• A SequentialAgent is instantiated containing two sub-agents: agent1 and agent2.
• agent1 utilizes the _StateTestingAgent mock to emit an event containing {"shared_key": "shared_val"} in its state_delta.
• agent2 defines a before_agent_callback that asserts ctx.state['shared_key'] == 'shared_val'.
• This ensures that context states propagate successfully across consecutive agent hooks inside a shared InvocationContext.

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

Please include a summary of passed pytest results.

Manual End-to-End (E2E) Tests:

Please provide instructions on how to manually test your changes, including any
necessary setup or configuration. Please provide logs or screenshots to help
reviewers better understand the fix.

Checklist

• I have read the CONTRIBUTING.md document.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end.
• Any dependent changes have been merged and published in downstream modules.

Additional context

Add any other context or screenshots about the feature request here.