@mixio-pro/kalaasetu-mcp 2.1.0 → 2.1.2-beta

package/src/tools/fal/dynamic-tools.ts

@@ -5,14 +5,13 @@
  * Each preset becomes a first-class MCP tool with its own Zod schema.
  */
 
-import jsonata from "jsonata";
-import { z } from "zod";
+import { z, type ZodTypeAny } from "zod";
+import { safeToolExecute } from "../../utils/tool-wrapper";
 import { resolveEnhancer } from "../../utils/prompt-enhancer-presets";
 import { sanitizeResponse } from "../../utils/sanitize";
-import { safeToolExecute } from "../../utils/tool-wrapper";
 import {
-  AUTHENTICATED_TIMEOUT,
   FAL_QUEUE_URL,
+  AUTHENTICATED_TIMEOUT,
   getApiKey,
   loadFalConfig,
   type FalPresetConfig,
@@ -31,7 +30,7 @@ async function wait(ms: number): Promise<void> {
 async function authenticatedRequest(
   url: string,
   method: "GET" | "POST" | "PUT" = "GET",
-  jsonData?: Record<string, any>
+  jsonData?: Record<string, any>,
 ): Promise<any> {
   const headers: Record<string, string> = {
     Authorization: `Key ${getApiKey()}`,
@@ -68,10 +67,10 @@ async function authenticatedRequest(
  * Sanitize parameters by removing null/undefined values.
  */
 function sanitizeParameters(
-  parameters: Record<string, any>
+  parameters: Record<string, any>,
 ): Record<string, any> {
   return Object.fromEntries(
-    Object.entries(parameters).filter(([_, v]) => v != null)
+    Object.entries(parameters).filter(([_, v]) => v != null),
   );
 }
 
@@ -84,7 +83,7 @@ interface ProgressContext {
     total: number;
   }) => Promise<void>;
   streamContent?: (
-    content: { type: "text"; text: string } | { type: "text"; text: string }[]
+    content: { type: "text"; text: string } | { type: "text"; text: string }[],
   ) => Promise<void>;
   log?: {
     info: (message: string, data?: any) => void;
@@ -92,13 +91,12 @@ interface ProgressContext {
   };
 }
 
-
 /**
  * Build a Zod object schema from preset input_schema using built-in z.fromJSONSchema.
  */
 function buildZodSchema(
   inputSchema: Record<string, any> | undefined,
-  defaultParams?: Record<string, any>
+  defaultParams?: Record<string, any>,
 ): z.ZodObject<any> {
   // Construct a properties object for JSON schema
   const properties: Record<string, any> = {};
@@ -129,13 +127,13 @@ function buildZodSchema(
       .optional()
       .describe(
         "If provided, resume polling for an existing request instead of starting a new one. " +
-          "Use the 'resume_endpoint' returned in an 'IN_PROGRESS' response."
+          "Use the 'resume_endpoint' returned in an 'IN_PROGRESS' response.",
       ),
     auto_enhance: z
       .boolean()
       .default(true) // Our custom default
       .describe(
-        "Whether to automatically enhance the prompt. Set to false to disable."
+        "Whether to automatically enhance the prompt. Set to false to disable.",
       ),
   });
 
@@ -148,7 +146,7 @@ function buildZodSchema(
 export function createToolFromPreset(preset: FalPresetConfig) {
   const zodSchema = buildZodSchema(
     preset.input_schema as Record<string, any>,
-    preset.defaultParams
+    preset.defaultParams,
   );
 
   const toolName = preset.presetName.startsWith("fal_")
@@ -164,266 +162,245 @@ export function createToolFromPreset(preset: FalPresetConfig) {
     parameters: zodSchema,
     timeoutMs: 90000, // 90 seconds MCP timeout
 
-    execute: async (
-      args: Record<string, any>,
-      context?: ProgressContext
-    ) => {
-      return safeToolExecute(async () => {
-        let statusUrl: string;
-        let responseUrl: string;
-        let requestId: string;
-
-        // Handle resume flow
-        if (args.resume_endpoint) {
-          if (args.resume_endpoint.startsWith("http")) {
-            statusUrl = args.resume_endpoint;
-            responseUrl = args.resume_endpoint.replace(/\/status$/, "");
-            const urlParts = args.resume_endpoint.split("/");
-            const lastPart = urlParts[urlParts.length - 1] || "";
-            requestId =
-              lastPart.replace("/status", "") ||
-              urlParts[urlParts.length - 2] ||
-              "unknown";
-            context?.log?.info(`Resuming with FAL URL: ${args.resume_endpoint}`);
-          } else {
-            // Legacy UUID format - reconstruct URL
-            requestId = args.resume_endpoint;
-            statusUrl = `${FAL_QUEUE_URL}/${preset.modelId}/requests/${requestId}/status`;
-            responseUrl = `${FAL_QUEUE_URL}/${preset.modelId}/requests/${requestId}`;
-            context?.log?.info(
-              `Resuming polling for ${preset.modelId} request: ${requestId}`
-            );
-          }
-        } else {
-          // New request - start generation
-          try {
-            const apiKey = getApiKey();
-            if (context?.streamContent) {
-              await context.streamContent({
-                type: "text" as const,
-                text: `[FAL] ✓ API key found (${apiKey.slice(0, 8)}...). Using model: ${preset.modelId}`,
-              });
+    execute: async (args: Record<string, any>, context?: ProgressContext) => {
+      return safeToolExecute(
+        async () => {
+          let statusUrl: string;
+          let responseUrl: string;
+          let requestId: string;
+
+          // Handle resume flow
+          if (args.resume_endpoint) {
+            if (args.resume_endpoint.startsWith("http")) {
+              statusUrl = args.resume_endpoint;
+              responseUrl = args.resume_endpoint.replace(/\/status$/, "");
+              const urlParts = args.resume_endpoint.split("/");
+              const lastPart = urlParts[urlParts.length - 1] || "";
+              requestId =
+                lastPart.replace("/status", "") ||
+                urlParts[urlParts.length - 2] ||
+                "unknown";
+              context?.log?.info(
+                `Resuming with FAL URL: ${args.resume_endpoint}`,
+              );
+            } else {
+              // Legacy UUID format - reconstruct URL
+              requestId = args.resume_endpoint;
+              statusUrl = `${FAL_QUEUE_URL}/${preset.modelId}/requests/${requestId}/status`;
+              responseUrl = `${FAL_QUEUE_URL}/${preset.modelId}/requests/${requestId}`;
+              context?.log?.info(
+                `Resuming polling for ${preset.modelId} request: ${requestId}`,
+              );
             }
-          } catch (keyError: any) {
-            throw keyError;
-          }
-
-          // Build parameters: input_schema defaults → defaultParams → user args
-          // Extract only the model parameters (exclude our internal fields)
-          const { resume_endpoint, auto_enhance, ...userParams } = args;
-
-          // Start with defaults from input_schema
-          const schemaDefaults: Record<string, any> = {};
-          if (preset.input_schema) {
-            for (const [key, paramSchema] of Object.entries(preset.input_schema)) {
-              if ((paramSchema as any).default !== undefined) {
-                schemaDefaults[key] = (paramSchema as any).default;
+          } else {
+            // New request - start generation
+            try {
+              const apiKey = getApiKey();
+              if (context?.streamContent) {
+                await context.streamContent({
+                  type: "text" as const,
+                  text: `[FAL] ✓ API key found (${apiKey.slice(0, 8)}...). Using model: ${preset.modelId}`,
+                });
               }
+            } catch (keyError: any) {
+              throw keyError;
             }
-          }
 
-          // Merge: schema defaults → preset defaultParams → user params
-          const mergedParams: Record<string, any> = {
-            ...schemaDefaults,
-            ...(preset.defaultParams || {}),
-            ...userParams,
-          };
-
-          // Apply transformations if defined
-          if (preset.transformers) {
-            // 1. JSONata transformation
-            if (preset.transformers.jsonata) {
-              try {
-                const expression = jsonata(preset.transformers.jsonata);
-                const transformed = await expression.evaluate(mergedParams);
-                if (transformed && typeof transformed === "object") {
-                  Object.assign(mergedParams, transformed);
+            // Build parameters: input_schema defaults → defaultParams → user args
+            // Extract only the model parameters (exclude our internal fields)
+            const { resume_endpoint, auto_enhance, ...userParams } = args;
+
+            // Start with defaults from input_schema
+            const schemaDefaults: Record<string, any> = {};
+            if (preset.input_schema) {
+              for (const [key, paramSchema] of Object.entries(
+                preset.input_schema,
+              )) {
+                if ((paramSchema as any).default !== undefined) {
+                  schemaDefaults[key] = (paramSchema as any).default;
                 }
-              } catch (err) {
-                context?.log?.info(`JSONata transformation failed: ${err}`);
               }
             }
 
-            // 2. Type-based coercion
-            if (preset.transformers.types) {
-              for (const [key, type] of Object.entries(preset.transformers.types)) {
-                if (mergedParams[key] !== undefined && mergedParams[key] !== null) {
-                  if (type === "number") {
-                    mergedParams[key] = Number(mergedParams[key]);
-                  } else if (type === "string") {
-                    mergedParams[key] = String(mergedParams[key]);
-                  } else if (type === "boolean") {
-                    mergedParams[key] = String(mergedParams[key]).toLowerCase() === "true";
+            // Merge: schema defaults → preset defaultParams → user params
+            const mergedParams = {
+              ...schemaDefaults,
+              ...(preset.defaultParams || {}),
+              ...userParams,
+            };
+
+            // Apply prompt enhancement if enabled
+            const shouldEnhance = auto_enhance !== false;
+            if (shouldEnhance && preset.promptEnhancer && mergedParams.prompt) {
+              const enhancerName =
+                typeof preset.promptEnhancer === "string"
+                  ? preset.promptEnhancer
+                  : null;
+
+              if (enhancerName === "ltx2") {
+                const { enhancePromptWithLLM, isLLMEnhancerAvailable } =
+                  await import("../../utils/llm-prompt-enhancer");
+                if (isLLMEnhancerAvailable()) {
+                  try {
+                    const originalPrompt = mergedParams.prompt;
+                    mergedParams.prompt = await enhancePromptWithLLM(
+                      mergedParams.prompt,
+                      "ltx2",
+                    );
+                    context?.log?.info(
+                      `LLM-enhanced prompt: "${originalPrompt}" → "${mergedParams.prompt}"`,
+                    );
+                  } catch (err) {
+                    context?.log?.info(
+                      `LLM enhancement failed, using original prompt`,
+                    );
+                  }
+                }
+              } else if (preset.promptEnhancer) {
+                const enhancer = resolveEnhancer(preset.promptEnhancer);
+                if (enhancer.hasTransformations()) {
+                  mergedParams.prompt = enhancer.enhance(mergedParams.prompt);
+                  const negatives = enhancer.getNegativeElements();
+                  if (negatives && !mergedParams.negative_prompt) {
+                    mergedParams.negative_prompt = negatives;
                   }
                 }
               }
             }
-          }
 
-          // Apply prompt enhancement if enabled
-          const shouldEnhance = auto_enhance !== false;
-          if (shouldEnhance && preset.promptEnhancer && mergedParams.prompt) {
-            const enhancerName =
-              typeof preset.promptEnhancer === "string"
-                ? preset.promptEnhancer
-                : null;
-
-            if (enhancerName === "ltx2") {
-              const { enhancePromptWithLLM, isLLMEnhancerAvailable } =
-                await import("../../utils/llm-prompt-enhancer");
-              if (isLLMEnhancerAvailable()) {
-                try {
-                  const originalPrompt = mergedParams.prompt;
-                  mergedParams.prompt = await enhancePromptWithLLM(
-                    mergedParams.prompt,
-                    "ltx2"
-                  );
-                  context?.log?.info(
-                    `LLM-enhanced prompt: "${originalPrompt}" → "${mergedParams.prompt}"`
-                  );
-                } catch (err) {
-                  context?.log?.info(
-                    `LLM enhancement failed, using original prompt`
-                  );
-                }
-              }
-            } else if (preset.promptEnhancer) {
-              const enhancer = resolveEnhancer(preset.promptEnhancer);
-              if (enhancer.hasTransformations()) {
-                mergedParams.prompt = enhancer.enhance(mergedParams.prompt);
-                const negatives = enhancer.getNegativeElements();
-                if (negatives && !mergedParams.negative_prompt) {
-                  mergedParams.negative_prompt = negatives;
-                }
-              }
+            const sanitizedParams = sanitizeParameters(mergedParams);
+            const url = `${FAL_QUEUE_URL}/${preset.modelId}`;
+
+            if (context?.streamContent) {
+              await context.streamContent({
+                type: "text" as const,
+                text: `[FAL] Submitting generation request to ${preset.modelId}...`,
+              });
             }
-          }
 
-          const sanitizedParams = sanitizeParameters(mergedParams);
-          const url = `${FAL_QUEUE_URL}/${preset.modelId}`;
+            const queueRes = await authenticatedRequest(
+              url,
+              "POST",
+              sanitizedParams,
+            );
 
-          if (context?.streamContent) {
-            await context.streamContent({
-              type: "text" as const,
-              text: `[FAL] Submitting generation request to ${preset.modelId}...`,
-            });
-          }
+            if (!queueRes.request_id && !queueRes.status_url) {
+              return JSON.stringify(sanitizeResponse(queueRes));
+            }
 
-          const queueRes = await authenticatedRequest(
-            url,
-            "POST",
-            sanitizedParams
-          );
+            requestId =
+              queueRes.request_id ||
+              queueRes.status_url?.split("/").pop() ||
+              "";
 
-          if (!queueRes.request_id && !queueRes.status_url) {
-            return JSON.stringify(sanitizeResponse(queueRes));
-          }
+            if (!requestId) {
+              throw new Error("Could not extract request ID from response");
+            }
 
-          requestId =
-            queueRes.request_id || queueRes.status_url?.split("/").pop() || "";
+            statusUrl =
+              queueRes.status_url ||
+              `${FAL_QUEUE_URL}/${preset.modelId}/requests/${requestId}/status`;
+            responseUrl =
+              queueRes.response_url ||
+              `${FAL_QUEUE_URL}/${preset.modelId}/requests/${requestId}`;
 
-          if (!requestId) {
-            throw new Error("Could not extract request ID from response");
+            if (context?.streamContent) {
+              await context.streamContent({
+                type: "text" as const,
+                text: `[FAL] Generation started. resume_endpoint: ${statusUrl}`,
+              });
+            }
           }
 
-          statusUrl =
-            queueRes.status_url ||
-            `${FAL_QUEUE_URL}/${preset.modelId}/requests/${requestId}/status`;
-          responseUrl =
-            queueRes.response_url ||
-            `${FAL_QUEUE_URL}/${preset.modelId}/requests/${requestId}`;
-
-          if (context?.streamContent) {
+          // Stream message for resume calls
+          if (args.resume_endpoint && context?.streamContent) {
             await context.streamContent({
               type: "text" as const,
-              text: `[FAL] Generation started. resume_endpoint: ${statusUrl}`,
+              text: `[FAL] Resuming status check for job: ${requestId}`,
             });
           }
-        }
 
-        // Stream message for resume calls
-        if (args.resume_endpoint && context?.streamContent) {
-          await context.streamContent({
-            type: "text" as const,
-            text: `[FAL] Resuming status check for job: ${requestId}`,
-          });
-        }
-
-        // Poll for completion
-        const startTime = Date.now();
-        const MAX_POLL_TIME = 60000; // 60 seconds internal timeout
-        let pollCount = 0;
-        const POLL_INTERVAL = 3000;
-
-        while (Date.now() - startTime < MAX_POLL_TIME) {
-          pollCount++;
-          let res;
-          try {
-            res = await authenticatedRequest(statusUrl, "GET");
-          } catch (e: any) {
-            if (`${e}`.includes("405")) {
-              context?.log?.info(
-                `Status check 405 on ${statusUrl}, trying fallback...`
-              );
-              res = await authenticatedRequest(responseUrl, "GET");
-              statusUrl = responseUrl;
-            } else {
-              throw e;
+          // Poll for completion
+          const startTime = Date.now();
+          const MAX_POLL_TIME = 60000; // 60 seconds internal timeout
+          let pollCount = 0;
+          const POLL_INTERVAL = 3000;
+
+          while (Date.now() - startTime < MAX_POLL_TIME) {
+            pollCount++;
+            let res;
+            try {
+              res = await authenticatedRequest(statusUrl, "GET");
+            } catch (e: any) {
+              if (`${e}`.includes("405")) {
+                context?.log?.info(
+                  `Status check 405 on ${statusUrl}, trying fallback...`,
+                );
+                res = await authenticatedRequest(responseUrl, "GET");
+                statusUrl = responseUrl;
+              } else {
+                throw e;
+              }
             }
-          }
 
-          if (res.status_url) statusUrl = res.status_url;
-          if (res.response_url) responseUrl = res.response_url;
+            if (res.status_url) statusUrl = res.status_url;
+            if (res.response_url) responseUrl = res.response_url;
 
-          if (context?.reportProgress) {
-            const elapsed = Date.now() - startTime;
-            const progressPercent = Math.min(
-              Math.round((elapsed / MAX_POLL_TIME) * 100),
-              99
-            );
-            await context.reportProgress({
-              progress: progressPercent,
-              total: 100,
-            });
-          }
+            if (context?.reportProgress) {
+              const elapsed = Date.now() - startTime;
+              const progressPercent = Math.min(
+                Math.round((elapsed / MAX_POLL_TIME) * 100),
+                99,
+              );
+              await context.reportProgress({
+                progress: progressPercent,
+                total: 100,
+              });
+            }
 
-          if (context?.streamContent && pollCount % 5 === 0) {
-            await context.streamContent({
-              type: "text" as const,
-              text: `[FAL] Still processing... (${Math.round(
-                (Date.now() - startTime) / 1000
-              )}s elapsed, status: ${res.status})`,
-            });
-          }
+            if (context?.streamContent && pollCount % 5 === 0) {
+              await context.streamContent({
+                type: "text" as const,
+                text: `[FAL] Still processing... (${Math.round(
+                  (Date.now() - startTime) / 1000,
+                )}s elapsed, status: ${res.status})`,
+              });
+            }
 
-          if (res.status === "COMPLETED") {
-            if (context?.reportProgress) {
-              await context.reportProgress({ progress: 100, total: 100 });
+            if (res.status === "COMPLETED") {
+              if (context?.reportProgress) {
+                await context.reportProgress({ progress: 100, total: 100 });
+              }
+              const finalResult = await authenticatedRequest(
+                responseUrl,
+                "GET",
+              );
+              return JSON.stringify(sanitizeResponse(finalResult));
             }
-            const finalResult = await authenticatedRequest(responseUrl, "GET");
-            return JSON.stringify(sanitizeResponse(finalResult));
-          }
 
-          if (res.status === "FAILED") {
-            throw new Error(
-              `Generation failed: ${JSON.stringify(res.error || res)}`
-            );
+            if (res.status === "FAILED") {
+              throw new Error(
+                `Generation failed: ${JSON.stringify(res.error || res)}`,
+              );
+            }
+
+            await wait(POLL_INTERVAL);
           }
 
-          await wait(POLL_INTERVAL);
-        }
-
-        // Timeout - return resume_endpoint
-        return JSON.stringify({
-          status: "IN_PROGRESS",
-          request_id: requestId,
-          resume_endpoint: statusUrl,
-          status_url: statusUrl,
-          response_url: responseUrl,
-          message:
-            "The generation is still in progress. Call this tool again with resume_endpoint to continue polling.",
-        });
-      }, preset.presetName);
+          // Timeout - return resume_endpoint
+          return JSON.stringify({
+            status: "IN_PROGRESS",
+            request_id: requestId,
+            resume_endpoint: statusUrl,
+            status_url: statusUrl,
+            response_url: responseUrl,
+            message:
+              "The generation is still in progress. Call this tool again with resume_endpoint to continue polling.",
+          });
+        },
+        preset.presetName,
+        { toolName: toolName },
+      );
     },
   };
 }