# 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.

"""Handles function calling for LLM flow."""

from __future__ import annotations

import asyncio
import base64
import binascii
from concurrent.futures import ThreadPoolExecutor
import contextvars
import copy
import inspect
import json
import logging
import threading
from typing import Any
from typing import AsyncGenerator
from typing import cast
from typing import Dict
from typing import Optional
from typing import TYPE_CHECKING

from google.adk.platform import uuid as platform_uuid
from google.adk.tools.computer_use.computer_use_tool import ComputerUseTool
from google.genai import types

from ...agents.active_streaming_tool import ActiveStreamingTool
from ...agents.live_request_queue import LiveRequestQueue
from ...auth.auth_tool import AuthConfig
from ...auth.auth_tool import AuthToolArguments
from ...events.event import Event
from ...events.event_actions import EventActions
from ...telemetry import _instrumentation
from ...telemetry.tracing import trace_merged_tool_calls
from ...telemetry.tracing import tracer
from ...tools.base_tool import BaseTool
from ...tools.tool_confirmation import ToolConfirmation
from ...tools.tool_context import ToolContext
from ...utils.context_utils import Aclosing

if TYPE_CHECKING:
  from ...agents.invocation_context import InvocationContext
  from ...agents.llm_agent import LlmAgent

AF_FUNCTION_CALL_ID_PREFIX = 'adk-'
REQUEST_EUC_FUNCTION_CALL_NAME = 'adk_request_credential'
REQUEST_CONFIRMATION_FUNCTION_CALL_NAME = 'adk_request_confirmation'
REQUEST_INPUT_FUNCTION_CALL_NAME = 'adk_request_input'

logger = logging.getLogger('google_adk.' + __name__)

# Global thread pool executors for running tools in background threads.
# This prevents blocking tools from blocking the event loop in Live API mode.
# Key is max_workers, value is the executor.
_TOOL_THREAD_POOLS: dict[int, ThreadPoolExecutor] = {}
_TOOL_THREAD_POOL_LOCK = threading.Lock()


def _detect_error_type_for_telemetry(
    tool: BaseTool,
    tool_context: ToolContext,
    function_response: Any,
) -> Optional[str]:
  """Detects an error type from a tool response for telemetry purposes.

  This does not modify the response. `_detect_error_in_response` is an optional
  per-tool hook accessed via `getattr` to avoid adding a public API on
  `BaseTool`. Any exception raised by the detector is logged and swallowed so
  that telemetry logic never breaks tool execution.

  Args:
    tool: The tool whose response is being inspected.
    tool_context: The tool context for the current invocation. Detection is
      skipped when the tool is requesting auth or confirmation.
    function_response: The raw response returned by the tool.

  Returns:
    The error type string reported by the tool's `_detect_error_in_response`
    hook, or `None` if no error was detected, no hook is defined, or the hook
    raised an exception.
  """
  try:
    if (
        tool_context.actions.requested_auth_configs
        or tool_context.actions.requested_tool_confirmations
    ):
      return None
    detector = getattr(tool, '_detect_error_in_response', None)
    if detector is None:
      return None
    return detector(function_response)
  except Exception:  # pylint: disable=broad-exception-caught
    # Never let telemetry logic break tool execution.
    logger.exception(
        'Error while detecting error type for telemetry from tool %r.',
        getattr(tool, 'name', tool),
    )
    return None


def _is_live_request_queue_annotation(param: inspect.Parameter) -> bool:
  """Check whether a parameter is annotated as LiveRequestQueue.

  Handles both the class itself and the string form produced by
  ``from __future__ import annotations``.
  """
  ann = param.annotation
  return ann is LiveRequestQueue or (
      isinstance(ann, str) and ann == 'LiveRequestQueue'
  )


def _get_tool_thread_pool(max_workers: int = 4) -> ThreadPoolExecutor:
  """Gets or creates a thread pool executor for tool execution.

  Args:
    max_workers: Maximum number of worker threads in the pool.

  Returns:
    A ThreadPoolExecutor with the specified max_workers.
  """
  if max_workers not in _TOOL_THREAD_POOLS:
    with _TOOL_THREAD_POOL_LOCK:
      if max_workers not in _TOOL_THREAD_POOLS:
        _TOOL_THREAD_POOLS[max_workers] = ThreadPoolExecutor(
            max_workers=max_workers, thread_name_prefix='adk_tool_executor'
        )
  return _TOOL_THREAD_POOLS[max_workers]


def _is_sync_tool(tool: BaseTool) -> bool:
  """Checks if a tool's underlying function is synchronous."""
  if not hasattr(tool, 'func'):
    return False
  func = tool.func
  return not (
      inspect.iscoroutinefunction(func)
      or inspect.isasyncgenfunction(func)
      or (
          hasattr(func, '__call__')
          and inspect.iscoroutinefunction(func.__call__)
      )
  )


