chore: import upstream snapshot with attribution

tests/unittests/flows/__init__.py

@@ -0,0 +1,13 @@
+# 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/unittests/flows/llm_flows/__init__.py

@@ -0,0 +1,13 @@
+# 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/unittests/flows/llm_flows/test_agent_transfer_system_instructions.py

@@ -0,0 +1,344 @@
+# 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.
+
+"""Behavioral tests for agent transfer system instructions.
+
+These tests verify the behavior of the agent transfer system by calling
+the request processor and checking the resulting system instructions not just
+implementation.
+"""
+
+from typing import AsyncGenerator
+
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import Agent
+from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows import agent_transfer
+from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
+from google.adk.models.llm_request import LlmRequest
+from google.adk.plugins.plugin_manager import PluginManager
+from google.adk.runners import RunConfig
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class _NonLlmAgent(BaseAgent):
+  """A minimal BaseAgent subclass that, like any non-LlmAgent, has no `mode`."""
+
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    yield Event(author=self.name, invocation_id=ctx.invocation_id)
+
+
+async def create_test_invocation_context(agent: Agent) -> InvocationContext:
+  """Helper to create constructed InvocationContext."""
+  session_service = InMemorySessionService()
+  memory_service = InMemoryMemoryService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+
+  return InvocationContext(
+      artifact_service=InMemoryArtifactService(),
+      session_service=session_service,
+      memory_service=memory_service,
+      plugin_manager=PluginManager(plugins=[]),
+      invocation_id='test_invocation_id',
+      agent=agent,
+      session=session,
+      user_content=types.Content(
+          role='user', parts=[types.Part.from_text(text='test')]
+      ),
+      run_config=RunConfig(),
+  )
+
+
+@pytest.mark.asyncio
+async def test_agent_transfer_includes_sorted_agent_names_in_system_instructions():
+  """Test that agent transfer adds NOTE with sorted agent names to system instructions."""
+  mockModel = testing_utils.MockModel.create(responses=[])
+
+  # Create agents with names that will test alphabetical sorting
+  z_agent = Agent(name='z_agent', model=mockModel, description='Last agent')
+  a_agent = Agent(name='a_agent', model=mockModel, description='First agent')
+  m_agent = Agent(name='m_agent', model=mockModel, description='Middle agent')
+  peer_agent = Agent(
+      name='peer_agent', model=mockModel, description='Peer agent'
+  )
+
+  # Create parent agent with a peer agent
+  parent_agent = Agent(
+      name='parent_agent',
+      model=mockModel,
+      sub_agents=[peer_agent],
+      description='Parent agent',
+  )
+
+  # Create main agent with sub-agents and parent (intentionally unsorted order)
+  main_agent = Agent(
+      name='main_agent',
+      model=mockModel,
+      sub_agents=[z_agent, a_agent, m_agent],  # Unsorted input
+      parent_agent=parent_agent,
+      description='Main coordinating agent',
+  )
+
+  # Create test context and LLM request
+  invocation_context = await create_test_invocation_context(main_agent)
+  llm_request = LlmRequest()
+
+  # Call the actual agent transfer request processor (this behavior we're testing)
+  async for _ in agent_transfer.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  # Check on the behavior: verify system instructions contain sorted agent names
+  instructions = llm_request.config.system_instruction
+
+  # The NOTE should contain agents in alphabetical order: sub-agents + parent + peers
+  expected_content = """\
+
+You have a list of other agents to transfer to:
+
+
+Agent name: z_agent
+Agent description: Last agent
+
+
+Agent name: a_agent
+Agent description: First agent
+
+
+Agent name: m_agent
+Agent description: Middle agent
+
+
+Agent name: parent_agent
+Agent description: Parent agent
+
+
+Agent name: peer_agent
+Agent description: Peer agent
+
+
+If you are the best to answer the question according to your description,
+you can answer it.
+
+If another agent is better for answering the question according to its
+description, call `transfer_to_agent` function to transfer the question to that
+agent. When transferring, do not generate any text other than the function
+call.
+
+**NOTE**: the only available agents for `transfer_to_agent` function are
+`a_agent`, `m_agent`, `parent_agent`, `peer_agent`, `z_agent`.
+
+If neither you nor the other agents are best for the question, transfer to your parent agent parent_agent."""
+
+  assert expected_content in instructions
+
+
+@pytest.mark.asyncio
+async def test_agent_transfer_system_instructions_without_parent():
+  """Test system instructions when agent has no parent."""
+  mockModel = testing_utils.MockModel.create(responses=[])
+
+  # Create agents without parent
+  sub_agent_1 = Agent(
+      name='agent1', model=mockModel, description='First sub-agent'
+  )
+  sub_agent_2 = Agent(
+      name='agent2', model=mockModel, description='Second sub-agent'
+  )
+
+  main_agent = Agent(
+      name='main_agent',
+      model=mockModel,
+      sub_agents=[sub_agent_1, sub_agent_2],
+      # No parent_agent
+      description='Main agent without parent',
+  )
+
+  # Create test context and LLM request
+  invocation_context = await create_test_invocation_context(main_agent)
+  llm_request = LlmRequest()
+
+  # Call the agent transfer request processor
+  async for _ in agent_transfer.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  # Assert behavior: should only include sub-agents in NOTE, no parent
+  instructions = llm_request.config.system_instruction
+
+  # Direct multiline string assertion showing the exact expected content
+  expected_content = """\
+
+You have a list of other agents to transfer to:
+
+
+Agent name: agent1
+Agent description: First sub-agent
+
+
+Agent name: agent2
+Agent description: Second sub-agent
+
+
+If you are the best to answer the question according to your description,
+you can answer it.
+
+If another agent is better for answering the question according to its
+description, call `transfer_to_agent` function to transfer the question to that
+agent. When transferring, do not generate any text other than the function
+call.
+
+**NOTE**: the only available agents for `transfer_to_agent` function are
+`agent1`, `agent2`."""
+
+  assert expected_content in instructions
+
+
+@pytest.mark.asyncio
+async def test_agent_transfer_simplified_parent_instructions():
+  """Test that parent agent instructions are simplified and not verbose."""
+  mockModel = testing_utils.MockModel.create(responses=[])
+
+  # Create agent with parent
+  sub_agent = Agent(name='sub_agent', model=mockModel, description='Sub agent')
+  parent_agent = Agent(
+      name='parent_agent', model=mockModel, description='Parent agent'
+  )
+
+  main_agent = Agent(
+      name='main_agent',
+      model=mockModel,
+      sub_agents=[sub_agent],
+      parent_agent=parent_agent,
+      description='Main agent with parent',
+  )
+
+  # Create test context and LLM request
+  invocation_context = await create_test_invocation_context(main_agent)
+  llm_request = LlmRequest()
+
+  # Call the agent transfer request processor
+  async for _ in agent_transfer.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  # Assert behavior: parent instructions should be simplified
+  instructions = llm_request.config.system_instruction
+
+  # Direct multiline string assertion showing the exact expected content
+  expected_content = """\
+
+You have a list of other agents to transfer to:
+
+
+Agent name: sub_agent
+Agent description: Sub agent
+
+
+Agent name: parent_agent
+Agent description: Parent agent
+
+
+If you are the best to answer the question according to your description,
+you can answer it.
+
+If another agent is better for answering the question according to its
+description, call `transfer_to_agent` function to transfer the question to that
+agent. When transferring, do not generate any text other than the function
+call.
+
+**NOTE**: the only available agents for `transfer_to_agent` function are
+`parent_agent`, `sub_agent`.
+
+If neither you nor the other agents are best for the question, transfer to your parent agent parent_agent."""
+
+  assert expected_content in instructions
+
+
+@pytest.mark.asyncio
+async def test_agent_transfer_no_instructions_when_no_transfer_targets():
+  """Test that no instructions are added when there are no transfer targets."""
+  mockModel = testing_utils.MockModel.create(responses=[])
+
+  # Create agent with no sub-agents and no parent
+  main_agent = Agent(
+      name='main_agent',
+      model=mockModel,
+      # No sub_agents, no parent_agent
+      description='Isolated agent',
+  )
+
+  # Create test context and LLM request
+  invocation_context = await create_test_invocation_context(main_agent)
+  llm_request = LlmRequest()
+  original_system_instruction = llm_request.config.system_instruction
+
+  # Call the agent transfer request processor
+  async for _ in agent_transfer.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  # Assert behavior: no instructions should be added
+  assert llm_request.config.system_instruction == original_system_instruction
+
+  instructions = llm_request.config.system_instruction or ''
+  assert '**NOTE**:' not in instructions
+  assert 'transfer_to_agent' not in instructions
+
+
+@pytest.mark.asyncio
+async def test_agent_transfer_with_non_llm_peer_agent():
+  """Peer agents that are not LlmAgents (no `mode`) must not break transfer."""
+  mockModel = testing_utils.MockModel.create(responses=[])
+
+  non_llm_peer = _NonLlmAgent(
+      name='non_llm_peer', description='A non-LlmAgent peer'
+  )
+  parent_agent = Agent(
+      name='parent_agent',
+      model=mockModel,
+      sub_agents=[non_llm_peer],
+      description='Parent agent',
+  )
+  main_agent = Agent(
+      name='main_agent',
+      model=mockModel,
+      parent_agent=parent_agent,
+      description='Main agent',
+  )
+
+  invocation_context = await create_test_invocation_context(main_agent)
+  llm_request = LlmRequest()
+
+  async for _ in agent_transfer.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  instructions = llm_request.config.system_instruction
+  assert 'non_llm_peer' in instructions

tests/unittests/flows/llm_flows/test_async_tool_callbacks.py

@@ -0,0 +1,301 @@
+# 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.
+
+from enum import Enum
+from functools import partial
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from unittest import mock
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.llm_agent import Agent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.functions import handle_function_calls_async
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class CallbackType(Enum):
+  SYNC = 1
+  ASYNC = 2
+
+
+class AsyncBeforeToolCallback:
+
+  def __init__(self, mock_response: Dict[str, Any]):
+    self.mock_response = mock_response
+
+  async def __call__(
+      self,
+      tool: FunctionTool,
+      args: Dict[str, Any],
+      tool_context: ToolContext,
+  ) -> Optional[Dict[str, Any]]:
+    return self.mock_response
+
+
+class AsyncAfterToolCallback:
+
+  def __init__(self, mock_response: Dict[str, Any]):
+    self.mock_response = mock_response
+
+  async def __call__(
+      self,
+      tool: FunctionTool,
+      args: Dict[str, Any],
+      tool_context: ToolContext,
+      tool_response: Dict[str, Any],
+  ) -> Optional[Dict[str, Any]]:
+    return self.mock_response
+
+
+async def invoke_tool_with_callbacks(
+    before_cb=None, after_cb=None
+) -> Optional[Event]:
+  def simple_fn(**kwargs) -> Dict[str, Any]:
+    return {"initial": "response"}
+
+  tool = FunctionTool(simple_fn)
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name="agent",
+      model=model,
+      tools=[tool],
+      before_tool_callback=before_cb,
+      after_tool_callback=after_cb,
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content=""
+  )
+  # Build function call event
+  function_call = types.FunctionCall(name=tool.name, args={})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {tool.name: tool}
+  return await handle_function_calls_async(
+      invocation_context,
+      event,
+      tools_dict,
+  )
+
+
+@pytest.mark.asyncio
+async def test_async_before_tool_callback():
+  mock_resp = {"test": "before_tool_callback"}
+  before_cb = AsyncBeforeToolCallback(mock_resp)
+  result_event = await invoke_tool_with_callbacks(before_cb=before_cb)
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_resp
+
+
+@pytest.mark.asyncio
+async def test_async_after_tool_callback():
+  mock_resp = {"test": "after_tool_callback"}
+  after_cb = AsyncAfterToolCallback(mock_resp)
+  result_event = await invoke_tool_with_callbacks(after_cb=after_cb)
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_resp
+
+
+def mock_async_before_cb_side_effect(
+    tool: FunctionTool,
+    args: Dict[str, Any],
+    tool_context: ToolContext,
+    ret_value: Optional[Dict[str, Any]] = None,
+):
+  if ret_value:
+    return ret_value
+  return None
+
+
+def mock_sync_before_cb_side_effect(
+    tool: FunctionTool,
+    args: Dict[str, Any],
+    tool_context: ToolContext,
+    ret_value: Optional[Dict[str, Any]] = None,
+):
+  if ret_value:
+    return ret_value
+  return None
+
+
+async def mock_async_after_cb_side_effect(
+    tool: FunctionTool,
+    args: Dict[str, Any],
+    tool_context: ToolContext,
+    tool_response: Dict[str, Any],
+    ret_value: Optional[Dict[str, Any]] = None,
+):
+  if ret_value:
+    return ret_value
+  return None
+
+
+def mock_sync_after_cb_side_effect(
+    tool: FunctionTool,
+    args: Dict[str, Any],
+    tool_context: ToolContext,
+    tool_response: Dict[str, Any],
+    ret_value: Optional[Dict[str, Any]] = None,
+):
+  if ret_value:
+    return ret_value
+  return None
+
+
+CALLBACK_PARAMS = [
+    pytest.param(
+        [
+            (None, CallbackType.SYNC),
+            ({"test": "callback_2_response"}, CallbackType.ASYNC),
+            ({"test": "callback_3_response"}, CallbackType.SYNC),
+            (None, CallbackType.ASYNC),
+        ],
+        {"test": "callback_2_response"},
+        [1, 1, 0, 0],
+        id="middle_async_callback_returns",
+    ),
+    pytest.param(
+        [
+            (None, CallbackType.SYNC),
+            (None, CallbackType.ASYNC),
+            (None, CallbackType.SYNC),
+            (None, CallbackType.ASYNC),
+        ],
+        {"initial": "response"},
+        [1, 1, 1, 1],
+        id="all_callbacks_return_none",
+    ),
+    pytest.param(
+        [
+            ({"test": "callback_1_response"}, CallbackType.SYNC),
+            ({"test": "callback_2_response"}, CallbackType.ASYNC),
+        ],
+        {"test": "callback_1_response"},
+        [1, 0],
+        id="first_sync_callback_returns",
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "callbacks, expected_response, expected_calls",
+    CALLBACK_PARAMS,
+)
+@pytest.mark.asyncio
+async def test_before_tool_callbacks_chain(
+    callbacks: List[tuple[Optional[Dict[str, Any]], int]],
+    expected_response: Dict[str, Any],
+    expected_calls: List[int],
+):
+  mock_before_cbs = []
+  for response, callback_type in callbacks:
+    if callback_type == CallbackType.ASYNC:
+      mock_cb = mock.AsyncMock(
+          side_effect=partial(
+              mock_async_before_cb_side_effect, ret_value=response
+          )
+      )
+    else:
+      mock_cb = mock.Mock(
+          side_effect=partial(
+              mock_sync_before_cb_side_effect, ret_value=response
+          )
+      )
+    mock_before_cbs.append(mock_cb)
+  result_event = await invoke_tool_with_callbacks(before_cb=mock_before_cbs)
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == expected_response
+
+  # Assert that the callbacks were called the expected number of times
+  for i, mock_cb in enumerate(mock_before_cbs):
+    expected_calls_count = expected_calls[i]
+    if expected_calls_count == 1:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited_once()
+      else:
+        mock_cb.assert_called_once()
+    elif expected_calls_count == 0:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_not_awaited()
+      else:
+        mock_cb.assert_not_called()
+    else:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited(expected_calls_count)
+      else:
+        mock_cb.assert_called(expected_calls_count)
+
+
+@pytest.mark.parametrize(
+    "callbacks, expected_response, expected_calls",
+    CALLBACK_PARAMS,
+)
+@pytest.mark.asyncio
+async def test_after_tool_callbacks_chain(
+    callbacks: List[tuple[Optional[Dict[str, Any]], int]],
+    expected_response: Dict[str, Any],
+    expected_calls: List[int],
+):
+  mock_after_cbs = []
+  for response, callback_type in callbacks:
+    if callback_type == CallbackType.ASYNC:
+      mock_cb = mock.AsyncMock(
+          side_effect=partial(
+              mock_async_after_cb_side_effect, ret_value=response
+          )
+      )
+    else:
+      mock_cb = mock.Mock(
+          side_effect=partial(
+              mock_sync_after_cb_side_effect, ret_value=response
+          )
+      )
+    mock_after_cbs.append(mock_cb)
+  result_event = await invoke_tool_with_callbacks(after_cb=mock_after_cbs)
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == expected_response
+
+  # Assert that the callbacks were called the expected number of times
+  for i, mock_cb in enumerate(mock_after_cbs):
+    expected_calls_count = expected_calls[i]
+    if expected_calls_count == 1:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited_once()
+      else:
+        mock_cb.assert_called_once()
+    elif expected_calls_count == 0:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_not_awaited()
+      else:
+        mock_cb.assert_not_called()
+    else:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited(expected_calls_count)
+      else:
+        mock_cb.assert_called(expected_calls_count)

tests/unittests/flows/llm_flows/test_audio_cache_manager.py

@@ -0,0 +1,440 @@
+# 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.
+
+import time
+from unittest.mock import AsyncMock
+from unittest.mock import Mock
+
+from google.adk.flows.llm_flows.audio_cache_manager import AudioCacheConfig
+from google.adk.flows.llm_flows.audio_cache_manager import AudioCacheManager
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class TestAudioCacheConfig:
+  """Test the AudioCacheConfig class."""
+
+  def test_default_values(self):
+    """Test that default configuration values are set correctly."""
+    config = AudioCacheConfig()
+    assert config.max_cache_size_bytes == 10 * 1024 * 1024  # 10MB
+    assert config.max_cache_duration_seconds == 300.0  # 5 minutes
+    assert config.auto_flush_threshold == 100
+
+  def test_custom_values(self):
+    """Test that custom configuration values are set correctly."""
+    config = AudioCacheConfig(
+        max_cache_size_bytes=5 * 1024 * 1024,
+        max_cache_duration_seconds=120.0,
+        auto_flush_threshold=50,
+    )
+    assert config.max_cache_size_bytes == 5 * 1024 * 1024
+    assert config.max_cache_duration_seconds == 120.0
+    assert config.auto_flush_threshold == 50
+
+
+class TestAudioCacheManager:
+  """Test the AudioCacheManager class."""
+
+  def setup_method(self):
+    """Set up test fixtures."""
+    self.config = AudioCacheConfig()
+    self.manager = AudioCacheManager(self.config)
+
+  @pytest.mark.asyncio
+  async def test_cache_input_audio(self):
+    """Test caching input audio data."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    audio_blob = types.Blob(data=b'test_audio_data', mime_type='audio/pcm')
+
+    # Initially no cache
+    assert invocation_context.input_realtime_cache is None
+
+    # Cache audio
+    self.manager.cache_audio(invocation_context, audio_blob, 'input')
+
+    # Verify cache is created and populated
+    assert invocation_context.input_realtime_cache is not None
+    assert len(invocation_context.input_realtime_cache) == 1
+
+    entry = invocation_context.input_realtime_cache[0]
+    assert entry.role == 'user'
+    assert entry.data == audio_blob
+    assert isinstance(entry.timestamp, float)
+
+  @pytest.mark.asyncio
+  async def test_cache_output_audio(self):
+    """Test caching output audio data."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    audio_blob = types.Blob(data=b'test_model_audio', mime_type='audio/wav')
+
+    # Initially no cache
+    assert invocation_context.output_realtime_cache is None
+
+    # Cache audio
+    self.manager.cache_audio(invocation_context, audio_blob, 'output')
+
+    # Verify cache is created and populated
+    assert invocation_context.output_realtime_cache is not None
+    assert len(invocation_context.output_realtime_cache) == 1
+
+    entry = invocation_context.output_realtime_cache[0]
+    assert entry.role == 'model'
+    assert entry.data == audio_blob
+    assert isinstance(entry.timestamp, float)
+
+  @pytest.mark.asyncio
+  async def test_multiple_audio_caching(self):
+    """Test caching multiple audio chunks."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Cache multiple input audio chunks
+    for i in range(3):
+      audio_blob = types.Blob(data=f'input_{i}'.encode(), mime_type='audio/pcm')
+      self.manager.cache_audio(invocation_context, audio_blob, 'input')
+
+    # Cache multiple output audio chunks
+    for i in range(2):
+      audio_blob = types.Blob(
+          data=f'output_{i}'.encode(), mime_type='audio/wav'
+      )
+      self.manager.cache_audio(invocation_context, audio_blob, 'output')
+
+    # Verify all chunks are cached
+    assert len(invocation_context.input_realtime_cache) == 3
+    assert len(invocation_context.output_realtime_cache) == 2
+
+  @pytest.mark.asyncio
+  async def test_flush_caches_both(self):
+    """Test flushing both input and output caches."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Set up mock artifact service
+    mock_artifact_service = AsyncMock()
+    mock_artifact_service.save_artifact.return_value = 123
+    invocation_context.artifact_service = mock_artifact_service
+
+    # Cache some audio
+    input_blob = types.Blob(data=b'input_data', mime_type='audio/pcm')
+    output_blob = types.Blob(data=b'output_data', mime_type='audio/wav')
+    self.manager.cache_audio(invocation_context, input_blob, 'input')
+    self.manager.cache_audio(invocation_context, output_blob, 'output')
+
+    # Flush caches
+    await self.manager.flush_caches(invocation_context)
+
+    # Verify caches are cleared
+    assert invocation_context.input_realtime_cache == []
+    assert invocation_context.output_realtime_cache == []
+
+    # Verify artifact service was called twice (once for each cache)
+    assert mock_artifact_service.save_artifact.call_count == 2
+
+  @pytest.mark.asyncio
+  async def test_flush_caches_selective(self):
+    """Test selectively flushing only one cache."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Set up mock artifact service
+    mock_artifact_service = AsyncMock()
+    mock_artifact_service.save_artifact.return_value = 123
+    invocation_context.artifact_service = mock_artifact_service
+
+    # Cache some audio
+    input_blob = types.Blob(data=b'input_data', mime_type='audio/pcm')
+    output_blob = types.Blob(data=b'output_data', mime_type='audio/wav')
+    self.manager.cache_audio(invocation_context, input_blob, 'input')
+    self.manager.cache_audio(invocation_context, output_blob, 'output')
+
+    # Flush only input cache
+    await self.manager.flush_caches(
+        invocation_context, flush_user_audio=True, flush_model_audio=False
+    )
+
+    # Verify only input cache is cleared
+    assert invocation_context.input_realtime_cache == []
+    assert len(invocation_context.output_realtime_cache) == 1
+
+    # Verify artifact service was called once
+    assert mock_artifact_service.save_artifact.call_count == 1
+
+  @pytest.mark.asyncio
+  async def test_flush_empty_caches(self):
+    """Test flushing when caches are empty."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Set up mock artifact service
+    mock_artifact_service = AsyncMock()
+    invocation_context.artifact_service = mock_artifact_service
+
+    # Flush empty caches (should not error)
+    await self.manager.flush_caches(invocation_context)
+
+    # Verify artifact service was not called
+    mock_artifact_service.save_artifact.assert_not_called()
+
+  @pytest.mark.asyncio
+  async def test_flush_without_artifact_service(self):
+    """Test flushing when no artifact service is available."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # No artifact service
+    invocation_context.artifact_service = None
+
+    # Cache some audio
+    input_blob = types.Blob(data=b'input_data', mime_type='audio/pcm')
+    self.manager.cache_audio(invocation_context, input_blob, 'input')
+
+    # Flush should not error but should not clear cache either
+    await self.manager.flush_caches(invocation_context)
+
+    # Cache should remain (no actual flushing happened)
+    assert len(invocation_context.input_realtime_cache) == 1
+
+  @pytest.mark.asyncio
+  async def test_flush_artifact_creation(self):
+    """Test that artifacts are created correctly during flush."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Set up mock services
+    mock_artifact_service = AsyncMock()
+    mock_artifact_service.save_artifact.return_value = 456
+    mock_session_service = AsyncMock()
+
+    invocation_context.artifact_service = mock_artifact_service
+    invocation_context.session_service = mock_session_service
+
+    # Cache audio with specific data
+    test_data = b'specific_test_audio_data'
+    audio_blob = types.Blob(data=test_data, mime_type='audio/pcm')
+    self.manager.cache_audio(invocation_context, audio_blob, 'input')
+
+    # Flush cache
+    await self.manager.flush_caches(invocation_context)
+
+    # Verify artifact was saved with correct data
+    mock_artifact_service.save_artifact.assert_called_once()
+    call_args = mock_artifact_service.save_artifact.call_args
+    saved_artifact = call_args.kwargs['artifact']
+    assert saved_artifact.inline_data.data == test_data
+    assert saved_artifact.inline_data.mime_type == 'audio/pcm'
+
+    # Verify session event was created
+    mock_session_service.append_event.assert_not_called()
+
+  def test_get_cache_stats_empty(self):
+    """Test getting statistics for empty caches."""
+    invocation_context = Mock()
+    invocation_context.input_realtime_cache = None
+    invocation_context.output_realtime_cache = None
+
+    stats = self.manager.get_cache_stats(invocation_context)
+
+    expected = {
+        'input_chunks': 0,
+        'output_chunks': 0,
+        'input_bytes': 0,
+        'output_bytes': 0,
+        'total_chunks': 0,
+        'total_bytes': 0,
+    }
+    assert stats == expected
+
+  @pytest.mark.asyncio
+  async def test_get_cache_stats_with_data(self):
+    """Test getting statistics for caches with data."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Cache some audio data of different sizes
+    input_blob1 = types.Blob(data=b'12345', mime_type='audio/pcm')  # 5 bytes
+    input_blob2 = types.Blob(
+        data=b'1234567890', mime_type='audio/pcm'
+    )  # 10 bytes
+    output_blob = types.Blob(data=b'abc', mime_type='audio/wav')  # 3 bytes
+
+    self.manager.cache_audio(invocation_context, input_blob1, 'input')
+    self.manager.cache_audio(invocation_context, input_blob2, 'input')
+    self.manager.cache_audio(invocation_context, output_blob, 'output')
+
+    stats = self.manager.get_cache_stats(invocation_context)
+
+    expected = {
+        'input_chunks': 2,
+        'output_chunks': 1,
+        'input_bytes': 15,  # 5 + 10
+        'output_bytes': 3,
+        'total_chunks': 3,
+        'total_bytes': 18,  # 15 + 3
+    }
+    assert stats == expected
+
+  @pytest.mark.asyncio
+  async def test_error_handling_in_flush(self):
+    """Test error handling during cache flush operations."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Set up mock artifact service that raises an error
+    mock_artifact_service = AsyncMock()
+    mock_artifact_service.save_artifact.side_effect = Exception(
+        'Artifact service error'
+    )
+    invocation_context.artifact_service = mock_artifact_service
+
+    # Cache some audio
+    audio_blob = types.Blob(data=b'test_data', mime_type='audio/pcm')
+    self.manager.cache_audio(invocation_context, audio_blob, 'input')
+
+    # Flush should not raise exception but should log error and retain cache
+    await self.manager.flush_caches(invocation_context)
+
+    # Cache should remain since flush failed
+    assert len(invocation_context.input_realtime_cache) == 1
+
+  @pytest.mark.asyncio
+  async def test_filename_uses_first_chunk_timestamp(self):
+    """Test that the filename timestamp comes from the first audio chunk, not flush time."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Set up mock services
+    mock_artifact_service = AsyncMock()
+    mock_artifact_service.save_artifact.return_value = 789
+    mock_session_service = AsyncMock()
+
+    invocation_context.artifact_service = mock_artifact_service
+    invocation_context.session_service = mock_session_service
+
+    # Cache multiple audio chunks with specific timestamps
+    first_timestamp = 1234567890.123  # First chunk timestamp
+    second_timestamp = 1234567891.456  # Second chunk timestamp (later)
+
+    # Manually create audio cache entries with specific timestamps
+    invocation_context.input_realtime_cache = []
+
+    from google.adk.agents.invocation_context import RealtimeCacheEntry
+
+    first_entry = RealtimeCacheEntry(
+        role='user',
+        data=types.Blob(data=b'first_chunk', mime_type='audio/pcm'),
+        timestamp=first_timestamp,
+    )
+
+    second_entry = RealtimeCacheEntry(
+        role='user',
+        data=types.Blob(data=b'second_chunk', mime_type='audio/pcm'),
+        timestamp=second_timestamp,
+    )
+
+    invocation_context.input_realtime_cache.extend([first_entry, second_entry])
+
+    # Sleep briefly to ensure current time is different from first timestamp
+    time.sleep(0.01)
+
+    # Flush cache
+    await self.manager.flush_caches(invocation_context)
+
+    # Verify artifact was saved
+    mock_artifact_service.save_artifact.assert_called_once()
+    call_args = mock_artifact_service.save_artifact.call_args
+    filename = call_args.kwargs['filename']
+
+    # Extract timestamp from filename (format: input_audio_{timestamp}.pcm)
+    expected_timestamp_ms = int(first_timestamp * 1000)
+    assert (
+        f'adk_live_audio_storage_input_audio_{expected_timestamp_ms}.pcm'
+        == filename
+    )
+
+    # Verify the timestamp in filename matches first chunk, not current time
+    current_timestamp_ms = int(time.time() * 1000)
+    assert expected_timestamp_ms != current_timestamp_ms  # Should be different
+    assert filename.startswith(
+        f'adk_live_audio_storage_input_audio_{expected_timestamp_ms}'
+    )
+
+  @pytest.mark.asyncio
+  async def test_flush_event_author_for_user_audio(self):
+    """Test that flushed user audio events have 'user' as author."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Set up mock artifact service
+    mock_artifact_service = AsyncMock()
+    mock_artifact_service.save_artifact.return_value = 123
+    invocation_context.artifact_service = mock_artifact_service
+
+    # Cache user input audio
+    input_blob = types.Blob(data=b'user_audio_data', mime_type='audio/pcm')
+    self.manager.cache_audio(invocation_context, input_blob, 'input')
+
+    # Flush cache and get events
+    events = await self.manager.flush_caches(
+        invocation_context, flush_user_audio=True, flush_model_audio=False
+    )
+
+    # Verify event author is 'user' for user audio
+    assert len(events) == 1
+    assert events[0].author == 'user'
+    assert events[0].content.role == 'user'
+
+  @pytest.mark.asyncio
+  async def test_flush_event_author_for_model_audio(self):
+    """Test that flushed model audio events have agent name as author, not 'model'."""
+    agent = testing_utils.create_test_agent(name='my_test_agent')
+    invocation_context = await testing_utils.create_invocation_context(agent)
+
+    # Set up mock artifact service
+    mock_artifact_service = AsyncMock()
+    mock_artifact_service.save_artifact.return_value = 123
+    invocation_context.artifact_service = mock_artifact_service
+
+    # Cache model output audio
+    output_blob = types.Blob(data=b'model_audio_data', mime_type='audio/wav')
+    self.manager.cache_audio(invocation_context, output_blob, 'output')
+
+    # Flush cache and get events
+    events = await self.manager.flush_caches(
+        invocation_context, flush_user_audio=False, flush_model_audio=True
+    )
+
+    # Verify event author is agent name (not 'model') for model audio
+    assert len(events) == 1
+    assert events[0].author == 'my_test_agent'  # Agent name, not 'model'
+    assert events[0].content.role == 'model'  # Role is still 'model'

tests/unittests/flows/llm_flows/test_base_llm_flow_partial_handling.py

@@ -0,0 +1,164 @@
+# 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.
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class BaseLlmFlowForTesting(BaseLlmFlow):
+  """Test implementation of BaseLlmFlow for testing purposes."""
+
+  pass
+
+
+@pytest.mark.asyncio
+async def test_run_async_breaks_on_partial_event():
+  """Test that run_async breaks when the last event is partial."""
+  # Create a mock model that returns partial responses
+  partial_response = LlmResponse(
+      content=types.Content(
+          role='model', parts=[types.Part.from_text(text='Partial response')]
+      ),
+      partial=True,
+  )
+
+  mock_model = testing_utils.MockModel.create(responses=[partial_response])
+
+  agent = Agent(name='test_agent', model=mock_model)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  flow = BaseLlmFlowForTesting()
+  events = []
+
+  # Collect events from the flow
+  async for event in flow.run_async(invocation_context):
+    events.append(event)
+
+  # Should have one event (the partial response)
+  assert len(events) == 1
+  assert events[0].partial is True
+  assert events[0].content.parts[0].text == 'Partial response'
+
+
+@pytest.mark.asyncio
+async def test_run_async_breaks_on_final_response():
+  """Test that run_async breaks when the last event is a final response."""
+  # Create a mock model that returns a final response
+  final_response = LlmResponse(
+      content=types.Content(
+          role='model', parts=[types.Part.from_text(text='Final response')]
+      ),
+      partial=False,
+      error_code=types.FinishReason.STOP,
+  )
+
+  mock_model = testing_utils.MockModel.create(responses=[final_response])
+
+  agent = Agent(name='test_agent', model=mock_model)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  flow = BaseLlmFlowForTesting()
+  events = []
+
+  # Collect events from the flow
+  async for event in flow.run_async(invocation_context):
+    events.append(event)
+
+  # Should have one event (the final response)
+  assert len(events) == 1
+  assert events[0].partial is False
+  assert events[0].content.parts[0].text == 'Final response'
+
+
+@pytest.mark.asyncio
+async def test_run_async_breaks_on_no_last_event():
+  """Test that run_async breaks when there is no last event."""
+  # Create a mock model that returns an empty response (no content)
+  empty_response = LlmResponse(content=None, partial=False)
+
+  mock_model = testing_utils.MockModel.create(responses=[empty_response])
+
+  agent = Agent(name='test_agent', model=mock_model)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  flow = BaseLlmFlowForTesting()
+  events = []
+
+  # Collect events from the flow
+  async for event in flow.run_async(invocation_context):
+    events.append(event)
+
+  # Should have no events because empty responses are filtered out
+  assert len(events) == 0
+
+
+@pytest.mark.asyncio
+async def test_run_async_breaks_on_first_partial_response():
+  """Test run_async breaks on the first partial response."""
+  # Create responses with mixed partial states
+  partial_response = LlmResponse(
+      content=types.Content(
+          role='model', parts=[types.Part.from_text(text='Partial response')]
+      ),
+      partial=True,
+  )
+
+  # These won't be reached because the flow breaks on the first partial
+  non_partial_response = LlmResponse(
+      content=types.Content(
+          role='model',
+          parts=[types.Part.from_text(text='Non-partial response')],
+      ),
+      partial=False,
+  )
+
+  final_partial_response = LlmResponse(
+      content=types.Content(
+          role='model',
+          parts=[types.Part.from_text(text='Final partial response')],
+      ),
+      partial=True,
+  )
+
+  mock_model = testing_utils.MockModel.create(
+      responses=[partial_response, non_partial_response, final_partial_response]
+  )
+
+  agent = Agent(name='test_agent', model=mock_model)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  flow = BaseLlmFlowForTesting()
+  events = []
+
+  # Collect events from the flow
+  async for event in flow.run_async(invocation_context):
+    events.append(event)
+
+  # Should have only one event, breaking on the first partial response
+  assert len(events) == 1
+  assert events[0].partial is True
+  assert events[0].content.parts[0].text == 'Partial response'

tests/unittests/flows/llm_flows/test_base_llm_flow_realtime.py

@@ -0,0 +1,232 @@
+# 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.
+
+from unittest import mock
+
+from google.adk.agents.live_request_queue import LiveRequest
+from google.adk.agents.live_request_queue import LiveRequestQueue
+from google.adk.agents.llm_agent import Agent
+from google.adk.agents.run_config import RunConfig
+from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
+from google.adk.models.llm_request import LlmRequest
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class TestBaseLlmFlow(BaseLlmFlow):
+  """Test implementation of BaseLlmFlow for testing purposes."""
+
+  pass
+
+
+@pytest.fixture
+def test_blob():
+  """Test blob for audio data."""
+  return types.Blob(data=b'\x00\xFF\x00\xFF', mime_type='audio/pcm')
+
+
+@pytest.fixture
+def mock_llm_connection():
+  """Mock LLM connection for testing."""
+  connection = mock.AsyncMock()
+  connection.send_realtime = mock.AsyncMock()
+  return connection
+
+
+@pytest.mark.asyncio
+async def test_send_to_model_with_disabled_vad(test_blob, mock_llm_connection):
+  """Test _send_to_model with automatic_activity_detection.disabled=True."""
+  # Create LlmRequest with disabled VAD
+  realtime_input_config = types.RealtimeInputConfig(
+      automatic_activity_detection=types.AutomaticActivityDetection(
+          disabled=True
+      )
+  )
+
+  # Create invocation context with live request queue
+  agent = Agent(name='test_agent', model='mock')
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent,
+      user_content='',
+      run_config=RunConfig(realtime_input_config=realtime_input_config),
+  )
+  invocation_context.live_request_queue = LiveRequestQueue()
+
+  # Create flow and start _send_to_model task
+  flow = TestBaseLlmFlow()
+
+  # Send a blob to the queue
+  live_request = LiveRequest(blob=test_blob)
+  invocation_context.live_request_queue.send(live_request)
+  invocation_context.live_request_queue.close()
+
+  # Run _send_to_model
+  await flow._send_to_model(mock_llm_connection, invocation_context)
+
+  mock_llm_connection.send_realtime.assert_called_once_with(test_blob)
+
+
+@pytest.mark.asyncio
+async def test_send_to_model_with_enabled_vad(test_blob, mock_llm_connection):
+  """Test _send_to_model with automatic_activity_detection.disabled=False.
+
+  Custom VAD activity signal is not supported so we should still disable it.
+  """
+  # Create LlmRequest with enabled VAD
+  realtime_input_config = types.RealtimeInputConfig(
+      automatic_activity_detection=types.AutomaticActivityDetection(
+          disabled=False
+      )
+  )
+
+  # Create invocation context with live request queue
+  agent = Agent(name='test_agent', model='mock')
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content=''
+  )
+  invocation_context.live_request_queue = LiveRequestQueue()
+
+  # Create flow and start _send_to_model task
+  flow = TestBaseLlmFlow()
+
+  # Send a blob to the queue
+  live_request = LiveRequest(blob=test_blob)
+  invocation_context.live_request_queue.send(live_request)
+  invocation_context.live_request_queue.close()
+
+  # Run _send_to_model
+  await flow._send_to_model(mock_llm_connection, invocation_context)
+
+  mock_llm_connection.send_realtime.assert_called_once_with(test_blob)
+
+
+@pytest.mark.asyncio
+async def test_send_to_model_without_realtime_config(
+    test_blob, mock_llm_connection
+):
+  """Test _send_to_model without realtime_input_config (default behavior)."""
+  # Create invocation context with live request queue
+  agent = Agent(name='test_agent', model='mock')
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content=''
+  )
+  invocation_context.live_request_queue = LiveRequestQueue()
+
+  # Create flow and start _send_to_model task
+  flow = TestBaseLlmFlow()
+
+  # Send a blob to the queue
+  live_request = LiveRequest(blob=test_blob)
+  invocation_context.live_request_queue.send(live_request)
+  invocation_context.live_request_queue.close()
+
+  # Run _send_to_model
+  await flow._send_to_model(mock_llm_connection, invocation_context)
+
+  mock_llm_connection.send_realtime.assert_called_once_with(test_blob)
+
+
+@pytest.mark.asyncio
+async def test_send_to_model_with_none_automatic_activity_detection(
+    test_blob, mock_llm_connection
+):
+  """Test _send_to_model with automatic_activity_detection=None."""
+  # Create LlmRequest with None automatic_activity_detection
+  realtime_input_config = types.RealtimeInputConfig(
+      automatic_activity_detection=None
+  )
+
+  # Create invocation context with live request queue
+  agent = Agent(name='test_agent', model='mock')
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent,
+      user_content='',
+      run_config=RunConfig(realtime_input_config=realtime_input_config),
+  )
+  invocation_context.live_request_queue = LiveRequestQueue()
+
+  # Create flow and start _send_to_model task
+  flow = TestBaseLlmFlow()
+
+  # Send a blob to the queue
+  live_request = LiveRequest(blob=test_blob)
+  invocation_context.live_request_queue.send(live_request)
+  invocation_context.live_request_queue.close()
+
+  # Run _send_to_model
+  await flow._send_to_model(mock_llm_connection, invocation_context)
+
+  mock_llm_connection.send_realtime.assert_called_once_with(test_blob)
+
+
+@pytest.mark.asyncio
+async def test_send_to_model_with_text_content(mock_llm_connection):
+  """Test _send_to_model with text content (not blob)."""
+  # Create invocation context with live request queue
+  agent = Agent(name='test_agent', model='mock')
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content=''
+  )
+  invocation_context.live_request_queue = LiveRequestQueue()
+
+  # Create flow and start _send_to_model task
+  flow = TestBaseLlmFlow()
+
+  # Send text content to the queue
+  content = types.Content(
+      role='user', parts=[types.Part.from_text(text='Hello')]
+  )
+  live_request = LiveRequest(content=content)
+  invocation_context.live_request_queue.send(live_request)
+  invocation_context.live_request_queue.close()
+
+  # Run _send_to_model
+  await flow._send_to_model(mock_llm_connection, invocation_context)
+
+  # Verify send_content was called instead of send_realtime
+  mock_llm_connection._send_content.assert_called_once_with(
+      content, partial=False
+  )
+  mock_llm_connection.send_realtime.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_send_to_model_with_intermediate_text_content(
+    mock_llm_connection,
+):
+  agent = Agent(name='test_agent', model='mock')
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content=''
+  )
+  invocation_context.live_request_queue = LiveRequestQueue()
+  invocation_context.session_service.append_event = mock.AsyncMock()
+
+  flow = TestBaseLlmFlow()
+
+  content = types.Content(
+      role='user', parts=[types.Part.from_text(text='progress')]
+  )
+  invocation_context.live_request_queue.send(
+      LiveRequest(content=content, partial=True)
+  )
+  invocation_context.live_request_queue.close()
+
+  await flow._send_to_model(mock_llm_connection, invocation_context)
+
+  mock_llm_connection._send_content.assert_called_once_with(
+      content, partial=True
+  )
+  invocation_context.session_service.append_event.assert_not_called()

