From 0c0153646d8a6bdde42b4322b295492b6a64e91a Mon Sep 17 00:00:00 2001
From: AshAnand34 <aashishanand2019@gmail.com>
Date: Sat, 30 Aug 2025 10:46:51 -0700
Subject: [PATCH] feat: add tool validation and filtering functionality

---
 app/routes/api.chat.ts          |  28 +++++++-
 src/api/utils/toolValidation.ts | 111 ++++++++++++++++++++++++++++++++
 2 files changed, 137 insertions(+), 2 deletions(-)
 create mode 100644 src/api/utils/toolValidation.ts

diff --git a/app/routes/api.chat.ts b/app/routes/api.chat.ts
index a05d272..9680e51 100644
--- a/app/routes/api.chat.ts
+++ b/app/routes/api.chat.ts
@@ -1,9 +1,9 @@
 import { getModel } from "~/chat/ai/providers.server";
 import type { modelID } from "~/chat/ai/providers.shared";
 import { streamText, type ToolSet, type UIMessage } from "ai";
-
 import type { StorageKey } from "~/chat/ai/providers.shared";
 import { MCPClientManager } from "agents/mcp/client";
+import { validateAndFilterTools } from "src/api/utils/toolValidation";
 
 // Allow streaming responses up to 30 seconds
 export const maxDuration = 60;
@@ -120,10 +120,34 @@ export async function action({
       const { id } = await mcp.connect(url);
       if (mcp.mcpConnections[id]?.connectionState === "ready") {
         const mcptools = await mcp.unstable_getAITools();
-        tools = { ...tools, ...mcptools };
+        // Validate and filter tools before merging
+        const validatedTools = validateAndFilterTools(mcptools);
+        const validToolCount = Object.keys(validatedTools).length;
+        const totalToolCount = Object.keys(mcptools).length;
+
+        if (validToolCount < totalToolCount) {
+          console.warn(
+            `[${url}] Some tools were invalid: ${totalToolCount - validToolCount} tools filtered out`,
+          );
+        }
+
+        if (validToolCount === 0) {
+          console.error(`[${url}] No valid tools found after validation`);
+        } else {
+          tools = { ...tools, ...validatedTools };
+        }
       }
     } catch (error) {
       console.error("Error getting tools for url", url, error);
+      // Log more detailed error information for debugging
+      if (error instanceof Error) {
+        console.error("Error details:", {
+          message: error.message,
+          stack: error.stack,
+          name: error.name,
+          url: url,
+        });
+      }
     }
   }
 
diff --git a/src/api/utils/toolValidation.ts b/src/api/utils/toolValidation.ts
new file mode 100644
index 0000000..ca09372
--- /dev/null
+++ b/src/api/utils/toolValidation.ts
@@ -0,0 +1,111 @@
+import { z } from "zod";
+
+// Define schema for parameter properties
+const parameterPropertySchema = z
+  .object({
+    type: z.string(),
+    description: z.string().optional(),
+    items: z
+      .object({
+        type: z.string(),
+      })
+      .optional(),
+    enum: z.array(z.string()).optional(),
+    format: z.string().optional(),
+  })
+  .catchall(z.any());
+
+// Define the parameter schema for tool inputs
+const parameterSchema = z
+  .object({
+    type: z.string(),
+    description: z.string().optional(),
+    properties: z.record(parameterPropertySchema).optional(),
+    required: z.array(z.string()).optional(),
+    additionalProperties: z.boolean().optional(),
+  })
+  .catchall(z.any());
+
+// Define the tool schema
+const toolSchema = z.object({
+  name: z.string(),
+  description: z.string(),
+  parameters: z
+    .object({
+      type: z.literal("object"),
+      properties: z.record(parameterPropertySchema),
+      required: z.array(z.string()).optional(),
+    })
+    .optional(),
+  paramsSchema: parameterSchema.optional(),
+});
+
+// Validate if input matches either schema format
+export function validateToolDefinition(tool: any): boolean {
+  try {
+    // Try validating against the toolSchema
+    const result = toolSchema.safeParse(tool);
+    if (result.success) {
+      // If tool has parameters property, ensure it's properly formatted
+      if (tool.parameters) {
+        if (tool.parameters.type !== "object") {
+          console.error(
+            `Invalid tool parameters type for tool: ${tool?.name}. Expected 'object', got '${tool.parameters.type}'`,
+          );
+          return false;
+        }
+        if (
+          !tool.parameters.properties ||
+          typeof tool.parameters.properties !== "object"
+        ) {
+          console.error(
+            `Invalid tool parameters properties for tool: ${tool?.name}`,
+          );
+          return false;
+        }
+      }
+      return true;
+    }
+
+    // Log validation errors for debugging
+    if (result.error) {
+      console.error(`Tool validation failed for tool: ${tool?.name}`, {
+        errors: result.error.errors,
+        invalidData: tool,
+      });
+    }
+
+    return false;
+  } catch (error) {
+    console.error(`Unexpected error validating tool: ${tool?.name}`, error);
+    return false;
+  }
+}
+
+export function validateAndFilterTools(
+  tools: Record<string, any>,
+): Record<string, any> {
+  const validTools: Record<string, any> = {};
+
+  for (const [name, tool] of Object.entries(tools)) {
+    // Skip if tool is not an object
+    if (!tool || typeof tool !== "object") {
+      console.error(`Invalid tool definition: ${name} is not an object`);
+      continue;
+    }
+
+    // Ensure tool has a name property that matches the key
+    const toolWithName = {
+      ...tool,
+      name: name,
+    };
+
+    if (validateToolDefinition(toolWithName)) {
+      validTools[name] = toolWithName;
+    } else {
+      console.error(`Tool validation failed for: ${name}`);
+    }
+  }
+
+  return validTools;
+}