async def _call_tool_in_thread_pool(
    tool: BaseTool,
    args: dict[str, Any],
    tool_context: ToolContext,
    max_workers: int = 4,
) -> Any:
  """Runs a tool in a thread pool to avoid blocking the event loop.

  For sync tools, this runs the tool's function directly in a background thread.
  For async tools, this creates a new event loop in the background thread and
  runs the async function there. This helps catch blocking I/O (like time.sleep,
  network calls, file I/O) that was mistakenly used inside async functions.

  Note: Due to Python's GIL, this does NOT help with pure Python CPU-bound code.
  Thread pool only helps when the GIL is released (blocking I/O, C extensions).

  Args:
    tool: The tool to execute.
    args: Arguments to pass to the tool.
    tool_context: The tool context.
    max_workers: Maximum number of worker threads in the pool.

  Returns:
    The result of running the tool.
  """
  from ...tools.function_tool import FunctionTool

  ctx = contextvars.copy_context()
  loop = asyncio.get_running_loop()
  executor = _get_tool_thread_pool(max_workers)

  if _is_sync_tool(tool):
    if isinstance(tool, FunctionTool):
      # For sync FunctionTool, call the underlying function directly.
      def run_sync_tool():
        args_to_call = tool._preprocess_args(args)
        signature = inspect.signature(tool.func)
        valid_params = {param for param in signature.parameters}
        if tool._context_param_name in valid_params:
          args_to_call[tool._context_param_name] = tool_context
        args_to_call = {
            k: v for k, v in args_to_call.items() if k in valid_params
        }
        mandatory_args = tool._get_mandatory_args()
        missing_mandatory_args = [
            arg for arg in mandatory_args if arg not in args_to_call
        ]
        if missing_mandatory_args:
          missing_mandatory_args_str = '\n'.join(missing_mandatory_args)
          error_str = (
              f'Invoking `{tool.name}()` failed as the following mandatory'
              ' input parameters are not present:\n'
              f'{missing_mandatory_args_str}\n'
              'You could retry calling this tool, but it is IMPORTANT for you'
              ' to provide all the mandatory parameters.'
          )
          return {'error': error_str}
        return tool.func(**args_to_call)

      return await loop.run_in_executor(
          executor, lambda: ctx.run(run_sync_tool)
      )
  else:
    # For async tools, run them in a new event loop in a background thread.
    # This helps when async functions contain blocking I/O (common user mistake)
    # that would otherwise block the main event loop.
    def run_async_tool_in_new_loop():
      # Create a new event loop for this thread
      return asyncio.run(tool.run_async(args=args, tool_context=tool_context))

    return await loop.run_in_executor(
        executor, lambda: ctx.run(run_async_tool_in_new_loop)
    )

  # Fall back to normal async execution for non-FunctionTool sync tools.
  return await tool.run_async(args=args, tool_context=tool_context)


def generate_client_function_call_id() -> str:
  return f'{AF_FUNCTION_CALL_ID_PREFIX}{platform_uuid.new_uuid()}'


def populate_client_function_call_id(model_response_event: Event) -> None:
  if not model_response_event.get_function_calls():
    return
  for function_call in model_response_event.get_function_calls():
    if not function_call.id:
      function_call.id = generate_client_function_call_id()


def remove_client_function_call_id(content: Optional[types.Content]) -> None:
  """Removes ADK-generated function call IDs from content before sending to LLM.

  Strips client-side function call/response IDs that start with 'adk-' prefix
  to avoid sending internal tracking IDs to the model.

  Args:
    content: Content containing function calls/responses to clean.
  """
  if content and content.parts:
    for part in content.parts:
      if (
          part.function_call
          and part.function_call.id
          and part.function_call.id.startswith(AF_FUNCTION_CALL_ID_PREFIX)
      ):
        part.function_call.id = None
      if (
          part.function_response
          and part.function_response.id
          and part.function_response.id.startswith(AF_FUNCTION_CALL_ID_PREFIX)
      ):
        part.function_response.id = None


def get_long_running_function_calls(
    function_calls: list[types.FunctionCall],
    tools_dict: dict[str, BaseTool],
) -> set[str]:
  long_running_tool_ids = set()
  for function_call in function_calls:
    if (
        function_call.name in tools_dict
        and tools_dict[function_call.name].is_long_running
    ):
      long_running_tool_ids.add(function_call.id)

  return long_running_tool_ids


def build_auth_request_event(
    invocation_context: InvocationContext,
    auth_requests: Dict[str, AuthConfig],
    *,
    author: Optional[str] = None,
    role: Optional[str] = None,
) -> Event:
  """Builds an auth request event with function calls for each auth request.

  This is a shared helper used by both tool-level auth (when a tool requests
  auth during execution) and toolset-level auth (before tool listing).

  Args:
    invocation_context: The invocation context.
    auth_requests: Dict mapping function_call_id to AuthConfig.
    author: The event author. Defaults to agent name.
    role: The content role. Defaults to None.

  Returns:
    Event with auth request function calls.
  """
  parts = []
  long_running_tool_ids = set()

  for function_call_id, auth_config in auth_requests.items():
    request_euc_function_call = types.FunctionCall(
        name=REQUEST_EUC_FUNCTION_CALL_NAME,
        id=generate_client_function_call_id(),
        args=AuthToolArguments(
            function_call_id=function_call_id,
            auth_config=auth_config,
        ).model_dump(mode='json', exclude_none=True, by_alias=True),
    )
    long_running_tool_ids.add(request_euc_function_call.id)
    parts.append(types.Part(function_call=request_euc_function_call))

  return Event(
      invocation_id=invocation_context.invocation_id,
      author=author or invocation_context.agent.name,
      branch=invocation_context.branch,
      content=types.Content(parts=parts, role=role),
      long_running_tool_ids=long_running_tool_ids,
  )