tests/unittests/flows/llm_flows/test_basic_processor.py

@@ -0,0 +1,375 @@
+# 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 basic LLM request processor."""
+
+from unittest import mock
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.run_config import RunConfig
+from google.adk.flows.llm_flows.basic import _BasicLlmRequestProcessor
+from google.adk.models.llm_request import LlmRequest
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.function_tool import FunctionTool
+from google.genai import types
+from pydantic import BaseModel
+from pydantic import Field
+import pytest
+
+
+class OutputSchema(BaseModel):
+  """Test schema for output."""
+
+  name: str = Field(description='A name')
+  value: int = Field(description='A value')
+
+
+def dummy_tool(query: str) -> str:
+  """A dummy tool for testing."""
+  return f'Result: {query}'
+
+
+async def _create_invocation_context(agent: LlmAgent) -> InvocationContext:
+  """Helper to create InvocationContext for testing."""
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  return InvocationContext(
+      invocation_id='test-id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+      run_config=RunConfig(),
+  )
+
+
+class TestBasicLlmRequestProcessor:
+  """Test class for _BasicLlmRequestProcessor."""
+
+  @pytest.mark.asyncio
+  async def test_sets_output_schema_when_no_tools(self):
+    """Test that processor sets output_schema when agent has no tools."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-2.5-flash',
+        output_schema=OutputSchema,
+        tools=[],  # No tools
+    )
+
+    invocation_context = await _create_invocation_context(agent)
+    llm_request = LlmRequest()
+    processor = _BasicLlmRequestProcessor()
+
+    # Process the request
+    events = []
+    async for event in processor.run_async(invocation_context, llm_request):
+      events.append(event)
+
+    # Should have set response_schema since agent has no tools
+    assert llm_request.config.response_schema == OutputSchema
+    assert llm_request.config.response_mime_type == 'application/json'
+
+  @pytest.mark.asyncio
+  async def test_skips_output_schema_when_tools_present(self, mocker):
+    """Test that processor skips output_schema when agent has tools."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-2.5-flash',
+        output_schema=OutputSchema,
+        tools=[FunctionTool(func=dummy_tool)],  # Has tools
+    )
+
+    invocation_context = await _create_invocation_context(agent)
+    llm_request = LlmRequest()
+    processor = _BasicLlmRequestProcessor()
+
+    can_use_output_schema_with_tools = mocker.patch(
+        'google.adk.flows.llm_flows.basic.can_use_output_schema_with_tools',
+        mock.MagicMock(return_value=False),
+    )
+
+    # Process the request
+    events = []
+    async for event in processor.run_async(invocation_context, llm_request):
+      events.append(event)
+
+    # Should NOT have set response_schema since agent has tools
+    assert llm_request.config.response_schema is None
+    assert llm_request.config.response_mime_type != 'application/json'
+
+    # Should have checked if output schema can be used with tools
+    can_use_output_schema_with_tools.assert_called_once_with(
+        agent.canonical_model
+    )
+
+  @pytest.mark.asyncio
+  async def test_sets_output_schema_when_tools_present(self, mocker):
+    """Test that processor skips output_schema when agent has tools."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-2.5-flash',
+        output_schema=OutputSchema,
+        tools=[FunctionTool(func=dummy_tool)],  # Has tools
+    )
+
+    invocation_context = await _create_invocation_context(agent)
+    llm_request = LlmRequest()
+    processor = _BasicLlmRequestProcessor()
+
+    can_use_output_schema_with_tools = mocker.patch(
+        'google.adk.flows.llm_flows.basic.can_use_output_schema_with_tools',
+        mock.MagicMock(return_value=True),
+    )
+
+    # Process the request
+    events = []
+    async for event in processor.run_async(invocation_context, llm_request):
+      events.append(event)
+
+    # Should have set response_schema since output schema can be used with tools
+    assert llm_request.config.response_schema == OutputSchema
+    assert llm_request.config.response_mime_type == 'application/json'
+
+    # Should have checked if output schema can be used with tools
+    can_use_output_schema_with_tools.assert_called_once_with(
+        agent.canonical_model
+    )
+
+  @pytest.mark.asyncio
+  async def test_no_output_schema_no_tools(self):
+    """Test that processor works normally when agent has no output_schema or tools."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-2.5-flash',
+        # No output_schema, no tools
+    )
+
+    invocation_context = await _create_invocation_context(agent)
+    llm_request = LlmRequest()
+    processor = _BasicLlmRequestProcessor()
+
+    # Process the request
+    events = []
+    async for event in processor.run_async(invocation_context, llm_request):
+      events.append(event)
+
+    # Should not have set anything
+    assert llm_request.config.response_schema is None
+    assert llm_request.config.response_mime_type != 'application/json'
+
+  @pytest.mark.asyncio
+  async def test_sets_model_name(self):
+    """Test that processor sets the model name correctly."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-2.5-flash',
+    )
+
+    invocation_context = await _create_invocation_context(agent)
+    llm_request = LlmRequest()
+    processor = _BasicLlmRequestProcessor()
+
+    # Process the request
+    events = []
+    async for event in processor.run_async(invocation_context, llm_request):
+      events.append(event)
+
+    # Should have set the model name
+    assert llm_request.model == 'gemini-2.5-flash'
+
+  @pytest.mark.asyncio
+  async def test_skips_output_schema_for_task_mode(self):
+    """Test that processor skips output_schema when agent is in task mode."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-2.5-flash',
+        mode='task',
+        output_schema=OutputSchema,
+    )
+
+    invocation_context = await _create_invocation_context(agent)
+    llm_request = LlmRequest()
+    processor = _BasicLlmRequestProcessor()
+
+    async for _ in processor.run_async(invocation_context, llm_request):
+      pass
+
+    assert llm_request.config.response_schema is None
+
+  @pytest.mark.asyncio
+  async def test_disables_affective_dialog_and_proactivity_for_gemini_3_x_live(
+      self,
+  ):
+    """Gemini 3.x Live does not support affective_dialog/proactivity."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-3.5-flash-lite-live-preview',
+    )
+    invocation_context = await _create_invocation_context(agent)
+    invocation_context.run_config = RunConfig(
+        enable_affective_dialog=True,
+        proactivity=types.ProactivityConfig(),
+    )
+    llm_request = LlmRequest()
+    processor = _BasicLlmRequestProcessor()
+
+    async for _ in processor.run_async(invocation_context, llm_request):
+      pass
+
+    assert llm_request.live_connect_config.enable_affective_dialog is None
+    assert llm_request.live_connect_config.proactivity is None
+
+  @pytest.mark.asyncio
+  async def test_keeps_affective_dialog_and_proactivity_for_non_gemini_3_x_live(
+      self,
+  ):
+    """Non-3.x live models keep the configured affective_dialog/proactivity."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-2.5-flash-live',
+    )
+    invocation_context = await _create_invocation_context(agent)
+    invocation_context.run_config = RunConfig(
+        enable_affective_dialog=True,
+        proactivity=types.ProactivityConfig(),
+    )
+    llm_request = LlmRequest()
+    processor = _BasicLlmRequestProcessor()
+
+    async for _ in processor.run_async(invocation_context, llm_request):
+      pass
+
+    assert llm_request.live_connect_config.enable_affective_dialog is True
+    assert llm_request.live_connect_config.proactivity is not None
+
+  @pytest.mark.asyncio
+  async def test_sets_translation_config(self):
+    """Translation config is forwarded to the live connect config."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-3.5-live-translate-preview',
+    )
+    invocation_context = await _create_invocation_context(agent)
+    invocation_context.run_config = RunConfig(
+        translation_config=types.TranslationConfig(
+            target_language_code='pl',
+            echo_target_language=True,
+        ),
+    )
+    llm_request = LlmRequest()
+    processor = _BasicLlmRequestProcessor()
+
+    async for _ in processor.run_async(invocation_context, llm_request):
+      pass
+
+    translation_config = llm_request.live_connect_config.translation_config
+    assert translation_config.target_language_code == 'pl'
+    assert translation_config.echo_target_language is True
+
+  @pytest.mark.asyncio
+  async def test_translation_config_defaults_to_none(self):
+    """Without a translation config the live connect field stays None."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-2.5-flash-live',
+    )
+    invocation_context = await _create_invocation_context(agent)
+    llm_request = LlmRequest()
+    processor = _BasicLlmRequestProcessor()
+
+    async for _ in processor.run_async(invocation_context, llm_request):
+      pass
+
+    assert llm_request.live_connect_config.translation_config is None
+
+  @pytest.mark.asyncio
+  async def test_preserves_merged_http_options(self):
+    """Test that processor preserves and merges existing http_options."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-1.5-flash',
+        generate_content_config=types.GenerateContentConfig(
+            http_options=types.HttpOptions(
+                timeout=1000,
+                headers={'Agent-Header': 'agent-val'},
+            )
+        ),
+    )
+
+    invocation_context = await _create_invocation_context(agent)
+    llm_request = LlmRequest()
+
+    # Simulate http_options propagated from RunConfig.
+    llm_request.config.http_options = types.HttpOptions(
+        timeout=500,  # Should override agent.
+        headers={
+            'RunConfig-Header': 'run-val',
+            'Agent-Header': 'run-val-override',
+        },
+    )
+
+    processor = _BasicLlmRequestProcessor()
+
+    async for _ in processor.run_async(invocation_context, llm_request):
+      pass
+
+    # RunConfig timeout wins.
+    assert llm_request.config.http_options.timeout == 500
+
+    # Headers merged, RunConfig wins on conflict.
+    assert (
+        llm_request.config.http_options.headers['RunConfig-Header'] == 'run-val'
+    )
+    assert (
+        llm_request.config.http_options.headers['Agent-Header']
+        == 'run-val-override'
+    )
+
+  @pytest.mark.asyncio
+  async def test_merges_http_options_without_headers(self):
+    """RunConfig timeout/extra_body merge even when no headers are set."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-1.5-flash',
+        generate_content_config=types.GenerateContentConfig(
+            http_options=types.HttpOptions(
+                timeout=1000,
+                headers={'Agent-Header': 'agent-val'},
+            )
+        ),
+    )
+
+    invocation_context = await _create_invocation_context(agent)
+    llm_request = LlmRequest()
+
+    # Propagated RunConfig http_options with no headers.
+    llm_request.config.http_options = types.HttpOptions(
+        timeout=500,
+        extra_body={'priority': 'high'},
+    )
+
+    processor = _BasicLlmRequestProcessor()
+
+    async for _ in processor.run_async(invocation_context, llm_request):
+      pass
+
+    # timeout and extra_body still merge despite empty headers.
+    assert llm_request.config.http_options.timeout == 500
+    assert llm_request.config.http_options.extra_body == {'priority': 'high'}
+    # Agent headers are untouched.
+    assert (
+        llm_request.config.http_options.headers['Agent-Header'] == 'agent-val'
+    )

tests/unittests/flows/llm_flows/test_code_execution.py

@@ -0,0 +1,337 @@
+# 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.
+
+"""Unit tests for Code Execution logic."""
+
+import ast
+import asyncio
+import datetime
+import threading
+from typing import Any
+from typing import Optional
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.code_executors.base_code_executor import BaseCodeExecutor
+from google.adk.code_executors.built_in_code_executor import BuiltInCodeExecutor
+from google.adk.code_executors.code_execution_utils import CodeExecutionInput
+from google.adk.code_executors.code_execution_utils import CodeExecutionResult
+from google.adk.code_executors.code_execution_utils import File
+from google.adk.flows.llm_flows._code_execution import _DATA_FILE_HELPER_LIB
+from google.adk.flows.llm_flows._code_execution import _get_data_file_preprocessing_code
+from google.adk.flows.llm_flows._code_execution import request_processor
+from google.adk.flows.llm_flows._code_execution import response_processor
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class _ExecutionRecord:
+  """Captures how the executor ran, for cross-thread inspection in tests."""
+
+  def __init__(self):
+    self.thread: Optional[Any] = None
+    self.released: Optional[bool] = None
+
+
+class _RecordingCodeExecutor(BaseCodeExecutor):
+  """A code executor that records the thread it runs on.
+
+  `execute_code` blocks on `release` so tests can verify it is offloaded from
+  the event loop: it records the running thread, signals `started`, then waits
+  for `release` before returning.
+  """
+
+  model_config = {'arbitrary_types_allowed': True}
+
+  started: threading.Event
+  release: threading.Event
+  record: _ExecutionRecord
+
+  def execute_code(
+      self,
+      invocation_context,
+      code_execution_input: CodeExecutionInput,
+  ) -> CodeExecutionResult:
+    self.record.thread = threading.current_thread()
+    self.started.set()
+    self.record.released = self.release.wait(timeout=2)
+    return CodeExecutionResult(stdout='ok')
+
+
+@pytest.mark.asyncio
+@patch('google.adk.flows.llm_flows._code_execution.datetime')
+async def test_builtin_code_executor_image_artifact_creation(mock_datetime):
+  """Test BuiltInCodeExecutor creates artifacts for images in response."""
+  mock_now = datetime.datetime(2025, 1, 1, 12, 0, 0)
+  mock_datetime.datetime.fromtimestamp.return_value.astimezone.return_value = (
+      mock_now
+  )
+  code_executor = BuiltInCodeExecutor()
+  agent = Agent(name='test_agent', code_executor=code_executor)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+  invocation_context.artifact_service = MagicMock()
+  invocation_context.artifact_service.save_artifact = AsyncMock(
+      return_value='v1'
+  )
+  llm_response = LlmResponse(
+      content=types.Content(
+          parts=[
+              types.Part(
+                  inline_data=types.Blob(
+                      mime_type='image/png',
+                      data=b'image1',
+                      display_name='image_1.png',
+                  )
+              ),
+              types.Part(text='this is text'),
+              types.Part(
+                  inline_data=types.Blob(mime_type='image/jpeg', data=b'image2')
+              ),
+          ]
+      )
+  )
+
+  events = []
+  async for event in response_processor.run_async(
+      invocation_context, llm_response
+  ):
+    events.append(event)
+
+  expected_timestamp = mock_now.strftime('%Y%m%d_%H%M%S')
+  expected_filename2 = f'{expected_timestamp}.jpeg'
+
+  assert invocation_context.artifact_service.save_artifact.call_count == 2
+  invocation_context.artifact_service.save_artifact.assert_any_call(
+      app_name=invocation_context.app_name,
+      user_id=invocation_context.user_id,
+      session_id=invocation_context.session.id,
+      filename='image_1.png',
+      artifact=types.Part.from_bytes(data=b'image1', mime_type='image/png'),
+  )
+  invocation_context.artifact_service.save_artifact.assert_any_call(
+      app_name=invocation_context.app_name,
+      user_id=invocation_context.user_id,
+      session_id=invocation_context.session.id,
+      filename=expected_filename2,
+      artifact=types.Part.from_bytes(data=b'image2', mime_type='image/jpeg'),
+  )
+
+  assert len(events) == 1
+  assert events[0].actions.artifact_delta == {
+      'image_1.png': 'v1',
+      expected_filename2: 'v1',
+  }
+  assert not events[0].content
+  assert llm_response.content is not None
+  assert len(llm_response.content.parts) == 3
+  assert (
+      llm_response.content.parts[0].text == 'Saved as artifact: image_1.png. '
+  )
+  assert not llm_response.content.parts[0].inline_data
+  assert llm_response.content.parts[1].text == 'this is text'
+  assert (
+      llm_response.content.parts[2].text
+      == f'Saved as artifact: {expected_filename2}. '
+  )
+  assert not llm_response.content.parts[2].inline_data
+
+
+@pytest.mark.asyncio
+@patch('google.adk.flows.llm_flows._code_execution.logger')
+async def test_logs_executed_code(mock_logger):
+  """Test that the response processor logs the code it executes."""
+  mock_code_executor = MagicMock(spec=BaseCodeExecutor)
+  mock_code_executor.code_block_delimiters = [('```python\n', '\n```')]
+  mock_code_executor.error_retry_attempts = 2
+  mock_code_executor.stateful = False
+  mock_code_executor.execute_code.return_value = CodeExecutionResult(
+      stdout='hello'
+  )
+
+  agent = Agent(name='test_agent', code_executor=mock_code_executor)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+  invocation_context.artifact_service = MagicMock()
+  invocation_context.artifact_service.save_artifact = AsyncMock()
+
+  llm_response = LlmResponse(
+      content=types.Content(
+          parts=[
+              types.Part(text='Here is some code:'),
+              types.Part(text='```python\nprint("hello")\n```'),
+          ]
+      )
+  )
+
+  _ = [
+      event
+      async for event in response_processor.run_async(
+          invocation_context, llm_response
+      )
+  ]
+
+  mock_code_executor.execute_code.assert_called_once()
+  mock_logger.debug.assert_called_once_with(
+      'Executed code:\n```\n%s\n```', 'print("hello")'
+  )
+
+
+def test_data_file_helper_lib_defines_crop():
+  """`explore_df` in the injected helper lib calls `crop`, which must exist."""
+  pd = pytest.importorskip('pandas')
+  namespace = {}
+  exec(_DATA_FILE_HELPER_LIB, namespace)  # pylint: disable=exec-used
+
+  crop = namespace['crop']
+  assert crop('short') == 'short'
+  assert crop('x' * 100, max_chars=10) == 'x' * 7 + '...'
+  assert crop('abcdef', max_chars=2) == 'ab'
+
+  # Regression for #4011: explore_df raised NameError when crop was undefined.
+  namespace['explore_df'](pd.DataFrame({'a': [1, 2], 'b': ['x', 'y']}))
+
+
+def test_get_data_file_preprocessing_code_injection_reproduction():
+  """Test that filenames with injection payloads are safely escaped."""
+  bad_filename = "'); print('PWNED')#"
+  file = File(name=bad_filename, mime_type='text/csv', content=b'')
+  code = _get_data_file_preprocessing_code(file)
+
+  tree = ast.parse(code)
+  for node in ast.walk(tree):
+    if isinstance(node, ast.Call):
+      if isinstance(node.func, ast.Name) and node.func.id == 'print':
+        if (
+            len(node.args) == 1
+            and isinstance(node.args[0], ast.Constant)
+            and node.args[0].value == 'PWNED'
+        ):
+          pytest.fail(
+              "Vulnerability reproduction: print('PWNED') was parsed as"
+              ' executable code!'
+          )
+
+  # Check that read_csv was called with bad_filename as a safe string literal.
+  read_csv_arg = None
+  for node in ast.walk(tree):
+    if (
+        isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and node.func.attr == 'read_csv'
+        and isinstance(node.func.value, ast.Name)
+        and node.func.value.id == 'pd'
+    ):
+      assert len(node.args) == 1
+      assert isinstance(node.args[0], ast.Constant)
+      read_csv_arg = node.args[0].value
+      break
+
+  assert read_csv_arg == bad_filename
+
+
+@pytest.mark.asyncio
+async def test_post_processor_does_not_block_event_loop():
+  """Response processor offloads blocking execute_code off the event loop."""
+  started = threading.Event()
+  release = threading.Event()
+  record = _ExecutionRecord()
+  loop_ran = False
+  code_executor = _RecordingCodeExecutor(
+      started=started, release=release, record=record
+  )
+  agent = Agent(name='test_agent', code_executor=code_executor)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+  invocation_context.artifact_service = MagicMock()
+  invocation_context.artifact_service.save_artifact = AsyncMock()
+
+  llm_response = LlmResponse(
+      content=types.Content(
+          parts=[types.Part(text='```python\nprint("hello")\n```')]
+      )
+  )
+
+  async def _release_when_started():
+    nonlocal loop_ran
+    while not started.is_set():
+      await asyncio.sleep(0.001)
+    loop_ran = True
+    release.set()
+
+  releaser = asyncio.create_task(_release_when_started())
+  _ = [
+      event
+      async for event in response_processor.run_async(
+          invocation_context, llm_response
+      )
+  ]
+  await releaser
+
+  assert record.thread is not threading.main_thread()
+  assert record.released is True
+  assert loop_ran is True
+
+
+@pytest.mark.asyncio
+async def test_pre_processor_runs_execute_code_off_the_loop():
+  """Request processor offloads blocking execute_code off the event loop."""
+  started = threading.Event()
+  release = threading.Event()
+  release.set()
+  record = _ExecutionRecord()
+  code_executor = _RecordingCodeExecutor(
+      started=started,
+      release=release,
+      record=record,
+      optimize_data_file=True,
+  )
+  agent = Agent(name='test_agent', code_executor=code_executor)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  llm_request = LlmRequest(
+      contents=[
+          types.Content(
+              role='user',
+              parts=[
+                  types.Part(
+                      inline_data=types.Blob(
+                          mime_type='text/csv',
+                          data=b'col1,col2\n1,2\n',
+                      )
+                  )
+              ],
+          )
+      ]
+  )
+
+  _ = [
+      event
+      async for event in request_processor.run_async(
+          invocation_context, llm_request
+      )
+  ]
+
+  assert record.thread is not threading.main_thread()

tests/unittests/flows/llm_flows/test_compaction_processor.py

@@ -0,0 +1,346 @@
+# 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 request-phase token compaction processor."""
+
+from unittest.mock import AsyncMock
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.apps.app import EventsCompactionConfig
+from google.adk.apps.llm_event_summarizer import LlmEventSummarizer
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows import compaction
+from google.adk.flows.llm_flows import contents
+from google.adk.flows.llm_flows.single_flow import SingleFlow
+from google.adk.models.llm_request import LlmRequest
+from google.adk.sessions.base_session_service import BaseSessionService
+from google.adk.sessions.session import Session
+from google.genai import types
+from google.genai.types import Content
+from google.genai.types import Part
+import pytest
+
+
+def _create_event(
+    *,
+    timestamp: float,
+    invocation_id: str,
+    text: str,
+    prompt_token_count: int | None = None,
+) -> Event:
+  usage_metadata = None
+  if prompt_token_count is not None:
+    usage_metadata = types.GenerateContentResponseUsageMetadata(
+        prompt_token_count=prompt_token_count
+    )
+  return Event(
+      timestamp=timestamp,
+      invocation_id=invocation_id,
+      author='user',
+      content=Content(role='user', parts=[Part(text=text)]),
+      usage_metadata=usage_metadata,
+  )
+
+
+def test_single_flow_includes_compaction_before_contents():
+  flow = SingleFlow()
+
+  compaction_index = flow.request_processors.index(compaction.request_processor)
+  contents_index = flow.request_processors.index(contents.request_processor)
+
+  assert compaction_index < contents_index
+
+
+@pytest.mark.asyncio
+async def test_compaction_request_processor_no_token_config():
+  session = Session(app_name='app', user_id='user', id='session', events=[])
+  session_service = AsyncMock(spec=BaseSessionService)
+  invocation_context = InvocationContext(
+      invocation_id='invocation',
+      agent=LlmAgent(name='agent'),
+      session=session,
+      session_service=session_service,
+      events_compaction_config=EventsCompactionConfig(
+          compaction_interval=2,
+          overlap_size=0,
+      ),
+  )
+
+  llm_request = LlmRequest()
+  processor = compaction.CompactionRequestProcessor()
+
+  events = []
+  async for event in processor.run_async(invocation_context, llm_request):
+    events.append(event)
+
+  assert not events
+  assert not invocation_context.token_compaction_checked
+  session_service.append_event.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_compaction_request_processor_runs_token_compaction():
+  session = Session(
+      app_name='app',
+      user_id='user',
+      id='session',
+      events=[
+          _create_event(timestamp=1.0, invocation_id='inv1', text='e1'),
+          _create_event(timestamp=2.0, invocation_id='inv2', text='e2'),
+          _create_event(
+              timestamp=3.0,
+              invocation_id='inv3',
+              text='e3',
+              prompt_token_count=100,
+          ),
+      ],
+  )
+  session_service = AsyncMock(spec=BaseSessionService)
+  mock_summarizer = AsyncMock(spec=LlmEventSummarizer)
+  compacted_event = Event(author='compactor', invocation_id=Event.new_id())
+  mock_summarizer.maybe_summarize_events.return_value = compacted_event
+
+  invocation_context = InvocationContext(
+      invocation_id='invocation',
+      agent=LlmAgent(name='agent'),
+      session=session,
+      session_service=session_service,
+      events_compaction_config=EventsCompactionConfig(
+          summarizer=mock_summarizer,
+          compaction_interval=999,
+          overlap_size=0,
+          token_threshold=50,
+          event_retention_size=1,
+      ),
+  )
+
+  llm_request = LlmRequest()
+  processor = compaction.CompactionRequestProcessor()
+
+  events = []
+  async for event in processor.run_async(invocation_context, llm_request):
+    events.append(event)
+
+  assert not events
+  assert invocation_context.token_compaction_checked
+
+  compacted_events_arg = mock_summarizer.maybe_summarize_events.call_args[1][
+      'events'
+  ]
+  assert [event.invocation_id for event in compacted_events_arg] == [
+      'inv1',
+      'inv2',
+  ]
+  session_service.append_event.assert_called_once_with(
+      session=session, event=compacted_event
+  )
+
+
+@pytest.mark.asyncio
+async def test_compaction_request_processor_compacts_with_latest_tool_response():
+  session = Session(
+      app_name='app',
+      user_id='user',
+      id='session',
+      events=[
+          _create_event(timestamp=1.0, invocation_id='inv1', text='e1'),
+          _create_event(timestamp=2.0, invocation_id='inv2', text='e2'),
+          Event(
+              timestamp=3.0,
+              invocation_id='current-inv',
+              author='agent',
+              content=Content(
+                  role='model',
+                  parts=[
+                      Part(
+                          function_call=types.FunctionCall(
+                              id='call-1', name='tool', args={}
+                          )
+                      )
+                  ],
+              ),
+          ),
+          Event(
+              timestamp=4.0,
+              invocation_id='current-inv',
+              author='agent',
+              content=Content(
+                  role='user',
+                  parts=[
+                      Part(
+                          function_response=types.FunctionResponse(
+                              id='call-1',
+                              name='tool',
+                              response={'result': 'ok'},
+                          )
+                      )
+                  ],
+              ),
+              usage_metadata=types.GenerateContentResponseUsageMetadata(
+                  prompt_token_count=100
+              ),
+          ),
+      ],
+  )
+  session_service = AsyncMock(spec=BaseSessionService)
+  mock_summarizer = AsyncMock(spec=LlmEventSummarizer)
+  compacted_event = Event(author='compactor', invocation_id=Event.new_id())
+  mock_summarizer.maybe_summarize_events.return_value = compacted_event
+
+  invocation_context = InvocationContext(
+      invocation_id='current-inv',
+      agent=LlmAgent(name='agent'),
+      session=session,
+      session_service=session_service,
+      events_compaction_config=EventsCompactionConfig(
+          summarizer=mock_summarizer,
+          compaction_interval=999,
+          overlap_size=0,
+          token_threshold=50,
+          event_retention_size=1,
+      ),
+  )
+
+  llm_request = LlmRequest()
+  processor = compaction.CompactionRequestProcessor()
+
+  events = []
+  async for event in processor.run_async(invocation_context, llm_request):
+    events.append(event)
+
+  assert not events
+  assert invocation_context.token_compaction_checked
+  compacted_events_arg = mock_summarizer.maybe_summarize_events.call_args[1][
+      'events'
+  ]
+  assert [event.invocation_id for event in compacted_events_arg] == [
+      'inv1',
+      'inv2',
+  ]
+  session_service.append_event.assert_called_once_with(
+      session=session, event=compacted_event
+  )
+
+
+@pytest.mark.asyncio
+async def test_compaction_request_processor_can_compact_current_user_event():
+  session = Session(
+      app_name='app',
+      user_id='user',
+      id='session',
+      events=[
+          _create_event(timestamp=1.0, invocation_id='inv1', text='e1'),
+          Event(
+              timestamp=2.0,
+              invocation_id='current-inv',
+              author='user',
+              content=Content(
+                  role='user',
+                  parts=[Part(text='latest user message')],
+              ),
+              usage_metadata=types.GenerateContentResponseUsageMetadata(
+                  prompt_token_count=100
+              ),
+          ),
+      ],
+  )
+  session_service = AsyncMock(spec=BaseSessionService)
+  mock_summarizer = AsyncMock(spec=LlmEventSummarizer)
+  compacted_event = Event(author='compactor', invocation_id=Event.new_id())
+  mock_summarizer.maybe_summarize_events.return_value = compacted_event
+
+  invocation_context = InvocationContext(
+      invocation_id='current-inv',
+      agent=LlmAgent(name='agent'),
+      session=session,
+      session_service=session_service,
+      events_compaction_config=EventsCompactionConfig(
+          summarizer=mock_summarizer,
+          compaction_interval=999,
+          overlap_size=0,
+          token_threshold=50,
+          event_retention_size=0,
+      ),
+  )
+
+  llm_request = LlmRequest()
+  processor = compaction.CompactionRequestProcessor()
+
+  events = []
+  async for event in processor.run_async(invocation_context, llm_request):
+    events.append(event)
+
+  assert not events
+  assert invocation_context.token_compaction_checked
+  compacted_events_arg = mock_summarizer.maybe_summarize_events.call_args[1][
+      'events'
+  ]
+  assert [event.invocation_id for event in compacted_events_arg] == [
+      'inv1',
+      'current-inv',
+  ]
+  session_service.append_event.assert_called_once_with(
+      session=session, event=compacted_event
+  )
+
+
+@pytest.mark.asyncio
+async def test_compaction_request_processor_not_marked_when_not_compacted():
+  session = Session(
+      app_name='app',
+      user_id='user',
+      id='session',
+      events=[
+          _create_event(timestamp=1.0, invocation_id='inv1', text='e1'),
+          _create_event(
+              timestamp=2.0,
+              invocation_id='inv2',
+              text='e2',
+              prompt_token_count=40,
+          ),
+      ],
+  )
+  session_service = AsyncMock(spec=BaseSessionService)
+  mock_summarizer = AsyncMock(spec=LlmEventSummarizer)
+  mock_summarizer.maybe_summarize_events.return_value = Event(
+      author='compactor',
+      invocation_id=Event.new_id(),
+  )
+
+  invocation_context = InvocationContext(
+      invocation_id='invocation',
+      agent=LlmAgent(name='agent'),
+      session=session,
+      session_service=session_service,
+      events_compaction_config=EventsCompactionConfig(
+          summarizer=mock_summarizer,
+          compaction_interval=999,
+          overlap_size=0,
+          token_threshold=50,
+          event_retention_size=1,
+      ),
+  )
+
+  llm_request = LlmRequest()
+  processor = compaction.CompactionRequestProcessor()
+
+  events = []
+  async for event in processor.run_async(invocation_context, llm_request):
+    events.append(event)
+
+  assert not events
+  assert not invocation_context.token_compaction_checked
+  mock_summarizer.maybe_summarize_events.assert_not_called()
+  session_service.append_event.assert_not_called()

