Skip to content
Open
88 changes: 88 additions & 0 deletions contributing/samples/adk_concurrent_agent_tool_call/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
# Concurrent Agent Tool Call Tests

This sample directory contains tests for concurrency issues that can occur when multiple agents or runners share toolsets and execute tools concurrently. The tests verify that closing one runner or completing one AgentTool call does not interrupt tools being executed by other runners or AgentTool calls that share the same toolset.

## Structure

- **`mock_tools.py`**: Common mock tools and toolsets used by all tests
- `MockTool`: A mock tool that waits for a `done_event` before completing
- `MockMcpToolset`: A mock MCP toolset with a closed event for testing concurrency

- **`runner_shared_toolset/`**: Tests concurrent runner behavior with shared toolsets
- Tests the scenario where two `InMemoryRunner` instances share the same agent and toolset
- Verifies that closing one runner doesn't interrupt tools being executed by the other runner

- **`agent_tool_parallel/`**: Tests parallel AgentTool call behavior
- Tests the scenario where a root agent calls a sub-agent via `AgentTool` multiple times in parallel
- Verifies that `AgentToolManager` properly handles parallel execution of `AgentTool` calls that share the same agent

## Problem Statement

Both test scenarios address similar concurrency issues:

1. **Runner Shared Toolset**: When multiple `Runner` instances share the same agent (and thus the same toolset), closing one runner should not interrupt tools being executed by other runners.

2. **AgentTool Parallel Calls**: When a root agent calls a sub-agent via `AgentTool` multiple times in parallel, each `AgentTool` call creates a `Runner` that uses the same sub-agent. When one `AgentTool` call completes and its runner closes, other parallel calls should not be interrupted.

## Running the Tests

### Runner Shared Toolset Test

```bash
# Run the test script directly
python -m contributing.samples.adk_concurrent_agent_tool_call.runner_shared_toolset.main

# Or use the ADK CLI
adk run contributing/samples/adk_concurrent_agent_tool_call/runner_shared_toolset
```

### AgentTool Parallel Call Test

```bash
# Run the test script directly
python -m contributing.samples.adk_concurrent_agent_tool_call.agent_tool_parallel.main

# Or use the ADK CLI
adk run contributing/samples/adk_concurrent_agent_tool_call/agent_tool_parallel
```

## Common Components

### MockTool

A mock tool that waits for a `done_event` before completing. It checks if the toolset's `closed_event` is set during execution and raises an error if interrupted.

### MockMcpToolset

A mock MCP toolset that simulates a stateful protocol. It creates a new `MockTool` instance on each `get_tools()` call (not cached), which is important for testing the concurrency scenario.

## Expected Behavior

Both tests should verify:

- Tools should start executing concurrently
- When one runner/AgentTool call completes, other parallel executions should continue
- All parallel executions should complete successfully without being interrupted
- No "interrupted" errors should appear in the events

## Key Testing Points

1. **Concurrent Tool Execution**: Verifies that multiple runners/AgentTool calls can execute tools from the same toolset simultaneously
2. **Toolset Closure Handling**: Ensures that closing one runner doesn't affect tools being executed by other runners
3. **State Management**: Tests that shared toolset state is properly managed across multiple runners/AgentTool calls
4. **Error Detection**: Checks for interruption errors in parallel executions

## Implementation Details

Both tests use monkey patching to track when tools are called:

- Patches `functions.__call_tool_async` to track running tools
- Uses `asyncio.Event` to synchronize tool execution
- Monitors events to detect any interruption errors

## Related Components

- **AgentTool**: The tool that wraps an agent and allows it to be called as a tool
- **AgentToolManager**: Manages runner registration and toolset cleanup for `AgentTool`
- **Runner**: The execution engine that orchestrates agent execution
- **BaseToolset**: Base class for toolsets that can be shared across multiple runners
13 changes: 13 additions & 0 deletions contributing/samples/adk_concurrent_agent_tool_call/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from . import agent

__all__ = ["agent"]
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# pylint: disable=g-importing-member

import os
import sys

SAMPLES_DIR = os.path.abspath(
os.path.join(os.path.dirname(__file__), "..", "..")
)
if SAMPLES_DIR not in sys.path:
sys.path.append(SAMPLES_DIR)

from adk_concurrent_agent_tool_call.mock_tools import MockMcpToolset
from google.adk import Agent
from google.adk.tools.agent_tool import AgentTool

# Create a MCP toolset for the sub-agent
sub_agent_mcp_toolset = MockMcpToolset()

sub_agent_system_prompt = """
You are a helpful sub-agent that can use tools to help users.
When asked to use the mcp_tool, you should call it.
"""

# Create a sub-agent that uses the MCP toolset
sub_agent = Agent(
model="gemini-2.5-flash",
name="sub_agent",
description=(
"A sub-agent that uses a MCP toolset for testing parallel AgentTool"
" calls."
),
instruction=sub_agent_system_prompt,
tools=[sub_agent_mcp_toolset],
)

# Create the root agent that uses AgentTool to call the sub-agent
root_agent_system_prompt = """
You are a helpful assistant that can call sub-agents as tools.
When asked to use the sub_agent tool, you should call it.
You can call multiple sub_agent tools in parallel if needed.
"""

root_agent = Agent(
model="gemini-2.5-flash",
name="root_agent",
description=(
"A root agent that calls sub-agents via AgentTool for testing parallel"
" execution."
),
instruction=root_agent_system_prompt,
tools=[AgentTool(agent=sub_agent)],
)
Loading