path: root/src/client/views/nodes/chatbot/tools/CreateNewTool.ts

import { Observation } from '../types/types';
import { Parameter, ParametersType, ToolInfo } from '../types/tool_types';
import { BaseTool } from './BaseTool';
import * as ts from 'typescript';
import { v4 as uuidv4 } from 'uuid';
import { Networking } from '../../../../Network';

// Forward declaration to avoid circular import
interface AgentLike {
    registerDynamicTool(toolName: string, toolInstance: BaseTool<ReadonlyArray<Parameter>>): void;
    notifyToolCreated(toolName: string, completeToolCode: string): void;
}

const createNewToolParams = [
    {
        name: 'toolName',
        type: 'string',
        description: 'The name of the new tool class (PascalCase) and filename. This will also be converted to lowercase for the action name.',
        required: true,
    },
    {
        name: 'toolCode',
        type: 'string',
        description:
            'The complete TypeScript code for the new tool class. IMPORTANT: Provide this as a single string without any XML formatting. Do not break it into multiple lines or add any XML tags. The tool must extend BaseTool, implement an async execute method, and have proper parameter definitions. Use CDATA format if needed: <![CDATA[your code here]]>',
        required: true,
    },
    {
        name: 'description',
        type: 'string',
        description: 'A brief description of what the tool does.',
        required: true,
    },
] as const;

type CreateNewToolParamsType = typeof createNewToolParams;

const createNewToolInfo: ToolInfo<CreateNewToolParamsType> = {
    name: 'createNewTool',
    description: `Creates a new tool for the agent to use based on research done with the codebase search, file content, and filenames tools. The new tool will be instantly available for use in the current session and saved as a proper TypeScript file.

IMPORTANT TOOL CREATION RULES:
1. Your tool will be created with proper imports adjusted for the dynamic subfolder location
2. Your tool MUST extend BaseTool with proper parameter type definition  
3. Your tool MUST implement an async execute method that returns Promise<Observation[]>
4. Your tool MUST call super() with the proper tool info configuration object
5. CRITICAL: The toolInfo.name property MUST be lowercase and should match the action name you want to use
6. Follow this EXACT pattern (imports will be added automatically):

\`\`\`typescript
const yourToolParams = [
  {
    name: 'inputParam',
    type: 'string', 
    description: 'Your parameter description',
    required: true
  }
] as const;

type YourToolParamsType = typeof yourToolParams;

const yourToolInfo: ToolInfo<YourToolParamsType> = {
  name: 'yourtoolname',
  description: 'Your tool description',
  citationRules: 'No citation needed.',
  parameterRules: yourToolParams
};

export class YourToolName extends BaseTool<YourToolParamsType> {
  constructor() {
    super(yourToolInfo);
  }

  async execute(args: ParametersType<YourToolParamsType>): Promise<Observation[]> {
    const { inputParam } = args;
    // Your implementation here
    return [{ type: 'text', text: 'Your result' }];
  }
}
\`\`\`

EXAMPLE - Character Count Tool:

\`\`\`typescript
const characterCountParams = [
  {
    name: 'text',
    type: 'string',
    description: 'The text to count characters in',
    required: true
  }
] as const;

type CharacterCountParamsType = typeof characterCountParams;

const characterCountInfo: ToolInfo<CharacterCountParamsType> = {
  name: 'charactercount',
  description: 'Counts characters in text, excluding spaces',
  citationRules: 'No citation needed.',
  parameterRules: characterCountParams
};

export class CharacterCountTool extends BaseTool<CharacterCountParamsType> {
  constructor() {
    super(characterCountInfo);
  }

  async execute(args: ParametersType<CharacterCountParamsType>): Promise<Observation[]> {
    const { text } = args;
    const count = text ? text.replace(/\\s/g, '').length : 0;
    return [{ type: 'text', text: \`Character count (excluding spaces): \${count}\` }];
  }
}
\`\`\``,
    citationRules: `No citation needed.`,
    parameterRules: createNewToolParams,
};

/**
 * This tool allows the agent to create new custom tools after researching the codebase.
 * It validates the provided code, dynamically compiles it, and registers it with the
 * Agent for immediate use.
 */
export class CreateNewTool extends BaseTool<CreateNewToolParamsType> {
    // Reference to the dynamic tool registry in the Agent class
    private dynamicToolRegistry: Map<string, BaseTool<ReadonlyArray<Parameter>>>;
    private existingTools: Record<string, BaseTool<ReadonlyArray<Parameter>>>;
    private agent?: AgentLike;