tests/unittests/flows/llm_flows/test_contents_branch.py

@@ -0,0 +1,300 @@
+# 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 branch filtering in contents module.
+
+Branch format: agent_1.agent_2.agent_3 (parent.child.grandchild)
+Child agents can see parent agents' events, but not sibling agents' events.
+"""
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.contents import request_processor
+from google.adk.models.llm_request import LlmRequest
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+@pytest.mark.asyncio
+async def test_branch_filtering_child_sees_parent():
+  """Test that child agents can see parent agents' events."""
+  agent = Agent(model="gemini-2.5-flash", name="child_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Set current branch as child of "parent_agent"
+  invocation_context.branch = "parent_agent.child_agent"
+
+  # Add events from parent and child levels
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("User message"),
+      ),
+      Event(
+          invocation_id="inv2",
+          author="parent_agent",
+          content=types.ModelContent("Parent agent response"),
+          branch="parent_agent",  # Parent branch - should be included
+      ),
+      Event(
+          invocation_id="inv3",
+          author="child_agent",
+          content=types.ModelContent("Child agent response"),
+          branch="parent_agent.child_agent",  # Current branch - should be included
+      ),
+      Event(
+          invocation_id="inv4",
+          author="child_agent",
+          content=types.ModelContent("Excluded response 1"),
+          branch="parent_agent.child_agent000",  # Prefix match BUT not itself/ancestor - should be excluded
+      ),
+      Event(
+          invocation_id="inv5",
+          author="child_agent",
+          content=types.ModelContent("Excluded response 2"),
+          branch="parent_agent.child",  # Prefix match BUT not itself/ancestor - should be excluded
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify child can see user message and parent events, but not sibling events
+  assert len(llm_request.contents) == 3
+  assert llm_request.contents[0] == types.UserContent("User message")
+  assert llm_request.contents[1].role == "user"
+  assert llm_request.contents[1].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[parent_agent] said: Parent agent response"),
+  ]
+  assert llm_request.contents[2] == types.ModelContent("Child agent response")
+
+
+@pytest.mark.asyncio
+async def test_branch_filtering_excludes_sibling_agents():
+  """Test that sibling agents cannot see each other's events."""
+  agent = Agent(model="gemini-2.5-flash", name="child_agent1")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Set current branch as first child
+  invocation_context.branch = "parent_agent.child_agent1"
+
+  # Add events from parent, current child, and sibling child
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("User message"),
+      ),
+      Event(
+          invocation_id="inv2",
+          author="parent_agent",
+          content=types.ModelContent("Parent response"),
+          branch="parent_agent",  # Parent - should be included
+      ),
+      Event(
+          invocation_id="inv3",
+          author="child_agent1",
+          content=types.ModelContent("Child1 response"),
+          branch="parent_agent.child_agent1",  # Current - should be included
+      ),
+      Event(
+          invocation_id="inv4",
+          author="child_agent2",
+          content=types.ModelContent("Sibling response"),
+          branch="parent_agent.child_agent2",  # Sibling - should be excluded
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify sibling events are excluded, but parent and current agent events included
+  assert len(llm_request.contents) == 3
+  assert llm_request.contents[0] == types.UserContent("User message")
+  assert llm_request.contents[1].role == "user"
+  assert llm_request.contents[1].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[parent_agent] said: Parent response"),
+  ]
+  assert llm_request.contents[2] == types.ModelContent("Child1 response")
+
+
+@pytest.mark.asyncio
+async def test_branch_filtering_no_branch_allows_all():
+  """Test that events are included when no branches are set."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # No current branch set (None)
+  invocation_context.branch = None
+
+  # Add events with and without branches
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("No branch message"),
+          branch=None,
+      ),
+      Event(
+          invocation_id="inv2",
+          author="agent1",
+          content=types.ModelContent("Agent with branch"),
+          branch="agent1",
+      ),
+      Event(
+          invocation_id="inv3",
+          author="user",
+          content=types.UserContent("Another no branch"),
+          branch=None,
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify all events are included when no current branch
+  assert len(llm_request.contents) == 3
+  assert llm_request.contents[0] == types.UserContent("No branch message")
+  assert llm_request.contents[1].role == "user"
+  assert llm_request.contents[1].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[agent1] said: Agent with branch"),
+  ]
+  assert llm_request.contents[2] == types.UserContent("Another no branch")
+
+
+@pytest.mark.asyncio
+async def test_branch_filtering_grandchild_sees_grandparent():
+  """Test that deeply nested child agents can see all ancestor events."""
+  agent = Agent(model="gemini-2.5-flash", name="grandchild_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Set deeply nested branch: grandparent.parent.grandchild
+  invocation_context.branch = "grandparent_agent.parent_agent.grandchild_agent"
+
+  # Add events from all levels of hierarchy
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="grandparent_agent",
+          content=types.ModelContent("Grandparent response"),
+          branch="grandparent_agent",
+      ),
+      Event(
+          invocation_id="inv2",
+          author="parent_agent",
+          content=types.ModelContent("Parent response"),
+          branch="grandparent_agent.parent_agent",
+      ),
+      Event(
+          invocation_id="inv3",
+          author="grandchild_agent",
+          content=types.ModelContent("Grandchild response"),
+          branch="grandparent_agent.parent_agent.grandchild_agent",
+      ),
+      Event(
+          invocation_id="inv4",
+          author="sibling_agent",
+          content=types.ModelContent("Sibling response"),
+          branch="grandparent_agent.parent_agent.sibling_agent",
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify only ancestors and current level are included
+  assert len(llm_request.contents) == 3
+  assert llm_request.contents[0].role == "user"
+  assert llm_request.contents[0].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[grandparent_agent] said: Grandparent response"),
+  ]
+  assert llm_request.contents[1].role == "user"
+  assert llm_request.contents[1].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[parent_agent] said: Parent response"),
+  ]
+  assert llm_request.contents[2] == types.ModelContent("Grandchild response")
+
+
+@pytest.mark.asyncio
+async def test_branch_filtering_parent_cannot_see_child():
+  """Test that parent agents cannot see child agents' events."""
+  agent = Agent(model="gemini-2.5-flash", name="parent_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Set current branch as parent
+  invocation_context.branch = "parent_agent"
+
+  # Add events from parent and its children
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("User message"),
+      ),
+      Event(
+          invocation_id="inv2",
+          author="parent_agent",
+          content=types.ModelContent("Parent response"),
+          branch="parent_agent",
+      ),
+      Event(
+          invocation_id="inv3",
+          author="child_agent",
+          content=types.ModelContent("Child response"),
+          branch="parent_agent.child_agent",
+      ),
+      Event(
+          invocation_id="inv4",
+          author="grandchild_agent",
+          content=types.ModelContent("Grandchild response"),
+          branch="parent_agent.child_agent.grandchild_agent",
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify parent cannot see child or grandchild events
+  assert llm_request.contents == [
+      types.UserContent("User message"),
+      types.ModelContent("Parent response"),
+  ]

tests/unittests/flows/llm_flows/test_contents_function.py

@@ -0,0 +1,592 @@
+# 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 function call/response rearrangement in contents module."""
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows import contents
+from google.adk.models.llm_request import LlmRequest
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+@pytest.mark.asyncio
+async def test_basic_function_call_response_processing():
+  """Test basic function call/response processing without rearrangement."""
+  agent = Agent(model="gemini-2.5-flash", name="test_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  function_call = types.FunctionCall(
+      id="call_123", name="search_tool", args={"query": "test"}
+  )
+  function_response = types.FunctionResponse(
+      id="call_123",
+      name="search_tool",
+      response={"results": ["item1", "item2"]},
+  )
+
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("Search for test"),
+      ),
+      Event(
+          invocation_id="inv2",
+          author="test_agent",
+          content=types.ModelContent([types.Part(function_call=function_call)]),
+      ),
+      Event(
+          invocation_id="inv3",
+          author="user",
+          content=types.UserContent(
+              [types.Part(function_response=function_response)]
+          ),
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in contents.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  # Verify no rearrangement occurred
+  assert llm_request.contents == [
+      types.UserContent("Search for test"),
+      types.ModelContent([types.Part(function_call=function_call)]),
+      types.UserContent([types.Part(function_response=function_response)]),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_rearrangement_with_intermediate_function_response():
+  """Test rearrangement when intermediate function response appears after call."""
+  agent = Agent(model="gemini-2.5-flash", name="test_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  function_call = types.FunctionCall(
+      id="long_call_123", name="long_running_tool", args={"task": "process"}
+  )
+  # First intermediate response
+  intermediate_response = types.FunctionResponse(
+      id="long_call_123",
+      name="long_running_tool",
+      response={"status": "processing", "progress": 50},
+  )
+  # Final response
+  final_response = types.FunctionResponse(
+      id="long_call_123",
+      name="long_running_tool",
+      response={"status": "completed", "result": "done"},
+  )
+
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("Run long process"),
+      ),
+      # Function call
+      Event(
+          invocation_id="inv2",
+          author="test_agent",
+          content=types.ModelContent([types.Part(function_call=function_call)]),
+      ),
+      # Intermediate function response appears right after call
+      Event(
+          invocation_id="inv3",
+          author="user",
+          content=types.UserContent(
+              [types.Part(function_response=intermediate_response)]
+          ),
+      ),
+      # Some conversation happens
+      Event(
+          invocation_id="inv4",
+          author="test_agent",
+          content=types.ModelContent("Still processing..."),
+      ),
+      # Final function response (this triggers rearrangement)
+      Event(
+          invocation_id="inv5",
+          author="user",
+          content=types.UserContent(
+              [types.Part(function_response=final_response)]
+          ),
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in contents.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  # Verify rearrangement: intermediate events removed, final response replaces intermediate
+  assert llm_request.contents == [
+      types.UserContent("Run long process"),
+      types.ModelContent([types.Part(function_call=function_call)]),
+      types.UserContent([types.Part(function_response=final_response)]),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_mixed_long_running_and_normal_function_calls():
+  """Test rearrangement with mixed long-running and normal function calls in same event."""
+  agent = Agent(model="gemini-2.5-flash", name="test_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  # Two function calls: one long-running, one normal
+  long_running_call = types.FunctionCall(
+      id="lro_call_456", name="long_running_tool", args={"task": "analyze"}
+  )
+  normal_call = types.FunctionCall(
+      id="normal_call_789", name="search_tool", args={"query": "test"}
+  )
+
+  # Intermediate response for long-running tool
+  lro_intermediate_response = types.FunctionResponse(
+      id="lro_call_456",
+      name="long_running_tool",
+      response={"status": "processing", "progress": 25},
+  )
+  # Response for normal tool (complete)
+  normal_response = types.FunctionResponse(
+      id="normal_call_789",
+      name="search_tool",
+      response={"results": ["item1", "item2"]},
+  )
+  # Final response for long-running tool
+  lro_final_response = types.FunctionResponse(
+      id="lro_call_456",
+      name="long_running_tool",
+      response={"status": "completed", "analysis": "done"},
+  )
+
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("Analyze data and search for info"),
+      ),
+      # Both function calls in same event
+      Event(
+          invocation_id="inv2",
+          author="test_agent",
+          content=types.ModelContent([
+              types.Part(function_call=long_running_call),
+              types.Part(function_call=normal_call),
+          ]),
+      ),
+      # Intermediate responses for both tools
+      Event(
+          invocation_id="inv3",
+          author="user",
+          content=types.UserContent([
+              types.Part(function_response=lro_intermediate_response),
+              types.Part(function_response=normal_response),
+          ]),
+      ),
+      # Some conversation
+      Event(
+          invocation_id="inv4",
+          author="test_agent",
+          content=types.ModelContent("Analysis in progress, search completed"),
+      ),
+      # Final response for long-running tool (triggers rearrangement)
+      Event(
+          invocation_id="inv5",
+          author="user",
+          content=types.UserContent(
+              [types.Part(function_response=lro_final_response)]
+          ),
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in contents.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  # Verify rearrangement: LRO intermediate replaced by final, normal tool preserved
+  assert llm_request.contents == [
+      types.UserContent("Analyze data and search for info"),
+      types.ModelContent([
+          types.Part(function_call=long_running_call),
+          types.Part(function_call=normal_call),
+      ]),
+      types.UserContent([
+          types.Part(function_response=lro_final_response),
+          types.Part(function_response=normal_response),
+      ]),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_completed_long_running_function_in_history():
+  """Test that completed long-running function calls in history.
+
+  Function call/response are properly rearranged and don't affect subsequent
+  conversation.
+  """
+  agent = Agent(model="gemini-2.5-flash", name="test_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  function_call = types.FunctionCall(
+      id="history_call_123", name="long_running_tool", args={"task": "process"}
+  )
+  intermediate_response = types.FunctionResponse(
+      id="history_call_123",
+      name="long_running_tool",
+      response={"status": "processing", "progress": 50},
+  )
+  final_response = types.FunctionResponse(
+      id="history_call_123",
+      name="long_running_tool",
+      response={"status": "completed", "result": "done"},
+  )
+
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("Start long process"),
+      ),
+      # Function call in history
+      Event(
+          invocation_id="inv2",
+          author="test_agent",
+          content=types.ModelContent([types.Part(function_call=function_call)]),
+      ),
+      # Intermediate response in history
+      Event(
+          invocation_id="inv3",
+          author="user",
+          content=types.UserContent(
+              [types.Part(function_response=intermediate_response)]
+          ),
+      ),
+      # Some conversation happens
+      Event(
+          invocation_id="inv4",
+          author="test_agent",
+          content=types.ModelContent("Still processing..."),
+      ),
+      # Final response completes the long-running function in history
+      Event(
+          invocation_id="inv5",
+          author="user",
+          content=types.UserContent(
+              [types.Part(function_response=final_response)]
+          ),
+      ),
+      # Agent acknowledges completion
+      Event(
+          invocation_id="inv6",
+          author="test_agent",
+          content=types.ModelContent("Process completed successfully!"),
+      ),
+      # Latest event is regular user message, not function response
+      Event(
+          invocation_id="inv7",
+          author="user",
+          content=types.UserContent("Great! What's next?"),
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in contents.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  # Verify the long-running function in history was rearranged correctly:
+  # - Intermediate response was replaced by final response
+  # - Non-function events (like "Still processing...") are preserved
+  # - No further rearrangement occurs since latest event is not function response
+  assert llm_request.contents == [
+      types.UserContent("Start long process"),
+      types.ModelContent([types.Part(function_call=function_call)]),
+      types.UserContent([types.Part(function_response=final_response)]),
+      types.ModelContent("Still processing..."),
+      types.ModelContent("Process completed successfully!"),
+      types.UserContent("Great! What's next?"),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_completed_mixed_function_calls_in_history():
+  """Test completed mixed long-running and normal function calls in history don't affect subsequent conversation."""
+  agent = Agent(model="gemini-2.5-flash", name="test_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  # Two function calls: one long-running, one normal
+  long_running_call = types.FunctionCall(
+      id="history_lro_123", name="long_running_tool", args={"task": "analyze"}
+  )
+  normal_call = types.FunctionCall(
+      id="history_normal_456", name="search_tool", args={"query": "data"}
+  )
+
+  # Intermediate response for long-running tool
+  lro_intermediate_response = types.FunctionResponse(
+      id="history_lro_123",
+      name="long_running_tool",
+      response={"status": "processing", "progress": 30},
+  )
+  # Complete response for normal tool
+  normal_response = types.FunctionResponse(
+      id="history_normal_456",
+      name="search_tool",
+      response={"results": ["result1", "result2"]},
+  )
+  # Final response for long-running tool
+  lro_final_response = types.FunctionResponse(
+      id="history_lro_123",
+      name="long_running_tool",
+      response={"status": "completed", "analysis": "finished"},
+  )
+
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("Analyze and search simultaneously"),
+      ),
+      # Both function calls in history
+      Event(
+          invocation_id="inv2",
+          author="test_agent",
+          content=types.ModelContent([
+              types.Part(function_call=long_running_call),
+              types.Part(function_call=normal_call),
+          ]),
+      ),
+      # Intermediate responses for both tools in history
+      Event(
+          invocation_id="inv3",
+          author="user",
+          content=types.UserContent([
+              types.Part(function_response=lro_intermediate_response),
+              types.Part(function_response=normal_response),
+          ]),
+      ),
+      # Some conversation in history
+      Event(
+          invocation_id="inv4",
+          author="test_agent",
+          content=types.ModelContent("Analysis continuing, search done"),
+      ),
+      # Final response completes the long-running function in history
+      Event(
+          invocation_id="inv5",
+          author="user",
+          content=types.UserContent(
+              [types.Part(function_response=lro_final_response)]
+          ),
+      ),
+      # Agent acknowledges completion
+      Event(
+          invocation_id="inv6",
+          author="test_agent",
+          content=types.ModelContent("Both tasks completed successfully!"),
+      ),
+      # Latest event is regular user message, not function response
+      Event(
+          invocation_id="inv7",
+          author="user",
+          content=types.UserContent("Perfect! What should we do next?"),
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in contents.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  # Verify mixed functions in history were rearranged correctly:
+  # - LRO intermediate was replaced by final response
+  # - Normal tool response was preserved
+  # - Non-function events preserved, no further rearrangement
+  assert llm_request.contents == [
+      types.UserContent("Analyze and search simultaneously"),
+      types.ModelContent([
+          types.Part(function_call=long_running_call),
+          types.Part(function_call=normal_call),
+      ]),
+      types.UserContent([
+          types.Part(function_response=lro_final_response),
+          types.Part(function_response=normal_response),
+      ]),
+      types.ModelContent("Analysis continuing, search done"),
+      types.ModelContent("Both tasks completed successfully!"),
+      types.UserContent("Perfect! What should we do next?"),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_function_rearrangement_preserves_other_content():
+  """Test that non-function content is preserved during rearrangement."""
+  agent = Agent(model="gemini-2.5-flash", name="test_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  function_call = types.FunctionCall(
+      id="preserve_test", name="long_running_tool", args={"test": "value"}
+  )
+  intermediate_response = types.FunctionResponse(
+      id="preserve_test",
+      name="long_running_tool",
+      response={"status": "processing"},
+  )
+  final_response = types.FunctionResponse(
+      id="preserve_test",
+      name="long_running_tool",
+      response={"output": "preserved"},
+  )
+
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("Before function call"),
+      ),
+      Event(
+          invocation_id="inv2",
+          author="test_agent",
+          content=types.ModelContent([
+              types.Part(text="I'll process this for you"),
+              types.Part(function_call=function_call),
+          ]),
+      ),
+      # Intermediate response with mixed content
+      Event(
+          invocation_id="inv3",
+          author="user",
+          content=types.UserContent([
+              types.Part(text="Intermediate prefix"),
+              types.Part(function_response=intermediate_response),
+              types.Part(text="Processing..."),
+          ]),
+      ),
+      # This should be removed during rearrangement
+      Event(
+          invocation_id="inv4",
+          author="test_agent",
+          content=types.ModelContent("Still working on it..."),
+      ),
+      # Final response with mixed content (triggers rearrangement)
+      Event(
+          invocation_id="inv5",
+          author="user",
+          content=types.UserContent([
+              types.Part(text="Final prefix"),
+              types.Part(function_response=final_response),
+              types.Part(text="Final suffix"),
+          ]),
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in contents.request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    pass
+
+  # Verify non-function content is preserved during rearrangement
+  # Intermediate response replaced by final, but ALL text content preserved
+  assert llm_request.contents == [
+      types.UserContent("Before function call"),
+      types.ModelContent([
+          types.Part(text="I'll process this for you"),
+          types.Part(function_call=function_call),
+      ]),
+      types.UserContent([
+          types.Part(text="Intermediate prefix"),
+          types.Part(function_response=final_response),
+          types.Part(text="Processing..."),
+          types.Part(text="Final prefix"),
+          types.Part(text="Final suffix"),
+      ]),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_error_when_function_response_without_matching_call():
+  """Test error when function response has no matching function call."""
+  agent = Agent(model="gemini-2.5-flash", name="test_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  # Function response without matching call
+  orphaned_response = types.FunctionResponse(
+      id="no_matching_call",
+      name="orphaned_tool",
+      response={"error": "no matching call"},
+  )
+
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("Regular message"),
+      ),
+      # Response without any prior matching function call
+      Event(
+          invocation_id="inv2",
+          author="user",
+          content=types.UserContent(
+              [types.Part(function_response=orphaned_response)]
+          ),
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # This should raise a ValueError during processing
+  with pytest.raises(ValueError, match="No function call event found"):
+    async for _ in contents.request_processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass

tests/unittests/flows/llm_flows/test_contents_other_agent.py

@@ -0,0 +1,494 @@
+# 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.
+
+"""Behavioral tests for other agent message processing in contents module."""
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.agents.run_config import RunConfig
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.contents import request_processor
+from google.adk.models.llm_request import LlmRequest
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+@pytest.mark.asyncio
+async def test_other_agent_message_appears_as_user_context():
+  """Test that messages from other agents appear as user context."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Add event from another agent
+  other_agent_event = Event(
+      invocation_id="test_inv",
+      author="other_agent",
+      content=types.ModelContent("Hello from other agent"),
+  )
+  invocation_context.session.events = [other_agent_event]
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify the other agent's message is presented as user context
+  assert llm_request.contents[0].role == "user"
+  assert llm_request.contents[0].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[other_agent] said: Hello from other agent"),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_other_agent_thoughts_are_excluded():
+  """Test that thoughts from other agents are excluded from context."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Add event from other agent with both regular text and thoughts
+  other_agent_event = Event(
+      invocation_id="test_inv",
+      author="other_agent",
+      content=types.ModelContent([
+          types.Part(text="Public message", thought=False),
+          types.Part(text="Private thought", thought=True),
+          types.Part(text="Another public message"),
+      ]),
+  )
+  invocation_context.session.events = [other_agent_event]
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify only non-thought parts are included (thoughts excluded)
+  assert llm_request.contents[0].role == "user"
+  assert llm_request.contents[0].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[other_agent] said: Public message"),
+      types.Part(text="[other_agent] said: Another public message"),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_other_agent_thoughts_can_be_included_as_context():
+  """Test opt-in inclusion of thoughts from other agents."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent,
+      run_config=RunConfig(include_thoughts_from_other_agents=True),
+  )
+  other_agent_event = Event(
+      invocation_id="test_inv",
+      author="other_agent",
+      content=types.ModelContent([
+          types.Part(text="Public message", thought=False),
+          types.Part(text="Private thought", thought=True),
+          types.Part(text="Another public message"),
+      ]),
+  )
+  invocation_context.session.events = [other_agent_event]
+
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  assert llm_request.contents[0].role == "user"
+  assert llm_request.contents[0].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[other_agent] said: Public message"),
+      types.Part(text="[other_agent] thought: Private thought"),
+      types.Part(text="[other_agent] said: Another public message"),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_other_agent_thought_only_message_can_be_included_as_context():
+  """Test opt-in inclusion of thought-only messages from other agents."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent,
+      run_config=RunConfig(include_thoughts_from_other_agents=True),
+  )
+  other_agent_event = Event(
+      invocation_id="test_inv",
+      author="other_agent",
+      content=types.ModelContent([
+          types.Part(text="First private thought", thought=True),
+          types.Part(text="Second private thought", thought=True),
+      ]),
+  )
+  invocation_context.session.events = [other_agent_event]
+
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  assert llm_request.contents[0].role == "user"
+  assert llm_request.contents[0].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[other_agent] thought: First private thought"),
+      types.Part(text="[other_agent] thought: Second private thought"),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_other_agent_thoughts_excluded_from_current_turn_only_context():
+  """Test include_contents='none' does not include other-agent thoughts."""
+  agent = Agent(
+      model="gemini-2.5-flash",
+      name="current_agent",
+      include_contents="none",
+  )
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent,
+      run_config=RunConfig(include_thoughts_from_other_agents=True),
+  )
+  invocation_context.session.events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("Earlier user message"),
+      ),
+      Event(
+          invocation_id="inv2",
+          author="other_agent",
+          content=types.ModelContent([
+              types.Part(text="Private thought", thought=True),
+              types.Part(text="Visible handoff"),
+          ]),
+      ),
+  ]
+
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  assert llm_request.contents == [
+      types.Content(
+          role="user",
+          parts=[
+              types.Part(text="For context:"),
+              types.Part(text="[other_agent] said: Visible handoff"),
+          ],
+      )
+  ]
+
+
+@pytest.mark.asyncio
+async def test_other_agent_function_calls():
+  """Test that function calls from other agents are preserved in context."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Add event from other agent with function call
+  function_call = types.FunctionCall(
+      id="func_123", name="search_tool", args={"query": "test query"}
+  )
+  other_agent_event = Event(
+      invocation_id="test_inv",
+      author="other_agent",
+      content=types.ModelContent([types.Part(function_call=function_call)]),
+  )
+  invocation_context.session.events = [other_agent_event]
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify function call is presented as context
+  assert llm_request.contents[0].role == "user"
+  assert llm_request.contents[0].parts == [
+      types.Part(text="For context:"),
+      types.Part(
+          text="""\
+[other_agent] called tool `search_tool` with parameters: {'query': 'test query'}"""
+      ),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_other_agent_function_responses():
+  """Test that function responses from other agents are properly formatted."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  # Add event from other agent with function response
+  function_response = types.FunctionResponse(
+      id="func_123",
+      name="search_tool",
+      response={"results": ["item1", "item2"]},
+  )
+  other_agent_event = Event(
+      invocation_id="test_inv",
+      author="other_agent",
+      content=types.Content(
+          role="user", parts=[types.Part(function_response=function_response)]
+      ),
+  )
+  invocation_context.session.events = [other_agent_event]
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify function response is presented as context
+  assert llm_request.contents[0].role == "user"
+  assert llm_request.contents[0].parts == [
+      types.Part(text="For context:"),
+      types.Part(
+          text=(
+              "[other_agent] `search_tool` tool returned result: {'results':"
+              " ['item1', 'item2']}"
+          )
+      ),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_other_agent_function_call_response():
+  """Test function call and response sequence from other agents."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Add function call event from other agent
+  function_call = types.FunctionCall(
+      id="func_123", name="calc_tool", args={"query": "6x7"}
+  )
+  call_event = Event(
+      invocation_id="test_inv1",
+      author="other_agent",
+      content=types.ModelContent([
+          types.Part(text="Let me calculate this"),
+          types.Part(function_call=function_call),
+      ]),
+  )
+  # Add function response event
+  function_response = types.FunctionResponse(
+      id="func_123", name="calc_tool", response={"result": 42}
+  )
+  response_event = Event(
+      invocation_id="test_inv2",
+      author="other_agent",
+      content=types.UserContent(
+          parts=[types.Part(function_response=function_response)]
+      ),
+  )
+  invocation_context.session.events = [call_event, response_event]
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify function call and response are properly formatted
+  assert len(llm_request.contents) == 2
+
+  # Function call from other agent
+  assert llm_request.contents[0].role == "user"
+  assert llm_request.contents[0].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[other_agent] said: Let me calculate this"),
+      types.Part(
+          text=(
+              "[other_agent] called tool `calc_tool` with parameters: {'query':"
+              " '6x7'}"
+          )
+      ),
+  ]
+  # Function response from other agent
+  assert llm_request.contents[1].role == "user"
+  assert llm_request.contents[1].parts == [
+      types.Part(text="For context:"),
+      types.Part(
+          text="[other_agent] `calc_tool` tool returned result: {'result': 42}"
+      ),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_other_agent_empty_content():
+  """Test that other agent messages with only thoughts or empty content are filtered out."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Add events: user message, other agents with empty content, user message
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("Hello"),
+      ),
+      # Other agent with only thoughts
+      Event(
+          invocation_id="inv2",
+          author="other_agent1",
+          content=types.ModelContent([
+              types.Part(text="This is a private thought", thought=True),
+              types.Part(text="Another private thought", thought=True),
+          ]),
+      ),
+      # Other agent with empty text and thoughts
+      Event(
+          invocation_id="inv3",
+          author="other_agent2",
+          content=types.ModelContent([
+              types.Part(text="", thought=False),
+              types.Part(text="Secret thought", thought=True),
+          ]),
+      ),
+      Event(
+          invocation_id="inv4",
+          author="user",
+          content=types.UserContent("World"),
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify empty content events are completely filtered out
+  assert llm_request.contents == [
+      types.UserContent("Hello"),
+      types.UserContent("World"),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_multiple_agents_in_conversation():
+  """Test handling multiple agents in a conversation flow."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  # Create a multi-agent conversation
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="user",
+          content=types.UserContent("Hello everyone"),
+      ),
+      Event(
+          invocation_id="inv2",
+          author="agent1",
+          content=types.ModelContent("Hi from agent1"),
+      ),
+      Event(
+          invocation_id="inv3",
+          author="agent2",
+          content=types.ModelContent("Hi from agent2"),
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify all messages are properly processed
+  assert len(llm_request.contents) == 3
+
+  # User message should remain as user
+  assert llm_request.contents[0] == types.UserContent("Hello everyone")
+  # Other agents' messages should be converted to user context
+  assert llm_request.contents[1].role == "user"
+  assert llm_request.contents[1].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[agent1] said: Hi from agent1"),
+  ]
+  assert llm_request.contents[2].role == "user"
+  assert llm_request.contents[2].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[agent2] said: Hi from agent2"),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_current_agent_messages_not_converted():
+  """Test that the current agent's own messages are not converted."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Add events from both current agent and other agent
+  events = [
+      Event(
+          invocation_id="inv1",
+          author="current_agent",
+          content=types.ModelContent("My own message"),
+      ),
+      Event(
+          invocation_id="inv2",
+          author="other_agent",
+          content=types.ModelContent("Other agent message"),
+      ),
+  ]
+  invocation_context.session.events = events
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify current agent's message stays as model role
+  # and other agent's message is converted to user context
+  assert len(llm_request.contents) == 2
+  assert llm_request.contents[0] == types.ModelContent("My own message")
+  assert llm_request.contents[1].role == "user"
+  assert llm_request.contents[1].parts == [
+      types.Part(text="For context:"),
+      types.Part(text="[other_agent] said: Other agent message"),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_user_messages_preserved():
+  """Test that user messages are preserved as-is."""
+  agent = Agent(model="gemini-2.5-flash", name="current_agent")
+  llm_request = LlmRequest(model="gemini-2.5-flash")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Add user message
+  user_event = Event(
+      invocation_id="inv1",
+      author="user",
+      content=types.UserContent("User message"),
+  )
+  invocation_context.session.events = [user_event]
+
+  # Process the request
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  # Verify user message is preserved exactly
+  assert len(llm_request.contents) == 1
+  assert llm_request.contents[0] == types.UserContent("User message")

tests/unittests/flows/llm_flows/test_context_cache_processor.py

@@ -0,0 +1,646 @@
+# 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 ContextCacheRequestProcessor."""
+
+import time
+from unittest.mock import MagicMock
+
+from google.adk.agents.context_cache_config import ContextCacheConfig
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.context_cache_processor import ContextCacheRequestProcessor
+from google.adk.models.cache_metadata import CacheMetadata
+from google.adk.models.llm_request import LlmRequest
+from google.adk.sessions.base_session_service import BaseSessionService
+from google.adk.sessions.session import Session
+from google.genai import types
+import pytest
+
+
+class TestContextCacheRequestProcessor:
+  """Test suite for ContextCacheRequestProcessor."""
+
+  def setup_method(self):
+    """Set up test fixtures."""
+    self.processor = ContextCacheRequestProcessor()
+    self.cache_config = ContextCacheConfig(
+        cache_intervals=10, ttl_seconds=1800, min_tokens=1024
+    )
+
+  def create_invocation_context(
+      self,
+      agent,
+      context_cache_config=None,
+      session_events=None,
+      invocation_id="test_invocation",
+  ):
+    """Helper to create InvocationContext."""
+    mock_session = Session(
+        id="test_session",
+        app_name="test_app",
+        user_id="test_user",
+        events=session_events or [],
+    )
+
+    mock_session_service = MagicMock(spec=BaseSessionService)
+
+    return InvocationContext(
+        agent=agent,
+        session=mock_session,
+        session_service=mock_session_service,
+        context_cache_config=context_cache_config,
+        invocation_id=invocation_id,
+    )
+
+  def create_cache_metadata(
+      self, invocations_used=1, cache_name="test-cache", contents_count=3
+  ):
+    """Helper to create CacheMetadata."""
+    return CacheMetadata(
+        cache_name=(
+            f"projects/test/locations/us-central1/cachedContents/{cache_name}"
+        ),
+        expire_time=time.time() + 1800,
+        fingerprint="test_fingerprint",
+        invocations_used=invocations_used,
+        contents_count=contents_count,
+        created_at=time.time() - 600,
+    )
+
+  async def test_no_cache_config(self):
+    """Test processor with no cache config."""
+    agent = LlmAgent(name="test_agent")
+    invocation_context = self.create_invocation_context(
+        agent, context_cache_config=None
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    # Process should complete without adding cache config
+    events = []
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      events.append(event)
+
+    assert len(events) == 0  # No events yielded
+    assert llm_request.cache_config is None
+
+  async def test_with_cache_config_no_session_events(self):
+    """Test processor with cache config but no session events."""
+    agent = LlmAgent(name="test_agent")
+    invocation_context = self.create_invocation_context(
+        agent, context_cache_config=self.cache_config
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    # Process should add cache config but no metadata
+    events = []
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      events.append(event)
+
+    assert len(events) == 0  # No events yielded
+    assert llm_request.cache_config == self.cache_config
+    assert llm_request.cache_metadata is None
+
+  async def test_with_cache_metadata_same_invocation(self):
+    """Test processor finds cache metadata from same invocation."""
+    agent = LlmAgent(name="test_agent")
+    cache_metadata = self.create_cache_metadata(invocations_used=5)
+
+    # Event with same invocation ID
+    events = [
+        Event(
+            author="test_agent",
+            cache_metadata=cache_metadata,
+            invocation_id="test_invocation",
+        )
+    ]
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=events,
+        invocation_id="test_invocation",
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    # Process should add cache config and metadata (same invocation, no increment)
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    assert llm_request.cache_config == self.cache_config
+    assert llm_request.cache_metadata == cache_metadata
+    assert llm_request.cache_metadata.invocations_used == 5  # No increment
+
+  async def test_with_cache_metadata_different_invocation(self):
+    """Test processor finds cache metadata from different invocation."""
+    agent = LlmAgent(name="test_agent")
+    cache_metadata = self.create_cache_metadata(invocations_used=5)
+
+    # Event with different invocation ID
+    events = [
+        Event(
+            author="test_agent",
+            cache_metadata=cache_metadata,
+            invocation_id="previous_invocation",
+        )
+    ]
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=events,
+        invocation_id="current_invocation",
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    # Process should add cache config and increment invocations_used
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    assert llm_request.cache_config == self.cache_config
+    assert llm_request.cache_metadata is not None
+    assert llm_request.cache_metadata.invocations_used == 6  # Incremented
+
+  async def test_cache_metadata_agent_filtering(self):
+    """Test that cache metadata is filtered by agent name."""
+    agent = LlmAgent(name="target_agent")
+    target_cache = self.create_cache_metadata(
+        invocations_used=3, cache_name="target"
+    )
+    other_cache = self.create_cache_metadata(
+        invocations_used=7, cache_name="other"
+    )
+
+    events = [
+        Event(
+            author="other_agent",
+            cache_metadata=other_cache,
+            invocation_id="other_invocation",
+        ),
+        Event(
+            author="target_agent",
+            cache_metadata=target_cache,
+            invocation_id="target_invocation",
+        ),
+    ]
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=events,
+        invocation_id="current_invocation",
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    # Should only use target_agent's cache metadata
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    assert llm_request.cache_metadata is not None
+    assert llm_request.cache_metadata.cache_name == target_cache.cache_name
+    assert llm_request.cache_metadata.invocations_used == 4  # target_cache + 1
+
+  async def test_latest_cache_metadata_selected(self):
+    """Test that the latest cache metadata is selected."""
+    agent = LlmAgent(name="test_agent")
+    older_cache = self.create_cache_metadata(
+        invocations_used=2, cache_name="older"
+    )
+    newer_cache = self.create_cache_metadata(
+        invocations_used=5, cache_name="newer"
+    )
+
+    # Events in chronological order (older first)
+    events = [
+        Event(
+            author="test_agent",
+            cache_metadata=older_cache,
+            invocation_id="older_invocation",
+        ),
+        Event(
+            author="test_agent",
+            cache_metadata=newer_cache,
+            invocation_id="newer_invocation",
+        ),
+    ]
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=events,
+        invocation_id="current_invocation",
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    # Should use the newer (latest) cache metadata
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    assert llm_request.cache_metadata is not None
+    assert llm_request.cache_metadata.cache_name == newer_cache.cache_name
+    assert llm_request.cache_metadata.invocations_used == 6  # newer_cache + 1
+
+  async def test_no_cache_metadata_events(self):
+    """Test when session has events but no cache metadata."""
+    agent = LlmAgent(name="test_agent")
+
+    events = [
+        Event(author="test_agent", cache_metadata=None),
+        Event(author="other_agent", cache_metadata=None),
+    ]
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=events,
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    # Should add cache config but no metadata
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    assert llm_request.cache_config == self.cache_config
+    assert llm_request.cache_metadata is None
+
+  async def test_empty_session(self):
+    """Test with empty session."""
+    agent = LlmAgent(name="test_agent")
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=[],
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    # Should add cache config but no metadata
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    assert llm_request.cache_config == self.cache_config
+    assert llm_request.cache_metadata is None
+
+  async def test_processor_yields_no_events(self):
+    """Test that processor yields no events."""
+    agent = LlmAgent(name="test_agent")
+
+    invocation_context = self.create_invocation_context(
+        agent, context_cache_config=self.cache_config
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    events = []
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      events.append(event)
+
+    # Processor should never yield events
+    assert len(events) == 0
+
+  async def test_mixed_events_scenario(self):
+    """Test complex scenario with mixed events."""
+    agent = LlmAgent(name="test_agent")
+    cache_metadata = self.create_cache_metadata(invocations_used=10)
+
+    events = [
+        Event(author="other_agent", cache_metadata=None),
+        Event(author="test_agent", cache_metadata=None),  # No cache metadata
+        Event(
+            author="different_agent", cache_metadata=cache_metadata
+        ),  # Wrong agent
+        Event(
+            author="test_agent",
+            cache_metadata=cache_metadata,
+            invocation_id="prev",
+        ),
+    ]
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=events,
+        invocation_id="current",
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    # Should find the test_agent's cache metadata and increment it
+    assert llm_request.cache_config == self.cache_config
+    assert llm_request.cache_metadata is not None
+    assert llm_request.cache_metadata.invocations_used == 11  # 10 + 1
+
+  async def test_cacheable_contents_token_count_extraction(self):
+    """Test that previous prompt token count is extracted and set."""
+    agent = LlmAgent(name="test_agent")
+
+    # Create event with usage metadata
+    event_with_tokens = Event(
+        author="test_agent",
+        usage_metadata=types.UsageMetadata(
+            prompt_token_count=1024,
+            response_token_count=256,
+            total_token_count=1280,
+        ),
+    )
+
+    events = [event_with_tokens]
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=events,
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    # Should extract token count from the event
+    assert llm_request.cacheable_contents_token_count == 1024
+
+  async def test_cacheable_contents_token_count_no_usage_metadata(self):
+    """Test when no usage metadata is available."""
+    agent = LlmAgent(name="test_agent")
+
+    events = [
+        Event(author="test_agent", usage_metadata=None),
+        Event(author="other_agent"),
+    ]
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=events,
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    # Should not set token count when no usage metadata
+    assert llm_request.cacheable_contents_token_count is None
+
+  async def test_cacheable_contents_token_count_agent_filtering(self):
+    """Test that token count is filtered by agent name."""
+    agent = LlmAgent(name="target_agent")
+
+    events = [
+        Event(
+            author="other_agent",
+            usage_metadata=types.UsageMetadata(prompt_token_count=2048),
+        ),
+        Event(
+            author="target_agent",
+            usage_metadata=types.UsageMetadata(prompt_token_count=1024),
+        ),
+    ]
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=events,
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    # Should use target_agent's token count, not other_agent's
+    assert llm_request.cacheable_contents_token_count == 1024
+
+  async def test_cacheable_contents_token_count_latest_selected(self):
+    """Test that the most recent token count is selected."""
+    agent = LlmAgent(name="test_agent")
+
+    events = [
+        Event(
+            author="test_agent",
+            usage_metadata=types.UsageMetadata(prompt_token_count=512),
+        ),
+        Event(
+            author="test_agent",
+            usage_metadata=types.UsageMetadata(prompt_token_count=1024),
+        ),
+    ]
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=events,
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    # Should use the latest (most recent) token count
+    assert llm_request.cacheable_contents_token_count == 1024
+
+  async def test_cache_metadata_and_token_count_both_found(self):
+    """Test that both cache metadata and token count are found in single pass."""
+    agent = LlmAgent(name="test_agent")
+    cache_metadata = self.create_cache_metadata(invocations_used=5)
+
+    events = [
+        Event(
+            author="test_agent",
+            cache_metadata=cache_metadata,
+            usage_metadata=types.UsageMetadata(prompt_token_count=1024),
+            invocation_id="previous_invocation",
+        ),
+    ]
+
+    invocation_context = self.create_invocation_context(
+        agent,
+        context_cache_config=self.cache_config,
+        session_events=events,
+        invocation_id="current_invocation",
+    )
+
+    llm_request = LlmRequest(
+        model="gemini-2.5-flash",
+        contents=[
+            types.Content(
+                role="user",
+                parts=[types.Part(text="Hello")],
+            )
+        ],
+    )
+
+    async for event in self.processor.run_async(
+        invocation_context, llm_request
+    ):
+      pass
+
+    # Should find both cache metadata and token count
+    assert llm_request.cache_metadata is not None
+    assert llm_request.cache_metadata.invocations_used == 6  # 5 + 1
+    assert llm_request.cacheable_contents_token_count == 1024

tests/unittests/flows/llm_flows/test_functions_error_messages.py

@@ -0,0 +1,89 @@
+# 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 enhanced error messages in function tool handling."""
+
+from google.adk.flows.llm_flows.functions import _get_tool
+from google.adk.tools import BaseTool
+from google.genai import types
+import pytest
+
+
+# Mock tool for testing error messages
+class MockTool(BaseTool):
+  """Mock tool for testing error messages."""
+
+  def __init__(self, name: str = 'mock_tool'):
+    super().__init__(name=name, description=f'Mock tool: {name}')
+
+  def call(self, *args, **kwargs):
+    return 'mock_response'
+
+
+def test_tool_not_found_enhanced_error():
+  """Verify enhanced error message for tool not found."""
+  function_call = types.FunctionCall(name='nonexistent_tool', args={})
+  tools_dict = {
+      'get_weather': MockTool(name='get_weather'),
+      'calculate_sum': MockTool(name='calculate_sum'),
+      'search_database': MockTool(name='search_database'),
+  }
+
+  with pytest.raises(ValueError) as exc_info:
+    _get_tool(function_call, tools_dict)
+
+  error_msg = str(exc_info.value)
+
+  # Verify error message components
+  assert 'nonexistent_tool' in error_msg
+  assert 'Available tools:' in error_msg
+  assert 'get_weather' in error_msg
+  assert 'Possible causes:' in error_msg
+  assert 'Suggested fixes:' in error_msg
+
+
+def test_tool_not_found_with_different_name():
+  """Verify error message contains basic information."""
+  function_call = types.FunctionCall(name='completely_different', args={})
+  tools_dict = {
+      'get_weather': MockTool(name='get_weather'),
+      'calculate_sum': MockTool(name='calculate_sum'),
+  }
+
+  with pytest.raises(ValueError) as exc_info:
+    _get_tool(function_call, tools_dict)
+
+  error_msg = str(exc_info.value)
+
+  # Verify error message contains basic information
+  assert 'completely_different' in error_msg
+  assert 'Available tools:' in error_msg
+
+
+def test_tool_not_found_shows_all_tools():
+  """Verify error message shows all tools (no truncation)."""
+  function_call = types.FunctionCall(name='nonexistent', args={})
+
+  # Create 100 tools
+  tools_dict = {f'tool_{i}': MockTool(name=f'tool_{i}') for i in range(100)}
+
+  with pytest.raises(ValueError) as exc_info:
+    _get_tool(function_call, tools_dict)
+
+  error_msg = str(exc_info.value)
+
+  # Verify all tools are shown (no truncation)
+  assert 'tool_0' in error_msg  # First tool shown
+  assert 'tool_99' in error_msg  # Last tool also shown
+  assert 'showing first 20 of' not in error_msg  # No truncation message

tests/unittests/flows/llm_flows/test_functions_long_running.py

@@ -0,0 +1,244 @@
+# 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.
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.tools.long_running_tool import LongRunningFunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai.types import Part
+
+from ... import testing_utils
+
+
+def test_async_function():
+  responses = [
+      Part.from_function_call(name='increase_by_one', args={'x': 1}),
+      'response1',
+      'response2',
+      'response3',
+      'response4',
+  ]
+  mockModel = testing_utils.MockModel.create(responses=responses)
+  function_called = 0
+
+  def increase_by_one(x: int, tool_context: ToolContext) -> int:
+    nonlocal function_called
+
+    function_called += 1
+    return {'status': 'pending'}
+
+  # Calls the first time.
+  agent = Agent(
+      name='root_agent',
+      model=mockModel,
+      tools=[LongRunningFunctionTool(func=increase_by_one)],
+  )
+  runner = testing_utils.InMemoryRunner(agent)
+  events = runner.run('test1')
+
+  # Asserts the requests.
+  assert len(mockModel.requests) == 2
+  # 1 item: user content
+  assert mockModel.requests[0].contents == [
+      testing_utils.UserContent('test1'),
+  ]
+  increase_by_one_call = Part.from_function_call(
+      name='increase_by_one', args={'x': 1}
+  )
+  pending_response = Part.from_function_response(
+      name='increase_by_one', response={'status': 'pending'}
+  )
+
+  assert testing_utils.simplify_contents(mockModel.requests[1].contents) == [
+      ('user', 'test1'),
+      ('model', increase_by_one_call),
+      ('user', pending_response),
+  ]
+
+  # Asserts the function calls.
+  assert function_called == 1
+
+  # Asserts the responses.
+  assert testing_utils.simplify_events(events) == [
+      (
+          'root_agent',
+          Part.from_function_call(name='increase_by_one', args={'x': 1}),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='increase_by_one', response={'status': 'pending'}
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+  assert events[0].long_running_tool_ids
+
+  # Updates with another pending progress.
+  still_waiting_response = Part.from_function_response(
+      name='increase_by_one', response={'status': 'still waiting'}
+  )
+  events = runner.run(testing_utils.UserContent(still_waiting_response))
+  # We have one new request.
+  assert len(mockModel.requests) == 3
+  assert testing_utils.simplify_contents(mockModel.requests[2].contents) == [
+      ('user', 'test1'),
+      ('model', increase_by_one_call),
+      ('user', still_waiting_response),
+  ]
+
+  assert testing_utils.simplify_events(events) == [('root_agent', 'response2')]
+
+  # Calls when the result is ready.
+  result_response = Part.from_function_response(
+      name='increase_by_one', response={'result': 2}
+  )
+  events = runner.run(testing_utils.UserContent(result_response))
+  # We have one new request.
+  assert len(mockModel.requests) == 4
+  assert testing_utils.simplify_contents(mockModel.requests[3].contents) == [
+      ('user', 'test1'),
+      ('model', increase_by_one_call),
+      ('user', result_response),
+  ]
+  assert testing_utils.simplify_events(events) == [('root_agent', 'response3')]
+
+  # Calls when the result is ready. Here we still accept the result and do
+  # another summarization. Whether this is the right behavior is TBD.
+  another_result_response = Part.from_function_response(
+      name='increase_by_one', response={'result': 3}
+  )
+  events = runner.run(testing_utils.UserContent(another_result_response))
+  # We have one new request.
+  assert len(mockModel.requests) == 5
+  assert testing_utils.simplify_contents(mockModel.requests[4].contents) == [
+      ('user', 'test1'),
+      ('model', increase_by_one_call),
+      ('user', another_result_response),
+  ]
+  assert testing_utils.simplify_events(events) == [('root_agent', 'response4')]
+
+  # At the end, function_called should still be 1.
+  assert function_called == 1
+
+
+def test_async_function_with_none_response():
+  responses = [
+      Part.from_function_call(name='increase_by_one', args={'x': 1}),
+      'response1',
+      'response2',
+      'response3',
+      'response4',
+  ]
+  mockModel = testing_utils.MockModel.create(responses=responses)
+  function_called = 0
+
+  def increase_by_one(x: int, tool_context: ToolContext) -> int:
+    nonlocal function_called
+    function_called += 1
+    return 'pending'
+
+  # Calls the first time.
+  agent = Agent(
+      name='root_agent',
+      model=mockModel,
+      tools=[LongRunningFunctionTool(func=increase_by_one)],
+  )
+  runner = testing_utils.InMemoryRunner(agent)
+  events = runner.run('test1')
+
+  # Asserts the requests.
+  assert len(mockModel.requests) == 2
+  # 1 item: user content
+  assert mockModel.requests[0].contents == [
+      testing_utils.UserContent('test1'),
+  ]
+  increase_by_one_call = Part.from_function_call(
+      name='increase_by_one', args={'x': 1}
+  )
+
+  assert testing_utils.simplify_contents(mockModel.requests[1].contents) == [
+      ('user', 'test1'),
+      ('model', increase_by_one_call),
+      (
+          'user',
+          Part.from_function_response(
+              name='increase_by_one', response={'result': 'pending'}
+          ),
+      ),
+  ]
+
+  # Asserts the function calls.
+  assert function_called == 1
+
+  # Asserts the responses.
+  assert testing_utils.simplify_events(events) == [
+      (
+          'root_agent',
+          Part.from_function_call(name='increase_by_one', args={'x': 1}),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='increase_by_one', response={'result': 'pending'}
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+  # Updates with another pending progress.
+  still_waiting_response = Part.from_function_response(
+      name='increase_by_one', response={'status': 'still waiting'}
+  )
+  events = runner.run(testing_utils.UserContent(still_waiting_response))
+  # We have one new request.
+  assert len(mockModel.requests) == 3
+  assert testing_utils.simplify_contents(mockModel.requests[2].contents) == [
+      ('user', 'test1'),
+      ('model', increase_by_one_call),
+      ('user', still_waiting_response),
+  ]
+
+  assert testing_utils.simplify_events(events) == [('root_agent', 'response2')]
+
+  # Calls when the result is ready.
+  result_response = Part.from_function_response(
+      name='increase_by_one', response={'result': 2}
+  )
+  events = runner.run(testing_utils.UserContent(result_response))
+  # We have one new request.
+  assert len(mockModel.requests) == 4
+  assert testing_utils.simplify_contents(mockModel.requests[3].contents) == [
+      ('user', 'test1'),
+      ('model', increase_by_one_call),
+      ('user', result_response),
+  ]
+  assert testing_utils.simplify_events(events) == [('root_agent', 'response3')]
+
+  # Calls when the result is ready. Here we still accept the result and do
+  # another summarization. Whether this is the right behavior is TBD.
+  another_result_response = Part.from_function_response(
+      name='increase_by_one', response={'result': 3}
+  )
+  events = runner.run(testing_utils.UserContent(another_result_response))
+  # We have one new request.
+  assert len(mockModel.requests) == 5
+  assert testing_utils.simplify_contents(mockModel.requests[4].contents) == [
+      ('user', 'test1'),
+      ('model', increase_by_one_call),
+      ('user', another_result_response),
+  ]
+  assert testing_utils.simplify_events(events) == [('root_agent', 'response4')]
+
+  # At the end, function_called should still be 1.
+  assert function_called == 1

tests/unittests/flows/llm_flows/test_functions_parallel.py

@@ -0,0 +1,107 @@
+# 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.
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.events.event_actions import EventActions
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+@pytest.mark.asyncio
+async def test_parallel_function_calls_with_state_change():
+  function_calls = [
+      types.Part.from_function_call(
+          name='update_session_state',
+          args={'key': 'test_key1', 'value': 'test_value1'},
+      ),
+      types.Part.from_function_call(
+          name='update_session_state',
+          args={'key': 'test_key2', 'value': 'test_value2'},
+      ),
+      types.Part.from_function_call(
+          name='transfer_to_agent', args={'agent_name': 'test_sub_agent'}
+      ),
+  ]
+  function_responses = [
+      types.Part.from_function_response(
+          name='update_session_state', response={'result': None}
+      ),
+      types.Part.from_function_response(
+          name='update_session_state', response={'result': None}
+      ),
+      types.Part.from_function_response(
+          name='transfer_to_agent', response={'result': None}
+      ),
+  ]
+
+  responses: list[types.Content] = [
+      function_calls,
+      'response1',
+  ]
+  function_called = 0
+  mock_model = testing_utils.MockModel.create(responses=responses)
+
+  async def update_session_state(
+      key: str, value: str, tool_context: ToolContext
+  ) -> None:
+    nonlocal function_called
+    function_called += 1
+    tool_context.state.update({key: value})
+    return
+
+  async def transfer_to_agent(
+      agent_name: str, tool_context: ToolContext
+  ) -> None:
+    nonlocal function_called
+    function_called += 1
+    tool_context.actions.transfer_to_agent = agent_name
+    return
+
+  test_sub_agent = Agent(
+      name='test_sub_agent',
+  )
+
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      tools=[update_session_state, transfer_to_agent],
+      sub_agents=[test_sub_agent],
+  )
+  runner = testing_utils.TestInMemoryRunner(agent)
+  events = await runner.run_async_with_new_session('test')
+
+  # Notice that the following assertion only checks the "contents" part of the events.
+  # The "actions" part will be checked later.
+  assert testing_utils.simplify_events(events) == [
+      ('root_agent', function_calls),
+      ('root_agent', function_responses),
+      ('test_sub_agent', 'response1'),
+  ]
+
+  # Asserts the function calls.
+  assert function_called == 3
+
+  # Asserts the actions in response event.
+  response_event = events[1]
+
+  assert response_event.actions == EventActions(
+      state_delta={
+          'test_key1': 'test_value1',
+          'test_key2': 'test_value2',
+      },
+      transfer_to_agent='test_sub_agent',
+  )

tests/unittests/flows/llm_flows/test_functions_parallel_error.py

@@ -0,0 +1,86 @@
+# 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.
+
+import asyncio
+from typing import Any
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.flows.llm_flows import functions
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+def function_call(function_call_id, name, args: dict[str, Any]) -> types.Part:
+  part = types.Part.from_function_call(name=name, args=args)
+  part.function_call.id = function_call_id
+  return part
+
+
+@pytest.mark.asyncio
+async def test_parallel_function_call_error_fail_fast():
+  id_1 = 'id_1'
+  id_2 = 'id_2'
+  responses = [
+      [
+          function_call(id_1, 'fail_tool', {}),
+          function_call(id_2, 'sleep_tool', {}),
+      ],
+      [
+          types.Part.from_text(text='final response'),
+      ],
+  ]
+
+  mock_model = testing_utils.MockModel.create(responses=responses)
+
+  fail_called = False
+  sleep_started = False
+  sleep_completed = False
+  sleep_cancelled = False
+
+  async def fail_tool(tool_context: ToolContext) -> str:
+    nonlocal fail_called
+    fail_called = True
+    raise ValueError('Tool failed intentionally')
+
+  async def sleep_tool(tool_context: ToolContext) -> str:
+    nonlocal sleep_started, sleep_completed, sleep_cancelled
+    sleep_started = True
+    try:
+      await asyncio.sleep(10)  # Sleep long enough to be cancelled
+      sleep_completed = True
+      return 'Tool succeeded'
+    except asyncio.CancelledError:
+      sleep_cancelled = True
+      raise
+
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      tools=[fail_tool, sleep_tool],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+
+  with pytest.raises(ValueError, match='Tool failed intentionally'):
+    await runner.run_async(
+        new_message=types.Content(parts=[types.Part(text='test')]),
+    )
+
+  assert fail_called
+  assert sleep_started
+  assert not sleep_completed
+  assert sleep_cancelled

tests/unittests/flows/llm_flows/test_functions_request_euc.py

@@ -0,0 +1,624 @@
+# 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.
+
+import json
+from typing import Any
+from typing import Optional
+
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode
+from fastapi.openapi.models import OAuthFlows
+from google.adk.agents.llm_agent import Agent
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.auth.auth_tool import AuthToolArguments
+from google.adk.flows.llm_flows import functions
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+
+from ... import testing_utils
+
+
+def function_call(function_call_id, name, args: dict[str, Any]) -> types.Part:
+  part = types.Part.from_function_call(name=name, args=args)
+  part.function_call.id = function_call_id
+  return part
+
+
+def test_function_request_euc():
+  responses = [
+      [
+          types.Part.from_function_call(name='call_external_api1', args={}),
+          types.Part.from_function_call(name='call_external_api2', args={}),
+      ],
+      [
+          types.Part.from_text(text='response1'),
+      ],
+  ]
+
+  auth_config1 = AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl='https://accounts.google.com/o/oauth2/auth',
+                  tokenUrl='https://oauth2.googleapis.com/token',
+                  scopes={
+                      'https://www.googleapis.com/auth/calendar': (
+                          'See, edit, share, and permanently delete all the'
+                          ' calendars you can access using Google Calendar'
+                      )
+                  },
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_1',
+              client_secret='oauth_client_secret1',
+          ),
+      ),
+  )
+  auth_config2 = AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl='https://accounts.google.com/o/oauth2/auth',
+                  tokenUrl='https://oauth2.googleapis.com/token',
+                  scopes={
+                      'https://www.googleapis.com/auth/calendar': (
+                          'See, edit, share, and permanently delete all the'
+                          ' calendars you can access using Google Calendar'
+                      )
+                  },
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_2',
+              client_secret='oauth_client_secret2',
+          ),
+      ),
+  )
+
+  mock_model = testing_utils.MockModel.create(responses=responses)
+
+  def call_external_api1(tool_context: ToolContext) -> Optional[int]:
+    tool_context.request_credential(auth_config1)
+
+  def call_external_api2(tool_context: ToolContext) -> Optional[int]:
+    tool_context.request_credential(auth_config2)
+
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      tools=[call_external_api1, call_external_api2],
+  )
+  runner = testing_utils.InMemoryRunner(agent)
+  events = runner.run('test')
+  assert events[0].content.parts[0].function_call is not None
+  assert events[0].content.parts[1].function_call is not None
+  auth_configs = list(events[2].actions.requested_auth_configs.values())
+  exchanged_auth_config1 = auth_configs[0]
+  exchanged_auth_config2 = auth_configs[1]
+  assert exchanged_auth_config1.auth_scheme == auth_config1.auth_scheme
+  assert (
+      exchanged_auth_config1.raw_auth_credential
+      == auth_config1.raw_auth_credential
+  )
+  assert (
+      exchanged_auth_config1.exchanged_auth_credential.oauth2.auth_uri
+      is not None
+  )
+  assert exchanged_auth_config2.auth_scheme == auth_config2.auth_scheme
+  assert (
+      exchanged_auth_config2.raw_auth_credential
+      == auth_config2.raw_auth_credential
+  )
+  assert (
+      exchanged_auth_config2.exchanged_auth_credential.oauth2.auth_uri
+      is not None
+  )
+  function_call_ids = list(events[2].actions.requested_auth_configs.keys())
+
+  for idx, part in enumerate(events[1].content.parts):
+    request_euc_function_call = part.function_call
+    assert request_euc_function_call is not None
+    assert (
+        request_euc_function_call.name
+        == functions.REQUEST_EUC_FUNCTION_CALL_NAME
+    )
+    args = AuthToolArguments.model_validate(request_euc_function_call.args)
+
+    assert args.function_call_id == function_call_ids[idx]
+    args.auth_config.auth_scheme.model_extra.clear()
+    assert args.auth_config.auth_scheme == auth_configs[idx].auth_scheme
+    assert (
+        args.auth_config.raw_auth_credential
+        == auth_configs[idx].raw_auth_credential
+    )
+
+  assert len(mock_model.requests) == 1
+
+
+def test_function_request_euc_args_are_json_serializable():
+  responses = [
+      [
+          types.Part.from_function_call(name='call_external_api', args={}),
+      ],
+      [
+          types.Part.from_text(text='response1'),
+      ],
+  ]
+
+  auth_config = AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl='https://accounts.google.com/o/oauth2/auth',
+                  tokenUrl='https://oauth2.googleapis.com/token',
+                  scopes={
+                      'https://www.googleapis.com/auth/calendar': (
+                          'See, edit, share, and permanently delete all the'
+                          ' calendars you can access using Google Calendar'
+                      )
+                  },
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id',
+              client_secret='oauth_client_secret',
+          ),
+      ),
+  )
+
+  mock_model = testing_utils.MockModel.create(responses=responses)
+
+  def call_external_api(tool_context: ToolContext) -> Optional[int]:
+    tool_context.request_credential(auth_config)
+
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      tools=[call_external_api],
+  )
+  runner = testing_utils.InMemoryRunner(agent)
+  events = runner.run('test')
+
+  request_euc_function_call = events[1].content.parts[0].function_call
+  assert (
+      request_euc_function_call.name == functions.REQUEST_EUC_FUNCTION_CALL_NAME
+  )
+
+  # python-mode dump leaves auth_scheme.type a live enum, breaking json.dumps
+  json.dumps(request_euc_function_call.args)
+  assert (
+      request_euc_function_call.args['authConfig']['authScheme']['type']
+      == 'oauth2'
+  )
+
+
+def test_function_get_auth_response():
+  id_1 = 'id_1'
+  id_2 = 'id_2'
+  responses = [
+      [
+          function_call(id_1, 'call_external_api1', {}),
+          function_call(id_2, 'call_external_api2', {}),
+      ],
+      [
+          types.Part.from_text(text='response1'),
+      ],
+      [
+          types.Part.from_text(text='response2'),
+      ],
+  ]
+
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  function_invoked = 0
+
+  auth_config1 = AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl='https://accounts.google.com/o/oauth2/auth',
+                  tokenUrl='https://oauth2.googleapis.com/token',
+                  scopes={
+                      'https://www.googleapis.com/auth/calendar': (
+                          'See, edit, share, and permanently delete all the'
+                          ' calendars you can access using Google Calendar'
+                      )
+                  },
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_1',
+              client_secret='oauth_client_secret1',
+          ),
+      ),
+  )
+  auth_config2 = AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl='https://accounts.google.com/o/oauth2/auth',
+                  tokenUrl='https://oauth2.googleapis.com/token',
+                  scopes={
+                      'https://www.googleapis.com/auth/calendar': (
+                          'See, edit, share, and permanently delete all the'
+                          ' calendars you can access using Google Calendar'
+                      )
+                  },
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_2',
+              client_secret='oauth_client_secret2',
+          ),
+      ),
+  )
+
+  auth_response1 = AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl='https://accounts.google.com/o/oauth2/auth',
+                  tokenUrl='https://oauth2.googleapis.com/token',
+                  scopes={
+                      'https://www.googleapis.com/auth/calendar': (
+                          'See, edit, share, and permanently delete all the'
+                          ' calendars you can access using Google Calendar'
+                      )
+                  },
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_1',
+              client_secret='oauth_client_secret1',
+          ),
+      ),
+      exchanged_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_1',
+              client_secret='oauth_client_secret1',
+              access_token='token1',
+          ),
+      ),
+  )
+  auth_response2 = AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl='https://accounts.google.com/o/oauth2/auth',
+                  tokenUrl='https://oauth2.googleapis.com/token',
+                  scopes={
+                      'https://www.googleapis.com/auth/calendar': (
+                          'See, edit, share, and permanently delete all the'
+                          ' calendars you can access using Google Calendar'
+                      )
+                  },
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_2',
+              client_secret='oauth_client_secret2',
+          ),
+      ),
+      exchanged_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_2',
+              client_secret='oauth_client_secret2',
+              access_token='token2',
+          ),
+      ),
+  )
+
+  def call_external_api1(tool_context: ToolContext) -> int:
+    nonlocal function_invoked
+    function_invoked += 1
+    auth_response = tool_context.get_auth_response(auth_config1)
+    if not auth_response:
+      tool_context.request_credential(auth_config1)
+      return
+    assert auth_response == auth_response1.exchanged_auth_credential
+    return 1
+
+  def call_external_api2(tool_context: ToolContext) -> int:
+    nonlocal function_invoked
+    function_invoked += 1
+    auth_response = tool_context.get_auth_response(auth_config2)
+    if not auth_response:
+      tool_context.request_credential(auth_config2)
+      return
+    assert auth_response == auth_response2.exchanged_auth_credential
+    return 2
+
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      tools=[call_external_api1, call_external_api2],
+  )
+  runner = testing_utils.InMemoryRunner(agent)
+  runner.run('test')
+  request_euc_function_call_event = runner.session.events[-2]
+  function_response1 = types.FunctionResponse(
+      name=request_euc_function_call_event.content.parts[0].function_call.name,
+      response=auth_response1.model_dump(),
+  )
+  function_response1.id = request_euc_function_call_event.content.parts[
+      0
+  ].function_call.id
+
+  function_response2 = types.FunctionResponse(
+      name=request_euc_function_call_event.content.parts[1].function_call.name,
+      response=auth_response2.model_dump(),
+  )
+  function_response2.id = request_euc_function_call_event.content.parts[
+      1
+  ].function_call.id
+  runner.run(
+      new_message=types.Content(
+          role='user',
+          parts=[
+              types.Part(function_response=function_response1),
+              types.Part(function_response=function_response2),
+          ],
+      ),
+  )
+
+  assert function_invoked == 4
+  request = mock_model.requests[-1]
+  content = request.contents[-1]
+  parts = content.parts
+  assert len(parts) == 2
+  assert parts[0].function_response.name == 'call_external_api1'
+  assert parts[0].function_response.response == {'result': 1}
+  assert parts[1].function_response.name == 'call_external_api2'
+  assert parts[1].function_response.response == {'result': 2}
+
+
+def test_function_get_auth_response_partial():
+  id_1 = 'id_1'
+  id_2 = 'id_2'
+  responses = [
+      [
+          function_call(id_1, 'call_external_api1', {}),
+          function_call(id_2, 'call_external_api2', {}),
+      ],
+      [
+          types.Part.from_text(text='response1'),
+      ],
+      [
+          types.Part.from_text(text='response2'),
+      ],
+      [
+          types.Part.from_text(text='final response'),
+      ],
+  ]
+
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  function_invoked = 0
+
+  auth_config1 = AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl='https://accounts.google.com/o/oauth2/auth',
+                  tokenUrl='https://oauth2.googleapis.com/token',
+                  scopes={
+                      'https://www.googleapis.com/auth/calendar': (
+                          'See, edit, share, and permanently delete all the'
+                          ' calendars you can access using Google Calendar'
+                      )
+                  },
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_1',
+              client_secret='oauth_client_secret1',
+          ),
+      ),
+  )
+  auth_config2 = AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl='https://accounts.google.com/o/oauth2/auth',
+                  tokenUrl='https://oauth2.googleapis.com/token',
+                  scopes={
+                      'https://www.googleapis.com/auth/calendar': (
+                          'See, edit, share, and permanently delete all the'
+                          ' calendars you can access using Google Calendar'
+                      )
+                  },
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_2',
+              client_secret='oauth_client_secret2',
+          ),
+      ),
+  )
+
+  auth_response1 = AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl='https://accounts.google.com/o/oauth2/auth',
+                  tokenUrl='https://oauth2.googleapis.com/token',
+                  scopes={
+                      'https://www.googleapis.com/auth/calendar': (
+                          'See, edit, share, and permanently delete all the'
+                          ' calendars you can access using Google Calendar'
+                      )
+                  },
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_1',
+              client_secret='oauth_client_secret1',
+          ),
+      ),
+      exchanged_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_1',
+              client_secret='oauth_client_secret1',
+              access_token='token1',
+          ),
+      ),
+  )
+  auth_response2 = AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl='https://accounts.google.com/o/oauth2/auth',
+                  tokenUrl='https://oauth2.googleapis.com/token',
+                  scopes={
+                      'https://www.googleapis.com/auth/calendar': (
+                          'See, edit, share, and permanently delete all the'
+                          ' calendars you can access using Google Calendar'
+                      )
+                  },
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_2',
+              client_secret='oauth_client_secret2',
+          ),
+      ),
+      exchanged_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id='oauth_client_id_2',
+              client_secret='oauth_client_secret2',
+              access_token='token2',
+          ),
+      ),
+  )
+
+  def call_external_api1(tool_context: ToolContext) -> int:
+    nonlocal function_invoked
+    function_invoked += 1
+    auth_response = tool_context.get_auth_response(auth_config1)
+    if not auth_response:
+      tool_context.request_credential(auth_config1)
+      return
+    assert auth_response == auth_response1.exchanged_auth_credential
+    return 1
+
+  def call_external_api2(tool_context: ToolContext) -> int:
+    nonlocal function_invoked
+    function_invoked += 1
+    auth_response = tool_context.get_auth_response(auth_config2)
+    if not auth_response:
+      tool_context.request_credential(auth_config2)
+      return
+    assert auth_response == auth_response2.exchanged_auth_credential
+    return 2
+
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      tools=[call_external_api1, call_external_api2],
+  )
+  runner = testing_utils.InMemoryRunner(agent)
+  runner.run('test')
+  request_euc_function_call_event = runner.session.events[-2]
+  function_response1 = types.FunctionResponse(
+      name=request_euc_function_call_event.content.parts[0].function_call.name,
+      response=auth_response1.model_dump(),
+  )
+  function_response1.id = request_euc_function_call_event.content.parts[
+      0
+  ].function_call.id
+
+  function_response2 = types.FunctionResponse(
+      name=request_euc_function_call_event.content.parts[1].function_call.name,
+      response=auth_response2.model_dump(),
+  )
+  function_response2.id = request_euc_function_call_event.content.parts[
+      1
+  ].function_call.id
+  runner.run(
+      new_message=types.Content(
+          role='user',
+          parts=[
+              types.Part(function_response=function_response1),
+          ],
+      ),
+  )
+
+  assert function_invoked == 3
+  assert len(mock_model.requests) == 2
+  request = mock_model.requests[-1]
+  content = request.contents[-1]
+  parts = content.parts
+  assert len(parts) == 2
+  assert parts[0].function_response.name == 'call_external_api1'
+  assert parts[0].function_response.response == {'result': 1}
+  assert parts[1].function_response.name == 'call_external_api2'
+  assert parts[1].function_response.response == {'result': None}
+
+  runner.run(
+      new_message=types.Content(
+          role='user',
+          parts=[
+              types.Part(function_response=function_response2),
+          ],
+      ),
+  )
+  assert function_invoked == 4
+  assert len(mock_model.requests) == 3
+  request = mock_model.requests[-1]
+  content = request.contents[-1]
+  parts = content.parts
+  assert len(parts) == 2
+  assert parts[0].function_response.name == 'call_external_api1'
+  assert parts[0].function_response.response == {'result': 1}
+  assert parts[1].function_response.name == 'call_external_api2'
+  assert parts[1].function_response.response == {'result': 2}

