ADK Streaming Grounding Callback Bug Report

Describe the bug

ADK callbacks that process grounding metadata and update session state fail to execute properly when streaming is enabled (/run_sse), but work correctly with non-streaming requests (/run). The callback executes but finds no grounding metadata in session.events, resulting in empty state updates.

To Reproduce

Steps to reproduce the behavior:

1. Create an agent with grounding callback:

from google.adk.agents import LlmAgent
from google.adk.tools import google_search
from google.adk.agents.callback_context import CallbackContext

def collect_research_sources_callback(callback_context: CallbackContext) -> None:
    """Collect sources from grounding metadata"""
    session = callback_context._invocation_context.session
    url_to_short_id = callback_context.state.get("url_to_short_id", {})
    sources = callback_context.state.get("sources", {})
    
    # Process grounding metadata from session events
    for event in session.events:
        if not (event.grounding_metadata and event.grounding_metadata.grounding_chunks):
            continue
        # ... process grounding data and update state
    
    callback_context.state["url_to_short_id"] = url_to_short_id
    callback_context.state["sources"] = sources

dummy_agent = LlmAgent(
    name="dummy_agent",
    model="gemini-2.5-flash",
    tools=[google_search],
    instruction="Use google_search tool to answer questions.",
    after_agent_callback=collect_research_sources_callback,
)

1. Test with streaming endpoint:

curl -X POST http://localhost:8000/run_sse \
-H "Content-Type: application/json" \
-d '{
"appName": "dummy_agent",
"userId": "u_123",
"sessionId": "s_streaming_test",
"newMessage": {
    "role": "user", 
    "parts": [{"text": "What is the weather in New York?"}]
}
}'

1. Observe the final state event:

{
  "actions": {
    "state_delta": {
      "url_to_short_id": {},  // Empty - callback found no sources
      "sources": {}           // Empty - no grounding metadata processed
    }
  }
}

1. Test with non-streaming endpoint:

curl -X POST http://localhost:8000/run \
-H "Content-Type: application/json" \
-d '{
"appName": "dummy_agent", 
"userId": "u_123",
"sessionId": "s_non_streaming_test",
"newMessage": {
    "role": "user",
    "parts": [{"text": "What is the weather in New York?"}]
}
}'

1. Observe callback works correctly - state contains actual grounding sources

Expected behavior

The after_agent_callback should execute with the same session.events data regardless of whether the request uses /run (non-streaming) or /run_sse (streaming). The callback should be able to access grounding metadata from Google Search tool calls in both scenarios.

Actual behavior

• Streaming (/run_sse): Callback executes but session.events contains no grounding metadata, resulting in empty state updates
• Non-streaming (/run): Callback executes correctly and finds grounding metadata in session.events

Screenshots

Final state event showing empty sources in streaming mode:

{
  "invocation_id": "e-7ce8dfbb-ef34-4351-9a67-6a68772404cb",
  "author": "dummy_agent", 
  "actions": {
    "state_delta": {
      "url_to_short_id": {},
      "sources": {}
    }
  }
}

Desktop:

• OS: macOS
• Python version: Python 3.10+
• ADK version: google-adk==1.4.2

Model Information:

• Model: gemini-2.5-flash
• Tool: google_search (from google.adk.tools)
• Callback: after_agent_callback

Additional context

This appears to be a timing or execution order issue where streaming requests don't properly populate session.events with grounding metadata before the callback executes. The callback function itself is correct (works in non-streaming mode), but the session state appears incomplete during streaming execution.

Impact: This prevents real-time source tracking and citation features in streaming applications, which is a critical use case for grounded AI agents.

Workaround: Currently using non-streaming /run endpoint, but this eliminates real-time user experience benefits.