@flink-app/flink 0.14.3 → 2.0.0-alpha.100

package/src/ai/ToolExecutor.ts

@@ -0,0 +1,485 @@
+import Ajv from "ajv";
+import addFormats from "ajv-formats";
+import { FlinkContext } from "../FlinkContext";
+import { forbidden } from "../FlinkErrors";
+import { FlinkLogFactory } from "../FlinkLogFactory";
+import { getRequestContext, getRequestUser } from "../FlinkRequestContext";
+import { FlinkTool, FlinkToolProps, ToolResult } from "./FlinkTool";
+import { FlinkToolSchema } from "./LLMAdapter";
+
+const toolLog = FlinkLogFactory.createLogger("flink.ai.tool");
+
+export class ToolExecutor<Ctx extends FlinkContext> {
+    private ajv: Ajv;
+    private compiledInputValidator?: ReturnType<Ajv["compile"]>;
+    private compiledOutputValidator?: ReturnType<Ajv["compile"]>;
+
+    constructor(
+        private toolProps: FlinkToolProps,
+        private toolFn: FlinkTool<Ctx, any, any>,
+        private ctx: Ctx,
+        private autoSchemas?: {
+            inputSchema?: any;
+            outputSchema?: any;
+            inputTypeHint?: 'void' | 'any' | 'named';
+            outputTypeHint?: 'void' | 'any' | 'named';
+        },
+        private allSchemas?: Record<string, any>
+    ) {
+        this.ajv = new Ajv({ allErrors: true });
+        addFormats(this.ajv);
+
+        // Pre-populate AJV with all schemas so $ref references resolve across schema boundaries
+        if (allSchemas) {
+            for (const schema of Object.values(allSchemas)) {
+                if (schema && schema.$id) {
+                    try {
+                        this.ajv.addSchema(schema);
+                    } catch {
+                        // Ignore duplicate schema errors (schema may already be added)
+                    }
+                }
+            }
+        }
+
+        // Pre-compile validators once at construction time (not per invocation)
+        if (toolProps.inputJsonSchema) {
+            this.compiledInputValidator = this.ajv.compile(toolProps.inputJsonSchema);
+        } else if (autoSchemas?.inputSchema) {
+            this.compiledInputValidator = this.ajv.compile(autoSchemas.inputSchema);
+        }
+
+        if (toolProps.outputJsonSchema) {
+            this.compiledOutputValidator = this.ajv.compile(toolProps.outputJsonSchema);
+        } else if (autoSchemas?.outputSchema) {
+            this.compiledOutputValidator = this.ajv.compile(autoSchemas.outputSchema);
+        }
+        // Log when using auto-schemas
+        if (autoSchemas?.inputSchema && !toolProps.inputSchema && !toolProps.inputJsonSchema) {
+            toolLog.debug(`Tool ${toolProps.id}: Using auto-generated schemas from type parameters`);
+        }
+
+        // Log when no schema is provided (valid for void/any input types)
+        const hasInputSchema = toolProps.inputSchema || toolProps.inputJsonSchema || autoSchemas?.inputSchema;
+        if (!hasInputSchema) {
+            toolLog.debug(`Tool ${toolProps.id}: No input schema provided (input type is void or any)`);
+        }
+
+        // Warn if both manual and auto provided (manual takes precedence)
+        if ((toolProps.inputSchema || toolProps.inputJsonSchema) && autoSchemas?.inputSchema) {
+            toolLog.debug(`Tool ${toolProps.id}: Manual schemas take precedence over auto-generated schemas`);
+        }
+
+        // Log warning if both manual Zod and JSON schemas are provided
+        if (toolProps.inputSchema && toolProps.inputJsonSchema) {
+            toolLog.warn(`Tool ${toolProps.id}: Both 'inputSchema' (Zod) and 'inputJsonSchema' are provided. ` + `Using 'inputSchema' for validation.`);
+        }
+
+        if (toolProps.outputSchema && toolProps.outputJsonSchema) {
+            toolLog.warn(`Tool ${toolProps.id}: Both 'outputSchema' (Zod) and 'outputJsonSchema' are provided. ` + `Using 'outputSchema' for validation.`);
+        }
+    }
+
+    /**
+     * Execute the tool with input
+     * @param input - Tool input data
+     * @param overrides - Optional overrides for user/permissions/conversationContext (for testing)
+     */
+    async execute(input: any, overrides?: { user?: any; permissions?: string[]; conversationContext?: any }): Promise<ToolResult<any>> {
+        // Get user, permissions, and conversationContext from AsyncLocalStorage or overrides
+        const user = overrides?.user ?? getRequestUser();
+        const userPermissions = overrides?.permissions ?? getRequestContext()?.userPermissions;
+        const conversationContext = overrides?.conversationContext;
+
+        // 1. Permission check
+        if (this.toolProps.permissions) {
+            const hasPermission = await this.checkPermissionsInternal(input, user, userPermissions);
+            if (!hasPermission) {
+                toolLog.debug("Tool invocator is missing required permission(s)", this.toolProps.permissions, "user has", userPermissions);
+                throw forbidden(`Permission denied for tool ${this.toolProps.id}`, "PERMISSION_DENIED");
+            }
+        }
+
+        // 2. Input validation (priority: Zod > JSON Schema > Auto-generated)
+        let validatedInput: any;
+        try {
+            if (this.toolProps.inputSchema) {
+                // Priority 1: Use Zod validation
+                validatedInput = this.toolProps.inputSchema.parse(input);
+            } else if (this.compiledInputValidator) {
+                // Priority 2 & 3: Use pre-compiled JSON Schema validator (manual or auto-generated)
+                const valid = this.compiledInputValidator(input);
+                if (!valid) {
+                    const errorDetails = this.formatAjvErrors(this.compiledInputValidator.errors || [], input);
+                    toolLog.warn(`Tool ${this.toolProps.id} input validation failed:`, errorDetails);
+                    return {
+                        success: false,
+                        error: `Invalid input for tool '${this.toolProps.id}': ${errorDetails}`,
+                        code: "VALIDATION_ERROR",
+                    };
+                }
+                validatedInput = input;
+            } else {
+                // No schema available - skip validation
+                validatedInput = input;
+            }
+        } catch (err: any) {
+            toolLog.warn(`Tool ${this.toolProps.id} input validation failed:`, err.message);
+
+            // Format Zod validation errors for better LLM understanding
+            const errorDetails = this.formatZodError(err, input);
+
+            return {
+                success: false,
+                error: `Invalid input for tool '${this.toolProps.id}': ${errorDetails}`,
+                code: "VALIDATION_ERROR",
+            };
+        }
+
+        // 3. Execute tool
+        toolLog.debug(`Executing tool ${this.toolProps.id}`);
+        let result: ToolResult<any>;
+        try {
+            toolLog.trace(this.toolFn.name + " input:", validatedInput);
+
+            result = await this.toolFn({
+                input: validatedInput,
+                ctx: this.ctx,
+                user,
+                permissions: userPermissions,
+                conversationCtx: conversationContext,
+            });
+        } catch (err: any) {
+            toolLog.error(`Tool ${this.toolProps.id} threw error:`, err.message);
+            return {
+                success: false,
+                error: `Tool execution failed: ${err.message}`,
+                code: "EXECUTION_ERROR",
+            };
+        }
+
+        // 4. Handle error results
+        if (!result.success) {
+            toolLog.warn(`Tool ${this.toolProps.id} returned error:`, result.error);
+            return result; // Return error result as-is
+        }
+
+        // 5. Output validation (priority: Zod > JSON Schema > Auto-generated)
+        if (this.toolProps.outputSchema) {
+            // Priority 1: Use Zod validation
+            try {
+                const validatedData = this.toolProps.outputSchema.parse(result.data);
+                return { success: true, data: validatedData };
+            } catch (err: any) {
+                toolLog.error(`Tool ${this.toolProps.id} output validation failed:`, err.message);
+                return {
+                    success: false,
+                    error: `Invalid output from tool ${this.toolProps.id}: ${err.message}`,
+                    code: "OUTPUT_VALIDATION_ERROR",
+                };
+            }
+        } else if (this.compiledOutputValidator) {
+            // Priority 2 & 3: Use pre-compiled JSON Schema validator (manual or auto-generated)
+            try {
+                const valid = this.compiledOutputValidator(result.data);
+                if (!valid) {
+                    const errorDetails = this.formatAjvErrors(this.compiledOutputValidator.errors || []);
+                    toolLog.error(`Tool ${this.toolProps.id} output validation failed:`, errorDetails);
+                    return {
+                        success: false,
+                        error: `Invalid output from tool ${this.toolProps.id}: ${errorDetails}`,
+                        code: "OUTPUT_VALIDATION_ERROR",
+                    };
+                }
+                return { success: true, data: result.data };
+            } catch (err: any) {
+                toolLog.error(`Tool ${this.toolProps.id} output validation failed:`, err.message);
+                return {
+                    success: false,
+                    error: `Invalid output from tool ${this.toolProps.id}: ${err.message}`,
+                    code: "OUTPUT_VALIDATION_ERROR",
+                };
+            }
+        }
+
+        // No output validation - return result as-is
+        return result;
+    }
+
+    getToolSchema(): FlinkToolSchema {
+        // Priority order: inputJsonSchema > inputSchema (Zod 4.x) > autoSchemas
+
+        // Priority 1: Manual JSON schema
+        if (this.toolProps.inputJsonSchema) {
+            return {
+                name: this.toolProps.id,
+                description: this.toolProps.description,
+                inputSchema: this.toolProps.inputJsonSchema,
+            };
+        }
+
+        // Priority 2: Zod schema (convert to JSON Schema using Zod 4.x)
+        if (this.toolProps.inputSchema) {
+            const z = require("zod") as any;
+
+            if (!z.toJSONSchema) {
+                throw new Error(
+                    `Tool ${this.toolProps.id}: Zod 4.x is required for automatic schema generation. ` +
+                        `Either upgrade to Zod 4.x or provide 'inputJsonSchema' in your tool definition, ` +
+                        `or use TypeScript type parameters for auto-generation.`
+                );
+            }
+
+            return {
+                name: this.toolProps.id,
+                description: this.toolProps.description,
+                inputSchema: z.toJSONSchema(this.toolProps.inputSchema),
+            };
+        }
+
+        // Priority 3: Auto-generated schema from type parameters
+        if (this.autoSchemas?.inputSchema) {
+            return {
+                name: this.toolProps.id,
+                description: this.toolProps.description,
+                inputSchema: this.resolveSchemaRefs(this.autoSchemas.inputSchema),
+            };
+        }
+
+        // No schema provided - return schema based on type hint
+        const typeHint = this.autoSchemas?.inputTypeHint;
+
+        if (typeHint === 'void') {
+            // void: Tool takes no input - reject any properties
+            return {
+                name: this.toolProps.id,
+                description: this.toolProps.description,
+                inputSchema: {
+                    type: "object",
+                    properties: {},
+                    additionalProperties: false,
+                },
+            };
+        } else if (typeHint === 'any') {
+            // any: Tool accepts any input - allow any properties
+            return {
+                name: this.toolProps.id,
+                description: this.toolProps.description,
+                inputSchema: {
+                    type: "object",
+                    additionalProperties: true,
+                },
+            };
+        } else {
+            // No type hint - default to void behavior (no input expected)
+            return {
+                name: this.toolProps.id,
+                description: this.toolProps.description,
+                inputSchema: {
+                    type: "object",
+                    properties: {},
+                    additionalProperties: false,
+                },
+            };
+        }
+    }
+
+    /**
+     * Resolve cross-schema $ref values into $defs so the schema is self-contained.
+     *
+     * Flink's schema manifest stores schemas as separate documents with IDs like
+     * "Canvas.ElementInput". These are valid for AJV (which uses a schema registry),
+     * but LLM providers like OpenAI require a single self-contained schema where all
+     * $ref values point to #/$defs/... entries at the top level.
+     */
+    private resolveSchemaRefs(schema: any): any {
+        if (!this.allSchemas) return schema;
+
+        const defs: Record<string, any> = {};
+
+        const collectRefs = (node: any, visited: Set<string>): void => {
+            if (!node || typeof node !== "object") return;
+
+            if (node.$ref && typeof node.$ref === "string" && !node.$ref.startsWith("#")) {
+                const refId = node.$ref;
+                if (!visited.has(refId) && this.allSchemas![refId]) {
+                    visited.add(refId);
+                    // Clone and strip top-level JSON Schema meta fields not valid inside $defs
+                    const def = JSON.parse(JSON.stringify(this.allSchemas![refId]));
+                    delete def.$id;
+                    delete def.$schema;
+                    defs[refId] = def;
+                    // Recurse into the referenced schema to collect its deps
+                    collectRefs(def, visited);
+                }
+            }
+
+            for (const value of Object.values(node)) {
+                collectRefs(value, visited);
+            }
+        };
+
+        const visited = new Set<string>();
+        collectRefs(schema, visited);
+
+        if (Object.keys(defs).length === 0) return schema;
+
+        // Deep clone and rewrite all non-standard $ref values to #/$defs/<id>
+        const resolved = JSON.parse(JSON.stringify(schema));
+        delete resolved.$id;
+        delete resolved.$schema;
+
+        const rewriteRefs = (node: any): void => {
+            if (!node || typeof node !== "object") return;
+
+            if (node.$ref && typeof node.$ref === "string" && !node.$ref.startsWith("#")) {
+                node.$ref = `#/$defs/${node.$ref}`;
+            }
+
+            for (const value of Object.values(node)) {
+                rewriteRefs(value);
+            }
+        };
+
+        rewriteRefs(resolved);
+        // Also rewrite refs inside the collected defs
+        for (const def of Object.values(defs)) {
+            rewriteRefs(def);
+        }
+
+        resolved.$defs = defs;
+        return resolved;
+    }
+
+    /**
+     * Get tool result for AI consumption
+     * Formats ToolResult into string for AI context
+     */
+    formatResultForAI(result: ToolResult<any>): string {
+        if (result.success) {
+            return JSON.stringify(result.data);
+        } else {
+            return `Error: ${result.error}${result.code ? ` (${result.code})` : ""}`;
+        }
+    }
+
+    /**
+     * Check if user has permission to use this tool
+     * Used by AgentRunner to filter tools before showing to LLM
+     *
+     * @param user - User object
+     * @param input - Tool input (for function-based permissions)
+     * @param userPermissions - Optional resolved permissions from auth plugin (preferred)
+     * @param conversationContext - Optional conversation context
+     */
+    async checkPermissions(user?: any, input?: any, userPermissions?: string[], conversationContext?: any): Promise<boolean> {
+        const perms = this.toolProps.permissions;
+        if (!perms) return true;
+        if (typeof perms === "function") return await perms(input ?? {}, user);
+
+        // Get effective permissions (prefer explicit userPermissions)
+        const effectivePerms = userPermissions || user?.permissions || [];
+
+        // If no user and no explicit permissions, deny access
+        if (!user && !userPermissions) return false;
+
+        const requiredPerms = Array.isArray(perms) ? perms : [perms];
+        return requiredPerms.every((p) => effectivePerms.includes(p));
+    }
+
+    private async checkPermissionsInternal(input: any, user?: any, userPermissions?: string[], conversationContext?: any): Promise<boolean> {
+        return this.checkPermissions(user, input, userPermissions, conversationContext);
+    }
+
+    /**
+     * Format Zod validation errors into LLM-friendly error messages
+     * Provides specific guidance on what's missing or incorrect
+     */
+    private formatZodError(err: any, receivedInput?: any): string {
+        // Check if it's a Zod error with issues array
+        if (err.issues && Array.isArray(err.issues)) {
+            const issues = err.issues.map((issue: any) => {
+                const path = issue.path.join(".");
+                const field = path || "input";
+
+                if (issue.code === "invalid_type") {
+                    if (issue.received === "undefined") {
+                        return `Missing required field '${field}' (expected ${issue.expected})`;
+                    }
+                    return `Field '${field}' has wrong type: expected ${issue.expected}, got ${issue.received}`;
+                }
+
+                if (issue.code === "too_small") {
+                    if (issue.type === "string") {
+                        return `Field '${field}' is too short (minimum ${issue.minimum} characters)`;
+                    }
+                    return `Field '${field}' is too small (minimum ${issue.minimum})`;
+                }
+
+                if (issue.code === "too_big") {
+                    if (issue.type === "string") {
+                        return `Field '${field}' is too long (maximum ${issue.maximum} characters)`;
+                    }
+                    return `Field '${field}' is too large (maximum ${issue.maximum})`;
+                }
+
+                // Generic fallback
+                return `${field}: ${issue.message}`;
+            });
+
+            if (receivedInput) {
+                const inputInfo =
+                    Object.keys(receivedInput || {}).length === 0 ? "You provided an empty object {}." : `You provided: ${JSON.stringify(receivedInput)}`;
+                return `${issues.join("; ")}. ${inputInfo}`;
+            }
+
+            return issues.join("; ");
+        }
+
+        // Fallback for non-Zod errors
+        return err.message || "Unknown validation error";
+    }
+
+    /**
+     * Format AJV validation errors into LLM-friendly error messages
+     */
+    private formatAjvErrors(errors: any[], receivedInput?: any): string {
+        const issues = errors.map((error) => {
+            const field = error.instancePath ? error.instancePath.substring(1).replace(/\//g, ".") : "input";
+
+            if (error.keyword === "required") {
+                return `Missing required field '${error.params.missingProperty}'`;
+            }
+
+            if (error.keyword === "type") {
+                return `Field '${field}' has wrong type: expected ${error.params.type}`;
+            }
+
+            if (error.keyword === "minLength") {
+                return `Field '${field}' is too short (minimum ${error.params.limit} characters)`;
+            }
+
+            if (error.keyword === "maxLength") {
+                return `Field '${field}' is too long (maximum ${error.params.limit} characters)`;
+            }
+
+            if (error.keyword === "minimum") {
+                return `Field '${field}' is too small (minimum ${error.params.limit})`;
+            }
+
+            if (error.keyword === "maximum") {
+                return `Field '${field}' is too large (maximum ${error.params.limit})`;
+            }
+
+            return `${field}: ${error.message}`;
+        });
+
+        if (receivedInput) {
+            const inputInfo =
+                Object.keys(receivedInput || {}).length === 0 ? "You provided an empty object {}." : `You provided: ${JSON.stringify(receivedInput)}`;
+            return `${issues.join("; ")}. ${inputInfo}`;
+        }
+
+        return issues.join("; ");
+    }
+}