google--adk-python/tests/unittests/tools/test_agent_tool.py

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

import asyncio
import json
from typing import Any
from typing import Optional

from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.llm_agent import Agent
from google.adk.agents.run_config import RunConfig
from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
from google.adk.events.event import Event
from google.adk.features import FeatureName
from google.adk.features._feature_registry import temporary_feature_override
from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.plugins.base_plugin import BasePlugin
from google.adk.plugins.plugin_manager import PluginManager
from google.adk.runners import Runner
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.tools.agent_tool import AgentTool
from google.adk.tools.tool_context import ToolContext
from google.adk.utils.variant_utils import GoogleLLMVariant
from google.genai import types
from google.genai.types import Part
from pydantic import BaseModel
import pytest
from pytest import mark

from .. import testing_utils

function_call_custom = Part.from_function_call(
    name='tool_agent', args={'custom_input': 'test1'}
)

function_call_no_schema = Part.from_function_call(
    name='tool_agent', args={'request': 'test1'}
)

function_response_custom = Part.from_function_response(
    name='tool_agent', response={'custom_output': 'response1'}
)

function_response_no_schema = Part.from_function_response(
    name='tool_agent', response={'result': 'response1'}
)


def change_state_callback(callback_context: CallbackContext):
  callback_context.state['state_1'] = 'changed_value'
  print('change_state_callback: ', callback_context.state)


@mark.asyncio
async def test_agent_tool_inherits_parent_app_name(monkeypatch):
  parent_app_name = 'parent_app'
  captured: dict[str, str] = {}

  class RecordingSessionService(InMemorySessionService):

    async def create_session(
        self,
        *,
        app_name: str,
        user_id: str,
        state: Optional[dict[str, Any]] = None,
        session_id: Optional[str] = None,
    ):
      captured['session_app_name'] = app_name
      return await super().create_session(
          app_name=app_name,
          user_id=user_id,
          state=state,
          session_id=session_id,
      )

  monkeypatch.setattr(
      'google.adk.sessions.in_memory_session_service.InMemorySessionService',
      RecordingSessionService,
  )

  async def _empty_async_generator():
    if False:
      yield None

  class StubRunner:

    def __init__(
        self,
        *,
        app_name: str,
        agent: Agent,
        artifact_service,
        session_service,
        memory_service,
        credential_service,
        plugins,
    ):
      del artifact_service, memory_service, credential_service
      captured['runner_app_name'] = app_name
      self.agent = agent
      self.session_service = session_service
      self.plugin_manager = PluginManager(plugins=plugins)
      self.app_name = app_name

    def run_async(
        self,
        *,
        user_id: str,
        session_id: str,
        invocation_id: Optional[str] = None,
        new_message: Optional[types.Content] = None,
        state_delta: Optional[dict[str, Any]] = None,
        run_config: Optional[RunConfig] = None,
    ):
      del (
          user_id,
          session_id,
          invocation_id,
          new_message,
          state_delta,
          run_config,
      )
      return _empty_async_generator()

    async def close(self):
      """Mock close method."""
      pass

  monkeypatch.setattr('google.adk.runners.Runner', StubRunner)

  tool_agent = Agent(
      name='tool_agent',
      model='test-model',
  )
  agent_tool = AgentTool(agent=tool_agent)
  root_agent = Agent(
      name='root_agent',
      model='test-model',
      tools=[agent_tool],
  )

  artifact_service = InMemoryArtifactService()
  parent_session_service = InMemorySessionService()
  parent_session = await parent_session_service.create_session(
      app_name=parent_app_name,
      user_id='user',
  )
  invocation_context = InvocationContext(
      artifact_service=artifact_service,
      session_service=parent_session_service,
      memory_service=InMemoryMemoryService(),
      plugin_manager=PluginManager(),
      invocation_id='invocation-id',
      agent=root_agent,
      session=parent_session,
      run_config=RunConfig(),
  )
  tool_context = ToolContext(invocation_context)

  assert tool_context._invocation_context.app_name == parent_app_name

  await agent_tool.run_async(
      args={'request': 'hello'},
      tool_context=tool_context,
  )

  assert captured['runner_app_name'] == parent_app_name
  assert captured['session_app_name'] == parent_app_name


def test_no_schema():
  mock_model = testing_utils.MockModel.create(
      responses=[
          function_call_no_schema,
          'response1',
          'response2',
      ]
  )

  tool_agent = Agent(
      name='tool_agent',
      model=mock_model,
  )

  root_agent = Agent(
      name='root_agent',
      model=mock_model,
      tools=[AgentTool(agent=tool_agent)],
  )

  runner = testing_utils.InMemoryRunner(root_agent)

  assert testing_utils.simplify_events(runner.run('test1')) == [
      ('root_agent', function_call_no_schema),
      ('root_agent', function_response_no_schema),
      ('root_agent', 'response2'),
  ]


