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

Describe the Bug:

When LlmAgent runs normally (no short-circuit), output_key correctly writes to event.actions.state_delta via __maybe_save_output_to_state(). The Runner commits this delta via append_event(). However, the value is not readable from state in:

1. The agent's own after_agent_callback
2. The next agent's before_agent_callback in a SequentialAgent

This means any callback that relies on reading an output_key value from state silently gets an empty/default value, while the LLM's actual output is lost.

Note: This is different from #4837, where output_key is never written because before_agent_callback short-circuits. In our case, the write succeeds (confirmed via instrumentation) — but the read fails.

Steps to Reproduce:

import asyncio
from google.adk.agents import LlmAgent
from google.adk.agents.callback_context import CallbackContext
from google.adk.runners import InMemoryRunner
from google.genai import types


async def check_output_key(callback_context: CallbackContext):
    """after_agent_callback: tries to read the output_key value."""
    result = callback_context.state.get("result", "NOT_FOUND")
    print(f"after_agent_callback: state['result'] = {repr(result)}")
    # Also check raw session.state
    raw = callback_context.session.state.get("result", "NOT_IN_RAW")
    print(f"after_agent_callback: session.state['result'] = {repr(raw)}")


agent = LlmAgent(
    name="my_agent",
    model="gemini-2.0-flash",
    instruction="Reply with a short greeting.",
    output_key="result",
    after_agent_callback=check_output_key,
)

runner = InMemoryRunner(agent=agent, app_name="test")


async def main():
    session = await runner.session_service.create_session(
        app_name="test", user_id="user1"
    )
    async for event in runner.run_async(
        user_id="user1",
        session_id=session.id,
        new_message=types.Content(parts=[types.Part(text="hello")]),
    ):
        pass

    # Post-workflow: this DOES work
    session = await runner.session_service.get_session(
        app_name="test", user_id="user1", session_id=session.id
    )
    print(f"post-workflow: state['result'] = {repr(session.state.get('result'))}")


asyncio.run(main())

Expected Behavior:

after_agent_callback: state['result'] = 'Hello! How can I help you?'
after_agent_callback: session.state['result'] = 'Hello! How can I help you?'
post-workflow: state['result'] = 'Hello! How can I help you?'

Observed Behavior:

after_agent_callback: state['result'] = 'NOT_FOUND'
after_agent_callback: session.state['result'] = 'NOT_IN_RAW'
post-workflow: state['result'] = 'Hello! How can I help you?'

The output_key value is available after the workflow completes (via get_session()), but not during after_agent_callback execution. This also affects the next agent's before_agent_callback in a SequentialAgent.

Debugging Evidence

We added print() instrumentation inside LlmAgent.__maybe_save_output_to_state():

DEBUG_output_key_WRITTEN agent=my_agent key=result result_len=5966

Immediately followed by our after_agent_callback:

after_agent_callback: state['result'] = ''        # empty
after_agent_callback: session.state['result'] = NOT_IN_RAW  # key doesn't exist

The output_key write succeeds (event.actions.state_delta["result"] is set, Runner calls append_event() which commits it), but the read fails in the callback.

Analysis

The issue appears to be in how _handle_after_agent_callback creates its CallbackContext:

# base_agent.py, line ~502
callback_context = CallbackContext(invocation_context)

The CallbackContext.__init__ creates a new State wrapper:

# callback_context.py, line ~51
self._state = State(
    value=invocation_context.session.state,
    delta=self._event_actions.state_delta,  # new, empty delta
)

Despite append_event() calling BaseSessionService._update_session_state() which does session.state.update({key: value}), the output_key value is not visible when the callback reads from state.get().

We verified that "result" not in callback_context.session.state returns True — the key genuinely does not exist in the session state dict at callback time.

Impact

This silently breaks any pattern where:

• after_agent_callback reads the agent's own output_key to post-process or extract structured data
• The next agent in a SequentialAgent reads a prior agent's output_key in its before_agent_callback

Workaround

Read the agent's output from callback_context.session.events instead of state:

async def check_output_key(callback_context):
    result = callback_context.state.get("result", "")

    # Fallback: read from session events when output_key state is not visible
    if not result:
        for event in reversed(callback_context.session.events):
            if (
                getattr(event, "author", "") == "my_agent"
                and event.content
                and event.content.parts
            ):
                for part in event.content.parts:
                    text = getattr(part, "text", "") or ""
                    if text.strip():
                        result = text.strip()
                        break
                if result:
                    break

Environment Details

• ADK Library Version: google-adk (latest, April 2026)
• Desktop OS: Windows Server 2019
• Python Version: 3.13.1
• Session Service: InMemorySessionService
• Tested with: Both Gemini and OpenAI (via LiteLlm) — same behavior on both providers

How often has this issue occurred?

• Always (100%) — reproducible on every request