Skip to content
Draft
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
28 changes: 26 additions & 2 deletions app/routes/api.chat.ts
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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,
});
}
}
}

Expand Down
111 changes: 111 additions & 0 deletions src/api/utils/toolValidation.ts
Original file line number Diff line number Diff line change
@@ -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;
}
Loading