# 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

import abc
import asyncio
import importlib
import inspect
import logging
from typing import Any
from typing import AsyncGenerator
from typing import Awaitable
from typing import Callable
from typing import ClassVar
from typing import Dict
from typing import Literal
from typing import Optional
from typing import Type
from typing import Union
import warnings

from google.genai import types
from pydantic import BaseModel
from pydantic import Field
from pydantic import field_validator
from pydantic import model_validator
from typing_extensions import override
from typing_extensions import TypeAlias

from ..code_executors.base_code_executor import BaseCodeExecutor
from ..events.event import Event
from ..features import experimental
from ..features import FeatureName
from ..flows.llm_flows.auto_flow import AutoFlow
from ..flows.llm_flows.base_llm_flow import BaseLlmFlow
from ..flows.llm_flows.single_flow import SingleFlow
from ..models.base_llm import BaseLlm
from ..models.llm_request import LlmRequest
from ..models.llm_response import LlmResponse
from ..models.registry import LLMRegistry
from ..planners.base_planner import BasePlanner
from ..tools.base_tool import BaseTool
from ..tools.base_toolset import BaseToolset
from ..tools.function_tool import FunctionTool
from ..tools.tool_configs import ToolConfig
from ..tools.tool_context import ToolContext
from ..utils._schema_utils import SchemaType
from ..utils._schema_utils import validate_schema
from ..utils.context_utils import Aclosing
from .base_agent import BaseAgent
from .base_agent import BaseAgentState
from .base_agent_config import BaseAgentConfig as BaseAgentConfig
from .callback_context import CallbackContext
from .context import Context
from .invocation_context import InvocationContext
from .llm_agent_config import LlmAgentConfig as LlmAgentConfig
from .readonly_context import ReadonlyContext

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

_SingleBeforeModelCallback: TypeAlias = Callable[
    [CallbackContext, LlmRequest],
    Union[Awaitable[Optional[LlmResponse]], Optional[LlmResponse]],
]

BeforeModelCallback: TypeAlias = Union[
    _SingleBeforeModelCallback,
    list[_SingleBeforeModelCallback],
]

_SingleAfterModelCallback: TypeAlias = Callable[
    [CallbackContext, LlmResponse],
    Union[Awaitable[Optional[LlmResponse]], Optional[LlmResponse]],
]

AfterModelCallback: TypeAlias = Union[
    _SingleAfterModelCallback,
    list[_SingleAfterModelCallback],
]

_SingleOnModelErrorCallback: TypeAlias = Callable[
    [CallbackContext, LlmRequest, Exception],
    Union[Awaitable[Optional[LlmResponse]], Optional[LlmResponse]],
]

OnModelErrorCallback: TypeAlias = Union[
    _SingleOnModelErrorCallback,
    list[_SingleOnModelErrorCallback],
]

_SingleBeforeToolCallback: TypeAlias = Callable[
    [BaseTool, dict[str, Any], ToolContext],
    Union[Awaitable[Optional[dict]], Optional[dict]],
]

BeforeToolCallback: TypeAlias = Union[
    _SingleBeforeToolCallback,
    list[_SingleBeforeToolCallback],
]

_SingleAfterToolCallback: TypeAlias = Callable[
    [BaseTool, dict[str, Any], ToolContext, dict],
    Union[Awaitable[Optional[dict]], Optional[dict]],
]

AfterToolCallback: TypeAlias = Union[
    _SingleAfterToolCallback,
    list[_SingleAfterToolCallback],
]

_SingleOnToolErrorCallback: TypeAlias = Callable[
    [BaseTool, dict[str, Any], ToolContext, Exception],
    Union[Awaitable[Optional[dict]], Optional[dict]],
]

OnToolErrorCallback: TypeAlias = Union[
    _SingleOnToolErrorCallback,
    list[_SingleOnToolErrorCallback],
]

InstructionProvider: TypeAlias = Callable[
    [ReadonlyContext], Union[str, Awaitable[str]]
]
ToolUnion: TypeAlias = Union[Callable, BaseTool, BaseToolset]


