# 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 __future__ import annotations

from abc import ABC
import asyncio
import inspect
import logging
from typing import AsyncGenerator
from typing import Optional
from typing import TYPE_CHECKING

from google.adk.platform import time as platform_time
from google.genai import types
from websockets.exceptions import ConnectionClosed
from websockets.exceptions import ConnectionClosedOK

from . import _output_schema_processor
from . import functions
from ...agents.base_agent import BaseAgent
from ...agents.callback_context import CallbackContext
from ...agents.invocation_context import InvocationContext
from ...agents.live_request_queue import LiveRequestQueue
from ...agents.readonly_context import ReadonlyContext
from ...agents.run_config import StreamingMode
from ...auth.auth_handler import AuthHandler
from ...auth.auth_tool import AuthConfig
from ...auth.credential_manager import CredentialManager
from ...events.event import Event
from ...models.base_llm_connection import BaseLlmConnection
from ...models.llm_request import LlmRequest
from ...models.llm_response import LlmResponse
from ...telemetry import tracing
from ...telemetry.tracing import trace_call_llm
from ...telemetry.tracing import trace_send_data
from ...telemetry.tracing import tracer
from ...tools.base_toolset import BaseToolset
from ...tools.tool_context import ToolContext
from ...utils.context_utils import Aclosing
from .audio_cache_manager import AudioCacheManager
from .functions import build_auth_request_event

# Prefix used by toolset auth credential IDs
TOOLSET_AUTH_CREDENTIAL_ID_PREFIX = '_adk_toolset_auth_'

if TYPE_CHECKING:
  from ...agents.llm_agent import LlmAgent
  from ...models.base_llm import BaseLlm
  from ._base_llm_processor import BaseLlmRequestProcessor
  from ._base_llm_processor import BaseLlmResponseProcessor

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

_ADK_AGENT_NAME_LABEL_KEY = 'adk_agent_name'

# Timing configuration
DEFAULT_TRANSFER_AGENT_DELAY = 1.0
DEFAULT_TASK_COMPLETION_DELAY = 1.0

# Statistics configuration
DEFAULT_ENABLE_CACHE_STATISTICS = False


def _finalize_model_response_event(
    llm_request: LlmRequest,
    llm_response: LlmResponse,
    model_response_event: Event,
) -> Event:
  """Finalize and build the model response event from LLM response.

  Merges the LLM response data into the model response event and
  populates function call IDs and long-running tool information.

  Args:
    llm_request: The original LLM request.
    llm_response: The LLM response from the model.
    model_response_event: The base event to populate.

  Returns:
    The finalized Event with LLM response data merged in.
  """
  finalized_event = Event.model_validate({
      **model_response_event.model_dump(exclude_none=True),
      **llm_response.model_dump(exclude_none=True),
  })

  if finalized_event.content:
    function_calls = finalized_event.get_function_calls()
    if function_calls:
      functions.populate_client_function_call_id(finalized_event)
      finalized_event.long_running_tool_ids = (
          functions.get_long_running_function_calls(
              function_calls, llm_request.tools_dict
          )
      )

  return finalized_event


async def _resolve_toolset_auth(
    invocation_context: InvocationContext,
    agent: LlmAgent,
) -> AsyncGenerator[Event, None]:
  """Resolves authentication for toolsets before tool listing.

  For each toolset with auth configured via get_auth_config():
  - If credential is available, populate auth_config.exchanged_auth_credential
  - If credential is not available, yield auth request event and interrupt

  Args:
    invocation_context: The invocation context.
    agent: The LLM agent.

  Yields:
    Auth request events if any toolset needs authentication.
  """
  if not agent.tools:
    return

  pending_auth_requests: dict[str, AuthConfig] = {}
  callback_context = CallbackContext(invocation_context)

  for tool_union in agent.tools:
    if not isinstance(tool_union, BaseToolset):
      continue

    auth_config = tool_union.get_auth_config()
    if not auth_config:
      continue

    try:
      credential = await CredentialManager(auth_config).get_auth_credential(
          callback_context
      )
    except ValueError as e:
      # Validation errors from CredentialManager should be logged but not
      # block the flow - the toolset may still work without auth
      logger.warning(
          'Failed to get auth credential for toolset %s: %s',
          type(tool_union).__name__,
          e,
      )
      credential = None

    if credential:
      # Populate in-place for toolset to use in get_tools()
      auth_config.exchanged_auth_credential = credential
    else:
      # Need auth - will interrupt
      toolset_id = (
          f'{TOOLSET_AUTH_CREDENTIAL_ID_PREFIX}{type(tool_union).__name__}'
      )
      pending_auth_requests[toolset_id] = auth_config

  if not pending_auth_requests:
    return

  # Build auth requests dict with generated auth requests
  auth_requests = {
      credential_id: AuthHandler(auth_config).generate_auth_request()
      for credential_id, auth_config in pending_auth_requests.items()
  }

  # Yield event with auth requests using the shared helper
  yield build_auth_request_event(
      invocation_context,
      auth_requests,
      author=agent.name,
  )

  # Interrupt invocation
  invocation_context.end_invocation = True