tests/unittests/flows/llm_flows/test_functions_sequential.py

@@ -0,0 +1,93 @@
+# 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.
+
+from typing import Any
+
+from google.adk.agents.llm_agent import Agent
+from google.genai import types
+
+from ... import testing_utils
+
+
+def function_call(args: dict[str, Any]) -> types.Part:
+  return types.Part.from_function_call(name='increase_by_one', args=args)
+
+
+def function_response(response: dict[str, Any]) -> types.Part:
+  return types.Part.from_function_response(
+      name='increase_by_one', response=response
+  )
+
+
+def test_sequential_calls():
+  responses = [
+      function_call({'x': 1}),
+      function_call({'x': 2}),
+      function_call({'x': 3}),
+      'response1',
+  ]
+  mockModel = testing_utils.MockModel.create(responses=responses)
+  function_called = 0
+
+  def increase_by_one(x: int) -> int:
+    nonlocal function_called
+    function_called += 1
+    return x + 1
+
+  agent = Agent(name='root_agent', model=mockModel, tools=[increase_by_one])
+  runner = testing_utils.InMemoryRunner(agent)
+  result = testing_utils.simplify_events(runner.run('test'))
+  assert result == [
+      ('root_agent', function_call({'x': 1})),
+      ('root_agent', function_response({'result': 2})),
+      ('root_agent', function_call({'x': 2})),
+      ('root_agent', function_response({'result': 3})),
+      ('root_agent', function_call({'x': 3})),
+      ('root_agent', function_response({'result': 4})),
+      ('root_agent', 'response1'),
+  ]
+
+  # Asserts the requests.
+  assert len(mockModel.requests) == 4
+  # 1 item: user content
+  assert testing_utils.simplify_contents(mockModel.requests[0].contents) == [
+      ('user', 'test')
+  ]
+  # 3 items: user content, function call / response for the 1st call
+  assert testing_utils.simplify_contents(mockModel.requests[1].contents) == [
+      ('user', 'test'),
+      ('model', function_call({'x': 1})),
+      ('user', function_response({'result': 2})),
+  ]
+  # 5 items: user content, function call / response for two calls
+  assert testing_utils.simplify_contents(mockModel.requests[2].contents) == [
+      ('user', 'test'),
+      ('model', function_call({'x': 1})),
+      ('user', function_response({'result': 2})),
+      ('model', function_call({'x': 2})),
+      ('user', function_response({'result': 3})),
+  ]
+  # 7 items: user content, function call / response for three calls
+  assert testing_utils.simplify_contents(mockModel.requests[3].contents) == [
+      ('user', 'test'),
+      ('model', function_call({'x': 1})),
+      ('user', function_response({'result': 2})),
+      ('model', function_call({'x': 2})),
+      ('user', function_response({'result': 3})),
+      ('model', function_call({'x': 3})),
+      ('user', function_response({'result': 4})),
+  ]
+
+  # Asserts the function calls.
+  assert function_called == 3