    constructor(toolRegistry: Map<string, BaseTool<ReadonlyArray<Parameter>>>, existingTools: Record<string, BaseTool<ReadonlyArray<Parameter>>> = {}, agent?: AgentLike) {
        super(createNewToolInfo);
        this.dynamicToolRegistry = toolRegistry;
        this.existingTools = existingTools;
        this.agent = agent;
    }

    /**
     * Validates TypeScript code for basic safety and correctness
     * @param code The TypeScript code to validate
     * @returns An object with validation result and any error messages
     */
    private validateToolCode(code: string, toolName: string): { valid: boolean; errors: string[] } {
        const errors: string[] = [];

        // Check for fundamental structure
        if (!code.includes('extends BaseTool')) {
            errors.push('Tool must extend BaseTool class');
        }

        if (!code.includes(`class ${toolName} extends`)) {
            errors.push(`Tool class name must match the provided toolName: ${toolName}`);
        }

        if (!code.includes('async execute(')) {
            errors.push('Tool must implement an async execute method');
        }

        if (!code.includes('super(')) {
            errors.push('Tool must call super() in constructor');
        }

        // Check if the tool exports the class correctly (should use export class)
        if (!code.includes(`export class ${toolName}`)) {
            errors.push(`Tools must export the class using: export class ${toolName}`);
        }

        // Check if tool info has name property in lowercase
        const nameMatch = code.match(/name\s*:\s*['"]([^'"]+)['"]/);
        if (nameMatch && nameMatch[1]) {
            const toolInfoName = nameMatch[1];
            if (toolInfoName !== toolInfoName.toLowerCase()) {
                errors.push(`Tool info name property must be lowercase. Found: "${toolInfoName}", should be "${toolInfoName.toLowerCase()}"`);
            }
        } else {
            errors.push('Tool info must have a name property');
        }

        // Check for type definition - make this more flexible
        const hasTypeDefinition = code.includes(`type ${toolName}ParamsType`) || code.includes(`type ${toolName.toLowerCase()}ParamsType`) || code.includes('ParamsType = typeof');
        if (!hasTypeDefinition) {
            errors.push(`Tool must define a type for parameters like: type ${toolName}ParamsType = typeof ${toolName.toLowerCase()}Params`);
        }

        // Check for ToolInfo type annotation - make this more flexible
        const hasToolInfoType = code.includes(`ToolInfo<${toolName}ParamsType>`) || code.includes(`ToolInfo<${toolName.toLowerCase()}ParamsType>`) || code.includes('ToolInfo<');
        if (!hasToolInfoType) {
            errors.push(`Tool info must be typed as ToolInfo<YourParamsType>`);
        }

        // Check for proper execute method typing - make this more flexible
        if (!code.includes(`ParametersType<${toolName}ParamsType>`) && !code.includes('args: ParametersType<')) {
            errors.push(`Execute method must have typed parameters: args: ParametersType<${toolName}ParamsType>`);
        }

        // Check for unsafe code patterns
        const unsafePatterns = [
            { pattern: /eval\s*\(/, message: 'eval() is not allowed' },
            { pattern: /Function\s*\(/, message: 'Function constructor is not allowed' },
            { pattern: /require\s*\(\s*['"]child_process['"]/, message: 'child_process module is not allowed' },
            { pattern: /require\s*\(\s*['"]fs['"]/, message: 'Direct fs module import is not allowed' },
            { pattern: /require\s*\(\s*['"]path['"]/, message: 'Direct path module import is not allowed' },
            { pattern: /process\.env/, message: 'Accessing process.env is not allowed' },
            { pattern: /import\s+.*['"]child_process['"]/, message: 'child_process module is not allowed' },
            { pattern: /import\s+.*['"]fs['"]/, message: 'Direct fs module import is not allowed' },
            { pattern: /import\s+.*['"]path['"]/, message: 'Direct path module import is not allowed' },
            { pattern: /\bnew\s+Function\b/, message: 'Function constructor is not allowed' },
            { pattern: /\bwindow\b/, message: 'Direct window access is not allowed' },
            { pattern: /\bdocument\b/, message: 'Direct document access is not allowed' },
            { pattern: /\blocation\b/, message: 'Direct location access is not allowed' },
            { pattern: /\bsessionStorage\b/, message: 'Direct sessionStorage access is not allowed' },
            { pattern: /\blocalStorage\b/, message: 'Direct localStorage access is not allowed' },
            { pattern: /fetch\s*\(/, message: 'Direct fetch calls are not allowed' },
            { pattern: /XMLHttpRequest/, message: 'Direct XMLHttpRequest use is not allowed' },
        ];

        for (const { pattern, message } of unsafePatterns) {
            if (pattern.test(code)) {
                errors.push(message);
            }
        }

        // Check if the tool name is already used by an existing tool
        const toolNameLower = toolName.toLowerCase();
        if (Object.keys(this.existingTools).some(key => key.toLowerCase() === toolNameLower) || Array.from(this.dynamicToolRegistry.keys()).some(key => key.toLowerCase() === toolNameLower)) {
            errors.push(`A tool with the name "${toolNameLower}" already exists. Please choose a different name.`);
        }

        // Use TypeScript compiler API to check for syntax errors
        try {
            const sourceFile = ts.createSourceFile(`${toolName}.ts`, code, ts.ScriptTarget.Latest, true);

            // Create a TypeScript program to check for type errors
            const options: ts.CompilerOptions = {
                target: ts.ScriptTarget.ES2020,
                module: ts.ModuleKind.ESNext,
                strict: true,
                esModuleInterop: true,
                skipLibCheck: true,
                forceConsistentCasingInFileNames: true,
            };

            // Perform additional static analysis on the AST
            const visitor = (node: ts.Node) => {
                // Check for potentially unsafe constructs
                if (ts.isCallExpression(node)) {
                    const expression = node.expression;
                    if (ts.isIdentifier(expression)) {
                        const name = expression.text;
                        if (name === 'eval' || name === 'Function') {
                            errors.push(`Use of ${name} is not allowed`);
                        }
                    }
                }

                // Recursively visit all child nodes
                ts.forEachChild(node, visitor);
            };

            visitor(sourceFile);
        } catch (error) {
            errors.push(`TypeScript syntax error: ${error}`);
        }

        return {
            valid: errors.length === 0,
            errors,
        };
    }

    /**
     * Extracts tool info name from the tool code
     * @param code The tool TypeScript code
     * @returns The tool info name or null if not found
     */
    private extractToolInfoName(code: string): string | null {
        const nameMatch = code.match(/name\s*:\s*['"]([^'"]+)['"]/);
        return nameMatch && nameMatch[1] ? nameMatch[1] : null;
    }

    /**
     * Extracts and parses parameter info from the tool code
     * @param code The tool TypeScript code
     * @returns An array of parameter objects
     */
    private extractToolParameters(code: string): Array<{ name: string; type: string; description: string; required: boolean }> {
        // Basic regex-based extraction - in a production environment, this should use the TypeScript AST
        const paramsMatch = code.match(/const\s+\w+Params\s*=\s*\[([\s\S]*?)\]\s*as\s*const/);
        if (!paramsMatch || !paramsMatch[1]) {
            return [];
        }

        const paramsText = paramsMatch[1];

        // Parse individual parameters
        const paramRegex = /{\s*name\s*:\s*['"]([^'"]+)['"]\s*,\s*type\s*:\s*['"]([^'"]+)['"]\s*,\s*description\s*:\s*['"]([^'"]+)['"]\s*,\s*required\s*:\s*(true|false)/g;
        const params = [];
        let match;

        while ((match = paramRegex.exec(paramsText)) !== null) {
            params.push({
                name: match[1],
                type: match[2],
                description: match[3],
                required: match[4] === 'true',
            });
        }

        return params;
    }

    /**
     * Generates the complete tool file content with proper imports for the dynamic subfolder
     * @param toolCode The user-provided tool code
     * @param toolName The name of the tool class
     * @returns The complete TypeScript file content
     */
    private generateCompleteToolFile(toolCode: string, toolName: string): string {
        // Add proper imports for the dynamic subfolder (one level deeper than regular tools)
        const imports = `import { Observation } from '../../types/types';
import { ParametersType, ToolInfo } from '../../types/tool_types';
import { BaseTool } from '../BaseTool';

`;

        // Clean the user code - remove any existing imports they might have added
        const cleanedCode = toolCode
            .replace(/import\s+[^;]+;?\s*/g, '') // Remove any import statements
            .trim();

        return imports + cleanedCode;
    }

    /**
     * Transpiles TypeScript code to JavaScript
     * @param code The TypeScript code to compile
     * @param filename The name of the file (for error reporting)
     * @returns The compiled JavaScript code
     */
    private transpileTypeScript(code: string, filename: string): { jsCode: string; errors: string[] } {
        try {
            const transpileOptions: ts.TranspileOptions = {
                compilerOptions: {
                    module: ts.ModuleKind.CommonJS, // Use CommonJS for dynamic imports
                    target: ts.ScriptTarget.ES2020,
                    moduleResolution: ts.ModuleResolutionKind.NodeJs,
                    esModuleInterop: true,
                    sourceMap: false,
                    strict: false, // Relax strict mode for dynamic compilation
                    noImplicitAny: false, // Allow implicit any types
                },
                reportDiagnostics: true,
                fileName: `${filename}.ts`,
            };

            const output = ts.transpileModule(code, transpileOptions);

            // Check for compilation errors
            const errors: string[] = [];
            if (output.diagnostics && output.diagnostics.length > 0) {
                for (const diagnostic of output.diagnostics) {
                    if (diagnostic.file && diagnostic.start !== undefined) {
                        const { line, character } = diagnostic.file.getLineAndCharacterOfPosition(diagnostic.start);
                        const message = ts.flattenDiagnosticMessageText(diagnostic.messageText, '\n');
                        errors.push(`Line ${line + 1}, Column ${character + 1}: ${message}`);
                    } else {
                        errors.push(ts.flattenDiagnosticMessageText(diagnostic.messageText, '\n'));
                    }
                }

                if (errors.length > 0) {
                    return { jsCode: '', errors };
                }
            }

            return { jsCode: output.outputText, errors: [] };
        } catch (error) {
            return {
                jsCode: '',
                errors: [`Transpilation failed: ${error instanceof Error ? error.message : String(error)}`],
            };
        }
    }

    /**
     * Dynamically evaluates and instantiates a tool from JavaScript code
     * @param jsCode The JavaScript code
     * @param toolName The name of the tool class
     * @returns An instance of the tool or null if instantiation failed
     */
    private async createDynamicTool(jsCode: string, toolName: string): Promise<BaseTool<ReadonlyArray<Parameter>> | null> {
        try {
            // Create a safe evaluation context with necessary globals
            const globalContext = {
                BaseTool: BaseTool, // Actual class reference
                exports: {},
                module: { exports: {} },
                require: (id: string) => {
                    // Mock require for the imports we know about
                    if (id.includes('types/types')) {
                        return { Observation: null };
                    }
                    if (id.includes('tool_types')) {
                        return { ParametersType: null, ToolInfo: null };
                    }
                    if (id.includes('BaseTool')) {
                        return { BaseTool: BaseTool };
                    }
                    return {};
                },
                console: console,
                // Add any other commonly needed globals
                JSON: JSON,
                Array: Array,
                Object: Object,
                String: String,
                Number: Number,
                Boolean: Boolean,
                Math: Math,
                Date: Date,
            };

            // Create function to evaluate in the proper context
            const evaluationFunction = new Function(
                ...Object.keys(globalContext),
                `"use strict";
                try {
                    ${jsCode}
                    // Get the exported class from the module
                    const ToolClass = exports.${toolName} || module.exports.${toolName} || module.exports;
                    if (ToolClass && typeof ToolClass === 'function') {
                        return new ToolClass();
                    } else {
                        console.error('Tool class not found in exports:', Object.keys(exports), Object.keys(module.exports));
                        return null;
                    }
                } catch (error) {
                    console.error('Error during tool evaluation:', error);
                    return null;
                }`
            );

            // Execute with our controlled globals
            const toolInstance = evaluationFunction(...Object.values(globalContext));

            if (!toolInstance) {
                console.error(`Failed to instantiate ${toolName} - no instance returned`);
                return null;
            }

            // Verify it's a proper BaseTool instance
            if (!(toolInstance instanceof BaseTool)) {
                console.error(`${toolName} is not a proper instance of BaseTool`);
                return null;
            }

            console.log(`Successfully created dynamic tool instance: ${toolName}`);
            return toolInstance;
        } catch (error) {
            console.error('Error creating dynamic tool:', error);
            return null;
        }
    }

    /**
     * Save the tool code to the server so it's available for future sessions
     * @param toolName The name of the tool
     * @param completeToolCode The complete TypeScript code for the tool with imports
     */
    private async saveToolToServer(toolName: string, completeToolCode: string): Promise<boolean> {
        try {
            // Create a server endpoint to save the tool
            const response = await Networking.PostToServer('/saveDynamicTool', {
                toolName: toolName,
                toolCode: completeToolCode,
            });

            // Type check the response to avoid property access errors
            return typeof response === 'object' && response !== null && 'success' in response && (response as { success: boolean }).success === true;
        } catch (error) {
            console.error('Failed to save tool to server:', error);
            return false;
        }
    }

    async execute(args: ParametersType<CreateNewToolParamsType>): Promise<Observation[]> {
        const { toolName, toolCode, description } = args;

        console.log(`Creating new tool: ${toolName}`);

        // Remove any markdown backticks that might be in the code
        const cleanedCode = (toolCode as string).replace(/```typescript|```/g, '').trim();

        if (!cleanedCode) {
            return [
                {
                    type: 'text',
                    text: 'Failed to extract tool code from the provided input. Please ensure the tool code is provided as valid TypeScript code.',
                },
            ];
        }

        // Validate the provided code
        const validation = this.validateToolCode(cleanedCode, toolName);
        if (!validation.valid) {
            return [
                {
                    type: 'text',
                    text: `Failed to create tool: Code validation failed with the following errors:\n- ${validation.errors.join('\n- ')}`,
                },
            ];
        }

        try {
            // Generate the complete tool file with proper imports
            const completeToolCode = this.generateCompleteToolFile(cleanedCode, toolName);

            // Extract tool info name from the code
            const toolInfoName = this.extractToolInfoName(cleanedCode);
            if (!toolInfoName) {
                return [
                    {
                        type: 'text',
                        text: 'Failed to extract tool info name from the code. Make sure the tool has a name property.',
                    },
                ];
            }

            // Extract parameters from the tool code
            const parameters = this.extractToolParameters(cleanedCode);

            // Transpile the TypeScript to JavaScript
            const { jsCode, errors } = this.transpileTypeScript(completeToolCode, toolName);
            if (errors.length > 0) {
                return [
                    {
                        type: 'text',
                        text: `Failed to transpile tool code with the following errors:\n- ${errors.join('\n- ')}`,
                    },
                ];
            }

            // Create a dynamic tool instance
            const toolInstance = await this.createDynamicTool(jsCode, toolName);
            if (!toolInstance) {
                return [
                    {
                        type: 'text',
                        text: 'Failed to instantiate the tool. Make sure it follows all the required patterns and properly extends BaseTool.',
                    },
                ];
            }

            // Register the tool in the dynamic registry
            // Use the name property from the tool info as the registry key
            this.dynamicToolRegistry.set(toolInfoName, toolInstance);

            // If we have a reference to the agent, tell it to register dynamic tool
            // This ensures the tool is properly loaded from the filesystem for the prompt system
            if (this.agent) {
                this.agent.registerDynamicTool(toolInfoName, toolInstance);
            }

            // Create the success message
            const successMessage = `Successfully created and registered new tool: ${toolName}\n\nThe tool is now available for use in the current session. You can call it using the action "${toolInfoName}".\n\nDescription: ${description}\n\nParameters: ${
                parameters.length > 0 ? parameters.map(p => `\n- ${p.name} (${p.type}${p.required ? ', required' : ''}): ${p.description}`).join('') : '\nNo parameters'
            }\n\nThe tool will be saved permanently after you confirm the page reload.`;

            // Notify the agent that a tool was created with the complete code for deferred saving
            // This will trigger the modal but NOT save to disk yet
            if (this.agent) {
                this.agent.notifyToolCreated(toolName, completeToolCode);
            }

            return [
                {
                    type: 'text',
                    text: successMessage,
                },
            ];
        } catch (error) {
            console.error(`Error creating new tool:`, error);
            return [
                {
                    type: 'text',
                    text: `Failed to create tool: ${(error as Error).message || 'Unknown error'}`,
                },
            ];
        }
    }

    /**
     * Public method to save tool to server (called by agent after user confirmation)
     * @param toolName The name of the tool
     * @param completeToolCode The complete TypeScript code for the tool with imports
     */
    public async saveToolToServerDeferred(toolName: string, completeToolCode: string): Promise<boolean> {
        return this.saveToolToServer(toolName, completeToolCode);
    }
}