def generate_auth_event(
    invocation_context: InvocationContext,
    function_response_event: Event,
) -> Optional[Event]:
  """Generates an auth request event from a function response event.

  This is used for tool-level auth where a tool requests credentials during
  execution.

  Args:
    invocation_context: The invocation context.
    function_response_event: The function response event with auth requests.

  Returns:
    Event with auth request function calls, or None if no auth requested.
  """
  if not function_response_event.actions.requested_auth_configs:
    return None

  return build_auth_request_event(
      invocation_context,
      function_response_event.actions.requested_auth_configs,
      role=function_response_event.content.role,
  )


def generate_request_confirmation_event(
    invocation_context: InvocationContext,
    function_call_event: Event,
    function_response_event: Event,
) -> Optional[Event]:
  """Generates a request confirmation event from a function response event."""
  if not function_response_event.actions.requested_tool_confirmations:
    return None
  parts = []
  long_running_tool_ids = set()
  function_calls = function_call_event.get_function_calls()
  for (
      function_call_id,
      tool_confirmation,
  ) in function_response_event.actions.requested_tool_confirmations.items():
    original_function_call = next(
        (fc for fc in function_calls if fc.id == function_call_id), None
    )
    if not original_function_call:
      continue
    request_confirmation_function_call = types.FunctionCall(
        name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
        args={
            'originalFunctionCall': original_function_call.model_dump(
                exclude_none=True, by_alias=True
            ),
            'toolConfirmation': tool_confirmation.model_dump(
                by_alias=True, exclude_none=True
            ),
        },
    )
    request_confirmation_function_call.id = generate_client_function_call_id()
    long_running_tool_ids.add(request_confirmation_function_call.id)
    parts.append(types.Part(function_call=request_confirmation_function_call))

  return Event(
      invocation_id=invocation_context.invocation_id,
      author=invocation_context.agent.name,
      branch=invocation_context.branch,
      content=types.Content(parts=parts, role='model'),
      long_running_tool_ids=long_running_tool_ids,
  )


async def handle_function_calls_async(
    invocation_context: InvocationContext,
    function_call_event: Event,
    tools_dict: dict[str, BaseTool],
    filters: Optional[set[str]] = None,
    tool_confirmation_dict: Optional[dict[str, ToolConfirmation]] = None,
) -> Optional[Event]:
  """Calls the functions and returns the function response event."""
  function_calls = function_call_event.get_function_calls()
  return await handle_function_call_list_async(
      invocation_context,
      function_calls,
      tools_dict,
      filters,
      tool_confirmation_dict,
  )


async def handle_function_call_list_async(
    invocation_context: InvocationContext,
    function_calls: list[types.FunctionCall],
    tools_dict: dict[str, BaseTool],
    filters: Optional[set[str]] = None,
    tool_confirmation_dict: Optional[dict[str, ToolConfirmation]] = None,
) -> Optional[Event]:
  """Calls the functions and returns the function response event."""

  agent = invocation_context.agent

  # Filter function calls
  filtered_calls = [
      fc for fc in function_calls if not filters or fc.id in filters
  ]

  if not filtered_calls:
    return None

  # Create tasks for parallel execution
  tasks = [
      asyncio.create_task(
          _execute_single_function_call_async(
              invocation_context,
              function_call,
              tools_dict,
              agent,
              tool_confirmation_dict[function_call.id]
              if tool_confirmation_dict
              else None,
          )
      )
      for function_call in filtered_calls
  ]

  # Wait for all tasks to complete
  try:
    function_response_events = await asyncio.gather(*tasks)
  except Exception:
    for t in tasks:
      if not t.done():
        t.cancel()
    await asyncio.gather(*tasks, return_exceptions=True)
    raise

  # Filter out None results
  function_response_events = [
      event for event in function_response_events if event is not None
  ]

  if not function_response_events:
    return None

  merged_event = merge_parallel_function_response_events(
      function_response_events
  )

  if len(function_response_events) > 1:
    # this is needed for debug traces of parallel calls
    # individual response with tool.name is traced in __build_response_event
    # (we drop tool.name from span name here as this is merged event)
    with tracer.start_as_current_span('execute_tool (merged)'):
      trace_merged_tool_calls(
          response_event_id=merged_event.id,
          function_response_event=merged_event,
          invocation_context=invocation_context,
      )
  return merged_event