tests/unittests/flows/llm_flows/test_functions_thread_pool.py

@@ -0,0 +1,621 @@
+# 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 thread pool execution of tools in Live API mode."""
+
+import asyncio
+import contextvars
+import threading
+import time
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.agents.run_config import RunConfig
+from google.adk.agents.run_config import ToolThreadPoolConfig
+from google.adk.flows.llm_flows.functions import _call_tool_in_thread_pool
+from google.adk.flows.llm_flows.functions import _get_tool_thread_pool
+from google.adk.flows.llm_flows.functions import _is_sync_tool
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.set_model_response_tool import SetModelResponseTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+from pydantic import BaseModel
+import pytest
+
+from ... import testing_utils
+
+
+@pytest.fixture(autouse=True)
+def cleanup_thread_pools():
+  yield
+  from google.adk.flows.llm_flows import functions
+
+  # Shutdown all pools
+  for pool in functions._TOOL_THREAD_POOLS.values():
+    pool.shutdown(wait=False)
+  functions._TOOL_THREAD_POOLS.clear()
+
+
+class TestIsSyncTool:
+  """Tests for the _is_sync_tool helper function."""
+
+  def test_sync_function_is_sync(self):
+    """Test that a synchronous function is detected as sync."""
+
+    def sync_func(x: int) -> int:
+      return x + 1
+
+    tool = FunctionTool(sync_func)
+    assert _is_sync_tool(tool) is True
+
+  def test_async_function_is_not_sync(self):
+    """Test that an async function is detected as not sync."""
+
+    async def async_func(x: int) -> int:
+      return x + 1
+
+    tool = FunctionTool(async_func)
+    assert _is_sync_tool(tool) is False
+
+  def test_async_generator_is_not_sync(self):
+    """Test that an async generator function is detected as not sync."""
+
+    async def async_gen_func(x: int):
+      yield x + 1
+
+    tool = FunctionTool(async_gen_func)
+    assert _is_sync_tool(tool) is False
+
+  def test_tool_without_func_returns_false(self):
+    """Test that a tool without func attribute returns False."""
+    tool = BaseTool(name='test', description='test tool')
+    assert _is_sync_tool(tool) is False
+
+
+class TestGetToolThreadPool:
+  """Tests for the _get_tool_thread_pool function."""
+
+  def test_returns_thread_pool_executor(self):
+    """Test that the function returns a ThreadPoolExecutor."""
+    from concurrent.futures import ThreadPoolExecutor
+
+    pool = _get_tool_thread_pool()
+    assert isinstance(pool, ThreadPoolExecutor)
+
+  def test_returns_same_pool_on_multiple_calls(self):
+    """Test that the same pool is returned on multiple calls (singleton)."""
+    pool1 = _get_tool_thread_pool()
+    pool2 = _get_tool_thread_pool()
+    assert pool1 is pool2
+
+  def test_different_max_workers_creates_different_pools(self):
+    """Test that different max_workers values create separate pools."""
+    pool_4 = _get_tool_thread_pool(max_workers=4)
+    pool_8 = _get_tool_thread_pool(max_workers=8)
+    assert pool_4 is not pool_8
+
+  def test_same_max_workers_returns_same_pool(self):
+    """Test that same max_workers returns the cached pool."""
+    pool1 = _get_tool_thread_pool(max_workers=16)
+    pool2 = _get_tool_thread_pool(max_workers=16)
+    assert pool1 is pool2
+
+
+class TestCallToolInThreadPool:
+  """Tests for the _call_tool_in_thread_pool function."""
+
+  @pytest.mark.asyncio
+  async def test_sync_tool_runs_in_thread_pool(self):
+    """Test that sync tools run in a separate thread."""
+    main_thread_id = threading.current_thread().ident
+    tool_thread_id = None
+
+    def sync_func() -> dict:
+      nonlocal tool_thread_id
+      tool_thread_id = threading.current_thread().ident
+      return {'result': 'success'}
+
+    tool = FunctionTool(sync_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {}, tool_context)
+
+    assert result == {'result': 'success'}
+    assert tool_thread_id is not None
+    assert tool_thread_id != main_thread_id
+
+  @pytest.mark.asyncio
+  async def test_async_tool_runs_in_thread_pool(self):
+    """Test that async tools run in a separate thread with new event loop."""
+    main_thread_id = threading.current_thread().ident
+    tool_thread_id = None
+
+    async def async_func() -> dict:
+      nonlocal tool_thread_id
+      tool_thread_id = threading.current_thread().ident
+      return {'result': 'async_success'}
+
+    tool = FunctionTool(async_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {}, tool_context)
+
+    assert result == {'result': 'async_success'}
+    assert tool_thread_id is not None
+    assert tool_thread_id != main_thread_id
+
+  @pytest.mark.asyncio
+  async def test_sync_tool_with_args(self):
+    """Test that sync tools receive arguments correctly."""
+
+    def sync_func(x: int, y: str) -> dict:
+      return {'sum': x, 'text': y}
+
+    tool = FunctionTool(sync_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(
+        tool, {'x': 42, 'y': 'hello'}, tool_context
+    )
+
+    assert result == {'sum': 42, 'text': 'hello'}
+
+  @pytest.mark.asyncio
+  async def test_sync_tool_missing_mandatory_args(self):
+    """Test sync tools return error dict when mandatory args are missing."""
+
+    def sync_func(x: int, y: str) -> dict:
+      return {'sum': x, 'text': y}
+
+    tool = FunctionTool(sync_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {'x': 42}, tool_context)
+
+    assert 'error' in result
+    assert 'mandatory input parameters are not present' in result['error']
+
+  @pytest.mark.asyncio
+  async def test_sync_tool_calling_asyncio_run(self):
+    """Test that sync tools can call asyncio.run internally."""
+
+    def sync_func_with_loop(x: int) -> dict:
+      async def inner_async():
+        return {'result': x * 2}
+
+      return asyncio.run(inner_async())
+
+    tool = FunctionTool(sync_func_with_loop)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {'x': 21}, tool_context)
+    assert result == {'result': 42}
+
+  @pytest.mark.asyncio
+  async def test_async_tool_with_args(self):
+    """Test that async tools receive arguments correctly."""
+
+    async def async_func(x: int, y: str) -> dict:
+      return {'sum': x, 'text': y}
+
+    tool = FunctionTool(async_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(
+        tool, {'x': 42, 'y': 'hello'}, tool_context
+    )
+
+    assert result == {'sum': 42, 'text': 'hello'}
+
+  @pytest.mark.asyncio
+  async def test_sync_tool_with_tool_context(self):
+    """Test that sync tools receive tool_context when requested."""
+
+    def sync_func_with_context(x: int, tool_context: ToolContext) -> dict:
+      return {'x': x, 'has_context': tool_context is not None}
+
+    tool = FunctionTool(sync_func_with_context)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {'x': 10}, tool_context)
+
+    assert result == {'x': 10, 'has_context': True}
+
+  @pytest.mark.asyncio
+  async def test_blocking_io_does_not_block_event_loop(self):
+    """Test that blocking I/O in thread pool doesn't block main event loop."""
+    event_loop_ticks = 0
+
+    async def ticker():
+      nonlocal event_loop_ticks
+      for _ in range(10):
+        await asyncio.sleep(0.01)
+        event_loop_ticks += 1
+
+    def blocking_sleep() -> dict:
+      time.sleep(0.15)  # Blocking sleep for 150ms
+      return {'result': 'done'}
+
+    tool = FunctionTool(blocking_sleep)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    # Run both ticker and blocking tool concurrently
+    ticker_task = asyncio.create_task(ticker())
+    result = await _call_tool_in_thread_pool(tool, {}, tool_context)
+    await ticker_task
+
+    assert result == {'result': 'done'}
+    # Ticker should have run multiple times while tool was sleeping
+    assert (
+        event_loop_ticks >= 5
+    ), f'Event loop should have ticked at least 5 times, got {event_loop_ticks}'
+
+  @pytest.mark.asyncio
+  @pytest.mark.parametrize(
+      'return_value,use_implicit_return',
+      [
+          (None, True),  # implicit None (no return statement)
+          (None, False),  # explicit `return None`
+          (0, False),  # falsy int
+          ('', False),  # falsy str
+          ({}, False),  # falsy dict
+          (False, False),  # falsy bool
+      ],
+  )
+  async def test_sync_tool_falsy_return_executes_exactly_once(
+      self, return_value, use_implicit_return
+  ):
+    """FunctionTools returning None or other falsy values must execute exactly once.
+
+    Previously, a None return was mistaken for the internal sentinel used to
+    signal 'non-FunctionTool, fall back to run_async', causing a second
+    invocation. The fix uses an identity-based sentinel so that None and other
+    falsy values (0, '', {}, False) are treated as valid results.
+    """
+    call_count = 0
+
+    def sync_func():
+      nonlocal call_count
+      call_count += 1
+      if not use_implicit_return:
+        return return_value
+      # implicit None — no return statement
+
+    tool = FunctionTool(sync_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {}, tool_context)
+
+    assert result == return_value
+    assert (
+        call_count == 1
+    ), f'Tool function executed {call_count} time(s); expected exactly 1.'
+
+  @pytest.mark.asyncio
+  async def test_sync_tool_exception_propagates(self):
+    """Test that exceptions from sync tools propagate correctly."""
+
+    def sync_func_raises() -> dict:
+      raise ValueError('Test error from sync tool')
+
+    tool = FunctionTool(sync_func_raises)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    with pytest.raises(ValueError, match='Test error from sync tool'):
+      await _call_tool_in_thread_pool(tool, {}, tool_context)
+
+  @pytest.mark.asyncio
+  async def test_async_tool_exception_propagates(self):
+    """Test that exceptions from async tools propagate correctly."""
+
+    async def async_func_raises() -> dict:
+      raise RuntimeError('Test error from async tool')
+
+    tool = FunctionTool(async_func_raises)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    with pytest.raises(RuntimeError, match='Test error from async tool'):
+      await _call_tool_in_thread_pool(tool, {}, tool_context)
+
+  @pytest.mark.asyncio
+  async def test_custom_max_workers_used(self):
+    """Test that custom max_workers parameter is passed to thread pool."""
+    pool_used = None
+
+    def sync_func() -> dict:
+      nonlocal pool_used
+      # The pool itself is global, so we just verify the call works
+      return {'result': 'success'}
+
+    tool = FunctionTool(sync_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    # Call with custom max_workers
+    result = await _call_tool_in_thread_pool(
+        tool, {}, tool_context, max_workers=12
+    )
+
+    assert result == {'result': 'success'}
+    # Verify the pool was created with custom max_workers
+    pool = _get_tool_thread_pool(max_workers=12)
+    assert pool is not None
+
+  @pytest.mark.asyncio
+  async def test_contextvars_propagation_sync_tool(self):
+    """Test that contextvars propagate to sync tools in thread pool."""
+    test_var = contextvars.ContextVar('test_var', default='default')
+    test_var.set('main_thread_value')
+
+    def sync_func() -> dict[str, str]:
+      return {'value': test_var.get()}
+
+    tool = FunctionTool(sync_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {}, tool_context)
+
+    assert result == {'value': 'main_thread_value'}
+
+  @pytest.mark.asyncio
+  async def test_contextvars_propagation_async_tool(self):
+    """Test that contextvars propagate to async tools in thread pool."""
+    test_var = contextvars.ContextVar('test_var', default='default')
+    test_var.set('main_thread_value')
+
+    async def async_func() -> dict[str, str]:
+      return {'value': test_var.get()}
+
+    tool = FunctionTool(async_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {}, tool_context)
+
+    assert result == {'value': 'main_thread_value'}
+
+  @pytest.mark.asyncio
+  async def test_sync_tool_returning_none_runs_exactly_once(self):
+    """Regression test for issue #5284.
+
+    A sync FunctionTool whose underlying function returns None must not
+    be re-invoked through the run_async fallback path.
+    """
+    call_count = 0
+
+    def side_effect_only_func() -> None:
+      nonlocal call_count
+      call_count += 1
+
+    tool = FunctionTool(side_effect_only_func)
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {}, tool_context)
+
+    assert result is None
+    assert call_count == 1
+
+  @pytest.mark.asyncio
+  async def test_non_function_tool_sync_falls_back_to_run_async(self):
+    """Sync tools that aren't FunctionTool subclasses go through run_async.
+
+    Covers the fall-through path used by tools like SetModelResponseTool
+    that have a sync ``func`` attribute but aren't FunctionTool instances.
+    """
+    run_async_call_count = 0
+
+    class _SyncNonFunctionTool(BaseTool):
+
+      def __init__(self):
+        super().__init__(name='custom_tool', description='desc')
+        # Sync attribute so _is_sync_tool returns True.
+        self.func = lambda: 'unused'
+
+      async def run_async(self, *, args, tool_context):
+        nonlocal run_async_call_count
+        run_async_call_count += 1
+        return {'via': 'run_async'}
+
+    tool = _SyncNonFunctionTool()
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(tool, {}, tool_context)
+
+    assert result == {'via': 'run_async'}
+    assert run_async_call_count == 1
+
+  @pytest.mark.asyncio
+  async def test_set_model_response_tool_falls_back_to_run_async(self):
+    """SetModelResponseTool — the real-world non-FunctionTool sync tool."""
+
+    class _Schema(BaseModel):
+      answer: str
+
+    tool = SetModelResponseTool(output_schema=_Schema)
+    # Precondition: this is the code path the bug report referenced.
+    assert _is_sync_tool(tool)
+
+    model = testing_utils.MockModel.create(responses=[])
+    agent = Agent(name='test_agent', model=model, tools=[tool])
+    invocation_context = await testing_utils.create_invocation_context(
+        agent=agent, user_content=''
+    )
+    tool_context = ToolContext(
+        invocation_context=invocation_context,
+        function_call_id='test_id',
+    )
+
+    result = await _call_tool_in_thread_pool(
+        tool, {'answer': 'hello'}, tool_context
+    )
+
+    assert result == {'answer': 'hello'}
+
+
+class TestToolThreadPoolConfig:
+  """Tests for the tool_thread_pool_config in RunConfig."""
+
+  def test_default_is_none(self):
+    """Test that tool_thread_pool_config defaults to None."""
+    config = RunConfig()
+    assert config.tool_thread_pool_config is None
+
+  def test_can_be_set_with_defaults(self):
+    """Test that tool_thread_pool_config can be set with default values."""
+    config = RunConfig(tool_thread_pool_config=ToolThreadPoolConfig())
+    assert config.tool_thread_pool_config is not None
+    assert config.tool_thread_pool_config.max_workers == 4
+
+  def test_can_set_custom_max_workers(self):
+    """Test that max_workers can be customized."""
+    config = RunConfig(
+        tool_thread_pool_config=ToolThreadPoolConfig(max_workers=8)
+    )
+    assert config.tool_thread_pool_config.max_workers == 8
+
+  def test_max_workers_must_be_positive(self):
+    """Test that max_workers must be >= 1."""
+    with pytest.raises(ValueError):
+      ToolThreadPoolConfig(max_workers=0)
+
+  def test_max_workers_rejects_negative(self):
+    """Test that negative max_workers is rejected."""
+    with pytest.raises(ValueError):
+      ToolThreadPoolConfig(max_workers=-1)

tests/unittests/flows/llm_flows/test_identity.py

@@ -0,0 +1,94 @@
+# 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.
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.flows.llm_flows import identity
+from google.adk.models.llm_request import LlmRequest
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+@pytest.mark.asyncio
+async def test_no_description():
+  request = LlmRequest(
+      model="gemini-2.5-flash",
+      config=types.GenerateContentConfig(system_instruction=""),
+  )
+  agent = Agent(model="gemini-2.5-flash", name="agent")
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  async for _ in identity.request_processor.run_async(
+      invocation_context,
+      request,
+  ):
+    pass
+
+  assert request.config.system_instruction == (
+      """You are an agent. Your internal name is "agent"."""
+  )
+
+
+@pytest.mark.asyncio
+async def test_with_description():
+  request = LlmRequest(
+      model="gemini-2.5-flash",
+      config=types.GenerateContentConfig(system_instruction=""),
+  )
+  agent = Agent(
+      model="gemini-2.5-flash",
+      name="agent",
+      description="test description",
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  async for _ in identity.request_processor.run_async(
+      invocation_context,
+      request,
+  ):
+    pass
+
+  assert (
+      request.config.system_instruction
+      == """\
+You are an agent. Your internal name is "agent". The description about you is "test description"."""
+  )
+
+
+@pytest.mark.asyncio
+async def test_single_turn_agent():
+  request = LlmRequest(
+      model="gemini-1.5-flash",
+      config=types.GenerateContentConfig(system_instruction=""),
+  )
+  agent = Agent(
+      name="agent",
+      mode="single_turn",
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+
+  async for _ in identity.request_processor.run_async(
+      invocation_context,
+      request,
+  ):
+    pass
+
+  assert request.config.system_instruction == ""

tests/unittests/flows/llm_flows/test_interactions_processor.py

@@ -0,0 +1,280 @@
+# 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 the interactions processor."""
+
+from unittest.mock import MagicMock
+
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows import interactions_processor
+from google.genai import types
+import pytest
+
+
+class TestInteractionsRequestProcessor:
+  """Tests for InteractionsRequestProcessor."""
+
+  def test_find_previous_interaction_id_empty_events(self):
+    """Test that None is returned when there are no events."""
+    processor = interactions_processor.InteractionsRequestProcessor()
+    invocation_context = MagicMock()
+    invocation_context.session.events = []
+    invocation_context.branch = None
+    invocation_context.agent.name = "test_agent"
+
+    result = processor._find_previous_interaction_id(invocation_context)
+    assert result is None
+
+  def test_find_previous_interaction_id_user_only_events(self):
+    """Test that None is returned when only user events exist."""
+    processor = interactions_processor.InteractionsRequestProcessor()
+    events = [
+        Event(
+            invocation_id="inv1",
+            author="user",
+            content=types.UserContent("Hello"),
+        ),
+        Event(
+            invocation_id="inv2",
+            author="user",
+            content=types.UserContent("World"),
+        ),
+    ]
+    invocation_context = MagicMock()
+    invocation_context.session.events = events
+    invocation_context.branch = None
+    invocation_context.agent.name = "test_agent"
+
+    result = processor._find_previous_interaction_id(invocation_context)
+    assert result is None
+
+  def test_find_previous_interaction_id_no_interaction_id(self):
+    """Test that None is returned when model events have no interaction_id."""
+    processor = interactions_processor.InteractionsRequestProcessor()
+    events = [
+        Event(
+            invocation_id="inv1",
+            author="user",
+            content=types.UserContent("Hello"),
+        ),
+        Event(
+            invocation_id="inv2",
+            author="test_agent",
+            content=types.ModelContent("Response without interaction_id"),
+        ),
+    ]
+    invocation_context = MagicMock()
+    invocation_context.session.events = events
+    invocation_context.branch = None
+    invocation_context.agent.name = "test_agent"
+
+    result = processor._find_previous_interaction_id(invocation_context)
+    assert result is None
+
+  def test_find_previous_interaction_id_from_model_event(self):
+    """Test that interaction_id is returned from model event."""
+    processor = interactions_processor.InteractionsRequestProcessor()
+    events = [
+        Event(
+            invocation_id="inv1",
+            author="user",
+            content=types.UserContent("Hello"),
+        ),
+        Event(
+            invocation_id="inv2",
+            author="test_agent",
+            content=types.ModelContent("Response"),
+            interaction_id="interaction_123",
+        ),
+    ]
+    invocation_context = MagicMock()
+    invocation_context.session.events = events
+    invocation_context.branch = None
+    invocation_context.agent.name = "test_agent"
+
+    result = processor._find_previous_interaction_id(invocation_context)
+    assert result == "interaction_123"
+
+  def test_find_previous_interaction_id_returns_most_recent(self):
+    """Test that the most recent interaction_id is returned."""
+    processor = interactions_processor.InteractionsRequestProcessor()
+    events = [
+        Event(
+            invocation_id="inv1",
+            author="user",
+            content=types.UserContent("Hello"),
+        ),
+        Event(
+            invocation_id="inv2",
+            author="test_agent",
+            content=types.ModelContent("First response"),
+            interaction_id="interaction_first",
+        ),
+        Event(
+            invocation_id="inv3",
+            author="user",
+            content=types.UserContent("Second message"),
+        ),
+        Event(
+            invocation_id="inv4",
+            author="test_agent",
+            content=types.ModelContent("Second response"),
+            interaction_id="interaction_second",
+        ),
+    ]
+    invocation_context = MagicMock()
+    invocation_context.session.events = events
+    invocation_context.branch = None
+    invocation_context.agent.name = "test_agent"
+
+    result = processor._find_previous_interaction_id(invocation_context)
+    assert result == "interaction_second"
+
+  def test_find_previous_interaction_id_skips_user_events(self):
+    """Test that user events with interaction_id are skipped."""
+    processor = interactions_processor.InteractionsRequestProcessor()
+    events = [
+        Event(
+            invocation_id="inv1",
+            author="test_agent",
+            content=types.ModelContent("Model response"),
+            interaction_id="interaction_model",
+        ),
+        Event(
+            invocation_id="inv2",
+            author="user",
+            content=types.UserContent("User message"),
+            interaction_id="interaction_user",  # This should be skipped
+        ),
+    ]
+    invocation_context = MagicMock()
+    invocation_context.session.events = events
+    invocation_context.branch = None
+    invocation_context.agent.name = "test_agent"
+
+    result = processor._find_previous_interaction_id(invocation_context)
+    assert result == "interaction_model"
+
+  def test_is_event_in_branch_no_branch(self):
+    """Test branch filtering with no current branch."""
+    # Event without branch should be included when no current branch
+    event = Event(
+        invocation_id="inv1",
+        author="test",
+        content=types.ModelContent("test"),
+    )
+    assert interactions_processor._is_event_in_branch(None, event) is True
+
+    # Event with branch should be excluded when no current branch
+    event_with_branch = Event(
+        invocation_id="inv2",
+        author="test",
+        content=types.ModelContent("test"),
+        branch="some_branch",
+    )
+    assert (
+        interactions_processor._is_event_in_branch(None, event_with_branch)
+        is False
+    )
+
+  def test_is_event_in_branch_same_branch(self):
+    """Test that events in the same branch are included."""
+    event = Event(
+        invocation_id="inv1",
+        author="test",
+        content=types.ModelContent("test"),
+        branch="root.child",
+    )
+    assert (
+        interactions_processor._is_event_in_branch("root.child", event) is True
+    )
+
+  def test_is_event_in_branch_different_branch(self):
+    """Test that events in different branches are excluded."""
+    event = Event(
+        invocation_id="inv1",
+        author="test",
+        content=types.ModelContent("test"),
+        branch="root.other",
+    )
+    assert (
+        interactions_processor._is_event_in_branch("root.child", event) is False
+    )
+
+  def test_is_event_in_branch_root_events_included(self):
+    """Test that root events (no branch) are included in child branches."""
+    event = Event(
+        invocation_id="inv1",
+        author="test",
+        content=types.ModelContent("test"),
+    )
+    assert (
+        interactions_processor._is_event_in_branch("root.child", event) is True
+    )
+
+
+def _evt(author: str, interaction_id: str | None, branch: str | None) -> Event:
+  return Event(author=author, interaction_id=interaction_id, branch=branch)
+
+
+def test_find_previous_interaction_id_returns_latest_for_agent():
+  events = [
+      _evt("my_agent", "int_1", None),
+      _evt("user", None, None),
+      _evt("my_agent", "int_2", None),
+      _evt("other_agent", "int_3", None),
+  ]
+
+  result = interactions_processor._find_previous_interaction_state(
+      events, agent_name="my_agent", current_branch=None
+  )
+
+  assert result[0] == "int_2"
+
+
+def test_find_previous_interaction_id_respects_branch():
+  events = [
+      _evt("my_agent", "int_main", None),
+      _evt("my_agent", "int_other_branch", "branch_b"),
+  ]
+
+  result = interactions_processor._find_previous_interaction_state(
+      events, agent_name="my_agent", current_branch="branch_a"
+  )
+
+  assert result[0] == "int_main"
+
+
+def test_find_previous_interaction_id_none_when_absent():
+  events = [_evt("user", None, None)]
+
+  result = interactions_processor._find_previous_interaction_state(
+      events, agent_name="my_agent", current_branch=None
+  )
+
+  assert result[0] is None
+
+
+def test_find_previous_interaction_state_returns_both_ids():
+  events = [
+      Event(author="my_agent", interaction_id="int_1", environment_id="env_1"),
+      Event(author="user"),
+      Event(author="my_agent", interaction_id="int_2", environment_id="env_2"),
+  ]
+
+  state = interactions_processor._find_previous_interaction_state(
+      events, agent_name="my_agent", current_branch=None
+  )
+
+  assert state == ("int_2", "env_2")

tests/unittests/flows/llm_flows/test_live_tool_callbacks.py

@@ -0,0 +1,465 @@
+# 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.
+
+from enum import Enum
+from functools import partial
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from unittest import mock
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.functions import handle_function_calls_live
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class CallbackType(Enum):
+  SYNC = 1
+  ASYNC = 2
+
+
+class AsyncBeforeToolCallback:
+
+  def __init__(self, mock_response: Dict[str, Any]):
+    self.mock_response = mock_response
+
+  async def __call__(
+      self,
+      tool: FunctionTool,
+      args: Dict[str, Any],
+      tool_context: ToolContext,
+  ) -> Optional[Dict[str, Any]]:
+    return self.mock_response
+
+
+class AsyncAfterToolCallback:
+
+  def __init__(self, mock_response: Dict[str, Any]):
+    self.mock_response = mock_response
+
+  async def __call__(
+      self,
+      tool: FunctionTool,
+      args: Dict[str, Any],
+      tool_context: ToolContext,
+      tool_response: Dict[str, Any],
+  ) -> Optional[Dict[str, Any]]:
+    return self.mock_response
+
+
+async def invoke_tool_with_callbacks_live(
+    before_cb=None, after_cb=None
+) -> Optional[Event]:
+  """Test helper to invoke a tool with callbacks using handle_function_calls_live."""
+
+  def simple_fn(**kwargs) -> Dict[str, Any]:
+    return {"initial": "response"}
+
+  tool = FunctionTool(simple_fn)
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name="agent",
+      model=model,
+      tools=[tool],
+      before_tool_callback=before_cb,
+      after_tool_callback=after_cb,
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content=""
+  )
+  # Build function call event
+  function_call = types.FunctionCall(name=tool.name, args={})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {tool.name: tool}
+  return await handle_function_calls_live(
+      invocation_context,
+      event,
+      tools_dict,
+  )
+
+
+def mock_sync_before_cb_side_effect(
+    tool, args, tool_context, ret_value=None
+) -> Optional[Dict[str, Any]]:
+  return ret_value
+
+
+async def mock_async_before_cb_side_effect(
+    tool, args, tool_context, ret_value=None
+) -> Optional[Dict[str, Any]]:
+  return ret_value
+
+
+def mock_sync_after_cb_side_effect(
+    tool, args, tool_context, tool_response, ret_value=None
+) -> Optional[Dict[str, Any]]:
+  return ret_value
+
+
+async def mock_async_after_cb_side_effect(
+    tool, args, tool_context, tool_response, ret_value=None
+) -> Optional[Dict[str, Any]]:
+  return ret_value
+
+
+@pytest.mark.asyncio
+async def test_live_async_before_tool_callback():
+  """Test that async before tool callbacks work in live mode."""
+  mock_resp = {"test": "before_tool_callback"}
+  before_cb = AsyncBeforeToolCallback(mock_resp)
+  result_event = await invoke_tool_with_callbacks_live(before_cb=before_cb)
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_resp
+
+
+@pytest.mark.asyncio
+async def test_live_async_after_tool_callback():
+  """Test that async after tool callbacks work in live mode."""
+  mock_resp = {"test": "after_tool_callback"}
+  after_cb = AsyncAfterToolCallback(mock_resp)
+  result_event = await invoke_tool_with_callbacks_live(after_cb=after_cb)
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_resp
+
+
+@pytest.mark.asyncio
+async def test_live_sync_before_tool_callback():
+  """Test that sync before tool callbacks work in live mode."""
+
+  def sync_before_cb(tool, args, tool_context):
+    return {"test": "sync_before_callback"}
+
+  result_event = await invoke_tool_with_callbacks_live(before_cb=sync_before_cb)
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == {"test": "sync_before_callback"}
+
+
+@pytest.mark.asyncio
+async def test_live_sync_after_tool_callback():
+  """Test that sync after tool callbacks work in live mode."""
+
+  def sync_after_cb(tool, args, tool_context, tool_response):
+    return {"test": "sync_after_callback"}
+
+  result_event = await invoke_tool_with_callbacks_live(after_cb=sync_after_cb)
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == {"test": "sync_after_callback"}
+
+
+# Test parameters for callback chains
+CALLBACK_PARAMS = [
+    # Test single sync callback returning None (should allow tool execution)
+    ([(None, CallbackType.SYNC)], {"initial": "response"}, [1]),
+    # Test single async callback returning None (should allow tool execution)
+    ([(None, CallbackType.ASYNC)], {"initial": "response"}, [1]),
+    # Test single sync callback returning response (should skip tool execution)
+    ([({}, CallbackType.SYNC)], {}, [1]),
+    # Test single async callback returning response (should skip tool execution)
+    ([({}, CallbackType.ASYNC)], {}, [1]),
+    # Test callback chain where an empty dict from the first callback doesn't
+    # stop the chain, allowing the second callback to execute.
+    (
+        [({}, CallbackType.SYNC), ({"second": "callback"}, CallbackType.ASYNC)],
+        {"second": "callback"},
+        [1, 1],
+    ),
+    # Test callback chain where first returns None, second returns response
+    (
+        [(None, CallbackType.SYNC), ({}, CallbackType.ASYNC)],
+        {},
+        [1, 1],
+    ),
+    # Test mixed sync/async chain where all return None
+    (
+        [(None, CallbackType.SYNC), (None, CallbackType.ASYNC)],
+        {"initial": "response"},
+        [1, 1],
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "callbacks, expected_response, expected_calls",
+    CALLBACK_PARAMS,
+)
+@pytest.mark.asyncio
+async def test_live_before_tool_callbacks_chain(
+    callbacks: List[tuple[Optional[Dict[str, Any]], int]],
+    expected_response: Dict[str, Any],
+    expected_calls: List[int],
+):
+  """Test that before tool callback chains work correctly in live mode."""
+  mock_before_cbs = []
+  for response, callback_type in callbacks:
+    if callback_type == CallbackType.ASYNC:
+      mock_cb = mock.AsyncMock(
+          side_effect=partial(
+              mock_async_before_cb_side_effect, ret_value=response
+          )
+      )
+    else:
+      mock_cb = mock.Mock(
+          side_effect=partial(
+              mock_sync_before_cb_side_effect, ret_value=response
+          )
+      )
+    mock_before_cbs.append(mock_cb)
+
+  result_event = await invoke_tool_with_callbacks_live(
+      before_cb=mock_before_cbs
+  )
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == expected_response
+
+  # Assert that the callbacks were called the expected number of times
+  for i, mock_cb in enumerate(mock_before_cbs):
+    expected_calls_count = expected_calls[i]
+    if expected_calls_count == 1:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited_once()
+      else:
+        mock_cb.assert_called_once()
+    elif expected_calls_count == 0:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_not_awaited()
+      else:
+        mock_cb.assert_not_called()
+    else:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited(expected_calls_count)
+      else:
+        mock_cb.assert_called(expected_calls_count)
+
+
+@pytest.mark.parametrize(
+    "callbacks, expected_response, expected_calls",
+    CALLBACK_PARAMS,
+)
+@pytest.mark.asyncio
+async def test_live_after_tool_callbacks_chain(
+    callbacks: List[tuple[Optional[Dict[str, Any]], int]],
+    expected_response: Dict[str, Any],
+    expected_calls: List[int],
+):
+  """Test that after tool callback chains work correctly in live mode."""
+  mock_after_cbs = []
+  for response, callback_type in callbacks:
+    if callback_type == CallbackType.ASYNC:
+      mock_cb = mock.AsyncMock(
+          side_effect=partial(
+              mock_async_after_cb_side_effect, ret_value=response
+          )
+      )
+    else:
+      mock_cb = mock.Mock(
+          side_effect=partial(
+              mock_sync_after_cb_side_effect, ret_value=response
+          )
+      )
+    mock_after_cbs.append(mock_cb)
+
+  result_event = await invoke_tool_with_callbacks_live(after_cb=mock_after_cbs)
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == expected_response
+
+  # Assert that the callbacks were called the expected number of times
+  for i, mock_cb in enumerate(mock_after_cbs):
+    expected_calls_count = expected_calls[i]
+    if expected_calls_count == 1:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited_once()
+      else:
+        mock_cb.assert_called_once()
+    elif expected_calls_count == 0:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_not_awaited()
+      else:
+        mock_cb.assert_not_called()
+    else:
+      if isinstance(mock_cb, mock.AsyncMock):
+        mock_cb.assert_awaited(expected_calls_count)
+      else:
+        mock_cb.assert_called(expected_calls_count)
+
+
+@pytest.mark.asyncio
+async def test_live_mixed_callbacks():
+  """Test that both before and after callbacks work together in live mode."""
+
+  def before_cb(tool, args, tool_context):
+    # Modify args and let tool run
+    args["modified_by_before"] = True
+    return None
+
+  def after_cb(tool, args, tool_context, tool_response):
+    # Modify response
+    tool_response["modified_by_after"] = True
+    return tool_response
+
+  result_event = await invoke_tool_with_callbacks_live(
+      before_cb=before_cb, after_cb=after_cb
+  )
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  response = part.function_response.response
+  assert response["modified_by_after"] is True
+  assert "initial" in response  # Original response should still be there
+
+
+@pytest.mark.asyncio
+async def test_live_callback_compatibility_with_async():
+  """Test that live callbacks have the same behavior as async callbacks."""
+  # This test ensures that the behavior between handle_function_calls_async
+  # and handle_function_calls_live is consistent for callbacks
+
+  def before_cb(tool, args, tool_context):
+    return {"bypassed": "by_before_callback"}
+
+  # Test with async version
+  from google.adk.flows.llm_flows.functions import handle_function_calls_async
+
+  def simple_fn(**kwargs) -> Dict[str, Any]:
+    return {"initial": "response"}
+
+  tool = FunctionTool(simple_fn)
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name="agent",
+      model=model,
+      tools=[tool],
+      before_tool_callback=before_cb,
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content=""
+  )
+  function_call = types.FunctionCall(name=tool.name, args={})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {tool.name: tool}
+
+  # Get result from async version
+  async_result = await handle_function_calls_async(
+      invocation_context, event, tools_dict
+  )
+
+  # Get result from live version
+  live_result = await handle_function_calls_live(
+      invocation_context, event, tools_dict
+  )
+
+  # Both should have the same response
+  assert async_result is not None
+  assert live_result is not None
+  async_response = async_result.content.parts[0].function_response.response
+  live_response = live_result.content.parts[0].function_response.response
+  assert async_response == live_response == {"bypassed": "by_before_callback"}
+
+
+@pytest.mark.asyncio
+async def test_live_on_tool_error_callback_tool_not_found_noop():
+  """Test that on_tool_error_callback is a no-op when the tool is not found."""
+
+  def noop_on_tool_error_callback(tool, args, tool_context, error):
+    return None
+
+  def simple_fn(**kwargs) -> Dict[str, Any]:
+    return {"initial": "response"}
+
+  tool = FunctionTool(simple_fn)
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name="agent",
+      model=model,
+      tools=[tool],
+      on_tool_error_callback=noop_on_tool_error_callback,
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content=""
+  )
+  function_call = types.FunctionCall(name="nonexistent_function", args={})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {tool.name: tool}
+
+  with pytest.raises(ValueError):
+    await handle_function_calls_live(invocation_context, event, tools_dict)
+
+
+@pytest.mark.asyncio
+async def test_live_on_tool_error_callback_tool_not_found_modify_tool_response():
+  """Test that on_tool_error_callback modifies tool response when tool is not found."""
+
+  def mock_on_tool_error_callback(tool, args, tool_context, error):
+    return {"result": "on_tool_error_callback_response"}
+
+  def simple_fn(**kwargs) -> Dict[str, Any]:
+    return {"initial": "response"}
+
+  tool = FunctionTool(simple_fn)
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name="agent",
+      model=model,
+      tools=[tool],
+      on_tool_error_callback=mock_on_tool_error_callback,
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content=""
+  )
+  function_call = types.FunctionCall(name="nonexistent_function", args={})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {tool.name: tool}
+
+  result_event = await handle_function_calls_live(
+      invocation_context,
+      event,
+      tools_dict,
+  )
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == {
+      "result": "on_tool_error_callback_response"
+  }

