chore: import upstream snapshot with attribution

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"
+  }