fix(agents): fix visibility of output_key state delta in callbacks

Co-authored-by: Amaad Martin <amaadmartin@google.com>
PiperOrigin-RevId: 916112779

Author: Amaad Martin

src/google/adk/runners.py

@@ -608,7 +608,7 @@ async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
         async with Aclosing(
             self._exec_with_plugin(
                 invocation_context=invocation_context,
-                session=session,
+                session=invocation_context.session,
                 execute_fn=execute,
                 is_live_call=False,
             )
@@ -622,7 +622,7 @@ async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
           logger.debug('Running event compactor.')
           await _run_compaction_for_sliding_window(
               self.app,
-              session,
+              invocation_context.session,
               self.session_service,
               skip_token_compaction=invocation_context.token_compaction_checked,
           )
@@ -841,7 +841,7 @@ async def _exec_with_plugin(
 
     Args:
       invocation_context: The invocation context
-      session: The current session
+      session: The current session (ignored, kept for backward compatibility)
       execute_fn: A callable that returns an AsyncGenerator of Events
       is_live_call: Whether this is a live call
 
@@ -866,7 +866,7 @@ async def _exec_with_plugin(
       )
       if self._should_append_event(early_exit_event, is_live_call):
         await self.session_service.append_event(
-            session=session,
+            session=invocation_context.session,
             event=early_exit_event,
         )
       yield early_exit_event
@@ -931,13 +931,13 @@ async def _exec_with_plugin(
                 )
                 if self._should_append_event(event, is_live_call):
                   await self.session_service.append_event(
-                      session=session, event=output_event
+                      session=invocation_context.session, event=output_event
                   )
 
                 for buffered_event in buffered_events:
                   logger.debug('Appending buffered event: %s', buffered_event)
                   await self.session_service.append_event(
-                      session=session, event=buffered_event
+                      session=invocation_context.session, event=buffered_event
                   )
                   yield buffered_event  # yield buffered events to caller
                 buffered_events = []
@@ -947,12 +947,12 @@ async def _exec_with_plugin(
                 if self._should_append_event(event, is_live_call):
                   logger.debug('Appending non-buffered event: %s', event)
                   await self.session_service.append_event(
-                      session=session, event=output_event
+                      session=invocation_context.session, event=output_event
                   )
           else:
             if event.partial is not True:
               await self.session_service.append_event(
-                  session=session, event=output_event
+                  session=invocation_context.session, event=output_event
               )
 
           yield output_event
@@ -1004,8 +1004,8 @@ async def _append_new_message_to_session(
         file_name = f'artifact_{invocation_context.invocation_id}_{i}'
         await self.artifact_service.save_artifact(
             app_name=self.app_name,
-            user_id=session.user_id,
-            session_id=session.id,
+            user_id=invocation_context.session.user_id,
+            session_id=invocation_context.session.id,
             filename=file_name,
             artifact=part,
         )
@@ -1032,7 +1032,9 @@ async def _append_new_message_to_session(
     if function_call := invocation_context._find_matching_function_call(event):
       event.branch = function_call.branch
 
-    await self.session_service.append_event(session=session, event=event)
+    await self.session_service.append_event(
+        session=invocation_context.session, event=event
+    )
 
   async def run_live(
       self,
@@ -1127,7 +1129,9 @@ async def run_live(
     )
 
     root_agent = self.agent
-    invocation_context.agent = self._find_agent_to_run(session, root_agent)
+    invocation_context.agent = self._find_agent_to_run(
+        invocation_context.session, root_agent
+    )
 
     async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
       async with Aclosing(ctx.agent.run_live(ctx)) as agen:
@@ -1137,7 +1141,7 @@ async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
     async with Aclosing(
         self._exec_with_plugin(
             invocation_context=invocation_context,
-            session=session,
+            session=invocation_context.session,
             execute_fn=execute,
             is_live_call=True,
         )
@@ -1355,14 +1359,16 @@ async def _setup_context_for_new_invocation(
     # Step 2: Handle new message, by running callbacks and appending to
     # session.
     await self._handle_new_message(
-        session=session,
+        session=invocation_context.session,
         new_message=new_message,
         invocation_context=invocation_context,
         run_config=run_config,
         state_delta=state_delta,
     )
     # Step 3: Set agent to run for the invocation.
-    invocation_context.agent = self._find_agent_to_run(session, self.agent)
+    invocation_context.agent = self._find_agent_to_run(
+        invocation_context.session, self.agent
+    )
     return invocation_context
 
   async def _setup_context_for_resumed_invocation(
@@ -1411,7 +1417,7 @@ async def _setup_context_for_resumed_invocation(
     # Step 3: Maybe handle new message.
     if new_message:
       await self._handle_new_message(
-          session=session,
+          session=invocation_context.session,
           new_message=user_message,
           invocation_context=invocation_context,
           run_config=run_config,
@@ -1425,7 +1431,9 @@ async def _setup_context_for_resumed_invocation(
     # started from a sub-agent and paused on a sub-agent.
     # We should find the appropriate agent to run to continue the invocation.
     if self.agent.name not in invocation_context.end_of_agents:
-      invocation_context.agent = self._find_agent_to_run(session, self.agent)
+      invocation_context.agent = self._find_agent_to_run(
+          invocation_context.session, self.agent
+      )
     return invocation_context
 
   def _find_user_message_for_invocation(
@@ -1559,7 +1567,7 @@ async def _handle_new_message(
       if 'save_input_blobs_as_artifacts' in run_config.model_fields_set:
         deprecated_save_blobs = run_config.save_input_blobs_as_artifacts
       await self._append_new_message_to_session(
-          session=session,
+          session=invocation_context.session,
           new_message=new_message,
           invocation_context=invocation_context,
           save_input_blobs_as_artifacts=deprecated_save_blobs,

tests/unittests/agents/test_output_key_visibility.py

@@ -0,0 +1,180 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for LlmAgent output_key visibility in callbacks."""
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.live_request_queue import LiveRequestQueue
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.sequential_agent import SequentialAgent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.auto_flow import AutoFlow
+from google.genai import types
+import pytest
+from pytest_mock import MockerFixture
+
+from .. import testing_utils
+
+# Standard MockModel will be used instead of SafeMockModel
+
+
+@pytest.mark.asyncio
+async def test_output_key_visibility_in_after_agent_callback():
+  """Test that output_key state delta is visible in after_agent_callback."""
+  mock_response = "Hello! How can I help you?"
+  mock_model = testing_utils.MockModel.create(responses=[mock_response])
+
+  callback_called = False
+  captured_state_value = None
+  captured_session_state_value = None
+
+  async def check_output_key(callback_context: CallbackContext):
+    nonlocal callback_called, captured_state_value, captured_session_state_value
+    callback_called = True
+    captured_state_value = callback_context.state.get("result", "NOT_FOUND")
+    captured_session_state_value = callback_context.session.state.get(
+        "result", "NOT_IN_RAW"
+    )
+
+  agent = LlmAgent(
+      name="my_agent",
+      model=mock_model,
+      instruction="Reply with a short greeting.",
+      output_key="result",
+      after_agent_callback=check_output_key,
+  )
+
+  runner = testing_utils.InMemoryRunner(root_agent=agent)
+
+  events = await runner.run_async(new_message="hello")
+
+  assert callback_called, "Callback was not called"
+
+  assert (
+      captured_state_value == mock_response
+  ), f"Expected {mock_response}, got {captured_state_value}"
+  assert (
+      captured_session_state_value == mock_response
+  ), f"Expected {mock_response}, got {captured_session_state_value}"
+
+
+@pytest.mark.asyncio
+async def test_output_key_visibility_in_run_live(mocker: MockerFixture):
+  """Test that output_key state delta is visible in after_agent_callback in run_live."""
+  mock_response = "Hello! How can I help you?"
+  mock_model = testing_utils.MockModel.create(responses=[mock_response])
+
+  callback_called = False
+  captured_state_value = None
+  captured_session_state_value = None
+
+  async def check_output_key(callback_context: CallbackContext):
+    nonlocal callback_called, captured_state_value, captured_session_state_value
+    callback_called = True
+    captured_state_value = callback_context.state.get("result", "NOT_FOUND")
+    captured_session_state_value = callback_context.session.state.get(
+        "result", "NOT_IN_RAW"
+    )
+
+  agent = LlmAgent(
+      name="my_agent",
+      model=mock_model,
+      instruction="Reply with a short greeting.",
+      output_key="result",
+      after_agent_callback=check_output_key,
+  )
+
+  async def mock_auto_flow_run_live(self, ctx):
+    yield Event(
+        id=Event.new_id(),
+        invocation_id=ctx.invocation_id,
+        author=ctx.agent.name,
+        content=types.Content(parts=[types.Part(text=mock_response)]),
+    )
+
+  mocker.patch.object(AutoFlow, "run_live", mock_auto_flow_run_live)
+
+  runner = testing_utils.InMemoryRunner(root_agent=agent)
+  live_queue = LiveRequestQueue()
+
+  agen = runner.runner.run_live(
+      user_id="test_user",
+      session_id=runner.session.id,
+      live_request_queue=live_queue,
+  )
+
+  # Send a message to trigger the agent
+  live_queue.send_content(
+      types.Content(role="user", parts=[types.Part(text="hello")])
+  )
+
+  live_queue.close()
+
+  async for event in agen:
+    pass
+
+  assert callback_called, "Callback was not called"
+  assert (
+      captured_state_value == mock_response
+  ), f"Expected {mock_response}, got {captured_state_value}"
+  assert (
+      captured_session_state_value == mock_response
+  ), f"Expected {mock_response}, got {captured_session_state_value}"
+
+
+@pytest.mark.asyncio
+async def test_output_key_visibility_in_sequential_agent():
+  """Test that output_key state delta is visible in next agent's before_agent_callback."""
+  mock_response = "Hello from agent 1!"
+  mock_model = testing_utils.MockModel.create(
+      responses=[mock_response, "Hello from agent 2!"]
+  )
+
+  callback_called = False
+  captured_session_state_value = None
+
+  async def check_before_agent(callback_context: CallbackContext):
+    nonlocal callback_called, captured_session_state_value
+    callback_called = True
+    captured_session_state_value = callback_context.session.state.get(
+        "result", "NOT_FOUND"
+    )
+
+  agent_1 = LlmAgent(
+      name="agent_1",
+      model=mock_model,
+      instruction="Reply with a short greeting.",
+      output_key="result",
+  )
+
+  agent_2 = LlmAgent(
+      name="agent_2",
+      model=mock_model,
+      instruction="Reply with a short greeting.",
+      before_agent_callback=check_before_agent,
+  )
+
+  sequential_agent = SequentialAgent(
+      name="seq_agent",
+      sub_agents=[agent_1, agent_2],
+  )
+
+  runner = testing_utils.InMemoryRunner(root_agent=sequential_agent)
+
+  events = await runner.run_async(new_message="hello")
+
+  assert callback_called, "Callback was not called"
+  assert (
+      captured_session_state_value == mock_response
+  ), f"Expected {mock_response}, got {captured_session_state_value}"