tests/unittests/flows/llm_flows/test_llm_callback_span_consistency.py

@@ -0,0 +1,384 @@
+# 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 that before/after/error model callbacks all observe the same call_llm span."""
+
+from typing import AsyncGenerator
+from typing import Optional
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.llm_agent import Agent
+from google.adk.agents.run_config import RunConfig
+from google.adk.agents.run_config import StreamingMode
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
+from google.adk.utils.context_utils import Aclosing
+from google.genai import types
+from google.genai.errors import ClientError
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+import pytest
+
+from ... import testing_utils
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+_SPAN_ID_INVALID = 0
+
+
+class _SpanCapture:
+  """Stores the span ID and trace ID observed from within a callback."""
+
+  def __init__(self):
+    self.span_id: int = _SPAN_ID_INVALID
+    self.trace_id: int = 0
+
+  def capture(self):
+    span = trace.get_current_span()
+    ctx = span.get_span_context()
+    if ctx and ctx.span_id != _SPAN_ID_INVALID:
+      self.span_id = ctx.span_id
+      self.trace_id = ctx.trace_id
+
+
+class SpanCapturingPlugin(BasePlugin):
+  """Plugin that records the active span ID in each callback."""
+
+  def __init__(self):
+    self.name = 'span_capturing_plugin'
+    self.before_capture = _SpanCapture()
+    self.after_capture = _SpanCapture()
+    self.error_capture = _SpanCapture()
+
+    self._short_circuit_before = False
+    self._short_circuit_response: Optional[LlmResponse] = None
+
+  async def before_model_callback(
+      self,
+      *,
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+  ) -> Optional[LlmResponse]:
+    self.before_capture.capture()
+    if self._short_circuit_before:
+      return self._short_circuit_response
+    return None
+
+  async def after_model_callback(
+      self,
+      *,
+      callback_context: CallbackContext,
+      llm_response: LlmResponse,
+  ) -> Optional[LlmResponse]:
+    self.after_capture.capture()
+    return None
+
+  async def on_model_error_callback(
+      self,
+      *,
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+      error: Exception,
+  ) -> Optional[LlmResponse]:
+    self.error_capture.capture()
+    # Return a response so the error doesn't propagate.
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text='error_handled')]
+        )
+    )
+
+
+# Install a real TracerProvider so spans are recorded (not NoOp).
+# This must happen at module level *before* any tracer is obtained,
+# because the OTel SDK only allows setting the provider once.
+_provider = TracerProvider()
+trace.set_tracer_provider(_provider)
+
+
+_MOCK_ERROR = ClientError(
+    code=500,
+    response_json={
+        'error': {
+            'code': 500,
+            'message': 'Model error.',
+            'status': 'INTERNAL',
+        }
+    },
+)
+
+
+# ---------------------------------------------------------------------------
+# Tests: non-CFC success path
+# ---------------------------------------------------------------------------
+
+
+def test_before_and_after_callbacks_share_same_span():
+  """before_model_callback and after_model_callback see the same span ID."""
+  plugin = SpanCapturingPlugin()
+  mock_model = testing_utils.MockModel.create(responses=['hello'])
+  agent = Agent(name='root_agent', model=mock_model)
+  runner = testing_utils.InMemoryRunner(agent, plugins=[plugin])
+
+  runner.run('test')
+
+  assert (
+      plugin.before_capture.span_id != _SPAN_ID_INVALID
+  ), 'before_model_callback did not observe a valid span'
+  assert (
+      plugin.after_capture.span_id != _SPAN_ID_INVALID
+  ), 'after_model_callback did not observe a valid span'
+  assert plugin.before_capture.span_id == plugin.after_capture.span_id, (
+      'before_model_callback and after_model_callback saw different spans:'
+      f' before={plugin.before_capture.span_id:#x},'
+      f' after={plugin.after_capture.span_id:#x}'
+  )
+
+
+def test_callbacks_same_trace_id():
+  """before and after callbacks are in the same trace."""
+  plugin = SpanCapturingPlugin()
+  mock_model = testing_utils.MockModel.create(responses=['hello'])
+  agent = Agent(name='root_agent', model=mock_model)
+  runner = testing_utils.InMemoryRunner(agent, plugins=[plugin])
+
+  runner.run('test')
+
+  assert plugin.before_capture.trace_id != 0
+  assert (
+      plugin.before_capture.trace_id == plugin.after_capture.trace_id
+  ), 'before and after callbacks are in different traces'
+
+
+# ---------------------------------------------------------------------------
+# Tests: non-CFC error path
+# ---------------------------------------------------------------------------
+
+
+def test_before_and_error_callbacks_share_same_span():
+  """before_model_callback and on_model_error_callback see the same span."""
+  plugin = SpanCapturingPlugin()
+  mock_model = testing_utils.MockModel.create(error=_MOCK_ERROR, responses=[])
+  agent = Agent(name='root_agent', model=mock_model)
+  runner = testing_utils.InMemoryRunner(agent, plugins=[plugin])
+
+  runner.run('test')
+
+  assert (
+      plugin.before_capture.span_id != _SPAN_ID_INVALID
+  ), 'before_model_callback did not observe a valid span'
+  assert (
+      plugin.error_capture.span_id != _SPAN_ID_INVALID
+  ), 'on_model_error_callback did not observe a valid span'
+  assert plugin.before_capture.span_id == plugin.error_capture.span_id, (
+      'before_model_callback and on_model_error_callback saw different'
+      f' spans: before={plugin.before_capture.span_id:#x},'
+      f' error={plugin.error_capture.span_id:#x}'
+  )
+
+
+# ---------------------------------------------------------------------------
+# Tests: short-circuit path (before_model_callback returns a response)
+# ---------------------------------------------------------------------------
+
+
+def test_short_circuit_before_callback_sees_valid_span():
+  """When before_model_callback short-circuits, it sees call_llm span."""
+  plugin = SpanCapturingPlugin()
+  plugin._short_circuit_before = True
+  plugin._short_circuit_response = LlmResponse(
+      content=testing_utils.ModelContent(
+          [types.Part.from_text(text='short_circuited')]
+      )
+  )
+  mock_model = testing_utils.MockModel.create(responses=['unused'])
+  agent = Agent(name='root_agent', model=mock_model)
+  runner = testing_utils.InMemoryRunner(agent, plugins=[plugin])
+
+  runner.run('test')
+
+  assert (
+      plugin.before_capture.span_id != _SPAN_ID_INVALID
+  ), 'before_model_callback did not observe a valid span on short-circuit'
+  # after_model_callback should NOT have been called.
+  assert plugin.after_capture.span_id == _SPAN_ID_INVALID
+
+
+# ---------------------------------------------------------------------------
+# Tests: all three callbacks share same span on error path
+# ---------------------------------------------------------------------------
+
+
+def test_all_three_callbacks_share_span_on_error():
+  """A plugin that implements all three callbacks sees the same span ID.
+
+  When the LLM errors and on_model_error_callback returns a recovery
+  response, after_model_callback also runs on that response.  All three
+  callbacks must observe the same call_llm span.
+  """
+  plugin = SpanCapturingPlugin()
+  mock_model = testing_utils.MockModel.create(error=_MOCK_ERROR, responses=[])
+  agent = Agent(name='root_agent', model=mock_model)
+  runner = testing_utils.InMemoryRunner(agent, plugins=[plugin])
+
+  runner.run('test')
+
+  # All three callbacks should have been called with valid spans.
+  assert plugin.before_capture.span_id != _SPAN_ID_INVALID
+  assert plugin.error_capture.span_id != _SPAN_ID_INVALID
+  assert plugin.after_capture.span_id != _SPAN_ID_INVALID
+  # And they should all share the same call_llm span.
+  assert (
+      plugin.before_capture.span_id == plugin.error_capture.span_id
+  ), 'before and error callbacks saw different spans'
+  assert (
+      plugin.before_capture.span_id == plugin.after_capture.span_id
+  ), 'before and after callbacks saw different spans on error recovery'
+
+
+# ---------------------------------------------------------------------------
+# Tests: CFC (Controlled Function Calling) / live path
+# ---------------------------------------------------------------------------
+
+
+class _CfcTestFlow(BaseLlmFlow):
+  """BaseLlmFlow subclass that stubs run_live for CFC testing."""
+
+  def __init__(self, live_responses: list[LlmResponse]):
+    self._live_responses = live_responses
+
+  async def run_live(
+      self, invocation_context
+  ) -> AsyncGenerator[LlmResponse, None]:
+    for resp in self._live_responses:
+      yield resp
+
+
+@pytest.mark.asyncio
+async def test_cfc_before_and_after_callbacks_share_same_span():
+  """CFC path: before_model_callback and after_model_callback share span."""
+  plugin = SpanCapturingPlugin()
+  mock_model = testing_utils.MockModel.create(responses=['unused'])
+  agent = Agent(name='root_agent', model=mock_model)
+
+  live_response = LlmResponse(
+      content=testing_utils.ModelContent(
+          [types.Part.from_text(text='live_hello')]
+      ),
+      turn_complete=True,
+  )
+  flow = _CfcTestFlow(live_responses=[live_response])
+
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent,
+      user_content='test',
+      run_config=RunConfig(
+          support_cfc=True,
+          streaming_mode=StreamingMode.SSE,
+      ),
+      plugins=[plugin],
+  )
+  model_response_event = Event(
+      id=Event.new_id(),
+      invocation_id=invocation_context.invocation_id,
+      author='root_agent',
+  )
+
+  responses = []
+  async with Aclosing(
+      flow._call_llm_async(
+          invocation_context,
+          LlmRequest(model='mock'),
+          model_response_event,
+      )
+  ) as agen:
+    async for resp in agen:
+      responses.append(resp)
+
+  assert len(responses) >= 1
+  assert (
+      plugin.before_capture.span_id != _SPAN_ID_INVALID
+  ), 'CFC: before_model_callback did not observe a valid span'
+  assert (
+      plugin.after_capture.span_id != _SPAN_ID_INVALID
+  ), 'CFC: after_model_callback did not observe a valid span'
+  assert plugin.before_capture.span_id == plugin.after_capture.span_id, (
+      'CFC: before_model_callback and after_model_callback saw different'
+      f' spans: before={plugin.before_capture.span_id:#x},'
+      f' after={plugin.after_capture.span_id:#x}'
+  )
+
+
+@pytest.mark.asyncio
+async def test_cfc_error_callback_shares_span():
+  """CFC path: on_model_error_callback shares span with before callback."""
+  plugin = SpanCapturingPlugin()
+  mock_model = testing_utils.MockModel.create(responses=['unused'])
+  agent = Agent(name='root_agent', model=mock_model)
+
+  # Flow whose run_live raises an error.
+  class _ErrorCfcFlow(BaseLlmFlow):
+
+    async def run_live(self, invocation_context):
+      # Make this a proper async generator that raises.
+      if False:
+        yield  # pragma: no cover — makes this an async generator
+      raise _MOCK_ERROR
+
+  flow = _ErrorCfcFlow()
+
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent,
+      user_content='test',
+      run_config=RunConfig(
+          support_cfc=True,
+          streaming_mode=StreamingMode.SSE,
+      ),
+      plugins=[plugin],
+  )
+  model_response_event = Event(
+      id=Event.new_id(),
+      invocation_id=invocation_context.invocation_id,
+      author='root_agent',
+  )
+
+  responses = []
+  async with Aclosing(
+      flow._call_llm_async(
+          invocation_context,
+          LlmRequest(model='mock'),
+          model_response_event,
+      )
+  ) as agen:
+    async for resp in agen:
+      responses.append(resp)
+
+  assert (
+      plugin.before_capture.span_id != _SPAN_ID_INVALID
+  ), 'CFC error: before_model_callback did not observe a valid span'
+  assert (
+      plugin.error_capture.span_id != _SPAN_ID_INVALID
+  ), 'CFC error: on_model_error_callback did not observe a valid span'
+  assert (
+      plugin.before_capture.span_id == plugin.error_capture.span_id
+  ), 'CFC error: before and error callbacks saw different spans'
+
+
+if __name__ == '__main__':
+  pytest.main([__file__])

tests/unittests/flows/llm_flows/test_model_callbacks.py

@@ -0,0 +1,195 @@
+# 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.
+
+from typing import Any
+from typing import Optional
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.llm_agent import Agent
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+from pydantic import BaseModel
+import pytest
+
+from ... import testing_utils
+
+
+class MockBeforeModelCallback(BaseModel):
+  mock_response: str
+
+  def __call__(
+      self,
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+  ) -> LlmResponse:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.mock_response)]
+        )
+    )
+
+
+class MockAfterModelCallback(BaseModel):
+  mock_response: str
+
+  def __call__(
+      self,
+      callback_context: CallbackContext,
+      llm_response: LlmResponse,
+  ) -> LlmResponse:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.mock_response)]
+        )
+    )
+
+
+class MockOnModelCallback(BaseModel):
+  mock_response: str
+
+  def __call__(
+      self,
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+      error: Exception,
+  ) -> LlmResponse:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.mock_response)]
+        )
+    )
+
+
+def noop_callback(**kwargs) -> Optional[LlmResponse]:
+  pass
+
+
+def test_before_model_callback():
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=MockBeforeModelCallback(
+          mock_response='before_model_callback'
+      ),
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', 'before_model_callback'),
+  ]
+
+
+def test_before_model_callback_noop():
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=noop_callback,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', 'model_response'),
+  ]
+
+
+def test_before_model_callback_end():
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=MockBeforeModelCallback(
+          mock_response='before_model_callback',
+      ),
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', 'before_model_callback'),
+  ]
+
+
+def test_after_model_callback():
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_model_callback=MockAfterModelCallback(
+          mock_response='after_model_callback'
+      ),
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', 'after_model_callback'),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_after_model_callback_noop():
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_model_callback=noop_callback,
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  assert testing_utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [('root_agent', 'model_response')]
+
+
+@pytest.mark.asyncio
+async def test_on_model_callback_model_error_noop():
+  """Test that the on_model_error_callback is a no-op when the model returns an error."""
+  mock_model = testing_utils.MockModel.create(
+      responses=[], error=SystemError('error')
+  )
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_model_error_callback=noop_callback,
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  with pytest.raises(SystemError):
+    await runner.run_async_with_new_session('test')
+
+
+@pytest.mark.asyncio
+async def test_on_model_callback_model_error_modify_model_response():
+  """Test that the on_model_error_callback can modify the model response."""
+  mock_model = testing_utils.MockModel.create(
+      responses=[], error=SystemError('error')
+  )
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_model_error_callback=MockOnModelCallback(
+          mock_response='on_model_error_callback_response'
+      ),
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  assert testing_utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [('root_agent', 'on_model_error_callback_response')]

tests/unittests/flows/llm_flows/test_nl_planning.py

@@ -0,0 +1,220 @@
+# 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.
+
+"""Unit tests for NL planning logic."""
+
+from typing import List
+from typing import Optional
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.llm_agent import Agent
+from google.adk.flows.llm_flows._nl_planning import request_processor
+from google.adk.flows.llm_flows._nl_planning import response_processor
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.adk.planners.built_in_planner import BuiltInPlanner
+from google.adk.planners.plan_re_act_planner import PlanReActPlanner
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+@pytest.mark.asyncio
+async def test_built_in_planner_content_list_unchanged():
+  """Test that BuiltInPlanner doesn't modify LlmRequest content list."""
+  planner = BuiltInPlanner(thinking_config=types.ThinkingConfig())
+  agent = Agent(name='test_agent', planner=planner)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+  # Create user/model/user conversation with thought in model response
+  llm_request = LlmRequest(
+      contents=[
+          types.UserContent(parts=[types.Part(text='Hello')]),
+          types.ModelContent(
+              parts=[
+                  types.Part(text='thinking...', thought=True),
+                  types.Part(text='Here is my response'),
+              ]
+          ),
+          types.UserContent(parts=[types.Part(text='Follow up')]),
+      ]
+  )
+  original_contents = llm_request.contents.copy()
+
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  assert llm_request.contents == original_contents
+
+
+@pytest.mark.asyncio
+async def test_built_in_planner_apply_thinking_config_called():
+  """Test that BuiltInPlanner.apply_thinking_config is called."""
+  planner = BuiltInPlanner(thinking_config=types.ThinkingConfig())
+  planner.apply_thinking_config = MagicMock()
+  agent = Agent(name='test_agent', planner=planner)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+  llm_request = LlmRequest()
+
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  planner.apply_thinking_config.assert_called_once_with(llm_request)
+
+
+@pytest.mark.asyncio
+async def test_plan_react_planner_instruction_appended():
+  """Test that PlanReActPlanner appends planning instruction."""
+  planner = PlanReActPlanner()
+  planner.build_planning_instruction = MagicMock(
+      return_value='Test instruction'
+  )
+  agent = Agent(name='test_agent', planner=planner)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  llm_request = LlmRequest()
+  llm_request.config.system_instruction = 'Original instruction'
+
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  assert llm_request.config.system_instruction == ("""\
+Original instruction
+
+Test instruction""")
+
+
+@pytest.mark.asyncio
+async def test_remove_thought_from_request_with_thoughts():
+  """Test that PlanReActPlanner removes thought flags from content parts."""
+  planner = PlanReActPlanner()
+  agent = Agent(name='test_agent', planner=planner)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+  llm_request = LlmRequest(
+      contents=[
+          types.UserContent(parts=[types.Part(text='initial query')]),
+          types.ModelContent(
+              parts=[
+                  types.Part(text='Text with thought', thought=True),
+                  types.Part(text='Regular text'),
+              ]
+          ),
+          types.UserContent(parts=[types.Part(text='follow up')]),
+      ]
+  )
+
+  async for _ in request_processor.run_async(invocation_context, llm_request):
+    pass
+
+  assert all(
+      part.thought is None
+      for content in llm_request.contents
+      for part in content.parts or []
+  )
+
+
+class OverriddenBuiltInPlanner(BuiltInPlanner):
+  """Subclass that overrides process_planning_response."""
+
+  def __init__(self, *, thinking_config: types.ThinkingConfig):
+    super().__init__(thinking_config=thinking_config)
+    self.process_planning_response_called = False
+    self.received_parts = None
+
+  def process_planning_response(
+      self,
+      callback_context: CallbackContext,
+      response_parts: List[types.Part],
+  ) -> Optional[List[types.Part]]:
+    self.process_planning_response_called = True
+    self.received_parts = response_parts
+    return response_parts
+
+
+class NonOverriddenBuiltInPlanner(BuiltInPlanner):
+  """Subclass that does NOT override process_planning_response."""
+
+  pass
+
+
+@pytest.mark.asyncio
+async def test_overridden_subclass_process_planning_response_called():
+  """Test that subclasses overriding process_planning_response have it called.
+
+  Regression test for issue #4133.
+  """
+  planner = OverriddenBuiltInPlanner(thinking_config=types.ThinkingConfig())
+  agent = Agent(name='test_agent', planner=planner)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  response_parts = [
+      types.Part(text='thinking...', thought=True),
+      types.Part(text='Here is my response'),
+  ]
+  llm_response = LlmResponse(
+      content=types.Content(role='model', parts=response_parts)
+  )
+
+  async for _ in response_processor.run_async(invocation_context, llm_response):
+    pass
+
+  assert planner.process_planning_response_called
+  assert planner.received_parts == response_parts
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'planner_class',
+    [BuiltInPlanner, NonOverriddenBuiltInPlanner],
+    ids=['base_class', 'non_overridden_subclass'],
+)
+async def test_process_planning_response_not_called_without_override(
+    planner_class,
+):
+  """Test that process_planning_response is not called for base or non-overridden subclasses."""
+  planner = planner_class(thinking_config=types.ThinkingConfig())
+  agent = Agent(name='test_agent', planner=planner)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+
+  response_parts = [
+      types.Part(text='thinking...', thought=True),
+      types.Part(text='Here is my response'),
+  ]
+  llm_response = LlmResponse(
+      content=types.Content(role='model', parts=response_parts)
+  )
+
+  with patch.object(
+      BuiltInPlanner,
+      'process_planning_response',
+  ) as mock_method:
+    async for _ in response_processor.run_async(
+        invocation_context, llm_response
+    ):
+      pass
+    mock_method.assert_not_called()

tests/unittests/flows/llm_flows/test_other_configs.py

@@ -0,0 +1,47 @@
+# 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.
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.tools.tool_context import ToolContext
+from google.genai.types import Part
+from pydantic import BaseModel
+
+from ... import testing_utils
+
+
+def test_output_schema():
+  class CustomOutput(BaseModel):
+    custom_field: str
+
+  response = [
+      'response1',
+  ]
+  mockModel = testing_utils.MockModel.create(responses=response)
+  root_agent = Agent(
+      name='root_agent',
+      model=mockModel,
+      output_schema=CustomOutput,
+      disallow_transfer_to_parent=True,
+      disallow_transfer_to_peers=True,
+  )
+
+  runner = testing_utils.InMemoryRunner(root_agent)
+
+  assert testing_utils.simplify_events(runner.run('test1')) == [
+      ('root_agent', 'response1'),
+  ]
+  assert len(mockModel.requests) == 1
+  assert mockModel.requests[0].config.response_schema == CustomOutput
+  assert mockModel.requests[0].config.response_mime_type == 'application/json'
+  assert mockModel.requests[0].config.labels == {'adk_agent_name': 'root_agent'}

tests/unittests/flows/llm_flows/test_output_schema_processor.py