def test_use_plugins():
  """The agent tool can use plugins from parent runner."""

  class ModelResponseCapturePlugin(BasePlugin):

    def __init__(self):
      super().__init__('plugin')
      self.model_responses = {}

    async def after_model_callback(
        self,
        *,
        callback_context: CallbackContext,
        llm_response: LlmResponse,
    ) -> Optional[LlmResponse]:
      response_text = []
      for part in llm_response.content.parts:
        if not part.text:
          continue
        response_text.append(part.text)
      if response_text:
        if callback_context.agent_name not in self.model_responses:
          self.model_responses[callback_context.agent_name] = []
        self.model_responses[callback_context.agent_name].append(
            ''.join(response_text)
        )

  mock_model = testing_utils.MockModel.create(
      responses=[
          function_call_no_schema,
          'response1',
          'response2',
      ]
  )

  tool_agent = Agent(
      name='tool_agent',
      model=mock_model,
  )

  root_agent = Agent(
      name='root_agent',
      model=mock_model,
      tools=[AgentTool(agent=tool_agent)],
  )

  model_response_capture = ModelResponseCapturePlugin()
  runner = testing_utils.InMemoryRunner(
      root_agent, plugins=[model_response_capture]
  )

  assert testing_utils.simplify_events(runner.run('test1')) == [
      ('root_agent', function_call_no_schema),
      ('root_agent', function_response_no_schema),
      ('root_agent', 'response2'),
  ]

  # should be able to capture response from both root and tool agent.
  assert model_response_capture.model_responses == {
      'tool_agent': ['response1'],
      'root_agent': ['response2'],
  }


def test_update_state():
  """The agent tool can read and change parent state."""

  mock_model = testing_utils.MockModel.create(
      responses=[
          function_call_no_schema,
          '{"custom_output": "response1"}',
          'response2',
      ]
  )

  tool_agent = Agent(
      name='tool_agent',
      model=mock_model,
      instruction='input: {state_1}',
      before_agent_callback=change_state_callback,
  )

  root_agent = Agent(
      name='root_agent',
      model=mock_model,
      tools=[AgentTool(agent=tool_agent)],
  )

  runner = testing_utils.InMemoryRunner(root_agent)
  runner.session.state['state_1'] = 'state1_value'

  runner.run('test1')
  assert (
      'input: changed_value' in mock_model.requests[1].config.system_instruction
  )
  assert runner.session.state['state_1'] == 'changed_value'


@mark.asyncio
async def test_update_artifacts():
  """The agent tool can read and write artifacts."""

  async def before_tool_agent(callback_context: CallbackContext):
    # Artifact 1 should be available in the tool agent.
    artifact = await callback_context.load_artifact('artifact_1')
    await callback_context.save_artifact(
        'artifact_2', Part.from_text(text=artifact.text + ' 2')
    )

  tool_agent = SequentialAgent(
      name='tool_agent',
      before_agent_callback=before_tool_agent,
  )

  async def before_main_agent(callback_context: CallbackContext):
    await callback_context.save_artifact(
        'artifact_1', Part.from_text(text='test')
    )

  async def after_main_agent(callback_context: CallbackContext):
    # Artifact 2 should be available after the tool agent.
    artifact_2 = await callback_context.load_artifact('artifact_2')
    await callback_context.save_artifact(
        'artifact_3', Part.from_text(text=artifact_2.text + ' 3')
    )

  mock_model = testing_utils.MockModel.create(
      responses=[function_call_no_schema, 'response2']
  )
  root_agent = Agent(
      name='root_agent',
      before_agent_callback=before_main_agent,
      after_agent_callback=after_main_agent,
      tools=[AgentTool(agent=tool_agent)],
      model=mock_model,
  )

  runner = testing_utils.InMemoryRunner(root_agent)
  runner.run('test1')

  async def load_artifact(filename: str):
    return await runner.runner.artifact_service.load_artifact(
        app_name='test_app',
        user_id='test_user',
        session_id=runner.session_id,
        filename=filename,
    )

  assert await runner.runner.artifact_service.list_artifact_keys(
      app_name='test_app', user_id='test_user', session_id=runner.session_id
  ) == ['artifact_1', 'artifact_2', 'artifact_3']

  assert await load_artifact('artifact_1') == Part.from_text(text='test')
  assert await load_artifact('artifact_2') == Part.from_text(text='test 2')
  assert await load_artifact('artifact_3') == Part.from_text(text='test 2 3')


@mark.parametrize(
    'env_variables',
    [
        'GOOGLE_AI',
        # TODO: re-enable after fix.
        # 'VERTEX',
    ],
    indirect=True,
)
def test_custom_schema(env_variables):
  class CustomInput(BaseModel):
    custom_input: str

  class CustomOutput(BaseModel):
    custom_output: str

  mock_model = testing_utils.MockModel.create(
      responses=[
          function_call_custom,
          '{"custom_output": "response1"}',
          'response2',
      ]
  )

  tool_agent = Agent(
      name='tool_agent',
      model=mock_model,
      input_schema=CustomInput,
      output_schema=CustomOutput,
      output_key='tool_output',
  )

  root_agent = Agent(
      name='root_agent',
      model=mock_model,
      tools=[AgentTool(agent=tool_agent)],
  )

  runner = testing_utils.InMemoryRunner(root_agent)
  runner.session.state['state_1'] = 'state1_value'

  assert testing_utils.simplify_events(runner.run('test1')) == [
      ('root_agent', function_call_custom),
      ('root_agent', function_response_custom),
      ('root_agent', 'response2'),
  ]

  assert runner.session.state['tool_output'] == {'custom_output': 'response1'}

  assert len(mock_model.requests) == 3
  # The second request is the tool agent request.
  assert mock_model.requests[1].config.response_schema == CustomOutput
  assert mock_model.requests[1].config.response_mime_type == 'application/json'


