feat: Add tool callback chaining support for live streaming

**Highlights:**

-   **Callback Chaining:** Now supports a list of before and after callbacks, as in the non-live version. Callbacks are executed in order, stopping when one returns a response.
-   **Unit Tests:** Added new unit tests for before and after callbacks, async and sync versions, callback chains and mixed callbacks.
-   **Sample Agent:** Introduced a new example agent with multiple callbacks showing various behaviors: audit, security, validation, and enhancement. This provides practical usage examples.

PiperOrigin-RevId: 783884562

contributing/samples/live_tool_callbacks_agent/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 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 . import agent

contributing/samples/live_tool_callbacks_agent/agent.py

@@ -0,0 +1,269 @@
+# Copyright 2025 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 datetime import datetime
+import random
+import time
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from google.adk import Agent
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+
+
+def get_weather(location: str, tool_context: ToolContext) -> Dict[str, Any]:
+  """Get weather information for a location.
+  Args:
+    location: The city or location to get weather for.
+  Returns:
+    A dictionary containing weather information.
+  """
+  # Simulate weather data
+  temperatures = [-10, -5, 0, 5, 10, 15, 20, 25, 30, 35]
+  conditions = ["sunny", "cloudy", "rainy", "snowy", "windy"]
+
+  return {
+      "location": location,
+      "temperature": random.choice(temperatures),
+      "condition": random.choice(conditions),
+      "humidity": random.randint(30, 90),
+      "timestamp": datetime.now().isoformat(),
+  }
+
+
+async def calculate_async(operation: str, x: float, y: float) -> Dict[str, Any]:
+  """Perform async mathematical calculations.
+  Args:
+    operation: The operation to perform (add, subtract, multiply, divide).
+    x: First number.
+    y: Second number.
+  Returns:
+    A dictionary containing the calculation result.
+  """
+  # Simulate some async work
+  await asyncio.sleep(0.1)
+
+  operations = {
+      "add": x + y,
+      "subtract": x - y,
+      "multiply": x * y,
+      "divide": x / y if y != 0 else float("inf"),
+  }
+
+  result = operations.get(operation.lower(), "Unknown operation")
+
+  return {
+      "operation": operation,
+      "x": x,
+      "y": y,
+      "result": result,
+      "timestamp": datetime.now().isoformat(),
+  }
+
+
+def log_activity(message: str, tool_context: ToolContext) -> Dict[str, str]:
+  """Log an activity message with timestamp.
+  Args:
+    message: The message to log.
+  Returns:
+    A dictionary confirming the log entry.
+  """
+  if "activity_log" not in tool_context.state:
+    tool_context.state["activity_log"] = []
+
+  log_entry = {"timestamp": datetime.now().isoformat(), "message": message}
+  tool_context.state["activity_log"].append(log_entry)
+
+  return {
+      "status": "logged",
+      "entry": log_entry,
+      "total_entries": len(tool_context.state["activity_log"]),
+  }
+
+
+# Before tool callbacks
+def before_tool_audit_callback(
+    tool, args: Dict[str, Any], tool_context: ToolContext
+) -> Optional[Dict[str, Any]]:
+  """Audit callback that logs all tool calls before execution."""
+  print(f"🔍 AUDIT: About to call tool '{tool.name}' with args: {args}")
+
+  # Add audit info to tool context state
+  if "audit_log" not in tool_context.state:
+    tool_context.state["audit_log"] = []
+
+  tool_context.state["audit_log"].append({
+      "type": "before_call",
+      "tool_name": tool.name,
+      "args": args,
+      "timestamp": datetime.now().isoformat(),
+  })
+
+  # Return None to allow normal tool execution
+  return None
+
+
+def before_tool_security_callback(
+    tool, args: Dict[str, Any], tool_context: ToolContext
+) -> Optional[Dict[str, Any]]:
+  """Security callback that can block certain tool calls."""
+  # Example: Block weather requests for restricted locations
+  if tool.name == "get_weather" and args.get("location", "").lower() in [
+      "classified",
+      "secret",
+  ]:
+    print(
+        "🚫 SECURITY: Blocked weather request for restricted location:"
+        f" {args.get('location')}"
+    )
+    return {
+        "error": "Access denied",
+        "reason": "Location access is restricted",
+        "requested_location": args.get("location"),
+    }
+
+  # Allow other calls to proceed
+  return None
+
+
+async def before_tool_async_callback(
+    tool, args: Dict[str, Any], tool_context: ToolContext
+) -> Optional[Dict[str, Any]]:
+  """Async before callback that can add preprocessing."""
+  print(f"⚡ ASYNC BEFORE: Processing tool '{tool.name}' asynchronously")
+
+  # Simulate some async preprocessing
+  await asyncio.sleep(0.05)
+
+  # For calculation tool, we could add validation
+  if (
+      tool.name == "calculate_async"
+      and args.get("operation") == "divide"
+      and args.get("y") == 0
+  ):
+    print("🚫 VALIDATION: Prevented division by zero")
+    return {
+        "error": "Division by zero",
+        "operation": args.get("operation"),
+        "x": args.get("x"),
+        "y": args.get("y"),
+    }
+
+  return None
+
+
+# After tool callbacks
+def after_tool_enhancement_callback(
+    tool,
+    args: Dict[str, Any],
+    tool_context: ToolContext,
+    tool_response: Dict[str, Any],
+) -> Optional[Dict[str, Any]]:
+  """Enhance tool responses with additional metadata."""
+  print(f"✨ ENHANCE: Adding metadata to response from '{tool.name}'")
+
+  # Add enhancement metadata
+  enhanced_response = tool_response.copy()
+  enhanced_response.update({
+      "enhanced": True,
+      "enhancement_timestamp": datetime.now().isoformat(),
+      "tool_name": tool.name,
+      "execution_context": "live_streaming",
+  })
+
+  return enhanced_response
+
+
+async def after_tool_async_callback(
+    tool,
+    args: Dict[str, Any],
+    tool_context: ToolContext,
+    tool_response: Dict[str, Any],
+) -> Optional[Dict[str, Any]]:
+  """Async after callback for post-processing."""
+  print(
+      f"🔄 ASYNC AFTER: Post-processing response from '{tool.name}'"
+      " asynchronously"
+  )
+
+  # Simulate async post-processing
+  await asyncio.sleep(0.05)
+
+  # Add async processing metadata
+  processed_response = tool_response.copy()
+  processed_response.update({
+      "async_processed": True,
+      "processing_time": "0.05s",
+      "processor": "async_after_callback",
+  })
+
+  return processed_response
+
+
+import asyncio
+
+# Create the agent with tool callbacks
+root_agent = Agent(
+    model="gemini-2.0-flash-live-preview-04-09",  # for Vertex project
+    # model="gemini-2.0-flash-live-001",  # for AI studio key
+    name="tool_callbacks_agent",
+    description=(
+        "Live streaming agent that demonstrates tool callbacks functionality. "
+        "It can get weather, perform calculations, and log activities while "
+        "showing how before and after tool callbacks work in live mode."
+    ),
+    instruction="""
+      You are a helpful assistant that can:
+      1. Get weather information for any location using the get_weather tool
+      2. Perform mathematical calculations using the calculate_async tool
+      3. Log activities using the log_activity tool
+
+      Important behavioral notes:
+      - You have several callbacks that will be triggered before and after tool calls
+      - Before callbacks can audit, validate, or even block tool calls
+      - After callbacks can enhance or modify tool responses
+      - Some locations like "classified" or "secret" are restricted for weather requests
+      - Division by zero will be prevented by validation callbacks
+      - All your tool responses will be enhanced with additional metadata
+
+      When users ask you to test callbacks, explain what's happening with the callback system.
+      Be conversational and explain the callback behavior you observe.
+    """,
+    tools=[
+        get_weather,
+        calculate_async,
+        log_activity,
+    ],
+    # Multiple before tool callbacks (will be processed in order until one returns a response)
+    before_tool_callback=[
+        before_tool_audit_callback,
+        before_tool_security_callback,
+        before_tool_async_callback,
+    ],
+    # Multiple after tool callbacks (will be processed in order until one returns a response)
+    after_tool_callback=[
+        after_tool_enhancement_callback,
+        after_tool_async_callback,
+    ],
+    generate_content_config=types.GenerateContentConfig(
+        safety_settings=[
+            types.SafetySetting(
+                category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+                threshold=types.HarmBlockThreshold.OFF,
+            ),
+        ]
+    ),
+)