@@ -0,0 +1,522 @@
+# 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 output schema processor functionality."""
+
+from unittest import mock
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.run_config import RunConfig
+from google.adk.flows.llm_flows.single_flow import SingleFlow
+from google.adk.models.llm_request import LlmRequest
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.function_tool import FunctionTool
+from pydantic import BaseModel
+from pydantic import Field
+import pytest
+
+
+class PersonSchema(BaseModel):
+  """Test schema for structured output."""
+
+  name: str = Field(description="A person's name")
+  age: int = Field(description="A person's age")
+  city: str = Field(description='The city they live in')
+
+
+def dummy_tool(query: str) -> str:
+  """A dummy tool for testing."""
+  return f'Searched for: {query}'
+
+
+async def _create_invocation_context(agent: LlmAgent) -> InvocationContext:
+  """Helper to create InvocationContext for testing."""
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  return InvocationContext(
+      invocation_id='test-id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+      run_config=RunConfig(),
+  )
+
+
+@pytest.mark.asyncio
+async def test_output_schema_with_tools_validation_removed():
+  """Test that LlmAgent now allows output_schema with tools."""
+  # This should not raise an error anymore
+  agent = LlmAgent(
+      name='test_agent',
+      model='gemini-2.5-flash',
+      output_schema=PersonSchema,
+      tools=[FunctionTool(func=dummy_tool)],
+  )
+
+  assert agent.output_schema == PersonSchema
+  assert len(agent.tools) == 1
+
+
+@pytest.mark.asyncio
+async def test_output_schema_with_sub_agents():
+  """Test that LlmAgent now allows output_schema with sub_agents."""
+  sub_agent = LlmAgent(
+      name='sub_agent',
+      model='gemini-2.5-flash',
+  )
+  agent = LlmAgent(
+      name='test_agent',
+      model='gemini-2.5-flash',
+      output_schema=PersonSchema,
+      sub_agents=[sub_agent],
+  )
+
+  assert agent.output_schema == PersonSchema
+  assert len(agent.sub_agents) == 1
+
+
+@pytest.mark.asyncio
+async def test_basic_processor_skips_output_schema_with_tools():
+  """Test that basic processor doesn't set output_schema when tools are present."""
+  from google.adk.flows.llm_flows.basic import _BasicLlmRequestProcessor
+
+  agent = LlmAgent(
+      name='test_agent',
+      model='gemini-2.5-flash',
+      output_schema=PersonSchema,
+      tools=[FunctionTool(func=dummy_tool)],
+  )
+
+  invocation_context = await _create_invocation_context(agent)
+
+  llm_request = LlmRequest()
+  processor = _BasicLlmRequestProcessor()
+
+  # Process the request
+  events = []
+  async for event in processor.run_async(invocation_context, llm_request):
+    events.append(event)
+
+  # Should not have set response_schema since agent has tools
+  assert llm_request.config.response_schema is None
+  assert llm_request.config.response_mime_type != 'application/json'
+
+
+@pytest.mark.asyncio
+async def test_basic_processor_sets_output_schema_without_tools():
+  """Test that basic processor still sets output_schema when no tools are present."""
+  from google.adk.flows.llm_flows.basic import _BasicLlmRequestProcessor
+
+  agent = LlmAgent(
+      name='test_agent',
+      model='gemini-2.5-flash',
+      output_schema=PersonSchema,
+      tools=[],  # No tools
+  )
+
+  invocation_context = await _create_invocation_context(agent)
+
+  llm_request = LlmRequest()
+  processor = _BasicLlmRequestProcessor()
+
+  # Process the request
+  events = []
+  async for event in processor.run_async(invocation_context, llm_request):
+    events.append(event)
+
+  # Should have set response_schema since agent has no tools
+  assert llm_request.config.response_schema == PersonSchema
+  assert llm_request.config.response_mime_type == 'application/json'
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'output_schema_with_tools_allowed',
+    [
+        False,
+        True,
+    ],
+)
+async def test_output_schema_request_processor(
+    output_schema_with_tools_allowed, mocker
+):
+  """Test that output schema processor adds set_model_response tool."""
+  from google.adk.flows.llm_flows._output_schema_processor import _OutputSchemaRequestProcessor
+
+  agent = LlmAgent(
+      name='test_agent',
+      model='gemini-2.5-flash',
+      output_schema=PersonSchema,
+      tools=[FunctionTool(func=dummy_tool)],
+  )
+
+  invocation_context = await _create_invocation_context(agent)
+
+  llm_request = LlmRequest()
+  processor = _OutputSchemaRequestProcessor()
+
+  can_use_output_schema_with_tools = mocker.patch(
+      'google.adk.flows.llm_flows._output_schema_processor.can_use_output_schema_with_tools',
+      mock.MagicMock(return_value=output_schema_with_tools_allowed),
+  )
+
+  # Process the request
+  events = []
+  async for event in processor.run_async(invocation_context, llm_request):
+    events.append(event)
+
+  if not output_schema_with_tools_allowed:
+    # Should have added set_model_response tool if output schema with tools is
+    # allowed
+    assert 'set_model_response' in llm_request.tools_dict
+    # Should have added instruction about using set_model_response
+    assert 'set_model_response' in llm_request.config.system_instruction
+  else:
+    # Should skip modifying LlmRequest
+    assert not llm_request.tools_dict
+    assert not llm_request.config.system_instruction
+
+  # Should have checked if output schema can be used with tools
+  can_use_output_schema_with_tools.assert_called_once_with(
+      agent.canonical_model
+  )
+
+
+@pytest.mark.asyncio
+async def test_set_model_response_tool():
+  """Test the set_model_response tool functionality."""
+  from google.adk.tools.set_model_response_tool import SetModelResponseTool
+  from google.adk.tools.tool_context import ToolContext
+
+  tool = SetModelResponseTool(PersonSchema)
+
+  agent = LlmAgent(name='test_agent', model='gemini-2.5-flash')
+  invocation_context = await _create_invocation_context(agent)
+  tool_context = ToolContext(invocation_context)
+
+  # Call the tool with valid data
+  result = await tool.run_async(
+      args={'name': 'John Doe', 'age': 30, 'city': 'New York'},
+      tool_context=tool_context,
+  )
+
+  # Verify the tool returns dict directly
+  assert result is not None
+  assert result['name'] == 'John Doe'
+  assert result['age'] == 30
+  assert result['city'] == 'New York'
+
+
+@pytest.mark.asyncio
+async def test_output_schema_helper_functions():
+  """Test the helper functions for handling set_model_response."""
+  from google.adk.events.event import Event
+  from google.adk.flows.llm_flows._output_schema_processor import create_final_model_response_event
+  from google.adk.flows.llm_flows._output_schema_processor import get_structured_model_response
+  from google.genai import types
+
+  agent = LlmAgent(
+      name='test_agent',
+      model='gemini-2.5-flash',
+      output_schema=PersonSchema,
+      tools=[FunctionTool(func=dummy_tool)],
+  )
+
+  invocation_context = await _create_invocation_context(agent)
+
+  # Test get_structured_model_response with a function response event
+  test_dict = {'name': 'Jane Smith', 'age': 25, 'city': 'Los Angeles'}
+  test_json = '{"name": "Jane Smith", "age": 25, "city": "Los Angeles"}'
+
+  # Create a function response event with set_model_response
+  function_response_event = Event(
+      author='test_agent',
+      content=types.Content(
+          role='user',
+          parts=[
+              types.Part(
+                  function_response=types.FunctionResponse(
+                      name='set_model_response', response=test_dict
+                  )
+              )
+          ],
+      ),
+  )
+
+  # Test get_structured_model_response function
+  extracted_json = get_structured_model_response(function_response_event)
+  assert extracted_json == test_json
+
+  # Test create_final_model_response_event function
+  final_event = create_final_model_response_event(invocation_context, test_json)
+  assert final_event.author == 'test_agent'
+  assert final_event.invocation_id == invocation_context.invocation_id
+  assert final_event.branch == invocation_context.branch
+  assert final_event.content.role == 'model'
+  assert final_event.content.parts[0].text == test_json
+
+  # Test get_structured_model_response with non-set_model_response function
+  other_function_response_event = Event(
+      author='test_agent',
+      content=types.Content(
+          role='user',
+          parts=[
+              types.Part(
+                  function_response=types.FunctionResponse(
+                      name='other_tool', response={'result': 'other response'}
+                  )
+              )
+          ],
+      ),
+  )
+
+  extracted_json = get_structured_model_response(other_function_response_event)
+  assert extracted_json is None
+
+
+@pytest.mark.asyncio
+async def test_get_structured_model_response_with_non_ascii():
+  """Test get_structured_model_response with non-ASCII characters."""
+  from google.adk.events.event import Event
+  from google.adk.flows.llm_flows._output_schema_processor import get_structured_model_response
+  from google.genai import types
+
+  # Test with a dictionary containing non-ASCII characters
+  test_dict = {'city': 'São Paulo'}
+  expected_json = '{"city": "São Paulo"}'
+
+  # Create a function response event
+  function_response_event = Event(
+      author='test_agent',
+      content=types.Content(
+          role='user',
+          parts=[
+              types.Part(
+                  function_response=types.FunctionResponse(
+                      name='set_model_response', response=test_dict
+                  )
+              )
+          ],
+      ),
+  )
+
+  # Get the structured response
+  extracted_json = get_structured_model_response(function_response_event)
+
+  # Assert that the output is the expected JSON string without escaped characters
+  assert extracted_json == expected_json
+
+
+@pytest.mark.asyncio
+async def test_get_structured_model_response_with_wrapped_result():
+  """Test get_structured_model_response with wrapped list result.
+
+  When a tool returns a non-dict (e.g., list), it gets wrapped as
+  {'result': [...]}.  This test ensures we correctly unwrap the result.
+  """
+  from google.adk.events.event import Event
+  from google.adk.flows.llm_flows._output_schema_processor import get_structured_model_response
+  from google.genai import types
+
+  # Simulate a list result wrapped by ADK's functions.py
+  wrapped_response = {
+      'result': [
+          {'name': 'Alice', 'age': 30},
+          {'name': 'Bob', 'age': 25},
+      ]
+  }
+  expected_json = '[{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}]'
+
+  # Create a function response event with wrapped result
+  function_response_event = Event(
+      author='test_agent',
+      content=types.Content(
+          role='user',
+          parts=[
+              types.Part(
+                  function_response=types.FunctionResponse(
+                      name='set_model_response', response=wrapped_response
+                  )
+              )
+          ],
+      ),
+  )
+
+  # Get the structured response
+  extracted_json = get_structured_model_response(function_response_event)
+
+  # Should extract the unwrapped list, not the wrapped dict
+  assert extracted_json == expected_json
+
+
+@pytest.mark.asyncio
+async def test_end_to_end_integration():
+  """Test the complete output schema with tools integration."""
+  agent = LlmAgent(
+      name='test_agent',
+      model='gemini-2.5-flash',
+      output_schema=PersonSchema,
+      tools=[FunctionTool(func=dummy_tool)],
+  )
+
+  invocation_context = await _create_invocation_context(agent)
+
+  # Create a flow and test the processors
+  flow = SingleFlow()
+  llm_request = LlmRequest()
+
+  # Run all request processors
+  async for event in flow._preprocess_async(invocation_context, llm_request):
+    pass
+
+  # Verify set_model_response tool was added
+  assert 'set_model_response' in llm_request.tools_dict
+
+  # Verify instruction was added
+  assert 'set_model_response' in llm_request.config.system_instruction
+
+  # Verify output_schema was NOT set on the model config
+  assert llm_request.config.response_schema is None
+
+
+@pytest.mark.asyncio
+async def test_flow_yields_both_events_for_set_model_response():
+  """Test that the flow yields both function response and final model response events."""
+  from google.adk.events.event import Event
+  from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
+  from google.adk.tools.set_model_response_tool import SetModelResponseTool
+  from google.genai import types
+
+  agent = LlmAgent(
+      name='test_agent',
+      model='gemini-2.5-flash',
+      output_schema=PersonSchema,
+      tools=[],
+  )
+
+  invocation_context = await _create_invocation_context(agent)
+  flow = BaseLlmFlow()
+
+  # Create a set_model_response tool and add it to the tools dict
+  set_response_tool = SetModelResponseTool(PersonSchema)
+  llm_request = LlmRequest()
+  llm_request.tools_dict['set_model_response'] = set_response_tool
+
+  # Create a function call event (model calling the function)
+  function_call_event = Event(
+      author='test_agent',
+      content=types.Content(
+          role='model',
+          parts=[
+              types.Part(
+                  function_call=types.FunctionCall(
+                      name='set_model_response',
+                      args={
+                          'name': 'Test User',
+                          'age': 30,
+                          'city': 'Test City',
+                      },
+                  )
+              )
+          ],
+      ),
+  )
+
+  # Test the postprocess function handling
+  events = []
+  async for event in flow._postprocess_handle_function_calls_async(
+      invocation_context, function_call_event, llm_request
+  ):
+    events.append(event)
+
+  # Should yield exactly 2 events: function response + final model response
+  assert len(events) == 2
+
+  # First event should be the function response
+  first_event = events[0]
+  assert first_event.get_function_responses()[0].name == 'set_model_response'
+  # The response should be the dict returned by the tool
+  assert first_event.get_function_responses()[0].response == {
+      'name': 'Test User',
+      'age': 30,
+      'city': 'Test City',
+  }
+
+  # Second event should be the final model response with JSON
+  second_event = events[1]
+  assert second_event.author == 'test_agent'
+  assert second_event.invocation_id == invocation_context.invocation_id
+  assert second_event.branch == invocation_context.branch
+  assert second_event.content.role == 'model'
+  assert (
+      second_event.content.parts[0].text
+      == '{"name": "Test User", "age": 30, "city": "Test City"}'
+  )
+
+
+@pytest.mark.asyncio
+async def test_flow_yields_only_function_response_for_normal_tools():
+  """Test that the flow yields only function response event for non-set_model_response tools."""
+  from google.adk.events.event import Event
+  from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
+  from google.genai import types
+
+  agent = LlmAgent(
+      name='test_agent',
+      model='gemini-2.5-flash',
+      tools=[FunctionTool(func=dummy_tool)],
+  )
+
+  invocation_context = await _create_invocation_context(agent)
+  flow = BaseLlmFlow()
+
+  # Create a dummy tool and add it to the tools dict
+  dummy_function_tool = FunctionTool(func=dummy_tool)
+  llm_request = LlmRequest()
+  llm_request.tools_dict['dummy_tool'] = dummy_function_tool
+
+  # Create a function call event (model calling the dummy tool)
+  function_call_event = Event(
+      author='test_agent',
+      content=types.Content(
+          role='model',
+          parts=[
+              types.Part(
+                  function_call=types.FunctionCall(
+                      name='dummy_tool', args={'query': 'test query'}
+                  )
+              )
+          ],
+      ),
+  )
+
+  # Test the postprocess function handling
+  events = []
+  async for event in flow._postprocess_handle_function_calls_async(
+      invocation_context, function_call_event, llm_request
+  ):
+    events.append(event)
+
+  # Should yield exactly 1 event: just the function response
+  assert len(events) == 1
+
+  # Should be the function response from dummy_tool
+  first_event = events[0]
+  assert first_event.get_function_responses()[0].name == 'dummy_tool'
+  assert first_event.get_function_responses()[0].response == {
+      'result': 'Searched for: test query'
+  }

tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py

@@ -0,0 +1,189 @@
+# 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.
+
+from typing import Optional
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.llm_agent import Agent
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
+from google.genai import types
+from google.genai.errors import ClientError
+import pytest
+
+from ... import testing_utils
+
+mock_error = ClientError(
+    code=429,
+    response_json={
+        'error': {
+            'code': 429,
+            'message': 'Quota exceeded.',
+            'status': 'RESOURCE_EXHAUSTED',
+        }
+    },
+)
+
+
+class MockPlugin(BasePlugin):
+  before_model_text = 'before_model_text from MockPlugin'
+  after_model_text = 'after_model_text from MockPlugin'
+  on_model_error_text = 'on_model_error_text from MockPlugin'
+
+  def __init__(self, name='mock_plugin'):
+    self.name = name
+    self.enable_before_model_callback = False
+    self.enable_after_model_callback = False
+    self.enable_on_model_error_callback = False
+    self.before_model_response = LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.before_model_text)]
+        )
+    )
+    self.after_model_response = LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.after_model_text)]
+        )
+    )
+    self.on_model_error_response = LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.on_model_error_text)]
+        )
+    )
+
+  async def before_model_callback(
+      self, *, callback_context: CallbackContext, llm_request: LlmRequest
+  ) -> Optional[LlmResponse]:
+    if not self.enable_before_model_callback:
+      return None
+    return self.before_model_response
+
+  async def after_model_callback(
+      self, *, callback_context: CallbackContext, llm_response: LlmResponse
+  ) -> Optional[LlmResponse]:
+    if not self.enable_after_model_callback:
+      return None
+    return self.after_model_response
+
+  async def on_model_error_callback(
+      self,
+      *,
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+      error: Exception,
+  ) -> Optional[LlmResponse]:
+    if not self.enable_on_model_error_callback:
+      return None
+    return self.on_model_error_response
+
+
+CANONICAL_MODEL_CALLBACK_CONTENT = 'canonical_model_callback_content'
+
+
+def canonical_agent_model_callback(**kwargs) -> Optional[LlmResponse]:
+  return LlmResponse(
+      content=testing_utils.ModelContent(
+          [types.Part.from_text(text=CANONICAL_MODEL_CALLBACK_CONTENT)]
+      )
+  )
+
+
+@pytest.fixture
+def mock_plugin():
+  return MockPlugin()
+
+
+def test_before_model_callback_with_plugin(mock_plugin):
+  """Tests that the model response is overridden by before_model_callback from the plugin."""
+  responses = ['model_response']
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  mock_plugin.enable_before_model_callback = True
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent, plugins=[mock_plugin])
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', mock_plugin.before_model_text),
+  ]
+
+
+def test_before_model_fallback_canonical_callback(mock_plugin):
+  """Tests that when plugin returns empty response, the model response is overridden by the canonical agent model callback."""
+  responses = ['model_response']
+  mock_plugin.enable_before_model_callback = False
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=canonical_agent_model_callback,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', CANONICAL_MODEL_CALLBACK_CONTENT),
+  ]
+
+
+def test_before_model_callback_fallback_model(mock_plugin):
+  """Tests that the model response is executed normally when both plugin and canonical agent model callback return empty response."""
+  responses = ['model_response']
+  mock_plugin.enable_before_model_callback = False
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent, plugins=[mock_plugin])
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', 'model_response'),
+  ]
+
+
+def test_on_model_error_callback_with_plugin(mock_plugin):
+  """Tests that the model error is handled by the plugin."""
+  mock_model = testing_utils.MockModel.create(error=mock_error, responses=[])
+  mock_plugin.enable_on_model_error_callback = True
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent, plugins=[mock_plugin])
+
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', mock_plugin.on_model_error_text),
+  ]
+
+
+def test_on_model_error_callback_fallback_to_runner(mock_plugin):
+  """Tests that the model error is not handled and falls back to raise from runner."""
+  mock_model = testing_utils.MockModel.create(error=mock_error, responses=[])
+  mock_plugin.enable_on_model_error_callback = False
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+  )
+
+  try:
+    testing_utils.InMemoryRunner(agent, plugins=[mock_plugin])
+  except Exception as e:
+    assert e == mock_error
+
+
+if __name__ == '__main__':
+  pytest.main([__file__])

tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py

@@ -0,0 +1,344 @@
+# 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.
+
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.functions import handle_function_calls_async
+from google.adk.flows.llm_flows.functions import handle_function_calls_live
+from google.adk.plugins.base_plugin import BasePlugin
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+from google.genai.errors import ClientError
+import pytest
+
+from ... import testing_utils
+
+mock_error = ClientError(
+    code=429,
+    response_json={
+        "error": {
+            "code": 429,
+            "message": "Quota exceeded.",
+            "status": "RESOURCE_EXHAUSTED",
+        }
+    },
+)
+
+
+class MockPlugin(BasePlugin):
+  before_tool_response = {"MockPlugin": "before_tool_response from MockPlugin"}
+  after_tool_response = {"MockPlugin": "after_tool_response from MockPlugin"}
+  on_tool_error_response = {
+      "MockPlugin": "on_tool_error_response from MockPlugin"
+  }
+
+  def __init__(self, name="mock_plugin"):
+    self.name = name
+    self.enable_before_tool_callback = False
+    self.enable_after_tool_callback = False
+    self.enable_on_tool_error_callback = False
+
+  async def before_tool_callback(
+      self,
+      *,
+      tool: BaseTool,
+      tool_args: dict[str, Any],
+      tool_context: ToolContext,
+  ) -> Optional[dict]:
+    if not self.enable_before_tool_callback:
+      return None
+    return self.before_tool_response
+
+  async def after_tool_callback(
+      self,
+      *,
+      tool: BaseTool,
+      tool_args: dict[str, Any],
+      tool_context: ToolContext,
+      result: dict,
+  ) -> Optional[dict]:
+    if not self.enable_after_tool_callback:
+      return None
+    return self.after_tool_response
+
+  async def on_tool_error_callback(
+      self,
+      *,
+      tool: BaseTool,
+      tool_args: dict[str, Any],
+      tool_context: ToolContext,
+      error: Exception,
+  ) -> Optional[dict]:
+    if not self.enable_on_tool_error_callback:
+      return None
+    return self.on_tool_error_response
+
+
+@pytest.fixture
+def mock_tool():
+  def simple_fn(**kwargs) -> Dict[str, Any]:
+    return {"initial": "response"}
+
+  return FunctionTool(simple_fn)
+
+
+@pytest.fixture
+def mock_error_tool():
+  def raise_error_fn(**kwargs) -> Dict[str, Any]:
+    raise mock_error
+
+  return FunctionTool(raise_error_fn)
+
+
+@pytest.fixture
+def mock_plugin():
+  return MockPlugin()
+
+
+async def invoke_tool_with_plugin(mock_tool, mock_plugin) -> Optional[Event]:
+  """Invokes a tool with a plugin."""
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name="agent",
+      model=model,
+      tools=[mock_tool],
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content="", plugins=[mock_plugin]
+  )
+  # Build function call event
+  function_call = types.FunctionCall(name=mock_tool.name, args={})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {mock_tool.name: mock_tool}
+  return await handle_function_calls_async(
+      invocation_context,
+      event,
+      tools_dict,
+  )
+
+
+@pytest.mark.asyncio
+async def test_async_before_tool_callback(mock_tool, mock_plugin):
+  mock_plugin.enable_before_tool_callback = True
+
+  result_event = await invoke_tool_with_plugin(mock_tool, mock_plugin)
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_plugin.before_tool_response
+
+
+@pytest.mark.asyncio
+async def test_async_after_tool_callback(mock_tool, mock_plugin):
+  mock_plugin.enable_after_tool_callback = True
+
+  result_event = await invoke_tool_with_plugin(mock_tool, mock_plugin)
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_plugin.after_tool_response
+
+
+@pytest.mark.asyncio
+async def test_async_on_tool_error_use_plugin_response(
+    mock_error_tool, mock_plugin
+):
+  mock_plugin.enable_on_tool_error_callback = True
+
+  result_event = await invoke_tool_with_plugin(mock_error_tool, mock_plugin)
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_plugin.on_tool_error_response
+
+
+@pytest.mark.asyncio
+async def test_async_on_tool_error_fallback_to_runner(
+    mock_error_tool, mock_plugin
+):
+  mock_plugin.enable_on_tool_error_callback = False
+
+  try:
+    await invoke_tool_with_plugin(mock_error_tool, mock_plugin)
+  except Exception as e:
+    assert e == mock_error
+
+
+async def invoke_tool_with_plugin_live(
+    mock_tool, mock_plugin
+) -> Optional[Event]:
+  """Invokes a tool with a plugin using the live path."""
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name="agent",
+      model=model,
+      tools=[mock_tool],
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content="", plugins=[mock_plugin]
+  )
+  # Build function call event
+  function_call = types.FunctionCall(name=mock_tool.name, args={})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {mock_tool.name: mock_tool}
+  return await handle_function_calls_live(
+      invocation_context,
+      event,
+      tools_dict,
+  )
+
+
+@pytest.mark.asyncio
+async def test_live_before_tool_callback(mock_tool, mock_plugin):
+  mock_plugin.enable_before_tool_callback = True
+
+  result_event = await invoke_tool_with_plugin_live(mock_tool, mock_plugin)
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_plugin.before_tool_response
+
+
+@pytest.mark.asyncio
+async def test_live_after_tool_callback(mock_tool, mock_plugin):
+  mock_plugin.enable_after_tool_callback = True
+
+  result_event = await invoke_tool_with_plugin_live(mock_tool, mock_plugin)
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_plugin.after_tool_response
+
+
+@pytest.mark.asyncio
+async def test_live_on_tool_error_use_plugin_response(
+    mock_error_tool, mock_plugin
+):
+  mock_plugin.enable_on_tool_error_callback = True
+
+  result_event = await invoke_tool_with_plugin_live(
+      mock_error_tool, mock_plugin
+  )
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_plugin.on_tool_error_response
+
+
+@pytest.mark.asyncio
+async def test_live_on_tool_error_fallback_to_runner(
+    mock_error_tool, mock_plugin
+):
+  mock_plugin.enable_on_tool_error_callback = False
+
+  try:
+    await invoke_tool_with_plugin_live(mock_error_tool, mock_plugin)
+  except Exception as e:
+    assert e == mock_error
+
+
+@pytest.mark.asyncio
+async def test_live_plugin_before_tool_callback_takes_priority(
+    mock_tool, mock_plugin
+):
+  """Plugin before_tool_callback should run before agent canonical callbacks."""
+  mock_plugin.enable_before_tool_callback = True
+
+  def agent_before_cb(tool, args, tool_context):
+    return {"agent": "should_not_be_called"}
+
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name="agent",
+      model=model,
+      tools=[mock_tool],
+      before_tool_callback=agent_before_cb,
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content="", plugins=[mock_plugin]
+  )
+  function_call = types.FunctionCall(name=mock_tool.name, args={})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {mock_tool.name: mock_tool}
+  result_event = await handle_function_calls_live(
+      invocation_context, event, tools_dict
+  )
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  # Plugin response should win, not the agent callback
+  assert part.function_response.response == mock_plugin.before_tool_response
+
+
+@pytest.mark.asyncio
+async def test_live_plugin_after_tool_callback_takes_priority(
+    mock_tool, mock_plugin
+):
+  """Plugin after_tool_callback should run before agent canonical callbacks."""
+  mock_plugin.enable_after_tool_callback = True
+
+  def agent_after_cb(tool, args, tool_context, tool_response):
+    return {"agent": "should_not_be_called"}
+
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name="agent",
+      model=model,
+      tools=[mock_tool],
+      after_tool_callback=agent_after_cb,
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content="", plugins=[mock_plugin]
+  )
+  function_call = types.FunctionCall(name=mock_tool.name, args={})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {mock_tool.name: mock_tool}
+  result_event = await handle_function_calls_live(
+      invocation_context, event, tools_dict
+  )
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  # Plugin response should win, not the agent callback
+  assert part.function_response.response == mock_plugin.after_tool_response
+
+
+if __name__ == "__main__":
+  pytest.main([__file__])

tests/unittests/flows/llm_flows/test_progressive_sse_streaming.py

@@ -0,0 +1,896 @@
+# 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 Progressive SSE Streaming Stage 1 implementation."""
+
+import asyncio
+from typing import Any
+from typing import AsyncGenerator
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.agents.run_config import RunConfig
+from google.adk.agents.run_config import StreamingMode
+from google.adk.models.base_llm import BaseLlm
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.adk.runners import InMemoryRunner
+from google.adk.utils.streaming_utils import StreamingResponseAggregator
+from google.genai import types
+import pytest
+
+
+def get_weather(location: str) -> dict[str, Any]:
+  """Mock weather function for testing.
+
+  Args:
+    location: The location to get the weather for.
+
+  Returns:
+    A dictionary containing the weather information.
+  """
+  return {
+      "temperature": 22,
+      "condition": "sunny",
+      "location": location,
+  }
+
+
+class StreamingMockModel(BaseLlm):
+  """A mock model that properly streams multiple chunks in a single call."""
+
+  model: str = "streaming-mock"
+  stream_chunks: list[LlmResponse] = []
+  call_count: int = 0
+
+  @classmethod
+  def supported_models(cls) -> list[str]:
+    return ["streaming-mock"]
+
+  async def generate_content_async(
+      self, llm_request: LlmRequest, stream: bool = False
+  ) -> AsyncGenerator[LlmResponse, None]:
+    """Yield all chunks in a single streaming call."""
+    self.call_count += 1
+
+    # Only stream on the first call
+    if self.call_count > 1:
+      # On subsequent calls, return a simple final response
+      yield LlmResponse(
+          content=types.Content(
+              role="model",
+              parts=[types.Part.from_text(text="Task completed.")],
+          ),
+          partial=False,
+      )
+      return
+
+    aggregator = StreamingResponseAggregator()
+
+    # Process each chunk through the aggregator
+    for chunk in self.stream_chunks:
+      # Convert LlmResponse to types.GenerateContentResponse
+      # Since we don't have the full response object, we'll simulate it
+      async for processed_chunk in aggregator.process_response(
+          self._llm_response_to_generate_content_response(chunk)
+      ):
+        yield processed_chunk
+
+    # Call close() to get the final aggregated response
+    if final_response := aggregator.close():
+      yield final_response
+
+  def _llm_response_to_generate_content_response(
+      self, llm_response: LlmResponse
+  ) -> types.GenerateContentResponse:
+    """Convert LlmResponse to GenerateContentResponse for aggregator."""
+    # Create a minimal GenerateContentResponse that the aggregator can process
+    candidates = []
+    if llm_response.content:
+      candidates.append(
+          types.Candidate(
+              content=llm_response.content,
+              finish_reason=llm_response.finish_reason,
+              finish_message=llm_response.error_message,
+          )
+      )
+
+    return types.GenerateContentResponse(
+        candidates=candidates,
+        usage_metadata=llm_response.usage_metadata,
+    )
+
+
+def test_progressive_sse_streaming_function_calls():
+  """Test that function calls are buffered and executed in parallel."""
+
+  # Setup: Create mock responses simulating streaming chunks
+  response1 = LlmResponse(
+      content=types.Content(
+          role="model", parts=[types.Part.from_text(text="Checking weather...")]
+      ),
+  )
+
+  response2 = LlmResponse(
+      content=types.Content(
+          role="model",
+          parts=[
+              types.Part.from_function_call(
+                  name="get_weather", args={"location": "Tokyo"}
+              )
+          ],
+      ),
+  )
+
+  response3 = LlmResponse(
+      content=types.Content(
+          role="model",
+          parts=[
+              types.Part.from_function_call(
+                  name="get_weather", args={"location": "New York"}
+              )
+          ],
+      ),
+      finish_reason=types.FinishReason.STOP,
+  )
+
+  # Create a streaming mock that yields all chunks in one call
+  mock_model = StreamingMockModel(
+      stream_chunks=[response1, response2, response3]
+  )
+
+  agent = Agent(
+      name="weather_agent",
+      model=mock_model,
+      tools=[get_weather],
+  )
+
+  run_config = RunConfig(streaming_mode=StreamingMode.SSE)
+
+  # Use the real InMemoryRunner to get access to run_config parameter
+  runner = InMemoryRunner(agent=agent)
+
+  # Create session manually
+  session = runner.session_service.create_session_sync(
+      app_name=runner.app_name, user_id="test_user"
+  )
+
+  events = []
+  for event in runner.run(
+      user_id="test_user",
+      session_id=session.id,
+      new_message=types.Content(
+          role="user",
+          parts=[types.Part.from_text(text="What is the weather?")],
+      ),
+      run_config=run_config,
+  ):
+    events.append(event)
+
+  # Verify event structure (Stage 1 expectations)
+  # Expected events:
+  # 0-2: Partial events (text + 2 FCs) - not executed
+  # 3: Final aggregated model event (text + 2 FCs) - partial=False
+  # 4: Aggregated function response (both get_weather results executed in
+  # parallel)
+  # 5: Final model response after FCs
+  assert len(events) == 6
+
+  assert events[0].partial
+  assert events[0].content.parts[0].text == "Checking weather..."
+
+  assert events[1].partial
+  assert events[1].content.parts[0].function_call.name == "get_weather"
+  assert events[1].content.parts[0].function_call.args["location"] == "Tokyo"
+
+  assert events[2].partial
+  assert events[2].content.parts[0].function_call.name == "get_weather"
+  assert events[2].content.parts[0].function_call.args["location"] == "New York"
+
+  assert not events[3].partial
+  assert events[3].content.parts[0].text == "Checking weather..."
+  assert events[3].content.parts[1].function_call.name == "get_weather"
+  assert events[3].content.parts[1].function_call.args["location"] == "Tokyo"
+  assert events[3].content.parts[2].function_call.name == "get_weather"
+  assert events[3].content.parts[2].function_call.args["location"] == "New York"
+
+  assert not events[4].partial
+  assert events[4].content.parts[0].function_response.name == "get_weather"
+  assert (
+      events[4].content.parts[0].function_response.response["location"]
+      == "Tokyo"
+  )
+  assert events[4].content.parts[1].function_response.name == "get_weather"
+  assert (
+      events[4].content.parts[1].function_response.response["location"]
+      == "New York"
+  )
+
+  assert not events[5].partial
+  assert events[5].content.parts[0].text == "Task completed."
+
+
+def test_progressive_sse_preserves_part_ordering():
+  """Test that part ordering is preserved, especially for thought parts.
+
+  This test verifies that when the model outputs:
+  - chunk1(thought1_1)
+  - chunk2(thought1_2)
+  - chunk3(text1_1)
+  - chunk4(text1_2)
+  - chunk5(FC1)
+  - chunk6(thought2_1)
+  - chunk7(thought2_2)
+  - chunk8(FC2)
+
+  The final aggregated output should be:
+  - Part(thought1)  # thought1_1 + thought1_2 merged
+  - Part(text1)     # text1_1 + text1_2 merged
+  - Part(FC1)
+  - Part(thought2)  # thought2_1 + thought2_2 merged
+  - Part(FC2)
+  """
+
+  # Create streaming chunks that test the ordering requirement
+  chunk1 = LlmResponse(
+      content=types.Content(
+          role="model",
+          parts=[types.Part(text="Initial thought part 1. ", thought=True)],
+      )
+  )
+
+  chunk2 = LlmResponse(
+      content=types.Content(
+          role="model",
+          parts=[types.Part(text="Initial thought part 2.", thought=True)],
+      )
+  )
+
+  chunk3 = LlmResponse(
+      content=types.Content(
+          role="model",
+          parts=[types.Part.from_text(text="Let me check Tokyo. ")],
+      )
+  )
+
+  chunk4 = LlmResponse(
+      content=types.Content(
+          role="model", parts=[types.Part.from_text(text="And New York too.")]
+      )
+  )
+
+  chunk5 = LlmResponse(
+      content=types.Content(
+          role="model",
+          parts=[
+              types.Part.from_function_call(
+                  name="get_weather", args={"location": "Tokyo"}
+              )
+          ],
+      )
+  )
+
+  chunk6 = LlmResponse(
+      content=types.Content(
+          role="model",
+          parts=[
+              types.Part(
+                  text="Now processing second thought part 1. ", thought=True
+              )
+          ],
+      )
+  )
+
+  chunk7 = LlmResponse(
+      content=types.Content(
+          role="model",
+          parts=[types.Part(text="Second thought part 2.", thought=True)],
+      )
+  )
+
+  chunk8 = LlmResponse(
+      content=types.Content(
+          role="model",
+          parts=[
+              types.Part.from_function_call(
+                  name="get_weather", args={"location": "New York"}
+              )
+          ],
+      ),
+      finish_reason=types.FinishReason.STOP,
+  )
+
+  mock_model = StreamingMockModel(
+      stream_chunks=[
+          chunk1,
+          chunk2,
+          chunk3,
+          chunk4,
+          chunk5,
+          chunk6,
+          chunk7,
+          chunk8,
+      ]
+  )
+
+  agent = Agent(
+      name="ordering_test_agent",
+      model=mock_model,
+      tools=[get_weather],
+  )
+
+  run_config = RunConfig(streaming_mode=StreamingMode.SSE)
+
+  # Use the real InMemoryRunner to get access to run_config parameter
+  runner = InMemoryRunner(agent=agent)
+
+  # Create session manually
+  session = runner.session_service.create_session_sync(
+      app_name=runner.app_name, user_id="test_user"
+  )
+
+  events = []
+  for event in runner.run(
+      user_id="test_user",
+      session_id=session.id,
+      new_message=types.Content(
+          role="user",
+          parts=[types.Part.from_text(text="What is the weather?")],
+      ),
+      run_config=run_config,
+  ):
+    events.append(event)
+
+  # Find the final aggregated model event (partial=False, from model)
+  aggregated_event = None
+  for event in events:
+    if (
+        not event.partial
+        and event.author == "ordering_test_agent"
+        and event.content
+        and len(event.content.parts) > 2
+    ):
+      aggregated_event = event
+      break
+
+  assert aggregated_event is not None, "Should find an aggregated model event"
+
+  # Verify the part ordering
+  parts = aggregated_event.content.parts
+  assert len(parts) == 5, f"Expected 5 parts, got {len(parts)}"
+
+  # Part 0: First thought (merged from chunk1 + chunk2)
+  assert parts[0].thought
+  assert parts[0].text == "Initial thought part 1. Initial thought part 2."
+
+  # Part 1: Regular text (merged from chunk3 + chunk4)
+  assert not parts[1].thought
+  assert parts[1].text == "Let me check Tokyo. And New York too."
+
+  # Part 2: First function call (from chunk5)
+  assert parts[2].function_call.name == "get_weather"
+  assert parts[2].function_call.args["location"] == "Tokyo"
+
+  # Part 3: Second thought (merged from chunk6 + chunk7)
+  assert parts[3].thought
+  assert (
+      parts[3].text
+      == "Now processing second thought part 1. Second thought part 2."
+  )
+
+  # Part 4: Second function call (from chunk8)
+  assert parts[4].function_call.name == "get_weather"
+  assert parts[4].function_call.args["location"] == "New York"
+
+
+def test_progressive_sse_streaming_function_call_arguments():
+  """Test streaming function call arguments feature.
+
+  This test simulates the streamFunctionCallArguments feature where a function
+  call's arguments are streamed incrementally across multiple chunks:
+
+  Chunk 1: FC name + partial location argument ("New ")
+  Chunk 2: Continue location argument ("York") -> concatenated to "New York"
+  Chunk 3: Add unit argument ("celsius"), willContinue=False -> FC complete
+
+  Expected result: FunctionCall(name="get_weather",
+                                 args={"location": "New York", "unit":
+                                 "celsius"},
+                                 id="fc_001")
+  """
+
+  aggregator = StreamingResponseAggregator()
+
+  # Chunk 1: FC name + partial location argument
+  chunk1_fc = types.FunctionCall(
+      name="get_weather",
+      id="fc_001",
+      partial_args=[
+          types.PartialArg(json_path="$.location", string_value="New ")
+      ],
+      will_continue=True,
+  )
+  chunk1 = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role="model", parts=[types.Part(function_call=chunk1_fc)]
+              )
+          )
+      ]
+  )
+
+  # Chunk 2: Continue streaming location argument
+  chunk2_fc = types.FunctionCall(
+      partial_args=[
+          types.PartialArg(json_path="$.location", string_value="York")
+      ],
+      will_continue=True,
+  )
+  chunk2 = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role="model", parts=[types.Part(function_call=chunk2_fc)]
+              )
+          )
+      ]
+  )
+
+  # Chunk 3: Add unit argument, FC complete
+  chunk3_fc = types.FunctionCall(
+      partial_args=[
+          types.PartialArg(json_path="$.unit", string_value="celsius")
+      ],
+      will_continue=False,  # FC complete
+  )
+  chunk3 = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role="model", parts=[types.Part(function_call=chunk3_fc)]
+              ),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ]
+  )
+
+  # Process all chunks through aggregator
+  processed_chunks = []
+  for chunk in [chunk1, chunk2, chunk3]:
+
+    async def process():
+      results = []
+      async for response in aggregator.process_response(chunk):
+        results.append(response)
+      return results
+
+    import asyncio
+
+    chunk_results = asyncio.run(process())
+    processed_chunks.extend(chunk_results)
+
+  # Get final aggregated response
+  final_response = aggregator.close()
+
+  # Verify final aggregated response has complete FC
+  assert final_response is not None
+  assert len(final_response.content.parts) == 1
+
+  fc_part = final_response.content.parts[0]
+  assert fc_part.function_call is not None
+  assert fc_part.function_call.name == "get_weather"
+  assert fc_part.function_call.id == "fc_001"
+
+  # Verify arguments were correctly assembled from streaming chunks
+  args = fc_part.function_call.args
+  assert args["location"] == "New York"  # "New " + "York" concatenated
+  assert args["unit"] == "celsius"
+
+
+def test_progressive_sse_preserves_thought_signature():
+  """Test that thought_signature is preserved when streaming FC arguments.
+
+  This test verifies that when a streaming function call has a thought_signature
+  in the Part, it is correctly preserved in the final aggregated FunctionCall.
+  """
+
+  aggregator = StreamingResponseAggregator()
+
+  # Create a thought signature (simulating what Gemini returns)
+  # thought_signature is bytes (base64 encoded)
+  test_thought_signature = b"test_signature_abc123"
+
+  # Chunk with streaming FC args and thought_signature
+  chunk_fc = types.FunctionCall(
+      name="add_5_numbers",
+      id="fc_003",
+      partial_args=[
+          types.PartialArg(json_path="$.num1", number_value=10),
+          types.PartialArg(json_path="$.num2", number_value=20),
+      ],
+      will_continue=False,
+  )
+
+  # Create Part with both function_call AND thought_signature
+  chunk_part = types.Part(
+      function_call=chunk_fc, thought_signature=test_thought_signature
+  )
+
+  chunk = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(role="model", parts=[chunk_part]),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ]
+  )
+
+  # Process chunk through aggregator
+  async def process():
+    results = []
+    async for response in aggregator.process_response(chunk):
+      results.append(response)
+    return results
+
+  import asyncio
+
+  asyncio.run(process())
+
+  # Get final aggregated response
+  final_response = aggregator.close()
+
+  # Verify thought_signature was preserved in the Part
+  assert final_response is not None
+  assert len(final_response.content.parts) == 1
+
+  fc_part = final_response.content.parts[0]
+  assert fc_part.function_call is not None
+  assert fc_part.function_call.name == "add_5_numbers"
+
+  assert fc_part.thought_signature == test_thought_signature
+
+
+def test_progressive_sse_handles_empty_function_call():
+  """Test that empty function calls are skipped.
+
+  When using streamFunctionCallArguments, Gemini may send an empty
+  functionCall: {} as the final chunk to signal streaming completion.
+  This test verifies that such empty function calls are properly skipped
+  and don't cause errors.
+  """
+
+  aggregator = StreamingResponseAggregator()
+
+  # Chunk 1: Streaming FC with partial args
+  chunk1_fc = types.FunctionCall(
+      name="concat_number_and_string",
+      id="fc_001",
+      partial_args=[
+          types.PartialArg(json_path="$.num", number_value=100),
+          types.PartialArg(json_path="$.s", string_value="ADK"),
+      ],
+      will_continue=False,
+  )
+  chunk1 = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role="model", parts=[types.Part(function_call=chunk1_fc)]
+              )
+          )
+      ]
+  )
+
+  # Chunk 2: Empty function call (streaming end marker)
+  chunk2_fc = types.FunctionCall()  # Empty function call
+  chunk2 = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role="model", parts=[types.Part(function_call=chunk2_fc)]
+              ),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ]
+  )
+
+  # Process all chunks through aggregator
+  async def process():
+    results = []
+    for chunk in [chunk1, chunk2]:
+      async for response in aggregator.process_response(chunk):
+        results.append(response)
+    return results
+
+  import asyncio
+
+  asyncio.run(process())
+
+  # Get final aggregated response
+  final_response = aggregator.close()
+
+  # Verify final response only has the real FC, not the empty one
+  assert final_response is not None
+  assert len(final_response.content.parts) == 1
+
+  fc_part = final_response.content.parts[0]
+  assert fc_part.function_call is not None
+  assert fc_part.function_call.name == "concat_number_and_string"
+  assert fc_part.function_call.id == "fc_001"
+
+  # Verify arguments
+  args = fc_part.function_call.args
+  assert args["num"] == 100
+  assert args["s"] == "ADK"
+
+
+@pytest.mark.parametrize(
+    "first_chunk_partial_args",
+    [
+        pytest.param(None, id="partial_args_none"),
+        pytest.param([], id="partial_args_empty_list"),
+    ],
+)
+def test_streaming_fc_chunk_with_will_continue_but_no_partial_args(
+    first_chunk_partial_args,
+):
+  """Test streaming function call with will_continue=True but no partial_args."""
+
+  aggregator = StreamingResponseAggregator()
+
+  # Chunk 1: FC name + will_continue=True, but NO partial_args (or empty list)
+  # This is the first chunk that Gemini 3 sends for streaming FC
+  chunk1_fc = types.FunctionCall(
+      name="my_tool",
+      id="fc_gemini3",
+      will_continue=True,
+      partial_args=first_chunk_partial_args,
+  )
+  chunk1_part = types.Part(
+      function_call=chunk1_fc,
+      thought_signature=b"test_sig_123",
+  )
+  chunk1 = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(role="model", parts=[chunk1_part])
+          )
+      ]
+  )
+
+  # Chunk 2: Middle chunk with partial_args, name is None
+  chunk2_fc = types.FunctionCall(
+      partial_args=[
+          types.PartialArg(json_path="$.document", string_value="Once upon ")
+      ],
+      will_continue=True,
+  )
+  chunk2 = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role="model", parts=[types.Part(function_call=chunk2_fc)]
+              )
+          )
+      ]
+  )
+
+  # Chunk 3: Another middle chunk continuing the string argument
+  chunk3_fc = types.FunctionCall(
+      partial_args=[
+          types.PartialArg(json_path="$.document", string_value="a time...")
+      ],
+      will_continue=True,
+  )
+  chunk3 = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role="model", parts=[types.Part(function_call=chunk3_fc)]
+              )
+          )
+      ]
+  )
+
+  # Chunk 4: Final chunk - no name, no partial_args, will_continue=False
+  # This signals the end of the streaming function call
+  chunk4_fc = types.FunctionCall(
+      will_continue=False,
+  )
+  chunk4 = types.GenerateContentResponse(
+      candidates=[
+          types.Candidate(
+              content=types.Content(
+                  role="model", parts=[types.Part(function_call=chunk4_fc)]
+              ),
+              finish_reason=types.FinishReason.STOP,
+          )
+      ]
+  )
+
+  # Process all chunks through aggregator
+  async def process():
+    results = []
+    for chunk in [chunk1, chunk2, chunk3, chunk4]:
+      async for response in aggregator.process_response(chunk):
+        results.append(response)
+    return results
+
+  processed_chunks = asyncio.run(process())
+
+  # All intermediate chunks should be marked as partial
+  assert all(chunk.partial for chunk in processed_chunks)
+
+  # Get final aggregated response
+  final_response = aggregator.close()
+
+  # Verify final aggregated response has the complete FC with accumulated args
+  assert final_response is not None
+  assert len(final_response.content.parts) == 1
+
+  fc_part = final_response.content.parts[0]
+  assert fc_part.function_call is not None
+  assert fc_part.function_call.name == "my_tool"
+  assert fc_part.function_call.id == "fc_gemini3"
+
+  # Verify the document argument was correctly accumulated
+  args = fc_part.function_call.args
+  assert "document" in args
+  assert (
+      args["document"] == "Once upon a time..."
+  )  # Concatenated from chunks 2 + 3
+
+  # Verify thought_signature was preserved from the first chunk
+  assert fc_part.thought_signature == b"test_sig_123"
+
+
+class PartialFunctionCallMockModel(BaseLlm):
+  """A mock model that yields partial function call events followed by final."""
+
+  model: str = "partial-fc-mock"
+  tool_call_count: int = 0
+
+  @classmethod
+  def supported_models(cls) -> list[str]:
+    return ["partial-fc-mock"]
+
+  async def generate_content_async(
+      self, llm_request: LlmRequest, stream: bool = False
+  ) -> AsyncGenerator[LlmResponse, None]:
+    """Yield partial FC events then final, simulating streaming behavior."""
+
+    # Check if this is a follow-up call (after function response)
+    has_function_response = False
+    for content in llm_request.contents:
+      for part in content.parts or []:
+        if part.function_response:
+          has_function_response = True
+          break
+
+    if has_function_response:
+      # Final response after function execution
+      yield LlmResponse(
+          content=types.Content(
+              role="model",
+              parts=[types.Part.from_text(text="Function executed once.")],
+          ),
+          partial=False,
+      )
+      return
+
+    # First call: yield partial FC events then final
+    # Partial event 1
+    yield LlmResponse(
+        content=types.Content(
+            role="model",
+            parts=[
+                types.Part.from_function_call(
+                    name="track_execution", args={"call_id": "partial_1"}
+                )
+            ],
+        ),
+        partial=True,
+    )
+
+    # Partial event 2
+    yield LlmResponse(
+        content=types.Content(
+            role="model",
+            parts=[
+                types.Part.from_function_call(
+                    name="track_execution", args={"call_id": "partial_2"}
+                )
+            ],
+        ),
+        partial=True,
+    )
+
+    # Final aggregated event (only this should trigger execution)
+    yield LlmResponse(
+        content=types.Content(
+            role="model",
+            parts=[
+                types.Part.from_function_call(
+                    name="track_execution", args={"call_id": "final"}
+                )
+            ],
+        ),
+        partial=False,
+        finish_reason=types.FinishReason.STOP,
+    )
+
+
+def test_partial_function_calls_not_executed_in_none_streaming_mode():
+  """Test that partial function call events are skipped regardless of mode."""
+  execution_log = []
+
+  def track_execution(call_id: str) -> str:
+    """A tool that logs each execution to verify call count."""
+    execution_log.append(call_id)
+    return f"Executed: {call_id}"
+
+  mock_model = PartialFunctionCallMockModel()
+
+  agent = Agent(
+      name="partial_fc_test_agent",
+      model=mock_model,
+      tools=[track_execution],
+  )
+
+  # Use StreamingMode.NONE to verify partial FCs are still skipped
+  run_config = RunConfig(streaming_mode=StreamingMode.NONE)
+
+  runner = InMemoryRunner(agent=agent)
+
+  session = runner.session_service.create_session_sync(
+      app_name=runner.app_name, user_id="test_user"
+  )
+
+  events = []
+  for event in runner.run(
+      user_id="test_user",
+      session_id=session.id,
+      new_message=types.Content(
+          role="user",
+          parts=[types.Part.from_text(text="Test partial FC handling")],
+      ),
+      run_config=run_config,
+  ):
+    events.append(event)
+
+  # Verify the tool was only executed once (from the final event)
+  assert (
+      len(execution_log) == 1
+  ), f"Expected 1 execution, got {len(execution_log)}: {execution_log}"
+  assert (
+      execution_log[0] == "final"
+  ), f"Expected 'final' execution, got: {execution_log[0]}"
+
+  # Verify partial events were yielded but not executed
+  partial_events = [e for e in events if e.partial]
+  assert (
+      len(partial_events) == 2
+  ), f"Expected 2 partial events, got {len(partial_events)}"
+
+  # Verify there's a function response event (from the final FC execution)
+  function_response_events = [
+      e
+      for e in events
+      if e.content
+      and e.content.parts
+      and any(p.function_response for p in e.content.parts)
+  ]
+  assert (
+      len(function_response_events) == 1
+  ), f"Expected 1 function response event, got {len(function_response_events)}"