@mark.parametrize(
    'env_variables',
    [
        'VERTEX',  # Test VERTEX_AI variant
    ],
    indirect=True,
)
def test_agent_tool_response_schema_no_output_schema_vertex_ai(
    env_variables,
):
  """Test AgentTool with no output schema has string response schema for VERTEX_AI."""
  tool_agent = Agent(
      name='tool_agent',
      model=testing_utils.MockModel.create(responses=['test response']),
  )

  agent_tool = AgentTool(agent=tool_agent)
  declaration = agent_tool._get_declaration()

  assert declaration.name == 'tool_agent'

  from google.adk.features import is_feature_enabled

  if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
    assert declaration.parameters_json_schema == {
        'type': 'object',
        'properties': {'request': {'type': 'string'}},
        'required': ['request'],
    }
    assert declaration.response_json_schema == {'type': 'string'}
  else:
    assert declaration.parameters.type == 'OBJECT'
    assert declaration.parameters.properties['request'].type == 'STRING'
    assert declaration.response is not None
    assert declaration.response.type == types.Type.STRING


@mark.parametrize(
    'env_variables',
    [
        'VERTEX',  # Test VERTEX_AI variant
    ],
    indirect=True,
)
def test_agent_tool_response_schema_with_output_schema_vertex_ai(
    env_variables,
):
  """Test AgentTool with output schema has object response schema for VERTEX_AI."""

  class CustomOutput(BaseModel):
    custom_output: str

  tool_agent = Agent(
      name='tool_agent',
      model=testing_utils.MockModel.create(responses=['test response']),
      output_schema=CustomOutput,
  )

  agent_tool = AgentTool(agent=tool_agent)
  declaration = agent_tool._get_declaration()

  assert declaration.name == 'tool_agent'
  # Should have object response schema for VERTEX_AI when output_schema exists
  from google.adk.features import is_feature_enabled

  if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
    assert declaration.response_json_schema == {'type': 'object'}
  else:
    assert declaration.response is not None
    assert declaration.response.type == types.Type.OBJECT


@mark.parametrize(
    'env_variables',
    [
        'GOOGLE_AI',  # Test GEMINI_API variant
    ],
    indirect=True,
)
def test_agent_tool_response_schema_gemini_api(
    env_variables,
):
  """Test AgentTool with GEMINI_API variant has no response schema."""

  class CustomOutput(BaseModel):
    custom_output: str

  tool_agent = Agent(
      name='tool_agent',
      model=testing_utils.MockModel.create(responses=['test response']),
      output_schema=CustomOutput,
  )

  agent_tool = AgentTool(agent=tool_agent)
  declaration = agent_tool._get_declaration()

  assert declaration.name == 'tool_agent'
  # GEMINI_API should not have response schema
  assert declaration.response is None


@mark.parametrize(
    'env_variables',
    [
        'VERTEX',  # Test VERTEX_AI variant
    ],
    indirect=True,
)
def test_agent_tool_response_schema_with_input_schema_vertex_ai(
    env_variables,
):
  """Test AgentTool with input and output schemas for VERTEX_AI."""

  class CustomInput(BaseModel):
    custom_input: str

  class CustomOutput(BaseModel):
    custom_output: str

  tool_agent = Agent(
      name='tool_agent',
      model=testing_utils.MockModel.create(responses=['test response']),
      input_schema=CustomInput,
      output_schema=CustomOutput,
  )

  agent_tool = AgentTool(agent=tool_agent)
  declaration = agent_tool._get_declaration()

  assert declaration.name == 'tool_agent'
  from google.adk.features import is_feature_enabled

  if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
    assert declaration.parameters_json_schema == {
        'title': 'CustomInput',
        'type': 'object',
        'properties': {
            'custom_input': {'title': 'Custom Input', 'type': 'string'}
        },
        'required': ['custom_input'],
    }
    assert declaration.response_json_schema == {'type': 'object'}
  else:
    assert declaration.parameters.type == 'OBJECT'
    assert declaration.parameters.properties['custom_input'].type == 'STRING'
    # Should have object response schema for VERTEX_AI when output_schema exists
    assert declaration.response is not None
    assert declaration.response.type == types.Type.OBJECT


@mark.parametrize(
    'env_variables',
    [
        'VERTEX',  # Test VERTEX_AI variant
    ],
    indirect=True,
)
def test_agent_tool_response_schema_with_input_schema_no_output_vertex_ai(
    env_variables,
):
  """Test AgentTool with input schema but no output schema for VERTEX_AI."""

  class CustomInput(BaseModel):
    custom_input: str

  tool_agent = Agent(
      name='tool_agent',
      model=testing_utils.MockModel.create(responses=['test response']),
      input_schema=CustomInput,
  )

  agent_tool = AgentTool(agent=tool_agent)
  declaration = agent_tool._get_declaration()

  assert declaration.name == 'tool_agent'
  from google.adk.features import is_feature_enabled

  if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
    assert declaration.parameters_json_schema == {
        'title': 'CustomInput',
        'type': 'object',
        'properties': {
            'custom_input': {'title': 'Custom Input', 'type': 'string'}
        },
        'required': ['custom_input'],
    }
    assert declaration.response_json_schema == {'type': 'string'}
  else:
    assert declaration.parameters.type == 'OBJECT'
    assert declaration.parameters.properties['custom_input'].type == 'STRING'
    # Should have string response schema for VERTEX_AI when no output_schema
    assert declaration.response is not None
    assert declaration.response.type == types.Type.STRING