async def _handle_before_model_callback(
    invocation_context: InvocationContext,
    llm_request: LlmRequest,
    model_response_event: Event,
) -> Optional[LlmResponse]:
  """Runs before-model callbacks (plugins then agent callbacks).

  Args:
    invocation_context: The invocation context.
    llm_request: The LLM request being built.
    model_response_event: The model response event for callback context.

  Returns:
    An LlmResponse if a callback short-circuits the LLM call, else None.
  """
  agent = invocation_context.agent

  callback_context = CallbackContext(
      invocation_context, event_actions=model_response_event.actions
  )

  # First run callbacks from the plugins.
  callback_response = (
      await invocation_context.plugin_manager.run_before_model_callback(
          callback_context=callback_context,
          llm_request=llm_request,
      )
  )
  if callback_response:
    return callback_response

  # If no overrides are provided from the plugins, further run the canonical
  # callbacks.
  if not agent.canonical_before_model_callbacks:
    return
  for callback in agent.canonical_before_model_callbacks:
    callback_response = callback(
        callback_context=callback_context, llm_request=llm_request
    )
    if inspect.isawaitable(callback_response):
      callback_response = await callback_response
    if callback_response:
      return callback_response


async def _handle_after_model_callback(
    invocation_context: InvocationContext,
    llm_response: LlmResponse,
    model_response_event: Event,
) -> Optional[LlmResponse]:
  """Runs after-model callbacks (plugins then agent callbacks).

  Also handles grounding metadata injection when google_search_agent is
  among the agent's tools.

  Args:
    invocation_context: The invocation context.
    llm_response: The LLM response to process.
    model_response_event: The model response event for callback context.

  Returns:
    An altered LlmResponse if a callback modifies it, else None.
  """
  agent = invocation_context.agent

  # Add grounding metadata to the response if needed.
  # TODO(b/448114567): Remove this function once the workaround is no longer needed.
  async def _maybe_add_grounding_metadata(
      response: Optional[LlmResponse] = None,
  ) -> Optional[LlmResponse]:
    readonly_context = ReadonlyContext(invocation_context)
    if (tools := invocation_context.canonical_tools_cache) is None:
      tools = await agent.canonical_tools(readonly_context)
      invocation_context.canonical_tools_cache = tools

    if not any(tool.name == 'google_search_agent' for tool in tools):
      return response
    ground_metadata = invocation_context.session.state.get(
        'temp:_adk_grounding_metadata', None
    )
    if not ground_metadata:
      return response

    if not response:
      response = llm_response
    response.grounding_metadata = ground_metadata
    return response

  callback_context = CallbackContext(
      invocation_context, event_actions=model_response_event.actions
  )

  # First run callbacks from the plugins.
  callback_response = (
      await invocation_context.plugin_manager.run_after_model_callback(
          callback_context=CallbackContext(invocation_context),
          llm_response=llm_response,
      )
  )
  if callback_response:
    return await _maybe_add_grounding_metadata(callback_response)

  # If no overrides are provided from the plugins, further run the canonical
  # callbacks.
  if not agent.canonical_after_model_callbacks:
    return await _maybe_add_grounding_metadata()
  for callback in agent.canonical_after_model_callbacks:
    callback_response = callback(
        callback_context=callback_context, llm_response=llm_response
    )
    if inspect.isawaitable(callback_response):
      callback_response = await callback_response
    if callback_response:
      return await _maybe_add_grounding_metadata(callback_response)
  return await _maybe_add_grounding_metadata()


