vibe reafactor tool loading logic from input (url and cli args)

Author: MQ37

src/mcp/utils.ts

@@ -2,8 +2,8 @@ import { createHash } from 'node:crypto';
 import { parse } from 'node:querystring';
 
 import { processInput } from '../input.js';
-import { addRemoveTools, betaTools, featureTools, getActorsAsTools } from '../tools/index.js';
-import type { FeatureToolKey, Input, ToolEntry } from '../types.js';
+import type { Input } from '../types.js';
+import { loadToolsFromInput } from '../utils/tools-loader.js';
 import { MAX_TOOL_NAME_LENGTH, SERVER_ID_LENGTH } from './const.js';
 
 /**
@@ -34,34 +34,14 @@ export function getProxyMCPServerToolName(url: string, toolName: string): string
 }
 
 /**
- * Process input parameters and get tools
+ * Process input parameters from URL and get tools
  * If URL contains query parameter `actors`, return tools from Actors otherwise return null.
  * @param url
  * @param apifyToken
  */
 export async function processParamsGetTools(url: string, apifyToken: string) {
     const input = parseInputParamsFromUrl(url);
-    let tools: ToolEntry[] = [];
-    if (input.actors) {
-        const actors = input.actors as string[];
-        // Normal Actors as a tool
-        tools = await getActorsAsTools(actors, apifyToken);
-    }
-    if (input.enableAddingActors) {
-        tools.push(...addRemoveTools);
-    }
-    if (input.beta) {
-        tools.push(...betaTools);
-    }
-    if (input.tools) {
-        for (const toolKey of input.tools) {
-            // Get tools by feature key
-            const keyTools = featureTools[toolKey as FeatureToolKey] || [];
-            // Push them into the tools array
-            tools.push(...keyTools);
-        }
-    }
-    return tools;
+    return await loadToolsFromInput(input, apifyToken);
 }
 
 export function parseInputParamsFromUrl(url: string): Input {

src/stdio.ts

@@ -22,10 +22,10 @@ import { hideBin } from 'yargs/helpers';
 
 import log from '@apify/log';
 
-import { defaults } from './const.js';
 import { ActorsMcpServer } from './mcp/server.js';
-import { featureTools, getActorsAsTools } from './tools/index.js';
-import type { FeatureToolKey } from './types.js';
+import { featureTools } from './tools/index.js';
+import type { FeatureToolKey, Input } from './types.js';
+import { loadToolsFromInput } from './utils/tools-loader.js';
 
 // Keeping this interface here and not types.ts since
 // it is only relevant to the CLI/STDIO transport in this file
@@ -108,14 +108,18 @@ if (!process.env.APIFY_TOKEN) {
 
 async function main() {
     const mcpServer = new ActorsMcpServer({ enableAddingActors, enableDefaultActors: false, enableBeta });
-    const tools = await getActorsAsTools(actorList.length ? actorList : defaults.actors, process.env.APIFY_TOKEN as string);
-    // Add feature tools based on the command line arguments
-    if (featureToolKeys.length > 0) {
-        for (const key of featureToolKeys) {
-            const featureToolsList = featureTools[key as FeatureToolKey] || [];
-            tools.push(...featureToolsList);
-        }
-    }
+
+    // Create an Input object from CLI arguments
+    const input: Input = {
+        actors: actorList.length ? actorList : [],
+        enableAddingActors,
+        beta: enableBeta,
+        tools: featureToolKeys as FeatureToolKey[],
+    };
+
+    // Use the shared tools loading logic
+    const tools = await loadToolsFromInput(input, process.env.APIFY_TOKEN as string, actorList.length === 0);
+
     mcpServer.upsertTools(tools);
 
     // Start server

src/utils/tools-loader.ts

@@ -0,0 +1,55 @@
+/**
+ * Shared logic for loading tools based on Input type.
+ * This eliminates duplication between stdio.ts and processParamsGetTools.
+ */
+
+import { defaults } from '../const.js';
+import { addRemoveTools, betaTools, featureTools, getActorsAsTools } from '../tools/index.js';
+import type { FeatureToolKey, Input, ToolEntry } from '../types.js';
+
+/**
+ * Load tools based on the provided Input object.
+ * This function is used by both the stdio.ts and the processParamsGetTools function.
+ *
+ * @param input The processed Input object
+ * @param apifyToken The Apify API token
+ * @param useDefaultActors Whether to use default actors if no actors are specified
+ * @returns An array of tool entries
+ */
+export async function loadToolsFromInput(
+    input: Input,
+    apifyToken: string,
+    useDefaultActors = false,
+): Promise<ToolEntry[]> {
+    let tools: ToolEntry[] = [];
+
+    // Load actors as tools
+    if (input.actors && (Array.isArray(input.actors) ? input.actors.length > 0 : input.actors)) {
+        const actors = Array.isArray(input.actors) ? input.actors : [input.actors];
+        tools = await getActorsAsTools(actors, apifyToken);
+    } else if (useDefaultActors) {
+        // Use default actors if no actors are specified and useDefaultActors is true
+        tools = await getActorsAsTools(defaults.actors, apifyToken);
+    }
+
+    // Add tools for adding/removing actors if enabled
+    if (input.enableAddingActors) {
+        tools.push(...addRemoveTools);
+    }
+
+    // Add beta tools if enabled
+    if (input.beta) {
+        tools.push(...betaTools);
+    }
+
+    // Add feature tools based on the input
+    if (input.tools) {
+        const toolKeys = Array.isArray(input.tools) ? input.tools : [input.tools];
+        for (const toolKey of toolKeys) {
+            const keyTools = featureTools[toolKey as FeatureToolKey] || [];
+            tools.push(...keyTools);
+        }
+    }
+
+    return tools;
+}