def test_include_plugins_default_true():
  """Test that plugins are propagated by default (include_plugins=True)."""

  # Create a test plugin that tracks callbacks
  class TrackingPlugin(BasePlugin):

    def __init__(self, name: str):
      super().__init__(name)
      self.before_agent_calls = 0

    async def before_agent_callback(self, **kwargs):
      self.before_agent_calls += 1

  tracking_plugin = TrackingPlugin(name='tracking')

  mock_model = testing_utils.MockModel.create(
      responses=[function_call_no_schema, 'response1', 'response2']
  )

  tool_agent = Agent(
      name='tool_agent',
      model=mock_model,
  )

  root_agent = Agent(
      name='root_agent',
      model=mock_model,
      tools=[AgentTool(agent=tool_agent)],  # Default include_plugins=True
  )

  runner = testing_utils.InMemoryRunner(root_agent, plugins=[tracking_plugin])
  runner.run('test1')

  # Plugin should be called for both root_agent and tool_agent.
  assert tracking_plugin.before_agent_calls == 2


def test_include_plugins_explicit_true():
  """Test that plugins are propagated when include_plugins=True."""

  class TrackingPlugin(BasePlugin):

    def __init__(self, name: str):
      super().__init__(name)
      self.before_agent_calls = 0

    async def before_agent_callback(self, **kwargs):
      self.before_agent_calls += 1

  tracking_plugin = TrackingPlugin(name='tracking')

  mock_model = testing_utils.MockModel.create(
      responses=[function_call_no_schema, 'response1', 'response2']
  )

  tool_agent = Agent(
      name='tool_agent',
      model=mock_model,
  )

  root_agent = Agent(
      name='root_agent',
      model=mock_model,
      tools=[AgentTool(agent=tool_agent, include_plugins=True)],
  )

  runner = testing_utils.InMemoryRunner(root_agent, plugins=[tracking_plugin])
  runner.run('test1')

  # Plugin should be called for both root_agent and tool_agent.
  assert tracking_plugin.before_agent_calls == 2


def test_include_plugins_false():
  """Test that plugins are NOT propagated when include_plugins=False."""

  class TrackingPlugin(BasePlugin):

    def __init__(self, name: str):
      super().__init__(name)
      self.before_agent_calls = 0

    async def before_agent_callback(self, **kwargs):
      self.before_agent_calls += 1

  tracking_plugin = TrackingPlugin(name='tracking')

  mock_model = testing_utils.MockModel.create(
      responses=[function_call_no_schema, 'response1', 'response2']
  )

  tool_agent = Agent(
      name='tool_agent',
      model=mock_model,
  )

  root_agent = Agent(
      name='root_agent',
      model=mock_model,
      tools=[AgentTool(agent=tool_agent, include_plugins=False)],
  )

  runner = testing_utils.InMemoryRunner(root_agent, plugins=[tracking_plugin])
  runner.run('test1')

  # Plugin should only be called for root_agent, not tool_agent.
  assert tracking_plugin.before_agent_calls == 1


@pytest.mark.asyncio
async def test_include_plugins_true_sub_runner_does_not_close_parent_plugins():
  """Sub-Runner must not close plugins owned by the parent runner."""

  class SlowClosePlugin(BasePlugin):

    def __init__(self, name: str):
      super().__init__(name)
      self.close_calls = 0

    async def close(self):
      self.close_calls += 1
      # Would otherwise blow past the sub-Runner's plugin_close_timeout.
      await asyncio.sleep(10)

  parent_plugin = SlowClosePlugin(name='parent_plugin')

  mock_model = testing_utils.MockModel.create(
      responses=[function_call_no_schema, 'response1', 'response2']
  )

  tool_agent = Agent(name='tool_agent', model=mock_model)
  root_agent = Agent(
      name='root_agent',
      model=mock_model,
      tools=[AgentTool(agent=tool_agent, include_plugins=True)],
  )

  runner = Runner(
      app_name='test_app',
      agent=root_agent,
      artifact_service=InMemoryArtifactService(),
      session_service=InMemorySessionService(),
      memory_service=InMemoryMemoryService(),
      plugins=[parent_plugin],
      # Tight timeout amplifies the bug if it regresses; with the fix, the
      # sub-Runner's close skips the parent's plugins entirely.
      plugin_close_timeout=0.01,
  )
  session = await runner.session_service.create_session(
      app_name='test_app', user_id='test_user'
  )
  # Must not raise RuntimeError("Failed to close plugins: ...") from the
  # sub-Runner closing the parent's slow-to-close plugin.
  async for _ in runner.run_async(
      user_id=session.user_id,
      session_id=session.id,
      new_message=testing_utils.get_user_content('test1'),
  ):
    pass

  # The sub-Runner must not have closed the parent's plugin.
  assert parent_plugin.close_calls == 0


def test_agent_tool_description_with_input_schema():
  """Test that agent description is propagated when using input_schema."""

  class CustomInput(BaseModel):
    """This is the Pydantic model docstring."""

    custom_input: str

  agent_description = 'This is the agent description that should be used'
  tool_agent = Agent(
      name='tool_agent',
      model=testing_utils.MockModel.create(responses=['test response']),
      description=agent_description,
      input_schema=CustomInput,
  )

  agent_tool = AgentTool(agent=tool_agent)
  declaration = agent_tool._get_declaration()

  # The description should come from the agent, not the Pydantic model
  assert declaration.description == agent_description


@pytest.fixture
def enable_json_schema_feature():
  """Fixture to enable JSON_SCHEMA_FOR_FUNC_DECL feature for a test."""
  with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True):
    yield