async def _execute_single_function_call_async(
    invocation_context: InvocationContext,
    function_call: types.FunctionCall,
    tools_dict: dict[str, BaseTool],
    agent: LlmAgent,
    tool_confirmation: Optional[ToolConfirmation] = None,
) -> Optional[Event]:
  """Execute a single function call with thread safety for state modifications."""

  async def _run_on_tool_error_callbacks(
      *,
      tool: BaseTool,
      tool_args: dict[str, Any],
      tool_context: ToolContext,
      error: Exception,
  ) -> Optional[dict[str, Any]]:
    """Runs the on_tool_error_callbacks for the given tool."""
    error_response = (
        await invocation_context.plugin_manager.run_on_tool_error_callback(
            tool=tool,
            tool_args=tool_args,
            tool_context=tool_context,
            error=error,
        )
    )
    if error_response is not None:
      return error_response

    for callback in agent.canonical_on_tool_error_callbacks:
      error_response = callback(
          tool=tool,
          args=tool_args,
          tool_context=tool_context,
          error=error,
      )
      if inspect.isawaitable(error_response):
        error_response = await error_response
      if error_response is not None:
        return error_response

    return None

  # Do not use "args" as the variable name, because it is a reserved keyword
  # in python debugger.
  # Make a deep copy to avoid being modified.
  function_args = (
      copy.deepcopy(function_call.args) if function_call.args else {}
  )
  detected_error_type: Optional[str] = None

  tool_context = _create_tool_context(
      invocation_context, function_call, tool_confirmation
  )

  try:
    tool = _get_tool(function_call, tools_dict)
  except ValueError as tool_error:
    tool = BaseTool(name=function_call.name, description='Tool not found')
    error_response = await _run_on_tool_error_callbacks(
        tool=tool,
        tool_args=function_args,
        tool_context=tool_context,
        error=tool_error,
    )
    if error_response is not None:
      return __build_response_event(
          tool, error_response, tool_context, invocation_context
      )
    else:
      raise tool_error

  async def _run_with_trace():
    nonlocal function_args, detected_error_type

    # Step 1: Check if plugin before_tool_callback overrides the function
    # response.
    function_response = (
        await invocation_context.plugin_manager.run_before_tool_callback(
            tool=tool, tool_args=function_args, tool_context=tool_context
        )
    )

    # Step 2: If no overrides are provided from the plugins, further run the
    # canonical callback.
    if function_response is None:
      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

    # Step 3: Otherwise, proceed calling the tool normally.
    if function_response is None:
      try:
        function_response = await __call_tool_async(
            tool, args=function_args, tool_context=tool_context
        )
      except Exception as tool_error:
        error_response = await _run_on_tool_error_callbacks(
            tool=tool,
            tool_args=function_args,
            tool_context=tool_context,
            error=tool_error,
        )
        if error_response is not None:
          function_response = error_response
        else:
          raise tool_error

    # Step 4: Check if plugin after_tool_callback overrides the function
    # response.
    altered_function_response = (
        await invocation_context.plugin_manager.run_after_tool_callback(
            tool=tool,
            tool_args=function_args,
            tool_context=tool_context,
            result=function_response,
        )
    )

    # Step 5: If no overrides are provided from the plugins, further run the
    # canonical after_tool_callbacks.
    if altered_function_response is None:
      for callback in agent.canonical_after_tool_callbacks:
        altered_function_response = callback(
            tool=tool,
            args=function_args,
            tool_context=tool_context,
            tool_response=function_response,
        )
        if inspect.isawaitable(altered_function_response):
          altered_function_response = await altered_function_response
        if altered_function_response:
          break

    # Step 6: If alternative response exists from after_tool_callback, use it
    # instead of the original function response.
    if altered_function_response is not None:
      function_response = altered_function_response

    if (
        tool.is_long_running or tool._defers_response
    ) and not function_response:
      # The tool either runs long (FR will arrive later via session
      # injection) or defers its response by design (e.g., the LlmAgent
      # wrapper for task delegation synthesizes the FR after the
      # sub-agent completes).  Either way, skip the auto-FR build when
      # the tool returned nothing.
      return None

    detected_error_type = _detect_error_type_for_telemetry(
        tool, tool_context, function_response
    )

    # Note: State deltas are not applied here - they are collected in
    # tool_context.actions.state_delta and applied later when the session
    # service processes the events

    # Builds the function response event.
    function_response_event = __build_response_event(
        tool, function_response, tool_context, invocation_context
    )
    return function_response_event

  async with _instrumentation.record_tool_execution(
      tool, agent, function_args, invocation_context=invocation_context
  ) as tel_ctx:
    tel_ctx.function_response_event = await _run_with_trace()
    tel_ctx.error_type = detected_error_type
    return tel_ctx.function_response_event