async def _run_and_handle_error(
    response_generator: AsyncGenerator[LlmResponse, None],
    invocation_context: InvocationContext,
    llm_request: LlmRequest,
    model_response_event: Event,
) -> AsyncGenerator[LlmResponse, None]:
  """Wraps an LLM response generator with error callback handling.

  Runs the response generator within a tracing span. If an error occurs,
  runs on-model-error callbacks (plugins then agent callbacks). If a
  callback returns a response, that response is yielded instead of
  re-raising the error.

  Args:
    response_generator: The async generator producing LLM responses.
    invocation_context: The invocation context.
    llm_request: The LLM request.
    model_response_event: The model response event.

  Yields:
    LlmResponse objects from the generator.

  Raises:
    The original model error if no error callback handles it.
  """
  agent = invocation_context.agent
  if not hasattr(agent, 'canonical_on_model_error_callbacks'):
    raise TypeError(
        'Expected agent to have canonical_on_model_error_callbacks'
        f' attribute, but got {type(agent)}'
    )

  async def _run_on_model_error_callbacks(
      *,
      callback_context: CallbackContext,
      llm_request: LlmRequest,
      error: Exception,
  ) -> Optional[LlmResponse]:
    error_response = (
        await invocation_context.plugin_manager.run_on_model_error_callback(
            callback_context=callback_context,
            llm_request=llm_request,
            error=error,
        )
    )
    if error_response is not None:
      return error_response

    for callback in agent.canonical_on_model_error_callbacks:
      error_response = callback(
          callback_context=callback_context,
          llm_request=llm_request,
          error=error,
      )
      if inspect.isawaitable(error_response):
        error_response = await error_response
      if error_response is not None:
        return error_response

    return None

  try:
    async with Aclosing(response_generator) as agen:
      async with tracing.use_inference_span(
          llm_request,
          invocation_context,
          model_response_event,
      ) as gc_span:
        async for llm_response in agen:
          if gc_span:
            tracing.trace_inference_result(
                gc_span,
                llm_response,
            )
          yield llm_response
  except Exception as model_error:
    callback_context = CallbackContext(
        invocation_context, event_actions=model_response_event.actions
    )
    error_response = await _run_on_model_error_callbacks(
        callback_context=callback_context,
        llm_request=llm_request,
        error=model_error,
    )
    if error_response is not None:
      yield error_response
    else:
      raise model_error


async def _process_agent_tools(
    invocation_context: InvocationContext,
    llm_request: LlmRequest,
) -> None:
  """Process the agent's tools and populate ``llm_request.tools_dict``.

  Iterates over the agent's ``tools`` list, converts each tool union
  (callable, BaseTool, or BaseToolset) into resolved ``BaseTool``
  instances, and calls ``process_llm_request`` on each to register
  tool declarations in the request.

  After this function returns, ``llm_request.tools_dict`` maps tool
  names to ``BaseTool`` instances ready for function call dispatch.

  Args:
    invocation_context: The invocation context (``agent`` is read
      from ``invocation_context.agent``).
    llm_request: The LLM request to populate with tool declarations.
  """
  agent = invocation_context.agent
  if not hasattr(agent, 'tools') or not agent.tools:
    return

  multiple_tools = len(agent.tools) > 1
  model = agent.canonical_model
  for tool_union in agent.tools:
    tool_context = ToolContext(invocation_context)

    # If it's a toolset, process it first
    if isinstance(tool_union, BaseToolset):
      await tool_union.process_llm_request(
          tool_context=tool_context, llm_request=llm_request
      )

    from ...agents.llm_agent import _convert_tool_union_to_tools

    # Then process all tools from this tool union
    tools = await _convert_tool_union_to_tools(
        tool_union,
        ReadonlyContext(invocation_context),
        model,
        multiple_tools,
    )
    for tool in tools:
      await tool.process_llm_request(
          tool_context=tool_context, llm_request=llm_request
      )