def test_agent_tool_no_schema_with_json_schema_feature(
    enable_json_schema_feature,
):
  """Test AgentTool without input_schema uses parameters_json_schema when feature enabled."""
  tool_agent = Agent(
      name='tool_agent',
      description='A tool agent for testing.',
      model=testing_utils.MockModel.create(responses=['test response']),
  )

  agent_tool = AgentTool(agent=tool_agent)
  declaration = agent_tool._get_declaration()

  assert declaration.model_dump(exclude_none=True) == {
      'name': 'tool_agent',
      'description': 'A tool agent for testing.',
      'parameters_json_schema': {
          'type': 'object',
          'properties': {
              'request': {'type': 'string'},
          },
          'required': ['request'],
      },
  }


@mark.parametrize(
    'env_variables',
    [
        'VERTEX',  # Test VERTEX_AI variant
    ],
    indirect=True,
)
def test_agent_tool_response_json_schema_no_output_schema_vertex_ai(
    env_variables,
    enable_json_schema_feature,
):
  """Test AgentTool with no output schema uses response_json_schema for VERTEX_AI when feature enabled."""
  tool_agent = Agent(
      name='tool_agent',
      description='A tool agent for testing.',
      model=testing_utils.MockModel.create(responses=['test response']),
  )

  agent_tool = AgentTool(agent=tool_agent)
  declaration = agent_tool._get_declaration()

  assert declaration.model_dump(exclude_none=True) == {
      'name': 'tool_agent',
      'description': 'A tool agent for testing.',
      'parameters_json_schema': {
          'type': 'object',
          'properties': {
              'request': {'type': 'string'},
          },
          'required': ['request'],
      },
      'response_json_schema': {'type': 'string'},
  }


@mark.parametrize(
    'env_variables',
    [
        'VERTEX',  # Test VERTEX_AI variant
    ],
    indirect=True,
)
def test_agent_tool_response_json_schema_with_output_schema_vertex_ai(
    env_variables,
    enable_json_schema_feature,
):
  """Test AgentTool with output schema uses response_json_schema for VERTEX_AI when feature enabled."""

  class CustomOutput(BaseModel):
    custom_output: str

  tool_agent = Agent(
      name='tool_agent',
      description='A tool agent for testing.',
      model=testing_utils.MockModel.create(responses=['test response']),
      output_schema=CustomOutput,
  )

  agent_tool = AgentTool(agent=tool_agent)
  declaration = agent_tool._get_declaration()

  assert declaration.model_dump(exclude_none=True) == {
      'name': 'tool_agent',
      'description': 'A tool agent for testing.',
      'parameters_json_schema': {
          'type': 'object',
          'properties': {
              'request': {'type': 'string'},
          },
          'required': ['request'],
      },
      'response_json_schema': {'type': 'object'},
  }


@mark.parametrize(
    'env_variables',
    [
        'GOOGLE_AI',  # Test GEMINI_API variant
    ],
    indirect=True,
)
def test_agent_tool_no_response_json_schema_gemini_api(
    env_variables,
    enable_json_schema_feature,
):
  """Test AgentTool with GEMINI_API variant has no response_json_schema when feature enabled."""

  class CustomOutput(BaseModel):
    custom_output: str

  tool_agent = Agent(
      name='tool_agent',
      description='A tool agent for testing.',
      model=testing_utils.MockModel.create(responses=['test response']),
      output_schema=CustomOutput,
  )

  agent_tool = AgentTool(agent=tool_agent)
  declaration = agent_tool._get_declaration()

  # GEMINI_API should not have response_json_schema
  assert declaration.model_dump(exclude_none=True) == {
      'name': 'tool_agent',
      'description': 'A tool agent for testing.',
      'parameters_json_schema': {
          'type': 'object',
          'properties': {
              'request': {'type': 'string'},
          },
          'required': ['request'],
      },
  }


@mark.parametrize(
    'env_variables',
    [
        'VERTEX',  # Test VERTEX_AI variant
    ],
    indirect=True,
)
def test_agent_tool_with_input_schema_uses_json_schema_feature(
    env_variables,
    enable_json_schema_feature,
):
  """Test AgentTool with input_schema uses parameters_json_schema when feature enabled."""

  class CustomInput(BaseModel):
    custom_input: str

  class CustomOutput(BaseModel):
    custom_output: str

  tool_agent = Agent(
      name='tool_agent',
      description='A tool agent for testing.',
      model=testing_utils.MockModel.create(responses=['test response']),
      input_schema=CustomInput,
      output_schema=CustomOutput,
  )

  agent_tool = AgentTool(agent=tool_agent)
  declaration = agent_tool._get_declaration()

  # When input_schema is provided, build_function_declaration uses Pydantic's
  # model_json_schema() which includes additional fields like 'title'
  assert declaration.model_dump(exclude_none=True) == {
      'name': 'tool_agent',
      'description': 'A tool agent for testing.',
      'parameters_json_schema': {
          'properties': {
              'custom_input': {'title': 'Custom Input', 'type': 'string'},
          },
          'required': ['custom_input'],
          'title': 'CustomInput',
          'type': 'object',
      },
      'response_json_schema': {'type': 'object'},
  }


@mark.asyncio
async def test_run_async_handles_none_parts_in_response():
  """Verify run_async handles None parts in response without raising TypeError."""

  # Mock model for the tool_agent that returns content with parts=None
  # This simulates the condition causing the TypeError
  tool_agent_model = testing_utils.MockModel.create(
      responses=[
          LlmResponse(
              content=types.Content(parts=None),
          )
      ]
  )

  tool_agent = Agent(
      name='tool_agent',
      model=tool_agent_model,
  )

  agent_tool = AgentTool(agent=tool_agent)

  session_service = InMemorySessionService()
  session = await session_service.create_session(
      app_name='test_app', user_id='test_user'
  )

  invocation_context = InvocationContext(
      invocation_id='invocation_id',
      agent=tool_agent,
      session=session,
      session_service=session_service,
  )
  tool_context = ToolContext(invocation_context=invocation_context)

  # This should not raise `TypeError: 'NoneType' object is not iterable`.
  tool_result = await agent_tool.run_async(
      args={'request': 'test request'}, tool_context=tool_context
  )

  assert tool_result == ''