async def handle_function_calls_live(
    invocation_context: InvocationContext,
    function_call_event: Event,
    tools_dict: dict[str, BaseTool],
) -> Event | None:
  """Calls the functions and returns the function response event."""
  from ...agents.llm_agent import LlmAgent

  agent = cast(LlmAgent, invocation_context.agent)
  function_calls = function_call_event.get_function_calls()

  if not function_calls:
    return None

  # Create async lock for active_streaming_tools and active_non_blocking_tool_tasks modifications
  active_tools_lock = asyncio.Lock()

  # Create tasks for parallel execution
  tasks = [
      asyncio.create_task(
          _execute_single_function_call_live(
              invocation_context,
              function_call,
              tools_dict,
              agent,
              active_tools_lock,
          )
      )
      for function_call in function_calls
  ]

  # Wait for all tasks to complete
  try:
    function_response_events = await asyncio.gather(*tasks)
  except Exception:
    for t in tasks:
      if not t.done():
        t.cancel()
    await asyncio.gather(*tasks, return_exceptions=True)
    raise

  # Filter out None results
  function_response_events = [
      event for event in function_response_events if event is not None
  ]

  for event in function_response_events:
    event.live_session_id = function_call_event.live_session_id

  if not function_response_events:
    return None

  merged_event = merge_parallel_function_response_events(
      function_response_events
  )
  if len(function_response_events) > 1:
    # this is needed for debug traces of parallel calls
    # individual response with tool.name is traced in __build_response_event
    # (we drop tool.name from span name here as this is merged event)
    with tracer.start_as_current_span('execute_tool (merged)'):
      trace_merged_tool_calls(
          response_event_id=merged_event.id,
          function_response_event=merged_event,
          invocation_context=invocation_context,
      )
  return merged_event


async def _execute_single_function_call_live(
    invocation_context: InvocationContext,
    function_call: types.FunctionCall,
    tools_dict: dict[str, BaseTool],
    agent: LlmAgent,
    active_tools_lock: asyncio.Lock,
) -> Optional[Event]:
  """Execute a single function call for live mode with thread safety."""

  async def _run_on_tool_error_callbacks(
      *,
      tool: BaseTool,
      tool_args: dict[str, Any],
      tool_context: ToolContext,
      error: Exception,
  ) -> Optional[dict[str, Any]]:
    """Runs the on_tool_error_callbacks for the given tool."""
    error_response = (
        await invocation_context.plugin_manager.run_on_tool_error_callback(
            tool=tool,
            tool_args=tool_args,
            tool_context=tool_context,
            error=error,
        )
    )
    if error_response is not None:
      return error_response

    for callback in agent.canonical_on_tool_error_callbacks:
      error_response = callback(
          tool=tool,
          args=tool_args,
          tool_context=tool_context,
          error=error,
      )
      if inspect.isawaitable(error_response):
        error_response = await error_response
      if error_response is not None:
        return error_response

    return None

  # Do not use "args" as the variable name, because it is a reserved keyword
  # in python debugger.
  # Make a deep copy to avoid being modified.
  function_args = (
      copy.deepcopy(function_call.args) if function_call.args else {}
  )
  detected_error_type: Optional[str] = None

  tool_context = _create_tool_context(invocation_context, function_call)

  try:
    tool = _get_tool(function_call, tools_dict)
  except ValueError as tool_error:
    tool = BaseTool(name=function_call.name, description='Tool not found')
    error_response = await _run_on_tool_error_callbacks(
        tool=tool,
        tool_args=function_args,
        tool_context=tool_context,
        error=tool_error,
    )
    if error_response is not None:
      return __build_response_event(
          tool, error_response, tool_context, invocation_context
      )
    raise tool_error

  async def _run_with_trace() -> Optional[Event]:
    """Executes the tool with full lifecycle management and telemetry.

    This function orchestrates the tool execution pipeline, including:
    1. Running plugin and canonical before-tool callbacks.
    2. Executing the actual tool logic.
    3. Running plugin and canonical after-tool callbacks.
    4. Detecting error types for telemetry.
    5. Building the final FunctionResponse Event to be returned.
    """
    nonlocal function_args, detected_error_type

    # Do not use "args" as the variable name, because it is a reserved keyword
    # in python debugger.
    # Make a deep copy to avoid being modified.
    function_response = None

    # Step 1: Check if plugin before_tool_callback overrides the function
    # response.
    function_response = (
        await invocation_context.plugin_manager.run_before_tool_callback(
            tool=tool, tool_args=function_args, tool_context=tool_context
        )
    )

    # Step 2: If no overrides are provided from the plugins, further run the
    # canonical callback.
    if function_response is None:
      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

    # Step 3: Otherwise, proceed calling the tool normally.
    if function_response is None:
      try:
        function_response = await _process_function_live_helper(
            tool,
            tool_context,
            function_call,
            function_args,
            invocation_context,
            active_tools_lock,
        )
      except Exception as tool_error:
        error_response = await _run_on_tool_error_callbacks(
            tool=tool,
            tool_args=function_args,
            tool_context=tool_context,
            error=tool_error,
        )
        if error_response is not None:
          function_response = error_response
        else:
          raise tool_error

    # Step 4: Check if plugin after_tool_callback overrides the function
    # response.
    altered_function_response = (
        await invocation_context.plugin_manager.run_after_tool_callback(
            tool=tool,
            tool_args=function_args,
            tool_context=tool_context,
            result=function_response,
        )
    )

    # Step 5: If no overrides are provided from the plugins, further run the
    # canonical after_tool_callbacks.
    if altered_function_response is None:
      for callback in agent.canonical_after_tool_callbacks:
        altered_function_response = callback(
            tool=tool,
            args=function_args,
            tool_context=tool_context,
            tool_response=function_response,
        )
        if inspect.isawaitable(altered_function_response):
          altered_function_response = await altered_function_response
        if altered_function_response:
          break

    # Step 6: If alternative response exists from after_tool_callback, use it
    # instead of the original function response.
    if altered_function_response is not None:
      function_response = altered_function_response

    if (
        tool.is_long_running or tool._defers_response
    ) and not function_response:
      # The tool either runs long (FR will arrive later via session
      # injection) or defers its response by design.  Skip the auto-FR
      # build when the tool returned nothing.
      return None

    detected_error_type = _detect_error_type_for_telemetry(
        tool, tool_context, function_response
    )

    # Note: State deltas are not applied here - they are collected in
    # tool_context.actions.state_delta and applied later when the session
    # service processes the events

    # Builds the function response event.
    function_response_event = __build_response_event(
        tool, function_response, tool_context, invocation_context
    )
    return function_response_event

  is_streaming = hasattr(tool, 'func') and inspect.isasyncgenfunction(tool.func)
  is_non_blocking = not is_streaming and tool.response_scheduling is not None
  if is_non_blocking:
    task_key = f'{tool.name}_{function_call.id}'

    async def _background_task() -> None:
      try:
        async with _instrumentation.record_tool_execution(
            tool, agent, function_args, invocation_context=invocation_context
        ) as tel_ctx:
          function_response_event = await _run_with_trace()
          tel_ctx.function_response_event = function_response_event
          tel_ctx.error_type = detected_error_type

          if function_response_event:
            if (
                invocation_context.session_service
                and invocation_context.session
            ):
              await invocation_context.session_service.append_event(
                  session=invocation_context.session,
                  event=function_response_event,
              )
            if (
                invocation_context.live_request_queue
                and function_response_event.content
            ):
              invocation_context.live_request_queue.send_content(
                  function_response_event.content
              )
      except Exception:
        logger.exception('Error running non-blocking tool %s', tool.name)
      finally:
        async with active_tools_lock:
          if (
              invocation_context.active_non_blocking_tool_tasks
              and task_key in invocation_context.active_non_blocking_tool_tasks
          ):
            del invocation_context.active_non_blocking_tool_tasks[task_key]

    task = asyncio.create_task(_background_task())
    async with active_tools_lock:
      if invocation_context.active_non_blocking_tool_tasks is None:
        invocation_context.active_non_blocking_tool_tasks = {}
      invocation_context.active_non_blocking_tool_tasks[task_key] = task
    return None
  else:
    async with _instrumentation.record_tool_execution(
        tool, agent, function_args, invocation_context=invocation_context
    ) as tel_ctx:
      tel_ctx.function_response_event = await _run_with_trace()
      tel_ctx.error_type = detected_error_type
      return tel_ctx.function_response_event


