Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions packages/core/src/agents/agentLoader.ts
Original file line number Diff line number Diff line change
Expand Up @@ -107,9 +107,11 @@ const localAgentSchema = z
display_name: z.string().optional(),
tools: z
.array(
z.string().refine((val) => isValidToolName(val), {
message: 'Invalid tool name',
}),
z
.string()
.refine((val) => isValidToolName(val, { allowWildcards: true }), {
message: 'Invalid tool name',
}),
)
.optional(),
model: z.string().optional(),
Expand Down
29 changes: 29 additions & 0 deletions packages/core/src/agents/local-executor.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ import {
type PartListUnion,
type Tool,
type CallableTool,
type FunctionDeclaration,
} from '@google/genai';
import type { Config } from '../config/config.js';
import { MockTool } from '../test-utils/mock-tool.js';
Expand Down Expand Up @@ -550,6 +551,34 @@ describe('LocalAgentExecutor', () => {

getToolSpy.mockRestore();
});

it('should not duplicate schemas when instantiated tools are provided in toolConfig', async () => {
// Create an instantiated mock tool
const instantiatedTool = new MockTool({ name: 'instantiated_tool' });

// Create an agent definition containing the instantiated tool
const definition = createTestDefinition([instantiatedTool]);

// Create the executor
const executor = await LocalAgentExecutor.create(
definition,
mockConfig,
onActivity,
);

// Extract the prepared tools list using the private method
const toolsList = (
executor as unknown as { prepareToolsList: () => FunctionDeclaration[] }
).prepareToolsList();

// Filter for the specific tool schema
const foundSchemas = (
toolsList as unknown as FunctionDeclaration[]
).filter((t: FunctionDeclaration) => t.name === 'instantiated_tool');

// Assert that there is exactly ONE schema for this tool
expect(foundSchemas).toHaveLength(1);
});
});