async def _run_agent_tool_with_parts(parts: list[types.Part]) -> Any:
  """Drives AgentTool with an inner agent whose final event content is `parts`."""

  class _StaticAgent(BaseAgent):

    async def _run_async_impl(self, ctx):
      yield Event(
          invocation_id=ctx.invocation_id,
          author=self.name,
          content=types.Content(role='model', parts=parts),
      )

  inner = _StaticAgent(name='inner_agent', description='static')
  agent_tool = AgentTool(agent=inner)

  session_service = InMemorySessionService()
  session = await session_service.create_session(
      app_name='test_app', user_id='test_user'
  )
  invocation_context = InvocationContext(
      invocation_id='invocation_id',
      agent=inner,
      session=session,
      session_service=session_service,
  )
  tool_context = ToolContext(invocation_context=invocation_context)

  return await agent_tool.run_async(
      args={'request': 'test request'}, tool_context=tool_context
  )


@mark.asyncio
async def test_run_async_extracts_text_only():
  """Plain text parts pass through unchanged."""
  result = await _run_agent_tool_with_parts([types.Part(text='hello world')])
  assert result == 'hello world'


@mark.asyncio
async def test_run_async_extracts_code_execution_result_only():
  """code_execution_result.output and executable_code.code are returned."""
  result = await _run_agent_tool_with_parts([
      types.Part(
          executable_code=types.ExecutableCode(
              language=types.Language.PYTHON, code='print(2 ** 10)'
          )
      ),
      types.Part(
          code_execution_result=types.CodeExecutionResult(
              outcome=types.Outcome.OUTCOME_OK, output='1024\n'
          )
      ),
  ])
  assert result == 'print(2 ** 10)\n1024'


@mark.asyncio
async def test_run_async_extracts_text_and_code_execution_result():
  """Mixed text + code parts are concatenated in order."""
  result = await _run_agent_tool_with_parts([
      types.Part(text='Here is the answer:'),
      types.Part(
          executable_code=types.ExecutableCode(
              language=types.Language.PYTHON, code='print(2 ** 10)'
          )
      ),
      types.Part(
          code_execution_result=types.CodeExecutionResult(
              outcome=types.Outcome.OUTCOME_OK, output='1024\n'
          )
      ),
  ])
  assert result == 'Here is the answer:\nprint(2 ** 10)\n1024'


@mark.asyncio
async def test_run_async_extracts_executable_code_only():
  """executable_code.code alone is returned when no result part follows."""
  result = await _run_agent_tool_with_parts([
      types.Part(
          executable_code=types.ExecutableCode(
              language=types.Language.PYTHON, code='print("hi")'
          )
      ),
  ])
  assert result == 'print("hi")'


@mark.asyncio
async def test_run_async_skips_thought_parts():
  """Parts marked thought=True are dropped regardless of kind."""
  result = await _run_agent_tool_with_parts([
      types.Part(text='thinking out loud', thought=True),
      types.Part(
          code_execution_result=types.CodeExecutionResult(
              outcome=types.Outcome.OUTCOME_OK, output='42\n'
          )
      ),
  ])
  assert result == '42'