async def _process_function_live_helper(
    tool,
    tool_context,
    function_call,
    function_args,
    invocation_context,
    active_tools_lock: asyncio.Lock,
):
  function_response = None
  # Check if this is a stop_streaming function call
  if (
      function_call.name == 'stop_streaming'
      and 'function_name' in function_args
  ):
    function_name = function_args['function_name']
    # Thread-safe access to active_streaming_tools
    async with active_tools_lock:
      active_tasks = invocation_context.active_streaming_tools
      if (
          active_tasks
          and function_name in active_tasks
          and active_tasks[function_name].task
          and not active_tasks[function_name].task.done()
      ):
        task = active_tasks[function_name].task
      else:
        task = None

    if task:
      task.cancel()
      try:
        # Wait for the task to be cancelled
        await asyncio.wait_for(task, timeout=1.0)
      except (asyncio.CancelledError, asyncio.TimeoutError):
        # Log the specific condition
        if task.cancelled():
          logging.info('Task %s was cancelled successfully', function_name)
        elif task.done():
          logging.info('Task %s completed during cancellation', function_name)
        else:
          logging.warning(
              'Task %s might still be running after cancellation timeout',
              function_name,
          )
          function_response = {
              'status': f'The task is not cancelled yet for {function_name}.'
          }
      if not function_response:
        # Clean up the reference under lock
        async with active_tools_lock:
          if (
              invocation_context.active_streaming_tools
              and function_name in invocation_context.active_streaming_tools
          ):
            invocation_context.active_streaming_tools[function_name].task = None
            invocation_context.active_streaming_tools[function_name].stream = (
                None
            )

        function_response = {
            'status': f'Successfully stopped streaming function {function_name}'
        }
    else:
      function_response = {
          'status': f'No active streaming function named {function_name} found'
      }
  elif hasattr(tool, 'func') and inspect.isasyncgenfunction(tool.func):
    # for streaming tool use case
    # we require the function to be an async generator function
    async def run_tool_and_update_queue(tool, function_args, tool_context):
      try:
        async with Aclosing(
            __call_tool_live(
                tool=tool,
                args=function_args,
                tool_context=tool_context,
                invocation_context=invocation_context,
            )
        ) as agen:
          async for result in agen:
            updated_content = _build_function_response_content(
                tool, result, tool_context.function_call_id
            )
            invocation_context.live_request_queue.send_content(
                updated_content, partial=True
            )
      except asyncio.CancelledError:
        raise  # Re-raise to properly propagate the cancellation

    task = asyncio.create_task(
        run_tool_and_update_queue(tool, function_args, tool_context)
    )

    async with active_tools_lock:

      if invocation_context.active_streaming_tools is None:
        invocation_context.active_streaming_tools = {}
      if tool.name in invocation_context.active_streaming_tools:
        invocation_context.active_streaming_tools[tool.name].task = task
      else:
        # Register the streaming tool lazily when the model calls it.
        invocation_context.active_streaming_tools[tool.name] = (
            ActiveStreamingTool(task=task)
        )
        logger.debug('Lazily registered streaming tool: %s', tool.name)

      # For input-streaming tools (those with `input_stream:
      # LiveRequestQueue`), create a dedicated LiveRequestQueue so
      # _send_to_model starts duplicating data to it. This also
      # handles re-invocation after stop_streaming reset .stream
      # to None.
      sig = inspect.signature(tool.func)
      if (
          'input_stream' in sig.parameters
          and _is_live_request_queue_annotation(sig.parameters['input_stream'])
      ):
        invocation_context.active_streaming_tools[tool.name].stream = (
            LiveRequestQueue()
        )

    # Immediately return a pending response.
    # This is required by current live model.
    function_response = {
        'status': (
            'The function is running asynchronously and the results are'
            ' pending.'
        )
    }
  else:
    # Check if we should run tools in thread pool to avoid blocking event loop
    thread_pool_config = invocation_context.run_config.tool_thread_pool_config
    if thread_pool_config is not None:
      function_response = await _call_tool_in_thread_pool(
          tool,
          args=function_args,
          tool_context=tool_context,
          max_workers=thread_pool_config.max_workers,
      )
    else:
      function_response = await __call_tool_async(
          tool, args=function_args, tool_context=tool_context
      )
  return function_response