describe('run (Execution Loop and Logic)', () => {
Expand Down
71 changes: 48 additions & 23 deletions packages/core/src/agents/local-executor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,13 @@ import {
type Schema,
} from '@google/genai';
import { ToolRegistry } from '../tools/tool-registry.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { type AnyDeclarativeTool } from '../tools/tools.js';
import {
DiscoveredMCPTool,
isMcpToolName,
parseMcpToolName,
MCP_TOOL_PREFIX,
} from '../tools/mcp-tool.js';
import { CompressionStatus } from '../core/turn.js';
import { type ToolCallRequestInfo } from '../scheduler/types.js';
import { type Message } from '../confirmation-bus/types.js';
Expand Down Expand Up @@ -141,28 +147,55 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
runtimeContext.getAgentRegistry().getAllAgentNames(),
);

const registerToolByName = (toolName: string) => {
const registerToolInstance = (tool: AnyDeclarativeTool) => {
// Check if the tool is a subagent to prevent recursion.
// We do not allow agents to call other agents.
if (allAgentNames.has(toolName)) {
if (allAgentNames.has(tool.name)) {
debugLogger.warn(
`[LocalAgentExecutor] Skipping subagent tool '${toolName}' for agent '${definition.name}' to prevent recursion.`,
`[LocalAgentExecutor] Skipping subagent tool '${tool.name}' for agent '${definition.name}' to prevent recursion.`,
);
return;
}

agentToolRegistry.registerTool(tool);
};

const registerToolByName = (toolName: string) => {
// Handle global wildcard
if (toolName === '*') {
for (const tool of parentToolRegistry.getAllTools()) {
registerToolInstance(tool);
}
return;
}

// Handle MCP wildcards
if (isMcpToolName(toolName)) {
if (toolName === `${MCP_TOOL_PREFIX}*`) {
for (const tool of parentToolRegistry.getAllTools()) {
if (tool instanceof DiscoveredMCPTool) {
registerToolInstance(tool);
}
}
return;
}

const parsed = parseMcpToolName(toolName);
if (parsed.serverName && parsed.toolName === '*') {
for (const tool of parentToolRegistry.getToolsByServer(
parsed.serverName,
)) {
registerToolInstance(tool);
}
return;
}
}

// If the tool is referenced by name, retrieve it from the parent
// registry and register it with the agent's isolated registry.
const tool = parentToolRegistry.getTool(toolName);
if (tool) {
if (tool instanceof DiscoveredMCPTool) {
// Subagents MUST use fully qualified names for MCP tools to ensure
// unambiguous tool calls and to comply with policy requirements.
// We automatically "upgrade" any MCP tool to its qualified version.
agentToolRegistry.registerTool(tool.asFullyQualifiedTool());
} else {
agentToolRegistry.registerTool(tool);
}
registerToolInstance(tool);
}
};

Expand Down Expand Up @@ -1171,22 +1204,14 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
const { toolConfig, outputConfig } = this.definition;

if (toolConfig) {
const toolNamesToLoad: string[] = [];
for (const toolRef of toolConfig.tools) {
if (typeof toolRef === 'string') {
toolNamesToLoad.push(toolRef);
} else if (typeof toolRef === 'object' && 'schema' in toolRef) {
// Tool instance with an explicit schema property.
toolsList.push(toolRef.schema);
} else {
if (typeof toolRef === 'object' && !('schema' in toolRef)) {
// Raw `FunctionDeclaration` object.
toolsList.push(toolRef);
}
}
// Add schemas from tools that were registered by name.
toolsList.push(
...this.toolRegistry.getFunctionDeclarationsFiltered(toolNamesToLoad),
);
// Add schemas from tools that were explicitly registered by name, wildcard, or instance.
toolsList.push(...this.toolRegistry.getFunctionDeclarations());
}

// Always inject complete_task.
Expand Down
20 changes: 1 addition & 19 deletions packages/core/src/tools/mcp-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ export function parseMcpToolName(name: string): {
// Remove the prefix
const withoutPrefix = name.slice(MCP_TOOL_PREFIX.length);
// The first segment is the server name, the rest is the tool name
// Must be strictly `server_tool` where neither are empty
const match = withoutPrefix.match(/^([^_]+)_(.+)$/);
if (match) {
return {
Expand Down Expand Up @@ -390,25 +391,6 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
`${this.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${this.serverToolName}`,
);
}

asFullyQualifiedTool(): DiscoveredMCPTool {
return new DiscoveredMCPTool(
this.mcpTool,
this.serverName,
this.serverToolName,
this.description,
this.parameterSchema,
this.messageBus,
this.trust,
this.isReadOnly,
this.getFullyQualifiedName(),
this.cliConfig,
this.extensionName,
this.extensionId,
this._toolAnnotations,
);
}

protected createInvocation(
params: ToolParams,
messageBus: MessageBus,
Expand Down
40 changes: 22 additions & 18 deletions packages/core/src/tools/tool-names.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,8 @@ vi.mock('./tool-names.js', async (importOriginal) => {
...actual,
TOOL_LEGACY_ALIASES: mockedAliases,
isValidToolName: vi.fn().mockImplementation((name: string, options) => {
if (mockedAliases[name]) return true;
if (Object.prototype.hasOwnProperty.call(mockedAliases, name))
return true;
return actual.isValidToolName(name, options);
}),
getToolAliases: vi.fn().mockImplementation((name: string) => {
Expand Down Expand Up @@ -55,11 +56,9 @@ describe('tool-names', () => {
expect(isValidToolName(`${DISCOVERED_TOOL_PREFIX}my_tool`)).toBe(true);
});

it('should validate MCP tool names (server__tool)', () => {
expect(isValidToolName('server__tool')).toBe(true);
expect(isValidToolName('my-server__my-tool')).toBe(true);
expect(isValidToolName('my.server__my:tool')).toBe(true);
expect(isValidToolName('my-server...truncated__tool')).toBe(true);
it('should validate modern MCP FQNs (mcp_server_tool)', () => {
expect(isValidToolName('mcp_server_tool')).toBe(true);
expect(isValidToolName('mcp_my-server_my-tool')).toBe(true);
});

it('should validate legacy tool aliases', async () => {
Expand All @@ -69,28 +68,33 @@ describe('tool-names', () => {
}
});

it('should reject invalid tool names', () => {
expect(isValidToolName('')).toBe(false);
expect(isValidToolName('invalid-name')).toBe(false);
expect(isValidToolName('server__')).toBe(false);
expect(isValidToolName('__tool')).toBe(false);
expect(isValidToolName('server__tool__extra')).toBe(false);
it('should return false for invalid tool names', () => {
expect(isValidToolName('invalid-tool-name')).toBe(false);
expect(isValidToolName('mcp_server')).toBe(false);
expect(isValidToolName('mcp__tool')).toBe(false);
expect(isValidToolName('mcp_invalid server_tool')).toBe(false);
expect(isValidToolName('mcp_server_invalid tool')).toBe(false);
expect(isValidToolName('mcp_server_')).toBe(false);
});

it('should handle wildcards when allowed', () => {
// Default: not allowed
expect(isValidToolName('*')).toBe(false);
expect(isValidToolName('server__*')).toBe(false);
expect(isValidToolName('mcp_*')).toBe(false);
expect(isValidToolName('mcp_server_*')).toBe(false);

// Explicitly allowed
expect(isValidToolName('*', { allowWildcards: true })).toBe(true);
expect(isValidToolName('server__*', { allowWildcards: true })).toBe(true);
expect(isValidToolName('mcp_*', { allowWildcards: true })).toBe(true);
expect(isValidToolName('mcp_server_*', { allowWildcards: true })).toBe(
true,
);

// Invalid wildcards
expect(isValidToolName('__*', { allowWildcards: true })).toBe(false);
expect(isValidToolName('server__tool*', { allowWildcards: true })).toBe(
false,
);
expect(isValidToolName('mcp__*', { allowWildcards: true })).toBe(false);
expect(
isValidToolName('mcp_server_tool*', { allowWildcards: true }),
).toBe(false);
});
});

Expand Down
51 changes: 38 additions & 13 deletions packages/core/src/tools/tool-names.ts
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,12 @@ export const DISCOVERED_TOOL_PREFIX = 'discovered_tool_';
/**
* List of all built-in tool names.
*/
import {
isMcpToolName,
parseMcpToolName,
MCP_TOOL_PREFIX,
} from './mcp-tool.js';

export const ALL_BUILTIN_TOOL_NAMES = [
GLOB_TOOL_NAME,
WRITE_TODOS_TOOL_NAME,
Expand Down Expand Up @@ -290,25 +296,44 @@ export function isValidToolName(
return true;
}

// MCP tools (format: server__tool)
if (name.includes('__')) {
const parts = name.split('__');
if (parts.length !== 2 || parts[0].length === 0 || parts[1].length === 0) {
// Handle standard MCP FQNs (mcp_server_tool or wildcards mcp_*, mcp_server_*)
if (isMcpToolName(name)) {
// Global wildcard: mcp_*
if (name === `${MCP_TOOL_PREFIX}*` && options.allowWildcards) {
return true;
}

// Explicitly reject names with empty server component (e.g. mcp__tool)
if (name.startsWith(`${MCP_TOOL_PREFIX}_`)) {
return false;
}

const server = parts[0];
const tool = parts[1];
const parsed = parseMcpToolName(name);
// Ensure that both components are populated. parseMcpToolName splits at the second _,
// so `mcp__tool` has serverName="", toolName="tool"
if (parsed.serverName && parsed.toolName) {
// Basic slug validation for server and tool names.
// We allow dots (.) and colons (:) as they are valid in function names and
// used for truncation markers.
const slugRegex = /^[a-z0-9_.:-]+$/i;

if (!slugRegex.test(parsed.serverName)) {
return false;
}

if (parsed.toolName === '*') {
return options.allowWildcards === true;
}

// A tool name consisting only of underscores is invalid.
if (/^_*$/.test(parsed.toolName)) {
return false;
}

if (tool === '*') {
return !!options.allowWildcards;
return slugRegex.test(parsed.toolName);
}

// Basic slug validation for server and tool names.
// We allow dots (.) and colons (:) as they are valid in function names and
// used for truncation markers.
const slugRegex = /^[a-z0-9_.:-]+$/i;
return slugRegex.test(server) && slugRegex.test(tool);
return false;
}

return false;
Expand Down
15 changes: 7 additions & 8 deletions packages/core/src/tools/tool-registry.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -310,13 +310,13 @@ describe('ToolRegistry', () => {
excludedTools: ['tool-a'],
},
{
name: 'should match simple MCP tool names, when qualified or unqualified',
tools: [mcpTool, mcpTool.asFullyQualifiedTool()],
name: 'should match simple MCP tool names',
tools: [mcpTool],
excludedTools: [mcpTool.name],
},
{
name: 'should match qualified MCP tool names when qualified or unqualified',
tools: [mcpTool, mcpTool.asFullyQualifiedTool()],
name: 'should match qualified MCP tool names',
tools: [mcpTool],
excludedTools: [mcpTool.name],
},
{
Expand Down Expand Up @@ -414,9 +414,9 @@ describe('ToolRegistry', () => {
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'desc');

// Register same MCP tool twice (one as alias, one as qualified)
// Register same MCP tool twice
toolRegistry.registerTool(mcpTool);
toolRegistry.registerTool(mcpTool);
toolRegistry.registerTool(mcpTool.asFullyQualifiedTool());

const toolNames = toolRegistry.getAllToolNames();
expect(toolNames).toEqual([`mcp_${serverName}_${toolName}`]);
Expand Down Expand Up @@ -698,9 +698,8 @@ describe('ToolRegistry', () => {
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'description');

// Register both alias and qualified
toolRegistry.registerTool(mcpTool);
toolRegistry.registerTool(mcpTool.asFullyQualifiedTool());
toolRegistry.registerTool(mcpTool);

const declarations = toolRegistry.getFunctionDeclarations();
expect(declarations).toHaveLength(1);
Expand Down
Loading
Loading