class BaseLlmFlow(ABC):
  """A basic flow that calls the LLM in a loop until a final response is generated.

  This flow ends when it transfers to another agent.
  """

  def __init__(self):
    self.request_processors: list[BaseLlmRequestProcessor] = []
    self.response_processors: list[BaseLlmResponseProcessor] = []

    # Initialize configuration and managers
    self.audio_cache_manager = AudioCacheManager()

  async def run_live(
      self,
      invocation_context: InvocationContext,
  ) -> AsyncGenerator[Event, None]:
    """Runs the flow using live api."""
    llm_request = LlmRequest()
    event_id = Event.new_id()

    # Preprocess before calling the LLM.
    async with Aclosing(
        self._preprocess_async(invocation_context, llm_request)
    ) as agen:
      async for event in agen:
        yield event
    if invocation_context.end_invocation:
      return

    llm = self.__get_llm(invocation_context)
    logger.debug(
        'Establishing live connection for agent: %s with llm request: %s',
        invocation_context.agent.name,
        llm_request,
    )

    attempt = 1
    while True:
      try:
        # On subsequent attempts, use the saved token to reconnect
        if invocation_context.live_session_resumption_handle:
          logger.info('Attempting to reconnect (Attempt %s)...', attempt)
          attempt += 1
          if not llm_request.live_connect_config:
            llm_request.live_connect_config = types.LiveConnectConfig()
          if not llm_request.live_connect_config.session_resumption:
            llm_request.live_connect_config.session_resumption = (
                types.SessionResumptionConfig()
            )
          llm_request.live_connect_config.session_resumption.handle = (
              invocation_context.live_session_resumption_handle
          )
          llm_request.live_connect_config.session_resumption.transparent = True

        logger.info(
            'Establishing live connection for agent: %s',
            invocation_context.agent.name,
        )
        async with llm.connect(llm_request) as llm_connection:
          if llm_request.contents:
            # Sends the conversation history to the model.
            with tracer.start_as_current_span('send_data'):
              # Combine regular contents with audio/transcription from session
              logger.debug('Sending history to model: %s', llm_request.contents)
              await llm_connection.send_history(llm_request.contents)
              trace_send_data(
                  invocation_context, event_id, llm_request.contents
              )

          send_task = asyncio.create_task(
              self._send_to_model(llm_connection, invocation_context)
          )

          try:
            async with Aclosing(
                self._receive_from_model(
                    llm_connection,
                    event_id,
                    invocation_context,
                    llm_request,
                )
            ) as agen:
              async for event in agen:
                # Empty event means the queue is closed.
                if not event:
                  break
                logger.debug('Receive new event: %s', event)
                yield event
                # send back the function response to models
                if event.get_function_responses():
                  logger.debug(
                      'Sending back last function response event: %s', event
                  )
                  invocation_context.live_request_queue.send_content(
                      event.content
                  )
                # We handle agent transfer here in `run_live` rather than
                # in `_postprocess_live` to prevent duplication of function
                # response processing. If agent transfer were handled in
                # `_postprocess_live`, events yielded from child agent's
                # `run_live` would bubble up to parent agent's `run_live`,
                # causing `event.get_function_responses()` to be true in both
                # child and parent, and `send_content()` to be called twice for
                # the same function response. By handling agent transfer here,
                # we ensure that only child agent processes its own function
                # responses after the transfer.
                if (
                    event.content
                    and event.content.parts
                    and event.content.parts[0].function_response
                    and event.content.parts[0].function_response.name
                    == 'transfer_to_agent'
                ):
                  await asyncio.sleep(DEFAULT_TRANSFER_AGENT_DELAY)
                  # cancel the tasks that belongs to the closed connection.
                  send_task.cancel()
                  logger.debug('Closing live connection')
                  await llm_connection.close()
                  logger.debug('Live connection closed.')
                  # transfer to the sub agent.
                  transfer_to_agent = event.actions.transfer_to_agent
                  if transfer_to_agent:
                    logger.debug('Transferring to agent: %s', transfer_to_agent)
                    agent_to_run = self._get_agent_to_run(
                        invocation_context, transfer_to_agent
                    )
                    async with Aclosing(
                        agent_to_run.run_live(invocation_context)
                    ) as agen:
                      async for item in agen:
                        yield item
                if (
                    event.content
                    and event.content.parts
                    and event.content.parts[0].function_response
                    and event.content.parts[0].function_response.name
                    == 'task_completed'
                ):
                  # this is used for sequential agent to signal the end of the agent.
                  await asyncio.sleep(DEFAULT_TASK_COMPLETION_DELAY)
                  # cancel the tasks that belongs to the closed connection.
                  send_task.cancel()
                  return
          finally:
            # Clean up
            if not send_task.done():
              send_task.cancel()
            try:
              await send_task
            except asyncio.CancelledError:
              pass
      except (ConnectionClosed, ConnectionClosedOK) as e:
        # when the session timeout, it will just close and not throw exception.
        # so this is for bad cases
        logger.error('Connection closed: %s.', e)
        raise
      except Exception as e:
        logger.error(
            'An unexpected error occurred in live flow: %s', e, exc_info=True
        )
        raise

  async def _send_to_model(
      self,
      llm_connection: BaseLlmConnection,
      invocation_context: InvocationContext,
  ):
    """Sends data to model."""
    while True:
      live_request_queue = invocation_context.live_request_queue
      live_request = await live_request_queue.get()
      # duplicate the live_request to all the active streams
      logger.debug(
          'Sending live request %s to active streams: %s',
          live_request,
          invocation_context.active_streaming_tools,
      )
      if invocation_context.active_streaming_tools:
        for active_streaming_tool in (
            invocation_context.active_streaming_tools
        ).values():
          if active_streaming_tool.stream:
            active_streaming_tool.stream.send(live_request)
      # Yield to event loop for cooperative multitasking
      await asyncio.sleep(0)

      if live_request.close:
        await llm_connection.close()
        return

      if live_request.activity_start:
        await llm_connection.send_realtime(types.ActivityStart())
      elif live_request.activity_end:
        await llm_connection.send_realtime(types.ActivityEnd())
      elif live_request.blob:
        # Cache input audio chunks before flushing
        self.audio_cache_manager.cache_audio(
            invocation_context, live_request.blob, cache_type='input'
        )

        await llm_connection.send_realtime(live_request.blob)

      if live_request.content:
        content = live_request.content
        # Persist user text content to session (similar to non-live mode)
        # Skip function responses - they are already handled separately
        is_function_response = content.parts and any(
            part.function_response for part in content.parts
        )
        if not is_function_response:
          if not content.role:
            content.role = 'user'
          user_content_event = Event(
              id=Event.new_id(),
              invocation_id=invocation_context.invocation_id,
              author='user',
              content=content,
          )
          await invocation_context.session_service.append_event(
              session=invocation_context.session,
              event=user_content_event,
          )
        await llm_connection.send_content(live_request.content)

  async def _receive_from_model(
      self,
      llm_connection: BaseLlmConnection,
      event_id: str,
      invocation_context: InvocationContext,
      llm_request: LlmRequest,
  ) -> AsyncGenerator[Event, None]:
    """Receive data from model and process events using BaseLlmConnection."""

    def get_author_for_event(llm_response):
      """Get the author of the event.

      When the model returns transcription, the author is "user". Otherwise, the
      author is the agent name(not 'model').

      Args:
        llm_response: The LLM response from the LLM call.
      """
      if (
          llm_response
          and llm_response.content
          and llm_response.content.role == 'user'
      ):
        return 'user'
      else:
        return invocation_context.agent.name

    try:
      while True:
        async with Aclosing(llm_connection.receive()) as agen:
          async for llm_response in agen:
            if llm_response.live_session_resumption_update:
              logger.info(
                  'Update session resumption handle:'
                  f' {llm_response.live_session_resumption_update}.'
              )
              invocation_context.live_session_resumption_handle = (
                  llm_response.live_session_resumption_update.new_handle
              )
            model_response_event = Event(
                id=Event.new_id(),
                invocation_id=invocation_context.invocation_id,
                author=get_author_for_event(llm_response),
            )

            async with Aclosing(
                self._postprocess_live(
                    invocation_context,
                    llm_request,
                    llm_response,
                    model_response_event,
                )
            ) as agen:
              async for event in agen:
                # Cache output audio chunks from model responses
                # TODO: support video data
                if (
                    invocation_context.run_config.save_live_blob
                    and event.content
                    and event.content.parts
                    and event.content.parts[0].inline_data
                    and event.content.parts[0].inline_data.mime_type.startswith(
                        'audio/'
                    )
                ):
                  audio_blob = types.Blob(
                      data=event.content.parts[0].inline_data.data,
                      mime_type=event.content.parts[0].inline_data.mime_type,
                  )
                  self.audio_cache_manager.cache_audio(
                      invocation_context, audio_blob, cache_type='output'
                  )

                yield event
        # Give opportunity for other tasks to run.
        await asyncio.sleep(0)
    except ConnectionClosedOK:
      pass

  async def run_async(
      self, invocation_context: InvocationContext
  ) -> AsyncGenerator[Event, None]:
    """Runs the flow."""
    while True:
      last_event = None
      async with Aclosing(self._run_one_step_async(invocation_context)) as agen:
        async for event in agen:
          last_event = event
          yield event
      if not last_event or last_event.is_final_response() or last_event.partial:
        if last_event and last_event.partial:
          logger.warning('The last event is partial, which is not expected.')
        break

  async def _run_one_step_async(
      self,
      invocation_context: InvocationContext,
  ) -> AsyncGenerator[Event, None]:
    """One step means one LLM call."""
    llm_request = LlmRequest()

    # Preprocess before calling the LLM.
    async with Aclosing(
        self._preprocess_async(invocation_context, llm_request)
    ) as agen:
      async for event in agen:
        yield event
    if invocation_context.end_invocation:
      return

    # Resume the LLM agent based on the last event from the current branch.
    # 1. User content: continue the normal flow
    # 2. Function call: call the tool and get the response event.
    events = invocation_context._get_events(
        current_invocation=True, current_branch=True
    )

    # Long running tool calls should have been handled before this point.
    # If there are still long running tool calls, it means the agent is paused
    # before, and its branch hasn't been resumed yet.
    if (
        invocation_context.is_resumable
        and events
        and len(events) > 1
        # TODO: here we are using the last 2 events to decide whether to pause
        # the invocation. But this is just being optimistic, we should find a
        # way to pause when the long running tool call is followed by more than
        # one text responses.
        and (
            invocation_context.should_pause_invocation(events[-1])
            or invocation_context.should_pause_invocation(events[-2])
        )
    ):
      return

    if (
        invocation_context.is_resumable
        and events
        and events[-1].get_function_calls()
    ):
      model_response_event = events[-1]
      async with Aclosing(
          self._postprocess_handle_function_calls_async(
              invocation_context, model_response_event, llm_request
          )
      ) as agen:
        async for event in agen:
          event.id = Event.new_id()
          yield event
        return

    # Calls the LLM.
    model_response_event = Event(
        id=Event.new_id(),
        invocation_id=invocation_context.invocation_id,
        author=invocation_context.agent.name,
        branch=invocation_context.branch,
    )
    async with Aclosing(
        self._call_llm_async(
            invocation_context, llm_request, model_response_event
        )
    ) as agen:
      async for llm_response in agen:
        # Postprocess after calling the LLM.
        async with Aclosing(
            self._postprocess_async(
                invocation_context,
                llm_request,
                llm_response,
                model_response_event,
            )
        ) as agen:
          async for event in agen:
            # Update the mutable event id to avoid conflict
            model_response_event.id = Event.new_id()
            model_response_event.timestamp = platform_time.get_time()
            yield event

  async def _preprocess_async(
      self, invocation_context: InvocationContext, llm_request: LlmRequest
  ) -> AsyncGenerator[Event, None]:
    agent = invocation_context.agent
    if not hasattr(agent, 'tools') or not hasattr(agent, 'canonical_model'):
      raise TypeError(
          'Expected agent to have tools and canonical_model attributes,'
          f' but got {type(agent)}'
      )

    # Runs processors.
    for processor in self.request_processors:
      async with Aclosing(
          processor.run_async(invocation_context, llm_request)
      ) as agen:
        async for event in agen:
          yield event

    # Resolve toolset authentication before tool listing.
    # This ensures credentials are ready before get_tools() is called.
    async with Aclosing(
        self._resolve_toolset_auth(invocation_context, agent)
    ) as agen:
      async for event in agen:
        yield event

    if invocation_context.end_invocation:
      return

    # Run processors for tools.
    await _process_agent_tools(invocation_context, llm_request)

  async def _postprocess_async(
      self,
      invocation_context: InvocationContext,
      llm_request: LlmRequest,
      llm_response: LlmResponse,
      model_response_event: Event,
  ) -> AsyncGenerator[Event, None]:
    """Postprocess after calling the LLM.

    Args:
      invocation_context: The invocation context.
      llm_request: The original LLM request.
      llm_response: The LLM response from the LLM call.
      model_response_event: A mutable event for the LLM response.

    Yields:
      A generator of events.
    """

    # Runs processors.
    async with Aclosing(
        self._postprocess_run_processors_async(invocation_context, llm_response)
    ) as agen:
      async for event in agen:
        yield event

    # Skip the model response event if there is no content and no error code.
    # This is needed for the code executor to trigger another loop.
    if (
        not llm_response.content
        and not llm_response.error_code
        and not llm_response.interrupted
    ):
      return

    # Builds the event.
    model_response_event = self._finalize_model_response_event(
        llm_request, llm_response, model_response_event
    )
    yield model_response_event

    # Handles function calls.
    if model_response_event.get_function_calls():

      # Skip partial function call events - they should not trigger execution
      # since partial events are not saved to session (see runners.py).
      # Only execute function calls in the non-partial events.
      if model_response_event.partial:
        return

      async with Aclosing(
          self._postprocess_handle_function_calls_async(
              invocation_context, model_response_event, llm_request
          )
      ) as agen:
        async for event in agen:
          yield event

  async def _postprocess_live(
      self,
      invocation_context: InvocationContext,
      llm_request: LlmRequest,
      llm_response: LlmResponse,
      model_response_event: Event,
  ) -> AsyncGenerator[Event, None]:
    """Postprocess after calling the LLM asynchronously.

    Args:
      invocation_context: The invocation context.
      llm_request: The original LLM request.
      llm_response: The LLM response from the LLM call.
      model_response_event: A mutable event for the LLM response.

    Yields:
      A generator of events.
    """

    # Runs processors.
    async with Aclosing(
        self._postprocess_run_processors_async(invocation_context, llm_response)
    ) as agen:
      async for event in agen:
        yield event

    # Skip the model response event if there is no content and no error code.
    # This is needed for the code executor to trigger another loop.
    # But don't skip control events like turn_complete or transcription events.
    if (
        not llm_response.content
        and not llm_response.error_code
        and not llm_response.interrupted
        and not llm_response.turn_complete
        and not llm_response.input_transcription
        and not llm_response.output_transcription
        and not llm_response.usage_metadata
    ):
      return

    # Handle transcription events ONCE per llm_response, outside the event loop
    if llm_response.input_transcription:
      model_response_event.input_transcription = (
          llm_response.input_transcription
      )
      model_response_event.partial = llm_response.partial
      yield model_response_event
      return

    if llm_response.output_transcription:
      model_response_event.output_transcription = (
          llm_response.output_transcription
      )
      model_response_event.partial = llm_response.partial
      yield model_response_event
      return

    # Flush audio caches based on control events using configurable settings
    if invocation_context.run_config.save_live_blob:
      flushed_events = await self._handle_control_event_flush(
          invocation_context, llm_response
      )
      for event in flushed_events:
        yield event
      if flushed_events:
        # NOTE below return is O.K. for now, because currently we only flush
        # events on interrupted or turn_complete. turn_complete is a pure
        # control event and interrupted is not with content but those content
        # is ignorable because model is already interrupted. If we have other
        # case to flush events in the future that are not pure control events,
        # we should not return here.
        return

    # Builds the event.
    model_response_event = self._finalize_model_response_event(
        llm_request, llm_response, model_response_event
    )
    yield model_response_event

    # Handles function calls.
    if model_response_event.get_function_calls():
      function_response_event = await functions.handle_function_calls_live(
          invocation_context, model_response_event, llm_request.tools_dict
      )
      # Always yield the function response event first
      yield function_response_event

      # Check if this is a set_model_response function response
      if json_response := _output_schema_processor.get_structured_model_response(
          function_response_event
      ):
        # Create and yield a final model response event
        final_event = (
            _output_schema_processor.create_final_model_response_event(
                invocation_context, json_response
            )
        )
        yield final_event

  async def _postprocess_run_processors_async(
      self, invocation_context: InvocationContext, llm_response: LlmResponse
  ) -> AsyncGenerator[Event, None]:
    for processor in self.response_processors:
      async with Aclosing(
          processor.run_async(invocation_context, llm_response)
      ) as agen:
        async for event in agen:
          yield event

  async def _postprocess_handle_function_calls_async(
      self,
      invocation_context: InvocationContext,
      function_call_event: Event,
      llm_request: LlmRequest,
  ) -> AsyncGenerator[Event, None]:
    if function_response_event := await functions.handle_function_calls_async(
        invocation_context, function_call_event, llm_request.tools_dict
    ):
      auth_event = functions.generate_auth_event(
          invocation_context, function_response_event
      )
      if auth_event:
        yield auth_event

      tool_confirmation_event = functions.generate_request_confirmation_event(
          invocation_context, function_call_event, function_response_event
      )
      if tool_confirmation_event:
        yield tool_confirmation_event

      # Always yield the function response event first
      yield function_response_event

      # Check if this is a set_model_response function response
      if json_response := _output_schema_processor.get_structured_model_response(
          function_response_event
      ):
        # Create and yield a final model response event
        final_event = (
            _output_schema_processor.create_final_model_response_event(
                invocation_context, json_response
            )
        )
        yield final_event
      transfer_to_agent = function_response_event.actions.transfer_to_agent
      if transfer_to_agent:
        agent_to_run = self._get_agent_to_run(
            invocation_context, transfer_to_agent
        )
        async with Aclosing(agent_to_run.run_async(invocation_context)) as agen:
          async for event in agen:
            yield event

  def _get_agent_to_run(
      self, invocation_context: InvocationContext, agent_name: str
  ) -> BaseAgent:
    root_agent = invocation_context.agent.root_agent
    agent_to_run = root_agent.find_agent(agent_name)
    if not agent_to_run:
      raise ValueError(f'Agent {agent_name} not found in the agent tree.')
    return agent_to_run

  async def _call_llm_async(
      self,
      invocation_context: InvocationContext,
      llm_request: LlmRequest,
      model_response_event: Event,
  ) -> AsyncGenerator[LlmResponse, None]:
    # Runs before_model_callback if it exists.
    if response := await self._handle_before_model_callback(
        invocation_context, llm_request, model_response_event
    ):
      yield response
      return

    llm_request.config = llm_request.config or types.GenerateContentConfig()
    llm_request.config.labels = llm_request.config.labels or {}

    # Add agent name as a label to the llm_request. This will help with slicing
    # the billing reports on a per-agent basis.
    if _ADK_AGENT_NAME_LABEL_KEY not in llm_request.config.labels:
      llm_request.config.labels[_ADK_AGENT_NAME_LABEL_KEY] = (
          invocation_context.agent.name
      )

    # Calls the LLM.
    llm = self.__get_llm(invocation_context)

    async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]:
      with tracer.start_as_current_span('call_llm') as span:
        if invocation_context.run_config.support_cfc:
          invocation_context.live_request_queue = LiveRequestQueue()
          responses_generator = self.run_live(invocation_context)
          async with Aclosing(
              self._run_and_handle_error(
                  responses_generator,
                  invocation_context,
                  llm_request,
                  model_response_event,
              )
          ) as agen:
            async for llm_response in agen:
              # Runs after_model_callback if it exists.
              if altered_llm_response := await self._handle_after_model_callback(
                  invocation_context, llm_response, model_response_event
              ):
                llm_response = altered_llm_response
              # only yield partial response in SSE streaming mode
              if (
                  invocation_context.run_config.streaming_mode
                  == StreamingMode.SSE
                  or not llm_response.partial
              ):
                yield llm_response
              if llm_response.turn_complete:
                invocation_context.live_request_queue.close()
        else:
          # Check if we can make this llm call or not. If the current call
          # pushes the counter beyond the max set value, then the execution is
          # stopped right here, and exception is thrown.
          invocation_context.increment_llm_call_count()
          responses_generator = llm.generate_content_async(
              llm_request,
              stream=invocation_context.run_config.streaming_mode
              == StreamingMode.SSE,
          )
          async with Aclosing(
              self._run_and_handle_error(
                  responses_generator,
                  invocation_context,
                  llm_request,
                  model_response_event,
              )
          ) as agen:
            async for llm_response in agen:
              trace_call_llm(
                  invocation_context,
                  model_response_event.id,
                  llm_request,
                  llm_response,
                  span,
              )
              # Runs after_model_callback if it exists.
              if altered_llm_response := await self._handle_after_model_callback(
                  invocation_context, llm_response, model_response_event
              ):
                llm_response = altered_llm_response

              yield llm_response

    async with Aclosing(_call_llm_with_tracing()) as agen:
      async for event in agen:
        yield event

  def _finalize_model_response_event(
      self,
      llm_request: LlmRequest,
      llm_response: LlmResponse,
      model_response_event: Event,
  ) -> Event:
    return _finalize_model_response_event(
        llm_request, llm_response, model_response_event
    )

  async def _resolve_toolset_auth(
      self,
      invocation_context: InvocationContext,
      agent: LlmAgent,
  ) -> AsyncGenerator[Event, None]:
    async with Aclosing(
        _resolve_toolset_auth(invocation_context, agent)
    ) as agen:
      async for event in agen:
        yield event

  async def _handle_before_model_callback(
      self,
      invocation_context: InvocationContext,
      llm_request: LlmRequest,
      model_response_event: Event,
  ) -> Optional[LlmResponse]:
    return await _handle_before_model_callback(
        invocation_context, llm_request, model_response_event
    )

  async def _handle_after_model_callback(
      self,
      invocation_context: InvocationContext,
      llm_response: LlmResponse,
      model_response_event: Event,
  ) -> Optional[LlmResponse]:
    return await _handle_after_model_callback(
        invocation_context, llm_response, model_response_event
    )

  async def _run_and_handle_error(
      self,
      response_generator: AsyncGenerator[LlmResponse, None],
      invocation_context: InvocationContext,
      llm_request: LlmRequest,
      model_response_event: Event,
  ) -> AsyncGenerator[LlmResponse, None]:
    async with Aclosing(
        _run_and_handle_error(
            response_generator,
            invocation_context,
            llm_request,
            model_response_event,
        )
    ) as agen:
      async for response in agen:
        yield response

  async def _handle_control_event_flush(
      self, invocation_context: InvocationContext, llm_response: LlmResponse
  ) -> list[Event]:
    """Handle audio cache flushing based on control events.

    Args:
      invocation_context: The invocation context containing audio caches.
      llm_response: The LLM response containing control event information.

    Returns:
      A list of Event objects created from the flushed caches.
    """

    # Log cache statistics if enabled
    if DEFAULT_ENABLE_CACHE_STATISTICS:
      stats = self.audio_cache_manager.get_cache_stats(invocation_context)
      logger.debug('Audio cache stats: %s', stats)

    if llm_response.interrupted:
      # user interrupts so the model will stop. we can flush model audio here
      return await self.audio_cache_manager.flush_caches(
          invocation_context,
          flush_user_audio=False,
          flush_model_audio=True,
      )
    elif llm_response.turn_complete:
      # turn completes so we can flush both user and model
      return await self.audio_cache_manager.flush_caches(
          invocation_context,
          flush_user_audio=True,
          flush_model_audio=True,
      )
    # TODO: Once generation_complete is surfaced on LlmResponse, we can flush
    # model audio here (flush_user_audio=False, flush_model_audio=True).
    return []

  def __get_llm(self, invocation_context: InvocationContext) -> BaseLlm:
    agent = invocation_context.agent

    # Check for conformance test replay mode
    if config := invocation_context.session.state.get('_adk_replay_config'):
      from ...cli.conformance._conformance_test_google_llm import _ConformanceTestGemini

      # Models are stateless, so the current replay state is cached in the
      # session state to maintain the state across model calls
      # key: (agent_name, user_message_index)
      # value: replay index
      user_message_index = config.get('user_message_index')
      replay_indexes = config.get('_adk_replay_indexes', {})
      if (agent.name, user_message_index) not in replay_indexes:
        replay_indexes[(agent.name, user_message_index)] = 0
      current_replay_index = replay_indexes[(agent.name, user_message_index)]

      config['current_replay_index'] = current_replay_index
      config['agent_name'] = agent.name
      model = _ConformanceTestGemini(
          config=config,
      )

      replay_indexes[(agent.name, user_message_index)] = (
          current_replay_index + 1
      )
      config['_adk_replay_indexes'] = replay_indexes
      return model

    if not hasattr(agent, 'canonical_model'):
      raise TypeError(
          'Expected agent to have canonical_model attribute,'
          f' but got {type(agent)}'
      )
    return agent.canonical_model