tests/unittests/flows/llm_flows/test_request_confirmation.py

@@ -0,0 +1,409 @@
+# 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.
+
+import json
+from unittest.mock import patch
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows import functions
+from google.adk.flows.llm_flows.request_confirmation import request_processor
+from google.adk.models.llm_request import LlmRequest
+from google.adk.tools.tool_confirmation import ToolConfirmation
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+MOCK_TOOL_NAME = "mock_tool"
+MOCK_FUNCTION_CALL_ID = "mock_function_call_id"
+MOCK_CONFIRMATION_FUNCTION_CALL_ID = "mock_confirmation_function_call_id"
+
+
+def mock_tool(param1: str):
+  """Mock tool function."""
+  return f"Mock tool result with {param1}"
+
+
+@pytest.mark.asyncio
+async def test_request_confirmation_processor_no_events():
+  """Test that the processor returns None when there are no events."""
+  agent = LlmAgent(name="test_agent", tools=[mock_tool])
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  llm_request = LlmRequest()
+
+  events = []
+  async for event in request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    events.append(event)
+
+  assert not events
+
+
+@pytest.mark.asyncio
+async def test_request_confirmation_processor_no_function_responses():
+  """Test that the processor returns None when the user event has no function responses."""
+  agent = LlmAgent(name="test_agent", tools=[mock_tool])
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  llm_request = LlmRequest()
+
+  invocation_context.session.events.append(
+      Event(author="user", content=types.Content())
+  )
+
+  events = []
+  async for event in request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    events.append(event)
+
+  assert not events
+
+
+@pytest.mark.asyncio
+async def test_request_confirmation_processor_no_confirmation_function_response():
+  """Test that the processor returns None when no confirmation function response is present."""
+  agent = LlmAgent(name="test_agent", tools=[mock_tool])
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  llm_request = LlmRequest()
+
+  invocation_context.session.events.append(
+      Event(
+          author="user",
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_response=types.FunctionResponse(
+                          name="other_function", response={}
+                      )
+                  )
+              ]
+          ),
+      )
+  )
+
+  events = []
+  async for event in request_processor.run_async(
+      invocation_context, llm_request
+  ):
+    events.append(event)
+
+  assert not events
+
+
+@pytest.mark.asyncio
+async def test_request_confirmation_processor_success():
+  """Test the successful processing of a tool confirmation."""
+  agent = LlmAgent(name="test_agent", tools=[mock_tool])
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  llm_request = LlmRequest()
+
+  original_function_call = types.FunctionCall(
+      name=MOCK_TOOL_NAME, args={"param1": "test"}, id=MOCK_FUNCTION_CALL_ID
+  )
+
+  tool_confirmation = ToolConfirmation(confirmed=False, hint="test hint")
+  tool_confirmation_args = {
+      "originalFunctionCall": original_function_call.model_dump(
+          exclude_none=True, by_alias=True
+      ),
+      "toolConfirmation": tool_confirmation.model_dump(
+          by_alias=True, exclude_none=True
+      ),
+  }
+
+  # Event with the request for confirmation
+  invocation_context.session.events.append(
+      Event(
+          author="agent",
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_call=types.FunctionCall(
+                          name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
+                          args=tool_confirmation_args,
+                          id=MOCK_CONFIRMATION_FUNCTION_CALL_ID,
+                      )
+                  )
+              ]
+          ),
+      )
+  )
+
+  # Event with the user's confirmation
+  user_confirmation = ToolConfirmation(confirmed=True)
+  invocation_context.session.events.append(
+      Event(
+          author="user",
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_response=types.FunctionResponse(
+                          name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
+                          id=MOCK_CONFIRMATION_FUNCTION_CALL_ID,
+                          response={
+                              "response": user_confirmation.model_dump_json()
+                          },
+                      )
+                  )
+              ]
+          ),
+      )
+  )
+
+  expected_event = Event(
+      author="agent",
+      content=types.Content(
+          parts=[
+              types.Part(
+                  function_response=types.FunctionResponse(
+                      name=MOCK_TOOL_NAME,
+                      id=MOCK_FUNCTION_CALL_ID,
+                      response={"result": "Mock tool result with test"},
+                  )
+              )
+          ]
+      ),
+  )
+
+  with patch(
+      "google.adk.flows.llm_flows.functions.handle_function_call_list_async"
+  ) as mock_handle_function_call_list_async:
+    mock_handle_function_call_list_async.return_value = expected_event
+
+    events = []
+    async for event in request_processor.run_async(
+        invocation_context, llm_request
+    ):
+      events.append(event)
+
+    assert len(events) == 1
+    assert events[0] == expected_event
+
+    mock_handle_function_call_list_async.assert_called_once()
+    args, _ = mock_handle_function_call_list_async.call_args
+
+    assert list(args[1]) == [original_function_call]  # function_calls
+    assert args[3] == {MOCK_FUNCTION_CALL_ID}  # tools_to_confirm
+    assert (
+        args[4][MOCK_FUNCTION_CALL_ID] == user_confirmation
+    )  # tool_confirmation_dict
+
+
+@pytest.mark.asyncio
+async def test_request_confirmation_processor_tool_not_confirmed():
+  """Test when the tool execution is not confirmed by the user."""
+  agent = LlmAgent(name="test_agent", tools=[mock_tool])
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  llm_request = LlmRequest()
+
+  original_function_call = types.FunctionCall(
+      name=MOCK_TOOL_NAME, args={"param1": "test"}, id=MOCK_FUNCTION_CALL_ID
+  )
+
+  tool_confirmation = ToolConfirmation(confirmed=False, hint="test hint")
+  tool_confirmation_args = {
+      "originalFunctionCall": original_function_call.model_dump(
+          exclude_none=True, by_alias=True
+      ),
+      "toolConfirmation": tool_confirmation.model_dump(
+          by_alias=True, exclude_none=True
+      ),
+  }
+
+  invocation_context.session.events.append(
+      Event(
+          author="agent",
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_call=types.FunctionCall(
+                          name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
+                          args=tool_confirmation_args,
+                          id=MOCK_CONFIRMATION_FUNCTION_CALL_ID,
+                      )
+                  )
+              ]
+          ),
+      )
+  )
+
+  user_confirmation = ToolConfirmation(confirmed=False)
+  invocation_context.session.events.append(
+      Event(
+          author="user",
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_response=types.FunctionResponse(
+                          name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
+                          id=MOCK_CONFIRMATION_FUNCTION_CALL_ID,
+                          response={
+                              "response": user_confirmation.model_dump_json()
+                          },
+                      )
+                  )
+              ]
+          ),
+      )
+  )
+
+  with patch(
+      "google.adk.flows.llm_flows.functions.handle_function_call_list_async"
+  ) as mock_handle_function_call_list_async:
+    mock_handle_function_call_list_async.return_value = Event(
+        author="agent",
+        content=types.Content(
+            parts=[
+                types.Part(
+                    function_response=types.FunctionResponse(
+                        name=MOCK_TOOL_NAME,
+                        id=MOCK_FUNCTION_CALL_ID,
+                        response={"error": "Tool execution not confirmed"},
+                    )
+                )
+            ]
+        ),
+    )
+
+    events = []
+    async for event in request_processor.run_async(
+        invocation_context, llm_request
+    ):
+      events.append(event)
+
+    assert len(events) == 1
+    mock_handle_function_call_list_async.assert_called_once()
+    args, _ = mock_handle_function_call_list_async.call_args
+    assert (
+        args[4][MOCK_FUNCTION_CALL_ID] == user_confirmation
+    )  # tool_confirmation_dict
+
+
+@pytest.mark.asyncio
+async def test_request_confirmation_processor_finds_user_confirmation_in_default_branch():
+  """Processor finds user confirmation in default branch when agent is in child branch.
+
+  Setup:
+    - Agent in 'child_branch'.
+    - RequestConfirmation event in 'child_branch'.
+    - User response event in default branch (None).
+  Act: Run request_processor.
+  Assert: Processor finds the response and triggers tool execution.
+  """
+  # Arrange
+  agent = LlmAgent(name="test_agent", tools=[mock_tool])
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  # Set branch for the agent context
+  invocation_context.branch = "child_branch"
+  llm_request = LlmRequest()
+
+  original_function_call = types.FunctionCall(
+      name=MOCK_TOOL_NAME, args={"param1": "test"}, id=MOCK_FUNCTION_CALL_ID
+  )
+
+  tool_confirmation = ToolConfirmation(confirmed=False, hint="test hint")
+  tool_confirmation_args = {
+      "originalFunctionCall": original_function_call.model_dump(
+          exclude_none=True, by_alias=True
+      ),
+      "toolConfirmation": tool_confirmation.model_dump(
+          by_alias=True, exclude_none=True
+      ),
+  }
+
+  # Event with the request for confirmation (in child branch)
+  invocation_context.session.events.append(
+      Event(
+          author="agent",
+          branch="child_branch",
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_call=types.FunctionCall(
+                          name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
+                          args=tool_confirmation_args,
+                          id=MOCK_CONFIRMATION_FUNCTION_CALL_ID,
+                      )
+                  )
+              ]
+          ),
+      )
+  )
+
+  # Event with the user's confirmation (in default branch, branch=None)
+  user_confirmation = ToolConfirmation(confirmed=True)
+  invocation_context.session.events.append(
+      Event(
+          author="user",
+          branch=None,
+          content=types.Content(
+              parts=[
+                  types.Part(
+                      function_response=types.FunctionResponse(
+                          name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
+                          id=MOCK_CONFIRMATION_FUNCTION_CALL_ID,
+                          response={
+                              "response": user_confirmation.model_dump_json()
+                          },
+                      )
+                  )
+              ]
+          ),
+      )
+  )
+
+  expected_event = Event(
+      author="agent",
+      branch="child_branch",
+      content=types.Content(
+          parts=[
+              types.Part(
+                  function_response=types.FunctionResponse(
+                      name=MOCK_TOOL_NAME,
+                      id=MOCK_FUNCTION_CALL_ID,
+                      response={"result": "Mock tool result with test"},
+                  )
+              )
+          ]
+      ),
+  )
+
+  # Act & Assert
+  with patch(
+      "google.adk.flows.llm_flows.functions.handle_function_call_list_async"
+  ) as mock_handle_function_call_list_async:
+    mock_handle_function_call_list_async.return_value = expected_event
+
+    events = []
+    async for event in request_processor.run_async(
+        invocation_context, llm_request
+    ):
+      events.append(event)
+
+    assert len(events) == 1
+    assert events[0] == expected_event

tests/unittests/flows/llm_flows/test_tool_callbacks.py

@@ -0,0 +1,434 @@
+# 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.
+
+from typing import Any
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+from google.genai.types import Part
+from pydantic import BaseModel
+import pytest
+
+from ... import testing_utils
+
+
+def simple_function(input_str: str) -> str:
+  return {'result': input_str}
+
+
+def simple_function_with_error() -> str:
+  raise SystemError('simple_function_with_error')
+
+
+class MockBeforeToolCallback(BaseModel):
+  """Mock before tool callback."""
+
+  mock_response: dict[str, object]
+  modify_tool_request: bool = False
+
+  def __call__(
+      self,
+      tool: BaseTool,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+  ) -> dict[str, object]:
+    if self.modify_tool_request:
+      args['input_str'] = 'modified_input'
+      return None
+    return self.mock_response
+
+
+class MockAfterToolCallback(BaseModel):
+  """Mock after tool callback."""
+
+  mock_response: dict[str, object]
+  modify_tool_request: bool = False
+  modify_tool_response: bool = False
+
+  def __call__(
+      self,
+      tool: BaseTool,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+      tool_response: dict[str, Any] = None,
+  ) -> dict[str, object]:
+    if self.modify_tool_request:
+      args['input_str'] = 'modified_input'
+      return None
+    if self.modify_tool_response:
+      tool_response['result'] = 'modified_output'
+      return tool_response
+    return self.mock_response
+
+
+class MockOnToolErrorCallback(BaseModel):
+  """Mock on tool error callback."""
+
+  mock_response: dict[str, object]
+  modify_tool_response: bool = False
+
+  def __call__(
+      self,
+      tool: BaseTool,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+      error: Exception,
+  ) -> dict[str, object]:
+    if self.modify_tool_response:
+      return self.mock_response
+    return None
+
+
+def noop_callback(
+    **kwargs,
+) -> dict[str, object]:
+  pass
+
+
+def test_before_tool_callback():
+  """Test that the before_tool_callback is called before the tool is called."""
+  responses = [
+      types.Part.from_function_call(name='simple_function', args={}),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_tool_callback=MockBeforeToolCallback(
+          mock_response={'test': 'before_tool_callback'}
+      ),
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', Part.from_function_call(name='simple_function', args={})),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function', response={'test': 'before_tool_callback'}
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+def test_before_tool_callback_noop():
+  """Test that the before_tool_callback is a no-op when not overridden."""
+  responses = [
+      types.Part.from_function_call(
+          name='simple_function', args={'input_str': 'simple_function_call'}
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_tool_callback=noop_callback,
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      (
+          'root_agent',
+          Part.from_function_call(
+              name='simple_function', args={'input_str': 'simple_function_call'}
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function',
+              response={'result': 'simple_function_call'},
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+def test_before_tool_callback_modify_tool_request():
+  """Test that the before_tool_callback modifies the tool request."""
+  responses = [
+      types.Part.from_function_call(name='simple_function', args={}),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_tool_callback=MockBeforeToolCallback(
+          mock_response={'test': 'before_tool_callback'},
+          modify_tool_request=True,
+      ),
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', Part.from_function_call(name='simple_function', args={})),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function',
+              response={'result': 'modified_input'},
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+def test_after_tool_callback():
+  """Test that the after_tool_callback is called after the tool is called."""
+  responses = [
+      types.Part.from_function_call(
+          name='simple_function', args={'input_str': 'simple_function_call'}
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_tool_callback=MockAfterToolCallback(
+          mock_response={'test': 'after_tool_callback'}
+      ),
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      (
+          'root_agent',
+          Part.from_function_call(
+              name='simple_function', args={'input_str': 'simple_function_call'}
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function', response={'test': 'after_tool_callback'}
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+def test_after_tool_callback_noop():
+  """Test that the after_tool_callback is a no-op when not overridden."""
+  responses = [
+      types.Part.from_function_call(
+          name='simple_function', args={'input_str': 'simple_function_call'}
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_tool_callback=noop_callback,
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      (
+          'root_agent',
+          Part.from_function_call(
+              name='simple_function', args={'input_str': 'simple_function_call'}
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function',
+              response={'result': 'simple_function_call'},
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+def test_after_tool_callback_modify_tool_response():
+  """Test that the after_tool_callback modifies the tool response."""
+  responses = [
+      types.Part.from_function_call(
+          name='simple_function', args={'input_str': 'simple_function_call'}
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_tool_callback=MockAfterToolCallback(
+          mock_response={'result': 'after_tool_callback'},
+          modify_tool_response=True,
+      ),
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      (
+          'root_agent',
+          Part.from_function_call(
+              name='simple_function', args={'input_str': 'simple_function_call'}
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function',
+              response={'result': 'modified_output'},
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+async def test_on_tool_error_callback_tool_not_found_noop():
+  """Test that the on_tool_error_callback is a no-op when the tool is not found."""
+  responses = [
+      types.Part.from_function_call(
+          name='nonexistent_function',
+          args={'input_str': 'simple_function_call'},
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_tool_error_callback=noop_callback,
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  with pytest.raises(ValueError):
+    await runner.run_async('test')
+
+
+def test_on_tool_error_callback_tool_not_found_modify_tool_response():
+  """Test that the on_tool_error_callback modifies the tool response when the tool is not found."""
+  responses = [
+      types.Part.from_function_call(
+          name='nonexistent_function',
+          args={'input_str': 'simple_function_call'},
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_tool_error_callback=MockOnToolErrorCallback(
+          mock_response={'result': 'on_tool_error_callback_response'},
+          modify_tool_response=True,
+      ),
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      (
+          'root_agent',
+          Part.from_function_call(
+              name='nonexistent_function',
+              args={'input_str': 'simple_function_call'},
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='nonexistent_function',
+              response={'result': 'on_tool_error_callback_response'},
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+async def test_on_tool_error_callback_tool_error_noop():
+  """Test that the on_tool_error_callback is a no-op when the tool returns an error."""
+  responses = [
+      types.Part.from_function_call(
+          name='simple_function_with_error',
+          args={},
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_tool_error_callback=noop_callback,
+      tools=[simple_function_with_error],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  with pytest.raises(SystemError):
+    await runner.run_async('test')
+
+
+def test_on_tool_error_callback_tool_error_modify_tool_response():
+  """Test that the on_tool_error_callback modifies the tool response when the tool returns an error."""
+
+  async def async_on_tool_error_callback(
+      tool: BaseTool,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+      error: Exception,
+  ) -> dict[str, object]:
+    if tool.name == 'simple_function_with_error':
+      return {'result': 'async_on_tool_error_callback_response'}
+    return None
+
+  responses = [
+      types.Part.from_function_call(
+          name='simple_function_with_error',
+          args={},
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_tool_error_callback=async_on_tool_error_callback,
+      tools=[simple_function_with_error],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      (
+          'root_agent',
+          Part.from_function_call(
+              name='simple_function_with_error',
+              args={},
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function_with_error',
+              response={'result': 'async_on_tool_error_callback_response'},
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]

tests/unittests/flows/llm_flows/test_tool_telemetry.py

@@ -0,0 +1,99 @@
+# 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.
+
+from typing import Any
+from typing import Dict
+from typing import Optional
+from unittest import mock
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.functions import handle_function_calls_async
+from google.adk.telemetry import tracing
+from google.adk.tools.function_tool import FunctionTool
+from google.genai import types
+
+from ... import testing_utils
+
+
+async def invoke_tool() -> Optional[Event]:
+  def simple_fn(**kwargs) -> Dict[str, Any]:
+    return {'result': 'test'}
+
+  tool = FunctionTool(simple_fn)
+  model = testing_utils.MockModel.create(responses=[])
+  agent = Agent(
+      name='agent',
+      model=model,
+      tools=[tool],
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content=''
+  )
+  function_call = types.FunctionCall(name=tool.name, args={'a': 1, 'b': 2})
+  content = types.Content(parts=[types.Part(function_call=function_call)])
+  event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+      content=content,
+  )
+  tools_dict = {tool.name: tool}
+  return await handle_function_calls_async(
+      invocation_context,
+      event,
+      tools_dict,
+  )
+
+
+async def test_simple_function_with_mocked_tracer(monkeypatch):
+  mock_start_as_current_span_func = mock.Mock()
+  returned_context_manager_mock = mock.MagicMock()
+  returned_context_manager_mock.__enter__.return_value = mock.Mock(
+      name='span_mock'
+  )
+  mock_start_as_current_span_func.return_value = returned_context_manager_mock
+
+  monkeypatch.setattr(
+      tracing.tracer, 'start_as_current_span', mock_start_as_current_span_func
+  )
+
+  mock_adk_trace_tool_call = mock.Mock()
+  monkeypatch.setattr(
+      'google.adk.telemetry.tracing.trace_tool_call',
+      mock_adk_trace_tool_call,
+  )
+
+  event = await invoke_tool()
+  assert event is not None
+
+  event = await invoke_tool()
+  assert event is not None
+
+  expected_span_name = 'execute_tool simple_fn'
+
+  assert mock_start_as_current_span_func.call_count == 2
+  mock_start_as_current_span_func.assert_any_call(expected_span_name)
+
+  assert returned_context_manager_mock.__enter__.call_count == 2
+  assert returned_context_manager_mock.__exit__.call_count == 2
+
+  assert mock_adk_trace_tool_call.call_count == 2
+  for call_args_item in mock_adk_trace_tool_call.call_args_list:
+    kwargs = call_args_item.kwargs
+    assert kwargs['tool'].name == 'simple_fn'
+    assert kwargs['args'] == {'a': 1, 'b': 2}
+    assert 'function_response_event' in kwargs
+    assert kwargs['function_response_event'].content.parts[
+        0
+    ].function_response.response == {'result': 'test'}

tests/unittests/flows/llm_flows/test_transcription_manager.py

@@ -0,0 +1,220 @@
+# 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.
+
+from unittest.mock import AsyncMock
+from unittest.mock import Mock
+
+from google.adk.flows.llm_flows.transcription_manager import TranscriptionManager
+from google.genai import types
+import pytest
+
+from ... import testing_utils
+
+
+class TestTranscriptionManager:
+  """Test the TranscriptionManager class."""
+
+  def setup_method(self):
+    """Set up test fixtures."""
+    self.manager = TranscriptionManager()
+
+  @pytest.mark.asyncio
+  async def test_handle_input_transcription(self):
+    """Test handling user input transcription events."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Set up mock session service
+    mock_session_service = AsyncMock()
+    invocation_context.session_service = mock_session_service
+
+    # Create test transcription
+    transcription = types.Transcription(text='Hello from user')
+
+    # Handle transcription
+    await self.manager.handle_input_transcription(
+        invocation_context, transcription
+    )
+
+    # Verify session service was called
+    mock_session_service.append_event.assert_not_called()
+
+  @pytest.mark.asyncio
+  async def test_handle_output_transcription(self):
+    """Test handling model output transcription events."""
+    agent = testing_utils.create_test_agent()
+    invocation_context = await testing_utils.create_invocation_context(agent)
+
+    # Set up mock session service
+    mock_session_service = AsyncMock()
+    invocation_context.session_service = mock_session_service
+
+    # Create test transcription
+    transcription = types.Transcription(text='Hello from model')
+
+    # Handle transcription
+    await self.manager.handle_output_transcription(
+        invocation_context, transcription
+    )
+
+    # Verify session service was called
+    mock_session_service.append_event.assert_not_called()
+
+  @pytest.mark.asyncio
+  async def test_handle_multiple_transcriptions(self):
+    """Test handling multiple transcription events."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Set up mock session service
+    mock_session_service = AsyncMock()
+    invocation_context.session_service = mock_session_service
+
+    # Handle multiple input transcriptions
+    for i in range(3):
+      transcription = types.Transcription(text=f'User message {i}')
+      await self.manager.handle_input_transcription(
+          invocation_context, transcription
+      )
+
+    # Handle multiple output transcriptions
+    for i in range(2):
+      transcription = types.Transcription(text=f'Model response {i}')
+      await self.manager.handle_output_transcription(
+          invocation_context, transcription
+      )
+
+    # Verify session service was called for each transcription
+    assert mock_session_service.append_event.call_count == 0
+
+  def test_get_transcription_stats_empty_session(self):
+    """Test getting transcription statistics for empty session."""
+    invocation_context = Mock()
+    invocation_context.session.events = []
+
+    stats = self.manager.get_transcription_stats(invocation_context)
+
+    expected = {
+        'input_transcriptions': 0,
+        'output_transcriptions': 0,
+        'total_transcriptions': 0,
+    }
+    assert stats == expected
+
+  def test_get_transcription_stats_with_events(self):
+    """Test getting transcription statistics for session with events."""
+    invocation_context = Mock()
+
+    # Create mock events
+    input_event1 = Mock()
+    input_event1.input_transcription = types.Transcription(text='User 1')
+    input_event1.output_transcription = None
+
+    input_event2 = Mock()
+    input_event2.input_transcription = types.Transcription(text='User 2')
+    input_event2.output_transcription = None
+
+    output_event = Mock()
+    output_event.input_transcription = None
+    output_event.output_transcription = types.Transcription(
+        text='Model response'
+    )
+
+    regular_event = Mock()
+    regular_event.input_transcription = None
+    regular_event.output_transcription = None
+
+    invocation_context.session.events = [
+        input_event1,
+        output_event,
+        input_event2,
+        regular_event,
+    ]
+
+    stats = self.manager.get_transcription_stats(invocation_context)
+
+    expected = {
+        'input_transcriptions': 2,
+        'output_transcriptions': 1,
+        'total_transcriptions': 3,
+    }
+    assert stats == expected
+
+  def test_get_transcription_stats_missing_attributes(self):
+    """Test getting transcription statistics when events don't have transcription attributes."""
+    invocation_context = Mock()
+
+    # Create mock events and explicitly set transcription attributes to None
+    event1 = Mock()
+    event1.input_transcription = None
+    event1.output_transcription = None
+
+    event2 = Mock()
+    event2.input_transcription = None
+    event2.output_transcription = None
+
+    invocation_context.session.events = [event1, event2]
+
+    stats = self.manager.get_transcription_stats(invocation_context)
+
+    expected = {
+        'input_transcriptions': 0,
+        'output_transcriptions': 0,
+        'total_transcriptions': 0,
+    }
+    assert stats == expected
+
+  @pytest.mark.asyncio
+  async def test_transcription_event_fields(self):
+    """Test that transcription events have correct field values."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Set up mock session service
+    mock_session_service = AsyncMock()
+    invocation_context.session_service = mock_session_service
+
+    # Create test transcription with specific content
+    transcription = types.Transcription(
+        text='Test transcription content', finished=True
+    )
+
+    # Handle input transcription
+    await self.manager.handle_input_transcription(
+        invocation_context, transcription
+    )
+
+  @pytest.mark.asyncio
+  async def test_transcription_with_different_data_types(self):
+    """Test handling transcriptions with different data types."""
+    invocation_context = await testing_utils.create_invocation_context(
+        testing_utils.create_test_agent()
+    )
+
+    # Set up mock session service
+    mock_session_service = AsyncMock()
+    invocation_context.session_service = mock_session_service
+
+    # Test with transcription that has basic fields only
+    transcription = types.Transcription(
+        text='Advanced transcription', finished=True
+    )
+
+    # Handle transcription
+    await self.manager.handle_input_transcription(
+        invocation_context, transcription
+    )