class TestAgentToolWithCompositeAgents:
  """Tests for AgentTool wrapping composite agents (SequentialAgent, etc.)."""

  def test_sequential_agent_with_first_sub_agent_input_schema(self):
    """Test that AgentTool exposes input_schema from first sub-agent of SequentialAgent."""

    class CustomInput(BaseModel):
      query: str
      language: str

    first_agent = Agent(
        name='first_agent',
        model=testing_utils.MockModel.create(responses=['response1']),
        input_schema=CustomInput,
    )

    second_agent = Agent(
        name='second_agent',
        model=testing_utils.MockModel.create(responses=['response2']),
    )

    sequence = SequentialAgent(
        name='sequence',
        description='Process the query through multiple steps',
        sub_agents=[first_agent, second_agent],
    )

    agent_tool = AgentTool(agent=sequence)
    declaration = agent_tool._get_declaration()

    # Should expose CustomInput schema, not fallback to 'request'
    assert declaration.name == 'sequence'
    assert declaration.description == 'Process the query through multiple steps'

    from google.adk.features import is_feature_enabled

    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
      assert declaration.parameters_json_schema == {
          'title': 'CustomInput',
          'type': 'object',
          'properties': {
              'query': {'title': 'Query', 'type': 'string'},
              'language': {'title': 'Language', 'type': 'string'},
          },
          'required': ['query', 'language'],
      }
    else:
      assert declaration.parameters.properties['query'].type == 'STRING'
      assert declaration.parameters.properties['language'].type == 'STRING'
      assert 'request' not in declaration.parameters.properties

  def test_sequential_agent_without_input_schema_falls_back_to_request(self):
    """Test that AgentTool falls back to 'request' when no sub-agent has input_schema."""

    first_agent = Agent(
        name='first_agent',
        model=testing_utils.MockModel.create(responses=['response1']),
    )

    second_agent = Agent(
        name='second_agent',
        model=testing_utils.MockModel.create(responses=['response2']),
    )

    sequence = SequentialAgent(
        name='sequence',
        description='Process the query through multiple steps',
        sub_agents=[first_agent, second_agent],
    )

    agent_tool = AgentTool(agent=sequence)
    declaration = agent_tool._get_declaration()

    # Should fall back to 'request' parameter
    assert declaration.name == 'sequence'

    from google.adk.features import is_feature_enabled

    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
      assert declaration.parameters_json_schema == {
          'type': 'object',
          'properties': {'request': {'type': 'string'}},
          'required': ['request'],
      }
    else:
      assert declaration.parameters.properties['request'].type == 'STRING'
      assert 'query' not in declaration.parameters.properties

  @mark.parametrize(
      'env_variables',
      [
          'VERTEX',
      ],
      indirect=True,
  )
  def test_sequential_agent_with_last_sub_agent_output_schema(
      self, env_variables
  ):
    """Test that AgentTool uses output_schema from last sub-agent of SequentialAgent."""

    class CustomOutput(BaseModel):
      result: str

    first_agent = Agent(
        name='first_agent',
        model=testing_utils.MockModel.create(responses=['response1']),
    )

    second_agent = Agent(
        name='second_agent',
        model=testing_utils.MockModel.create(responses=['response2']),
        output_schema=CustomOutput,
    )

    sequence = SequentialAgent(
        name='sequence',
        description='Process the query',
        sub_agents=[first_agent, second_agent],
    )

    agent_tool = AgentTool(agent=sequence)
    declaration = agent_tool._get_declaration()

    # Should have object response schema from last sub-agent
    from google.adk.features import is_feature_enabled

    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
      assert declaration.response_json_schema == {'type': 'object'}
    else:
      assert declaration.response is not None
      assert declaration.response.type == types.Type.OBJECT

  def test_nested_sequential_agent_input_schema(self):
    """Test that AgentTool recursively finds input_schema in nested composite agents."""

    class CustomInput(BaseModel):
      deep_query: str

    inner_agent = Agent(
        name='inner_agent',
        model=testing_utils.MockModel.create(responses=['response1']),
        input_schema=CustomInput,
    )

    inner_sequence = SequentialAgent(
        name='inner_sequence',
        sub_agents=[inner_agent],
    )

    outer_sequence = SequentialAgent(
        name='outer_sequence',
        description='Nested sequence',
        sub_agents=[inner_sequence],
    )

    agent_tool = AgentTool(agent=outer_sequence)
    declaration = agent_tool._get_declaration()

    # Should recursively find CustomInput from inner_agent
    assert declaration.name == 'outer_sequence'

    from google.adk.features import is_feature_enabled

    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
      assert declaration.parameters_json_schema == {
          'title': 'CustomInput',
          'type': 'object',
          'properties': {
              'deep_query': {'title': 'Deep Query', 'type': 'string'}
          },
          'required': ['deep_query'],
      }
    else:
      assert 'deep_query' in declaration.parameters.properties
      assert declaration.parameters.properties['deep_query'].type == 'STRING'
      assert 'request' not in declaration.parameters.properties

  @mark.parametrize(
      'env_variables',
      [
          'GOOGLE_AI',
          'VERTEX',
      ],
      indirect=True,
  )
  def test_sequential_agent_custom_schema_end_to_end(self, env_variables):
    """Test end-to-end flow with SequentialAgent using custom input/output schema."""

    class CustomInput(BaseModel):
      custom_input: str

    class CustomOutput(BaseModel):
      custom_output: str

    function_call_seq = Part.from_function_call(
        name='sequence', args={'custom_input': 'test_input'}
    )

    mock_model = testing_utils.MockModel.create(
        responses=[
            function_call_seq,
            '{"custom_output": "step1_response"}',
            '{"custom_output": "final_response"}',
            'root_response',
        ]
    )

    first_agent = Agent(
        name='first_agent',
        model=mock_model,
        input_schema=CustomInput,
    )

    second_agent = Agent(
        name='second_agent',
        model=mock_model,
        output_schema=CustomOutput,
        output_key='seq_output',
    )

    sequence = SequentialAgent(
        name='sequence',
        description='A sequential pipeline',
        sub_agents=[first_agent, second_agent],
    )

    root_agent = Agent(
        name='root_agent',
        model=mock_model,
        tools=[AgentTool(agent=sequence)],
    )

    runner = testing_utils.InMemoryRunner(root_agent)
    runner.run('test1')

    # Verify the tool declaration sent to LLM has the correct schema
    # The first request is from root_agent, which should have the tool declaration
    first_request = mock_model.requests[0]
    tool_declarations = first_request.config.tools
    assert len(tool_declarations) == 1

    sequence_tool = tool_declarations[0].function_declarations[0]
    assert sequence_tool.name == 'sequence'

    from google.adk.features import is_feature_enabled

    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
      assert sequence_tool.parameters_json_schema == {
          'title': 'CustomInput',
          'type': 'object',
          'properties': {
              'custom_input': {'title': 'Custom Input', 'type': 'string'}
          },
          'required': ['custom_input'],
      }
    else:
      # Should have 'custom_input' parameter from first sub-agent's input_schema
      assert 'custom_input' in sequence_tool.parameters.properties
      # Should NOT have the fallback 'request' parameter
      assert 'request' not in sequence_tool.parameters.properties

  def test_empty_sequential_agent_falls_back_to_request(self):
    """Test that AgentTool with empty SequentialAgent falls back to 'request'."""

    sequence = SequentialAgent(
        name='empty_sequence',
        description='An empty sequence',
        sub_agents=[],
    )

    agent_tool = AgentTool(agent=sequence)
    declaration = agent_tool._get_declaration()

    # Should fall back to 'request' parameter
    from google.adk.features import is_feature_enabled

    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
      assert declaration.parameters_json_schema == {
          'type': 'object',
          'properties': {'request': {'type': 'string'}},
          'required': ['request'],
      }
    else:
      assert declaration.parameters.properties['request'].type == 'STRING'