async def _convert_tool_union_to_tools(
    tool_union: ToolUnion,
    ctx: Optional[ReadonlyContext],
    model: Union[str, BaseLlm],
    multiple_tools: bool = False,
) -> list[BaseTool]:
  from ..tools.google_search_tool import GoogleSearchTool
  from ..tools.vertex_ai_search_tool import VertexAiSearchTool

  # Wrap google_search tool with AgentTool if there are multiple tools because
  # the built-in tools cannot be used together with other tools.
  # TODO: Remove once the workaround is no longer needed.
  if multiple_tools and isinstance(tool_union, GoogleSearchTool):
    from ..tools.google_search_agent_tool import create_google_search_agent
    from ..tools.google_search_agent_tool import GoogleSearchAgentTool

    search_tool = tool_union
    if search_tool.bypass_multi_tools_limit:
      return [GoogleSearchAgentTool(create_google_search_agent(model))]

  # Replace VertexAiSearchTool with DiscoveryEngineSearchTool if there are
  # multiple tools because the built-in tools cannot be used together with
  # other tools.
  # TODO: Remove once the workaround is no longer needed.
  if multiple_tools and isinstance(tool_union, VertexAiSearchTool):
    from ..tools.discovery_engine_search_tool import DiscoveryEngineSearchTool

    vais_tool = tool_union
    if vais_tool.bypass_multi_tools_limit:
      return [
          DiscoveryEngineSearchTool(
              data_store_id=vais_tool.data_store_id,
              data_store_specs=vais_tool.data_store_specs,
              search_engine_id=vais_tool.search_engine_id,
              filter=vais_tool.filter,
              max_results=vais_tool.max_results,
          )
      ]
  from ..workflow._base_node import BaseNode

  if isinstance(tool_union, BaseNode):
    from ..tools._node_tool import NodeTool
    from .base_agent import BaseAgent

    if isinstance(tool_union, BaseAgent):
      raise ValueError(
          f"Agent '{tool_union.name}' cannot be wrapped as a NodeTool. Agents"
          ' should be invoked as sub-agents.'
      )

    description = tool_union.description
    if not description:
      raise ValueError(
          f"Workflow/Node '{tool_union.name}' must have a description to be"
          ' wrapped as a tool.'
      )

    return [
        NodeTool(
            node=tool_union,
            name=tool_union.name,
            description=description,
        )
    ]

  if isinstance(tool_union, BaseTool):
    return [tool_union]
  if callable(tool_union):
    return [FunctionTool(func=tool_union)]

  # At this point, tool_union must be a BaseToolset
  try:
    return await tool_union.get_tools_with_prefix(ctx)
  except Exception as e:
    logger.warning(
        'Failed to get tools from toolset %s: %s',
        type(tool_union).__name__,
        e,
    )
    return []