contributing/samples/live_tool_callbacks_agent/readme.md

@@ -0,0 +1,94 @@
+# Live Tool Callbacks Agent
+
+This sample demonstrates how tool callbacks work in live (bidirectional streaming) mode. It showcases both `before_tool_callback` and `after_tool_callback` functionality with multiple callback chains, async callbacks, and various callback behaviors.
+
+## Features Demonstrated
+
+### Before Tool Callbacks
+1. **Audit Callback**: Logs all tool calls before execution
+2. **Security Callback**: Can block tool calls based on security rules (e.g., restricted locations)
+3. **Async Validation Callback**: Performs async validation and can prevent invalid operations
+
+### After Tool Callbacks
+1. **Enhancement Callback**: Adds metadata to tool responses
+2. **Async Post-processing Callback**: Performs async post-processing of responses
+
+### Tools Available
+- `get_weather`: Get weather information for any location
+- `calculate_async`: Perform mathematical calculations asynchronously
+- `log_activity`: Log activities with timestamps
+
+## Testing Scenarios
+
+### 1. Basic Callback Flow
+```
+"What's the weather in New York?"
+```
+Watch the console output to see:
+- Audit logging before the tool call
+- Security check (will pass for New York)
+- Response enhancement after the tool call
+
+### 2. Security Blocking
+```
+"What's the weather in classified?"
+```
+The security callback will block this request and return an error response.
+
+### 3. Validation Prevention
+```
+"Calculate 10 divided by 0"
+```
+The async validation callback will prevent division by zero.
+
+### 4. Multiple Tool Calls
+```
+"Get weather for London and calculate 5 + 3"
+```
+See how callbacks work with multiple parallel tool calls.
+
+### 5. Callback Chain Testing
+```
+"Log this activity: Testing callback chains"
+```
+Observe how multiple callbacks in the chain are processed.
+
+## Getting Started
+
+1. **Start the ADK Web Server**
+   ```bash
+   adk web
+   ```
+
+2. **Access the ADK Web UI**
+   Navigate to `http://localhost:8000`
+
+3. **Select the Agent**
+   Choose "tool_callbacks_agent" from the dropdown in the top-left corner
+
+4. **Start Streaming**
+   Click the **Audio** or **Video** icon to begin streaming
+
+5. **Test Callbacks**
+   Try the testing scenarios above and watch both the chat responses and the console output to see callbacks in action
+
+## What to Observe
+
+- **Console Output**: Watch for callback logs with emojis:
+  - 🔍 AUDIT: Audit callback logging
+  - 🚫 SECURITY: Security callback blocking
+  - ⚡ ASYNC BEFORE: Async preprocessing
+  - ✨ ENHANCE: Response enhancement
+  - 🔄 ASYNC AFTER: Async post-processing
+
+- **Enhanced Responses**: Tool responses will include additional metadata added by after callbacks
+
+- **Error Handling**: Security blocks and validation errors will be returned as proper error responses
+
+## Technical Notes
+
+- This sample demonstrates that tool callbacks now work identically in both regular and live streaming modes
+- Multiple callbacks are supported and processed in order
+- Both sync and async callbacks are supported
+- Callbacks can modify, enhance, or block tool execution
+- The callback system provides full control over the tool execution pipeline 