def _get_tool(
    function_call: types.FunctionCall, tools_dict: dict[str, BaseTool]
):
  """Returns the tool corresponding to the function call."""
  if function_call.name not in tools_dict:
    available = list(tools_dict.keys())
    error_msg = (
        f"Tool '{function_call.name}' not found.\nAvailable tools:"
        f" {', '.join(available)}\n\nPossible causes:\n  1. LLM hallucinated"
        ' the function name - review agent instruction clarity\n  2. Tool not'
        ' registered - verify agent.tools list\n  3. Name mismatch - check for'
        ' typos\n\nSuggested fixes:\n  - Review agent instruction to ensure'
        ' tool usage is clear\n  - Verify tool is included in agent.tools'
        ' list\n  - Check for typos in function name'
    )
    raise ValueError(error_msg)

  return tools_dict[function_call.name]


def _create_tool_context(
    invocation_context: InvocationContext,
    function_call: types.FunctionCall,
    tool_confirmation: Optional[ToolConfirmation] = None,
):
  """Creates a ToolContext object."""
  return ToolContext(
      invocation_context=invocation_context,
      function_call_id=function_call.id,
      tool_confirmation=tool_confirmation,
  )


def _get_tool_and_context(
    invocation_context: InvocationContext,
    function_call: types.FunctionCall,
    tools_dict: dict[str, BaseTool],
    tool_confirmation: Optional[ToolConfirmation] = None,
):
  """Returns the tool and tool context corresponding to the function call."""
  tool = _get_tool(function_call, tools_dict)
  tool_context = _create_tool_context(
      invocation_context,
      function_call,
      tool_confirmation,
  )

  return (tool, tool_context)


def _try_decode_computer_use_image(
    tool: BaseTool,
    function_result: dict[str, object],
) -> Optional[list[types.FunctionResponsePart]]:
  """Decodes the image from the function result for a computer use tool.

  Args:
    tool: The tool that produced the function result.
    function_result: The dictionary containing the function's result. This
      dictionary may be modified in-place to remove the 'image' key if an image
      is successfully decoded.

  Returns:
    A list containing a `types.FunctionResponsePart` with the decoded image
    data, or None if no image was found or decoding failed.
  """

  if not isinstance(tool, ComputerUseTool) or not isinstance(
      function_result, dict
  ):
    return None

  if (
      'image' not in function_result
      or 'data' not in function_result['image']
      or 'mimetype' not in function_result['image']
  ):
    return None

  try:
    image_data = base64.b64decode(function_result['image']['data'])
    mime_type = function_result['image']['mimetype']

    part = types.FunctionResponsePart.from_bytes(
        data=image_data, mime_type=mime_type
    )

    del function_result['image']
    return [part]
  except (binascii.Error, ValueError):
    logger.exception('Failed to decode image from computer use tool')
    return None


async def __call_tool_live(
    tool: BaseTool,
    args: dict[str, object],
    tool_context: ToolContext,
    invocation_context: InvocationContext,
) -> AsyncGenerator[Event, None]:
  """Calls the tool asynchronously (awaiting the coroutine)."""
  async with Aclosing(
      tool._call_live(
          args=args,
          tool_context=tool_context,
          invocation_context=invocation_context,
      )
  ) as agen:
    async for item in agen:
      yield item


async def __call_tool_async(
    tool: BaseTool,
    args: dict[str, Any],
    tool_context: ToolContext,
) -> Any:
  """Calls the tool."""
  return await tool.run_async(args=args, tool_context=tool_context)