# TODO: drop the explicit abc.ABC base once BaseNode surfaces ABCMeta to
# static type checkers.
class LlmAgent(BaseAgent, abc.ABC):
  """LLM-based Agent."""

  DEFAULT_MODEL: ClassVar[str] = 'gemini-3.5-flash'
  """System default model used when no model is set on an agent."""

  DEFAULT_LIVE_MODEL: ClassVar[str] = 'gemini-live-2.5-flash-native-audio'
  """System default model used for live mode when no model is set on an agent."""

  _default_model: ClassVar[Union[str, BaseLlm]] = DEFAULT_MODEL
  """Current default model used when an agent has no model set."""

  _default_live_model: ClassVar[Union[str, BaseLlm]] = DEFAULT_LIVE_MODEL
  """Current default model used for live mode when an agent has no model set."""

  model: Union[str, BaseLlm] = ''
  """The model to use for the agent.

  When not set, the agent will inherit the model from its ancestor. If no
  ancestor provides a model, the agent uses the default model configured via
  LlmAgent.set_default_model. The built-in default is gemini-3.5-flash.
  """

  config_type: ClassVar[Type[BaseAgentConfig]] = LlmAgentConfig
  """The config type for this agent.

  DEPRECATED: This attribute is deprecated and will be removed in a future
  version, along with the AgentConfig YAML loader.
  """

  instruction: Union[str, InstructionProvider] = ''
  """Dynamic instructions for the LLM model, guiding the agent's behavior.

  These instructions can contain placeholders like {variable_name} that will be
  resolved at runtime using session state and context.

  **Behavior depends on static_instruction:**
  - If static_instruction is None: instruction goes to system_instruction
  - If static_instruction is set: instruction goes to user content in the request

  This allows for context caching optimization where static content (static_instruction)
  comes first in the prompt, followed by dynamic content (instruction).
  """

  global_instruction: Union[str, InstructionProvider] = ''
  """Instructions for all the agents in the entire agent tree.

  DEPRECATED: This field is deprecated and will be removed in a future version.
  Use GlobalInstructionPlugin instead, which provides the same functionality
  at the App level. See migration guide for details.

  ONLY the global_instruction in root agent will take effect.

  For example: use global_instruction to make all agents have a stable identity
  or personality.
  """

  static_instruction: Optional[types.ContentUnion] = None
  """Static instruction content sent literally as system instruction at the beginning.

  This field is for content that never changes and doesn't contain placeholders.
  It's sent directly to the model without any processing or variable substitution.

  This field is primarily for context caching optimization. Static instructions
  are sent as system instruction at the beginning of the request, allowing
  for improved performance when the static portion remains unchanged. Live API
  has its own cache mechanism, thus this field doesn't work with Live API.

  **Impact on instruction field:**
  - When static_instruction is None: instruction → system_instruction
  - When static_instruction is set: instruction → user content (after static content)

  **Context Caching:**
  - **Implicit Cache**: Automatic caching by model providers (no config needed)
  - **Explicit Cache**: Cache explicitly created by user for instructions, tools and contents

  See below for more information of Implicit Cache and Explicit Cache
  Gemini API: https://ai.google.dev/gemini-api/docs/caching?lang=python
  Vertex API: https://cloud.google.com/vertex-ai/generative-ai/docs/context-cache/context-cache-overview

  Setting static_instruction alone does NOT enable caching automatically.
  For explicit caching control, configure context_cache_config at App level.

  **Content Support:**
  Accepts types.ContentUnion which includes:
  - str: Simple text instruction
  - types.Content: Rich content object
  - types.Part: Single part (text, inline_data, file_data, etc.)
  - PIL.Image.Image: Image object
  - types.File: File reference
  - list[PartUnion]: List of parts

  **Examples:**
  ```python
  # Simple string instruction
  static_instruction = "You are a helpful assistant."

  # Rich content with files
  static_instruction = types.Content(
      role='user',
      parts=[
          types.Part(text='You are a helpful assistant.'),
          types.Part(file_data=types.FileData(...))
      ]
  )
  ```
  """

  tools: list[ToolUnion] = Field(default_factory=list)
  """Tools available to this agent."""

  generate_content_config: Optional[types.GenerateContentConfig] = None
  """The additional content generation configurations.

  NOTE: not all fields are usable, e.g. tools must be configured via `tools`,
  thinking_config can be configured here or via the `planner`. If both are set, the planner's configuration takes precedence.

  For example: use this config to adjust model temperature, configure safety
  settings, etc.
  """

  mode: Literal['chat', 'task', 'single_turn'] | None = None
  """The delegation mode for this agent.

  Options:
    chat: Standard chat agent reachable via transfer_to_agent.
    task: Task agent that chats with the user to accomplish a task.
    single_turn: Agents that complete a task without chatting with the user.

  Default value is chat as a sub-agent, single_turn as a node in a workflow.
  """

  parallel_worker: bool | None = None
  """Whether to run the agent in parallel worker mode."""

  # LLM-based agent transfer configs - Start
  disallow_transfer_to_parent: bool = False
  """Disallows LLM-controlled transferring to the parent agent.

  NOTE: Setting this as True also prevents this agent from continuing to reply
  to the end-user, and will transfer control back to the parent agent in the
  next turn. This behavior prevents one-way transfer, in which end-user may be
  stuck with one agent that cannot transfer to other agents in the agent tree.
  """
  disallow_transfer_to_peers: bool = False
  """Disallows LLM-controlled transferring to the peer agents."""
  # LLM-based agent transfer configs - End

  include_contents: Literal['default', 'none'] = 'default'
  """Controls content inclusion in model requests.

  Options:
    default: Model receives relevant conversation history
    none: Model receives no prior history, operates solely on current
    instruction and input
  """

  # Controlled input/output configurations - Start
  input_schema: Optional[type[BaseModel]] = None
  """The input schema when agent is used as a tool."""
  output_schema: Optional[SchemaType] = None
  """The output schema when agent replies.

  Supports all schema types that the underlying Google GenAI API supports:
    - type[BaseModel]: e.g., MySchema
    - list[type[BaseModel]]: e.g., list[MySchema]
    - list[primitive]: e.g., list[str], list[int]
    - dict: Raw dict schemas
    - Schema: Google's Schema type

  NOTE:
    The ADK supports using `output_schema` and `tools` together. It works by
    exposing tools during the thought loop and enforcing structure only on the
    final output.
  """
  output_key: Optional[str] = None
  """The key in session state to store the output of the agent.

  Typically use cases:
  - Extracts agent reply for later use, such as in tools, callbacks, etc.
  - Connects agents to coordinate with each other.
  """
  # Controlled input/output configurations - End

  # Advance features - Start
  planner: Optional[BasePlanner] = None
  """Instructs the agent to make a plan and execute it step by step.

  NOTE:
    To use model's built-in thinking features, set the `thinking_config`
    field in `google.adk.planners.built_in_planner`.
  """

  code_executor: Optional[BaseCodeExecutor] = None
  """Allow agent to execute code blocks from model responses using the provided
  CodeExecutor.

  Check out available code executions in `google.adk.code_executor` package.

  NOTE:
    To use model's built-in code executor, use the `BuiltInCodeExecutor`.
  """
  # Advance features - End

  # Callbacks - Start
  before_model_callback: Optional[BeforeModelCallback] = None
  """Callback or list of callbacks to be called before calling the LLM.

  When a list of callbacks is provided, the callbacks will be called in the
  order they are listed until a callback does not return None.

  Args:
    callback_context: CallbackContext,
    llm_request: LlmRequest, The raw model request. Callback can mutate the
    request.

  Returns:
    The content to return to the user. When present, the model call will be
    skipped and the provided content will be returned to user.
  """
  after_model_callback: Optional[AfterModelCallback] = None
  """Callback or list of callbacks to be called after calling the LLM.

  When a list of callbacks is provided, the callbacks will be called in the
  order they are listed until a callback does not return None.

  Args:
    callback_context: CallbackContext,
    llm_response: LlmResponse, the actual model response.

  Returns:
    The content to return to the user. When present, the actual model response
    will be ignored and the provided content will be returned to user.
  """
  on_model_error_callback: Optional[OnModelErrorCallback] = None
  """Callback or list of callbacks to be called when a model call encounters an error.

  When a list of callbacks is provided, the callbacks will be called in the
  order they are listed until a callback does not return None.

  Args:
    callback_context: CallbackContext,
    llm_request: LlmRequest, The raw model request.
    error: The error from the model call.

  Returns:
    The content to return to the user. When present, the error will be
    ignored and the provided content will be returned to user.
  """
  before_tool_callback: Optional[BeforeToolCallback] = None
  """Callback or list of callbacks to be called before calling the tool.

  When a list of callbacks is provided, the callbacks will be called in the
  order they are listed until a callback does not return None.

  Args:
    tool: The tool to be called.
    args: The arguments to the tool.
    tool_context: ToolContext,

  Returns:
    The tool response. When present, the returned tool response will be used and
    the framework will skip calling the actual tool.
  """
  after_tool_callback: Optional[AfterToolCallback] = None
  """Callback or list of callbacks to be called after calling the tool.

  When a list of callbacks is provided, the callbacks will be called in the
  order they are listed until a callback does not return None.

  Args:
    tool: The tool to be called.
    args: The arguments to the tool.
    tool_context: ToolContext,
    tool_response: The response from the tool.

  Returns:
    When present, the returned dict will be used as tool result.
  """
  on_tool_error_callback: Optional[OnToolErrorCallback] = None
  """Callback or list of callbacks to be called when a tool call encounters an error.

  When a list of callbacks is provided, the callbacks will be called in the
  order they are listed until a callback does not return None.

  Args:
    tool: The tool to be called.
    args: The arguments to the tool.
    tool_context: ToolContext,
    error: The error from the tool call.

  Returns:
    When present, the returned dict will be used as tool result.
  """
  # Callbacks - End

  @override
  async def _handle_before_agent_callback(
      self, ctx: InvocationContext
  ) -> Optional[Event]:
    event = await super()._handle_before_agent_callback(ctx)
    if event is not None:
      self.__maybe_save_output_to_state(event)
    return event

  @override
  async def _run_async_impl(
      self, ctx: InvocationContext
  ) -> AsyncGenerator[Event, None]:
    agent_state = self._load_agent_state(ctx, BaseAgentState)

    # If there is a sub-agent to resume, run it and then end the current
    # agent.
    if agent_state is not None and (
        agent_to_transfer := self._get_subagent_to_resume(ctx)
    ):
      async with Aclosing(agent_to_transfer.run_async(ctx)) as agen:
        async for event in agen:
          yield event

      ctx.set_agent_state(self.name, end_of_agent=True)
      yield self._create_agent_state_event(ctx)
      return

    should_pause = False
    output_accumulator = ''
    async with Aclosing(self._llm_flow.run_async(ctx)) as agen:
      async for event in agen:
        self.__maybe_save_output_to_state(event)
        output_accumulator = self.__maybe_accumulate_streaming_output(
            event, output_accumulator
        )
        yield event
        if ctx.should_pause_invocation(event):
          # Do not pause immediately, wait until the long-running tool call is
          # executed.
          should_pause = True
    if should_pause:
      return

    if ctx.is_resumable:
      events = ctx._get_events(current_invocation=True, current_branch=True)
      if events and any(ctx.should_pause_invocation(e) for e in events[-2:]):
        return
      # Only yield an end state if the last event is no longer a long-running
      # tool call.
      ctx.set_agent_state(self.name, end_of_agent=True)
      yield self._create_agent_state_event(ctx)

  @override
  async def _run_live_impl(
      self, ctx: InvocationContext
  ) -> AsyncGenerator[Event, None]:
    output_accumulator = ''
    async with Aclosing(self._llm_flow.run_live(ctx)) as agen:
      async for event in agen:
        self.__maybe_save_output_to_state(event)
        output_accumulator = self.__maybe_accumulate_streaming_output(
            event, output_accumulator
        )
        yield event
      if ctx.end_invocation:
        return

  @override
  async def _run_impl(
      self,
      *,
      ctx: Context,
      node_input: Any,
  ) -> AsyncGenerator[Any, None]:
    """Runs the agent as a node in a workflow graph."""
    from ..utils.context_utils import Aclosing
    from ..workflow._llm_agent_wrapper import run_llm_agent_as_node

    async with Aclosing(
        run_llm_agent_as_node(self, ctx=ctx, node_input=node_input)
    ) as agen:
      async for event in agen:
        # Keep the agent's true event author so the outer NodeRunner does
        # not overwrite it with the parent workflow's event_author.
        if event.author:
          ctx.event_author = event.author
        yield event

  @property
  def canonical_model(self) -> BaseLlm:
    """The resolved self.model field as BaseLlm.

    This method is only for use by Agent Development Kit.
    """
    if isinstance(self.model, BaseLlm):
      return self.model
    elif self.model:  # model is non-empty str
      return LLMRegistry.new_llm(self.model)
    else:  # find model from ancestors.
      ancestor_agent = self.parent_agent
      while ancestor_agent is not None:
        if isinstance(ancestor_agent, LlmAgent):
          return ancestor_agent.canonical_model
        ancestor_agent = ancestor_agent.parent_agent
      return self._resolve_default_model()

  @property
  def canonical_live_model(self) -> BaseLlm:
    """The resolved self.model field as BaseLlm for live mode.

    This method is only for use by Agent Development Kit.
    """
    if isinstance(self.model, BaseLlm):
      return self.model
    elif self.model:  # model is non-empty str
      return LLMRegistry.new_llm(self.model)
    else:  # find model from ancestors.
      ancestor_agent = self.parent_agent
      while ancestor_agent is not None:
        if isinstance(ancestor_agent, LlmAgent):
          return ancestor_agent.canonical_live_model
        ancestor_agent = ancestor_agent.parent_agent
      return self._resolve_default_live_model()

  @classmethod
  def set_default_model(cls, model: Union[str, BaseLlm]) -> None:
    """Overrides the default model used when an agent has no model set."""
    if not isinstance(model, (str, BaseLlm)):
      raise TypeError('Default model must be a model name or BaseLlm.')
    if isinstance(model, str) and not model:
      raise ValueError('Default model must be a non-empty string.')
    cls._default_model = model

  @classmethod
  def _resolve_default_model(cls) -> BaseLlm:
    """Resolves the current default model to a BaseLlm instance."""
    default_model = cls._default_model
    if isinstance(default_model, BaseLlm):
      return default_model
    return LLMRegistry.new_llm(default_model)

  @classmethod
  def set_default_live_model(cls, model: Union[str, BaseLlm]) -> None:
    """Overrides the default model used for live mode when an agent has no model set."""
    if not isinstance(model, (str, BaseLlm)):
      raise TypeError('Default live model must be a model name or BaseLlm.')
    if isinstance(model, str) and not model:
      raise ValueError('Default live model must be a non-empty string.')
    cls._default_live_model = model

  @classmethod
  def _resolve_default_live_model(cls) -> BaseLlm:
    """Resolves the current default live model to a BaseLlm instance."""
    default_live_model = cls._default_live_model
    if isinstance(default_live_model, BaseLlm):
      return default_live_model
    return LLMRegistry.new_llm(default_live_model)

  async def canonical_instruction(
      self, ctx: ReadonlyContext
  ) -> tuple[str, bool]:
    """The resolved self.instruction field to construct instruction for this agent.

    This method is only for use by Agent Development Kit.

    Args:
      ctx: The context to retrieve the session state.

    Returns:
      A tuple of (instruction, bypass_state_injection).
      instruction: The resolved self.instruction field.
      bypass_state_injection: Whether the instruction is based on
      InstructionProvider.
    """
    if isinstance(self.instruction, str):
      return self.instruction, False
    else:
      instruction = self.instruction(ctx)
      if inspect.isawaitable(instruction):
        instruction = await instruction
      return instruction, True

  async def canonical_global_instruction(
      self, ctx: ReadonlyContext
  ) -> tuple[str, bool]:
    """The resolved self.instruction field to construct global instruction.

    This method is only for use by Agent Development Kit.

    Args:
      ctx: The context to retrieve the session state.

    Returns:
      A tuple of (instruction, bypass_state_injection).
      instruction: The resolved self.global_instruction field.
      bypass_state_injection: Whether the instruction is based on
      InstructionProvider.
    """
    # Issue deprecation warning if global_instruction is being used
    if self.global_instruction:
      warnings.warn(
          'global_instruction field is deprecated and will be removed in a'
          ' future version. Use GlobalInstructionPlugin instead for the same'
          ' functionality at the App level. See migration guide for details.',
          DeprecationWarning,
          stacklevel=2,
      )

    if isinstance(self.global_instruction, str):
      return self.global_instruction, False
    else:
      global_instruction = self.global_instruction(ctx)
      if inspect.isawaitable(global_instruction):
        global_instruction = await global_instruction
      return global_instruction, True

  async def canonical_tools(
      self, ctx: Optional[ReadonlyContext] = None
  ) -> list[BaseTool]:
    """The resolved self.tools field as a list of BaseTool based on the context.

    This method is only for use by Agent Development Kit.
    """
    # We may need to wrap some built-in tools if there are other tools
    # because the built-in tools cannot be used together with other tools.
    # TODO: Remove once the workaround is no longer needed.
    multiple_tools = len(self.tools) > 1
    model = self.canonical_model

    results = await asyncio.gather(*(
        _convert_tool_union_to_tools(tool_union, ctx, model, multiple_tools)
        for tool_union in self.tools
    ))

    resolved_tools = []
    for tools in results:
      resolved_tools.extend(tools)

    return resolved_tools

  @property
  def canonical_before_model_callbacks(
      self,
  ) -> list[_SingleBeforeModelCallback]:
    """The resolved self.before_model_callback field as a list of _SingleBeforeModelCallback.

    This method is only for use by Agent Development Kit.
    """
    if not self.before_model_callback:
      return []
    if isinstance(self.before_model_callback, list):
      return self.before_model_callback
    return [self.before_model_callback]

  @property
  def canonical_after_model_callbacks(self) -> list[_SingleAfterModelCallback]:
    """The resolved self.after_model_callback field as a list of _SingleAfterModelCallback.

    This method is only for use by Agent Development Kit.
    """
    if not self.after_model_callback:
      return []
    if isinstance(self.after_model_callback, list):
      return self.after_model_callback
    return [self.after_model_callback]

  @property
  def canonical_on_model_error_callbacks(
      self,
  ) -> list[_SingleOnModelErrorCallback]:
    """The resolved self.on_model_error_callback field as a list of _SingleOnModelErrorCallback.

    This method is only for use by Agent Development Kit.
    """
    if not self.on_model_error_callback:
      return []
    if isinstance(self.on_model_error_callback, list):
      return self.on_model_error_callback
    return [self.on_model_error_callback]

  @property
  def canonical_before_tool_callbacks(
      self,
  ) -> list[BeforeToolCallback]:
    """The resolved self.before_tool_callback field as a list of BeforeToolCallback.

    This method is only for use by Agent Development Kit.
    """
    if not self.before_tool_callback:
      return []
    if isinstance(self.before_tool_callback, list):
      return self.before_tool_callback
    return [self.before_tool_callback]

  @property
  def canonical_after_tool_callbacks(
      self,
  ) -> list[AfterToolCallback]:
    """The resolved self.after_tool_callback field as a list of AfterToolCallback.

    This method is only for use by Agent Development Kit.
    """
    if not self.after_tool_callback:
      return []
    if isinstance(self.after_tool_callback, list):
      return self.after_tool_callback
    return [self.after_tool_callback]

  @property
  def canonical_on_tool_error_callbacks(
      self,
  ) -> list[OnToolErrorCallback]:
    """The resolved self.on_tool_error_callback field as a list of OnToolErrorCallback.

    This method is only for use by Agent Development Kit.
    """
    if not self.on_tool_error_callback:
      return []
    if isinstance(self.on_tool_error_callback, list):
      return self.on_tool_error_callback
    return [self.on_tool_error_callback]

  @property
  def _llm_flow(self) -> BaseLlmFlow:
    if (
        self.disallow_transfer_to_parent
        and self.disallow_transfer_to_peers
        and not self.sub_agents
    ):
      return SingleFlow()
    else:
      return AutoFlow()

  def _get_subagent_to_resume(
      self, ctx: InvocationContext
  ) -> Optional[BaseAgent]:
    """Returns the sub-agent in the llm tree to resume if it exists.

    There are 2 cases where we need to transfer to and resume a sub-agent:
    1. The last event is a transfer to agent response from the current agent.
       In this case, we need to return the agent specified in the response.

    2. The last event's author isn't the current agent, or the user is
       responding to another agent's tool call.
       In this case, we need to return the LAST agent being transferred to
       from the current agent.
    """
    events = ctx._get_events(current_invocation=True, current_branch=True)
    if not events:
      return None

    last_event = events[-1]
    if last_event.author == self.name:
      # Last event is from current agent. Return transfer_to_agent in the event
      # if it exists, or None.
      return self.__get_transfer_to_agent_or_none(last_event, self.name)

    # Last event is from user or another agent.
    if last_event.author == 'user':
      function_call_event = ctx._find_matching_function_call(last_event)
      if not function_call_event:
        raise ValueError(
            'No agent to transfer to for resuming agent from function response'
            f' {self.name}'
        )
      if function_call_event.author == self.name:
        # User is responding to a tool call from the current agent.
        # Current agent should continue, so no sub-agent to resume.
        return None

    # Last event is from another agent, or from user for another agent's tool
    # call. We need to find the last agent we transferred to.
    for event in reversed(events):
      if agent := self.__get_transfer_to_agent_or_none(event, self.name):
        return agent

    return None

  def __get_agent_to_run(self, agent_name: str) -> BaseAgent:
    """Find the agent to run under the root agent by name."""
    agent_to_run = self.root_agent.find_agent(agent_name)
    if not agent_to_run:
      available = self._get_available_agent_names()
      error_msg = (
          f"Agent '{agent_name}' not found.\n"
          f"Available agents: {', '.join(available)}\n\n"
          'Possible causes:\n'
          '  1. Agent not registered before being referenced\n'
          '  2. Agent name mismatch (typo or case sensitivity)\n'
          '  3. Timing issue (agent referenced before creation)\n\n'
          'Suggested fixes:\n'
          '  - Verify agent is registered with root agent\n'
          '  - Check agent name spelling and case\n'
          '  - Ensure agents are created before being referenced'
      )
      raise ValueError(error_msg)
    return agent_to_run

  def _get_available_agent_names(self) -> list[str]:
    """Helper to get all agent names in the tree for error reporting.

    This is a private helper method used only for error message formatting.
    Traverses the agent tree starting from root_agent and collects all
    agent names for display in error messages.

    Returns:
      List of all agent names in the agent tree.
    """
    agents = []

    def collect_agents(agent: BaseAgent) -> None:
      agents.append(agent.name)
      if hasattr(agent, 'sub_agents') and agent.sub_agents:
        for sub_agent in agent.sub_agents:
          collect_agents(sub_agent)

    collect_agents(self.root_agent)
    return agents

  def __get_transfer_to_agent_or_none(
      self, event: Event, from_agent: str
  ) -> Optional[BaseAgent]:
    """Returns the agent to run if the event is a transfer to agent response."""
    function_responses = event.get_function_responses()
    if not function_responses:
      return None
    for function_response in function_responses:
      if (
          function_response.name == 'transfer_to_agent'
          and event.author == from_agent
          and event.actions.transfer_to_agent != from_agent
      ):
        return self.__get_agent_to_run(event.actions.transfer_to_agent)
    return None

  def __maybe_save_output_to_state(self, event: Event) -> None:
    """Saves the model output to state if needed."""
    # skip if the event was authored by some other agent (e.g. current agent
    # transferred to another agent)
    if event.author != self.name:
      logger.debug(
          'Skipping output save for agent %s: event authored by %s',
          self.name,
          event.author,
      )
      return

    if not self.output_key:
      return

    # Handle text responses
    if event.is_final_response() and event.content and event.content.parts:

      # Skip if no text parts at all to avoid overwriting state_delta values
      # already set (e.g. after_tool_callback with skip_summarization
      # on function_response-only events).
      has_text_part = any(
          part.text is not None and not part.thought
          for part in event.content.parts
      )

      if not has_text_part:
        return

      result = ''.join(
          part.text
          for part in event.content.parts
          if part.text and not part.thought
      )
      if self.output_schema:
        # If the result from the final chunk is just whitespace or empty,
        # it means this is an empty final chunk of a stream.
        # Do not attempt to parse it as JSON.
        if not result.strip():
          return
        result = validate_schema(self.output_schema, result)
      event.actions.state_delta[self.output_key] = result

  def __maybe_accumulate_streaming_output(
      self, event: Event, accumulator: str
  ) -> str:
    """Accumulates output_key text across a streaming model turn.

    Streaming with tool calls produces non-partial events that carry text
    alongside a function_call. is_final_response() rejects those, so
    __maybe_save_output_to_state skips them and the text on those events
    is dropped from output_key. Accumulate every non-partial text-bearing
    event from this agent across the model turn so the segments survive
    in session state. See issue #5590.

    No-op when accumulation doesn't apply (different author, no
    output_key, output_schema set, partial event, no content, no text).
    For applicable events, appends the event's text to ``accumulator``
    and writes the running value to state_delta[output_key], overwriting
    any value __maybe_save_output_to_state set on the same event.
    Returns the new accumulator value.
    """
    if (
        not self.output_key
        or self.output_schema
        or event.author != self.name
        or event.partial
        or not event.content
        or not event.content.parts
    ):
      return accumulator

    text = ''.join(
        part.text
        for part in event.content.parts
        if part.text and not part.thought
    )
    if not text:
      return accumulator

    accumulator += text
    event.actions.state_delta[self.output_key] = accumulator
    return accumulator

  @model_validator(mode='before')
  @classmethod
  def _pre_validate_tools(cls, data: Any) -> Any:
    if isinstance(data, dict) and 'tools' in data and data['tools']:
      from google.adk.agents.base_agent import BaseAgent
      from google.adk.tools._node_tool import NodeTool
      from google.adk.workflow._base_node import BaseNode

      new_tools = []
      for t in data['tools']:
        if isinstance(t, BaseAgent):
          raise ValueError(
              f"Agent '{t.name}' cannot be wrapped as a NodeTool. Agents should"
              ' be invoked as sub-agents.'
          )
        elif isinstance(t, BaseNode):
          description = t.description
          if not description:
            raise ValueError(
                f"Workflow/Node '{t.name}' must have a description to be"
                ' wrapped as a tool.'
            )
          new_tools.append(NodeTool(node=t, description=description))
        else:
          new_tools.append(t)
      data['tools'] = new_tools
    return data

  @model_validator(mode='after')
  def __model_validator_after(self) -> LlmAgent:
    return self

  @field_validator('generate_content_config', mode='after')
  @classmethod
  def validate_generate_content_config(
      cls, generate_content_config: Optional[types.GenerateContentConfig]
  ) -> types.GenerateContentConfig:
    if not generate_content_config:
      return types.GenerateContentConfig()
    if generate_content_config.tools:
      raise ValueError('All tools must be set via LlmAgent.tools.')
    if generate_content_config.system_instruction:
      raise ValueError(
          'System instruction must be set via LlmAgent.instruction.'
      )
    if generate_content_config.response_schema:
      raise ValueError(
          'Response schema must be set via LlmAgent.output_schema.'
      )
    return generate_content_config

  @override
  def model_post_init(self, __context: Any) -> None:
    """Provides a warning if multiple thinking configurations are found."""
    super().model_post_init(__context)

    from ..planners.built_in_planner import BuiltInPlanner

    if (
        self.generate_content_config is not None
        and self.generate_content_config.thinking_config is not None
        and isinstance(self.planner, BuiltInPlanner)
        and self.planner.thinking_config is not None
    ):
      warnings.warn(
          'Both `thinking_config` in `generate_content_config` and a '
          'planner with `thinking_config` are provided. The '
          "planner's configuration will take precedence.",
          UserWarning,
          stacklevel=3,
      )

    if self.mode == 'task':
      from .llm.task._finish_task_tool import FinishTaskTool

      self.tools.append(FinishTaskTool(self))

    # Add sub-agents as tools based on their mode
    from ..tools.agent_tool import _SingleTurnAgentTool
    from ..tools.agent_tool import _TaskAgentTool

    if self.sub_agents:
      for sub_agent in self.sub_agents:
        # `mode` is defined by whichever agent classes declare the field; any
        # agent that defines `mode` participates here. A sub-agent that does not
        # declare `mode` returns None and is never wrapped (it stays an
        # LLM-transfer target).
        mode = getattr(sub_agent, 'mode', None)
        # LlmAgent sub-agents default to chat mode (unchanged behavior).
        if isinstance(sub_agent, LlmAgent) and mode is None:
          sub_agent.mode = 'chat'
          mode = 'chat'
        if mode == 'single_turn':
          self.tools.append(_SingleTurnAgentTool(sub_agent))
        elif mode == 'task':
          self.tools.append(_TaskAgentTool(sub_agent))

  @classmethod
  @experimental(FeatureName.AGENT_CONFIG)
  def _resolve_tools(
      cls, tool_configs: list[ToolConfig], config_abs_path: str
  ) -> list[Any]:
    """Resolve tools from configuration.

    Args:
      tool_configs: List of tool configurations (ToolConfig objects).
      config_abs_path: The absolute path to the agent config file.

    Returns:
      List of resolved tool objects.
    """

    resolved_tools = []
    for tool_config in tool_configs:
      if '.' not in tool_config.name:
        # ADK built-in tools
        module = importlib.import_module('google.adk.tools')
        obj = getattr(module, tool_config.name)
      else:
        # User-defined tools
        from .config_agent_utils import _validate_module_reference

        _validate_module_reference(tool_config.name)
        module_path, obj_name = tool_config.name.rsplit('.', 1)
        module = importlib.import_module(module_path)
        obj = getattr(module, obj_name)

      if isinstance(obj, BaseTool) or isinstance(obj, BaseToolset):
        logger.debug(
            'Tool %s is an instance of BaseTool/BaseToolset.', tool_config.name
        )
        resolved_tools.append(obj)
      elif inspect.isclass(obj) and (
          issubclass(obj, BaseTool) or issubclass(obj, BaseToolset)
      ):
        logger.debug(
            'Tool %s is a sub-class of BaseTool/BaseToolset.', tool_config.name
        )
        resolved_tools.append(
            obj.from_config(tool_config.args, config_abs_path)
        )
      elif callable(obj):
        if tool_config.args:
          logger.debug(
              'Tool %s is a user-defined tool-generating function.',
              tool_config.name,
          )
          resolved_tools.append(obj(tool_config.args))
        else:
          logger.debug(
              'Tool %s is a user-defined function tool.', tool_config.name
          )
          resolved_tools.append(obj)
      else:
        raise ValueError(f'Invalid tool YAML config: {tool_config}.')

    return resolved_tools

  @override
  @classmethod
  @experimental(FeatureName.AGENT_CONFIG)
  def _parse_config(
      cls: Type[LlmAgent],
      config: LlmAgentConfig,
      config_abs_path: str,
      kwargs: Dict[str, Any],
  ) -> Dict[str, Any]:
    from .config_agent_utils import resolve_callbacks
    from .config_agent_utils import resolve_code_reference

    if config.model_code:
      kwargs['model'] = resolve_code_reference(config.model_code)
    elif config.model:
      kwargs['model'] = config.model
    if config.instruction:
      kwargs['instruction'] = config.instruction
    if config.static_instruction:
      kwargs['static_instruction'] = config.static_instruction
    if config.disallow_transfer_to_parent:
      kwargs['disallow_transfer_to_parent'] = config.disallow_transfer_to_parent
    if config.disallow_transfer_to_peers:
      kwargs['disallow_transfer_to_peers'] = config.disallow_transfer_to_peers
    if config.include_contents != 'default':
      kwargs['include_contents'] = config.include_contents
    if config.input_schema:
      kwargs['input_schema'] = resolve_code_reference(config.input_schema)
    if config.output_schema:
      kwargs['output_schema'] = resolve_code_reference(config.output_schema)
    if config.output_key:
      kwargs['output_key'] = config.output_key
    if config.tools:
      kwargs['tools'] = cls._resolve_tools(config.tools, config_abs_path)
    if config.before_model_callbacks:
      kwargs['before_model_callback'] = resolve_callbacks(
          config.before_model_callbacks
      )
    if config.after_model_callbacks:
      kwargs['after_model_callback'] = resolve_callbacks(
          config.after_model_callbacks
      )
    if config.before_tool_callbacks:
      kwargs['before_tool_callback'] = resolve_callbacks(
          config.before_tool_callbacks
      )
    if config.after_tool_callbacks:
      kwargs['after_tool_callback'] = resolve_callbacks(
          config.after_tool_callbacks
      )
    if config.generate_content_config:
      kwargs['generate_content_config'] = config.generate_content_config

    return kwargs


Agent: TypeAlias = LlmAgent
