Improve Function Call argument validation and typing (#2881)

Co-authored-by: N. Taylor Mullen <[email protected]>

1 files changed, 6 insertions, 15 deletions

diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts
index 6dca8cea..f2b7458b 100644
--- a/packages/core/src/tools/mcp-client.ts
+++ b/packages/core/src/tools/mcp-client.ts
@@ -14,7 +14,7 @@ import {
 import { parse } from 'shell-quote';
 import { MCPServerConfig } from '../config/config.js';
 import { DiscoveredMCPTool } from './mcp-tool.js';
-import { CallableTool, FunctionDeclaration, mcpToTool } from '@google/genai';
+import { Type, mcpToTool } from '@google/genai';
 import { sanitizeParameters, ToolRegistry } from './tool-registry.js';
 
 export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
@@ -275,13 +275,10 @@ async function connectAndDiscover(
   }
 
   try {
-    const mcpCallableTool: CallableTool = mcpToTool(mcpClient);
-    const discoveredToolFunctions = await mcpCallableTool.tool();
+    const mcpCallableTool = mcpToTool(mcpClient);
+    const tool = await mcpCallableTool.tool();
 
-    if (
-      !discoveredToolFunctions ||
-      !Array.isArray(discoveredToolFunctions.functionDeclarations)
-    ) {
+    if (!tool || !Array.isArray(tool.functionDeclarations)) {
       console.error(
         `MCP server '${mcpServerName}' did not return valid tool function declarations. Skipping.`,
       );
@@ -297,7 +294,7 @@ async function connectAndDiscover(
       return;
     }
 
-    for (const funcDecl of discoveredToolFunctions.functionDeclarations) {
+    for (const funcDecl of tool.functionDeclarations) {
       if (!funcDecl.name) {
         console.warn(
           `Discovered a function declaration without a name from MCP server '${mcpServerName}'. Skipping.`,
@@ -344,19 +341,13 @@ async function connectAndDiscover(
 
       sanitizeParameters(funcDecl.parameters);
 
-      // Ensure parameters is a valid JSON schema object, default to empty if not.
-      const parameterSchema: Record<string, unknown> =
-        funcDecl.parameters && typeof funcDecl.parameters === 'object'
-          ? { ...(funcDecl.parameters as FunctionDeclaration) }
-          : { type: 'object', properties: {} };
-
       toolRegistry.registerTool(
         new DiscoveredMCPTool(
           mcpCallableTool,
           mcpServerName,
           toolNameForModel,
           funcDecl.description ?? '',
-          parameterSchema,
+          funcDecl.parameters ?? { type: Type.OBJECT, properties: {} },
           funcDecl.name,
           mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
           mcpServerConfig.trust,