src/google/adk/flows/llm_flows/functions.py

@@ -267,36 +267,27 @@ async def handle_function_calls_live(
       # in python debugger.
       function_args = function_call.args or {}
       function_response = None
-      # # Calls the tool if before_tool_callback does not exist or returns None.
-      # if agent.before_tool_callback:
-      #   function_response = agent.before_tool_callback(
-      #       tool, function_args, tool_context
-      #   )
-      if agent.before_tool_callback:
-        function_response = agent.before_tool_callback(
+
+      # Handle before_tool_callbacks - iterate through the canonical callback
+      # list
+      for callback in agent.canonical_before_tool_callbacks:
+        function_response = callback(
             tool=tool, args=function_args, tool_context=tool_context
         )
         if inspect.isawaitable(function_response):
           function_response = await function_response
+        if function_response:
+          break
 
-      if not function_response:
+      if function_response is None:
         function_response = await _process_function_live_helper(
             tool, tool_context, function_call, function_args, invocation_context
         )
 
       # Calls after_tool_callback if it exists.
-      # if agent.after_tool_callback:
-      #   new_response = agent.after_tool_callback(
-      #       tool,
-      #       function_args,
-      #       tool_context,
-      #       function_response,
-      #   )
-      #   if new_response:
-      #     function_response = new_response
       altered_function_response = None
-      if agent.after_tool_callback:
-        altered_function_response = agent.after_tool_callback(
+      for callback in agent.canonical_after_tool_callbacks:
+        altered_function_response = callback(
             tool=tool,
             args=function_args,
             tool_context=tool_context,
@@ -304,6 +295,9 @@ async def handle_function_calls_live(
         )
         if inspect.isawaitable(altered_function_response):
           altered_function_response = await altered_function_response
+        if altered_function_response:
+          break
+
       if altered_function_response is not None:
         function_response = altered_function_response
 

tests/unittests/flows/llm_flows/test_live_tool_callbacks.py

@@ -0,0 +1,388 @@
+# Copyright 2025 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 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"}