def __build_response_event(
    tool: BaseTool,
    function_result: dict[str, object],
    tool_context: ToolContext,
    invocation_context: InvocationContext,
) -> Event:
  # Capture the raw result for display purposes before any normalization.
  display_result = function_result

  # Specs requires the result to be a dict.
  if not isinstance(function_result, dict):
    function_result = {'result': function_result}

  function_response_parts = None
  if isinstance(tool, ComputerUseTool):
    function_response_parts = _try_decode_computer_use_image(
        tool, function_result
    )

  content = _build_function_response_content(
      tool,
      function_result,
      tool_context.function_call_id,
      function_response_parts,
  )

  # When summarization is skipped, ensure a displayable text part is added so
  # the tool's output is not lost in UIs that don't render function responses.
  # Control-flow tools (e.g. exit_loop) set skip_summarization but return no
  # meaningful output; their None result is normalized to {'result': None}, so
  # skip those to avoid emitting a noisy "null" text part.
  has_displayable_result = display_result is not None and display_result != {
      'result': None
  }
  if (
      tool_context.actions.skip_summarization
      and 'error' not in function_result
      and has_displayable_result
  ):
    if isinstance(display_result, str):
      result_text = display_result
    else:
      result_text = json.dumps(display_result, ensure_ascii=False, default=str)
    content.parts.append(types.Part.from_text(text=result_text))

  function_response_event = Event(
      invocation_id=invocation_context.invocation_id,
      author=invocation_context.agent.name,
      content=content,
      actions=tool_context.actions,
      branch=invocation_context.branch,
  )

  return function_response_event


def _build_function_response_content(
    tool: BaseTool,
    function_result: object,
    function_call_id: Optional[str],
    function_response_parts: Optional[list[types.FunctionResponsePart]] = None,
) -> types.Content:
  """Builds the content carrying a tool result as a FunctionResponse."""
  # Specs requires the result to be a dict.
  if not isinstance(function_result, dict):
    function_result = {'result': function_result}

  part_function_response = types.Part.from_function_response(
      name=tool.name,
      response=function_result,
      parts=function_response_parts,
  )
  part_function_response.function_response.id = function_call_id
  if tool.response_scheduling is not None:
    part_function_response.function_response.scheduling = (
        tool.response_scheduling
    )

  return types.Content(role='user', parts=[part_function_response])


def deep_merge_dicts(d1: dict, d2: dict) -> dict:
  """Recursively merges d2 into d1."""
  for key, value in d2.items():
    if key in d1 and isinstance(d1[key], dict) and isinstance(value, dict):
      d1[key] = deep_merge_dicts(d1[key], value)
    else:
      d1[key] = value
  return d1


def merge_parallel_function_response_events(
    function_response_events: list['Event'],
) -> 'Event':
  if not function_response_events:
    raise ValueError('No function response events provided.')

  if len(function_response_events) == 1:
    return function_response_events[0]
  merged_parts = []
  for event in function_response_events:
    if event.content:
      for part in event.content.parts or []:
        merged_parts.append(part)

  # Use the first event as the "base" for common attributes
  base_event = function_response_events[0]

  # Merge actions from all events
  merged_actions_data: dict[str, Any] = {}
  aggregated_ui_widgets = []
  for event in function_response_events:
    if event.actions:
      actions_dict = event.actions.model_dump(exclude_none=True, by_alias=True)
      ui_widgets = actions_dict.pop(
          'renderUiWidgets', None
      ) or actions_dict.pop('render_ui_widgets', None)
      if ui_widgets:
        aggregated_ui_widgets.extend(ui_widgets)

      # Use `by_alias=True` because it converts the model to a dictionary while respecting field aliases, ensuring that the enum fields are correctly handled without creating a duplicate.
      merged_actions_data = deep_merge_dicts(
          merged_actions_data,
          actions_dict,
      )

  if aggregated_ui_widgets:
    merged_actions_data['renderUiWidgets'] = aggregated_ui_widgets

  merged_actions = EventActions.model_validate(merged_actions_data)

  # Create the new merged event
  merged_event = Event(
      invocation_id=base_event.invocation_id,
      author=base_event.author,
      branch=base_event.branch,
      content=types.Content(role='user', parts=merged_parts),
      actions=merged_actions,  # Aggregated from all parallel events
      live_session_id=base_event.live_session_id,
  )

  # Use the base_event as the timestamp
  merged_event.timestamp = base_event.timestamp
  return merged_event


def find_event_by_function_call_id(
    events: list[Event],
    function_call_id: str,
) -> Optional[Event]:
  """Finds the function call event that matches the function call id."""
  for event in reversed(events):
    for function_call in event.get_function_calls():
      if function_call.id == function_call_id:
        return event
  return None


def find_matching_function_call(
    events: list[Event],
) -> Optional[Event]:
  """Finds the function call event that matches the function response id of the last event."""
  if not events:
    return None

  last_event = events[-1]
  function_responses = last_event.get_function_responses()
  if not function_responses:
    return None

  return find_event_by_function_call_id(events[:-1], function_responses[0].id)