@mark.parametrize(
    'args,expected_text',
    [
        (
            {'brand': 'Nike', 'product': 'running shoes'},
            '{"brand": "Nike", "product": "running shoes"}',
        ),
        (
            {'request': 'find me Nike running shoes'},
            'find me Nike running shoes',
        ),
        (
            {'request': ''},
            '',
        ),
    ],
)
@mark.asyncio
async def test_no_schema_args_handling(monkeypatch, args, expected_text):
  """AgentTool.run_async handles fallback schema cases properly.

  - Non-'request' args are serialized as JSON.
  - 'request' key is kept as plain text (backward compatibility).
  - Empty string 'request' is correctly preserved instead of evaluating to
  false.
  """
  captured = {}

  async def _empty_async_generator():
    if False:
      yield None

  class StubRunner:

    def __init__(
        self,
        *,
        app_name: str,
        agent,
        artifact_service,
        session_service,
        memory_service,
        credential_service,
        plugins,
    ):
      del artifact_service, memory_service, credential_service
      self.agent = agent
      self.session_service = session_service
      self.plugin_manager = PluginManager(plugins=plugins)
      self.app_name = app_name

    def run_async(
        self,
        *,
        user_id: str,
        session_id: str,
        invocation_id=None,
        new_message=None,
        state_delta=None,
        run_config=None,
    ):
      captured['new_message'] = new_message
      return _empty_async_generator()

    async def close(self):
      pass

  monkeypatch.setattr('google.adk.runners.Runner', StubRunner)

  tool_agent = Agent(name='tool_agent', model='test-model')
  agent_tool = AgentTool(agent=tool_agent)
  root_agent = Agent(name='root_agent', model='test-model', tools=[agent_tool])

  session_service = InMemorySessionService()
  session = await session_service.create_session(
      app_name='test_app', user_id='user'
  )
  invocation_context = InvocationContext(
      artifact_service=InMemoryArtifactService(),
      session_service=session_service,
      memory_service=InMemoryMemoryService(),
      plugin_manager=PluginManager(),
      invocation_id='test-invocation',
      agent=root_agent,
      session=session,
      run_config=RunConfig(),
  )
  tool_context = ToolContext(invocation_context)

  await agent_tool.run_async(
      args=args,
      tool_context=tool_context,
  )

  assert captured['new_message'] is not None
  text = captured['new_message'].parts[0].text
  assert text == expected_text


@pytest.fixture
def setup_skip_summarization_runner():
  def _setup_runner(tool_agent_model_responses, tool_agent_output_schema=None):
    tool_agent_model = testing_utils.MockModel.create(
        responses=tool_agent_model_responses
    )
    tool_agent = Agent(
        name='tool_agent',
        model=tool_agent_model,
        output_schema=tool_agent_output_schema,
    )

    agent_tool = AgentTool(agent=tool_agent, skip_summarization=True)

    root_agent_model = testing_utils.MockModel.create(
        responses=[
            function_call_no_schema,
            'final_summary_text_that_should_not_be_reached',
        ]
    )

    root_agent = Agent(
        name='root_agent',
        model=root_agent_model,
        tools=[agent_tool],
    )
    return testing_utils.InMemoryRunner(root_agent)

  return _setup_runner


def test_agent_tool_skip_summarization_has_text_output(
    setup_skip_summarization_runner,
):
  """Tests that when skip_summarization is True, the final event contains text content."""
  runner = setup_skip_summarization_runner(
      tool_agent_model_responses=['tool_response_text']
  )
  events = runner.run('start')

  final_events = [e for e in events if e.is_final_response()]
  assert final_events
  last_event = final_events[-1]
  assert last_event.is_final_response()

  assert any(p.function_response for p in last_event.content.parts)

  assert [p.text for p in last_event.content.parts if p.text] == [
      'tool_response_text'
  ]


def test_agent_tool_skip_summarization_preserves_json_string_output(
    setup_skip_summarization_runner,
):
  """Tests that structured output string is preserved as text when skipping summarization."""
  runner = setup_skip_summarization_runner(
      tool_agent_model_responses=['{"field": "value"}']
  )
  events = runner.run('start')

  final_events = [e for e in events if e.is_final_response()]
  assert final_events
  last_event = final_events[-1]
  assert last_event.is_final_response()

  text_parts = [p.text for p in last_event.content.parts if p.text]

  # Check that the JSON string content is preserved exactly
  assert text_parts == ['{"field": "value"}']


def test_agent_tool_skip_summarization_handles_non_string_result(
    setup_skip_summarization_runner,
):
  """Tests that non-string (dict) output is correctly serialized as JSON text."""

  class CustomOutput(BaseModel):
    value: int

  runner = setup_skip_summarization_runner(
      tool_agent_model_responses=['{"value": 123}'],
      tool_agent_output_schema=CustomOutput,
  )
  events = runner.run('start')

  final_events = [e for e in events if e.is_final_response()]
  assert final_events
  last_event = final_events[-1]

  text_parts = [p.text for p in last_event.content.parts if p.text]

  assert text_parts == ['{"value": 123}']