feat: add tool selection by feature param (#172)

* add tools param to select tools by feature

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

* rename to tool categories, include beta tools as separte preview category, remove beta flag param

* simplify building the internal tool map

Author: MQ37

src/actor/server.ts

@@ -13,26 +13,9 @@ import log from '@apify/log';
 
 import { ActorsMcpServer } from '../mcp/server.js';
 import { parseInputParamsFromUrl } from '../mcp/utils.js';
-import { getActorsAsTools } from '../tools/actor.js';
 import { getHelpMessage, HEADER_READINESS_PROBE, Routes } from './const.js';
 import { getActorRunData } from './utils.js';
 
-/**
- * Helper function to load tools and actors based on input parameters
- * @param mcpServer The MCP server instance
- * @param url The request URL to parse parameters from
- * @param apifyToken The Apify token for authentication
- */
-async function loadToolsAndActors(mcpServer: ActorsMcpServer, url: string, apifyToken: string): Promise<void> {
-    const input = parseInputParamsFromUrl(url);
-    if (input.actors || input.enableAddingActors) {
-        await mcpServer.loadToolsFromUrl(url, apifyToken);
-    }
-    if (!input.actors) {
-        await mcpServer.loadDefaultActors(apifyToken);
-    }
-}
-
 export function createExpressApp(
     host: string,
     mcpServerOptions: {
@@ -85,13 +68,21 @@ export function createExpressApp(
         try {
             log.info(`Received GET message at: ${Routes.SSE}`);
             const mcpServer = new ActorsMcpServer(mcpServerOptions, false);
-            // Load tools from Actor input for backwards compatibility
-            if (mcpServerOptions.actors && mcpServerOptions.actors.length > 0) {
-                const tools = await getActorsAsTools(mcpServerOptions.actors, process.env.APIFY_TOKEN as string);
-                mcpServer.upsertTools(tools);
-            }
-            await loadToolsAndActors(mcpServer, req.url, process.env.APIFY_TOKEN as string);
             const transport = new SSEServerTransport(Routes.MESSAGE, res);
+
+            // Load MCP server tools
+            const apifyToken = process.env.APIFY_TOKEN as string;
+            const input = parseInputParamsFromUrl(req.url);
+            if (input.actors || input.enableAddingActors || input.tools) {
+                log.debug('[SSE] Loading tools from URL', { sessionId: transport.sessionId });
+                await mcpServer.loadToolsFromUrl(req.url, apifyToken);
+            }
+            // Load default tools if no actors are specified
+            if (!input.actors) {
+                log.debug('[SSE] Loading default tools', { sessionId: transport.sessionId });
+                await mcpServer.loadDefaultActors(apifyToken);
+            }
+
             transportsSSE[transport.sessionId] = transport;
             mcpServers[transport.sessionId] = mcpServer;
             await mcpServer.connect(transport);
@@ -164,13 +155,20 @@ export function createExpressApp(
                     enableJsonResponse: false, // Use SSE response mode
                 });
                 const mcpServer = new ActorsMcpServer(mcpServerOptions, false);
-                // Load tools from Actor input for backwards compatibility
-                if (mcpServerOptions.actors && mcpServerOptions.actors.length > 0) {
-                    const tools = await getActorsAsTools(mcpServerOptions.actors, process.env.APIFY_TOKEN as string);
-                    mcpServer.upsertTools(tools);
-                }
+
                 // Load MCP server tools
-                await loadToolsAndActors(mcpServer, req.url, process.env.APIFY_TOKEN as string);
+                const apifyToken = process.env.APIFY_TOKEN as string;
+                const input = parseInputParamsFromUrl(req.url);
+                if (input.actors || input.enableAddingActors || input.tools) {
+                    log.debug('[Streamable] Loading tools from URL', { sessionId: transport.sessionId });
+                    await mcpServer.loadToolsFromUrl(req.url, apifyToken);
+                }
+                // Load default tools if no actors are specified
+                if (!input.actors) {
+                    log.debug('[Streamable] Loading default tools', { sessionId: transport.sessionId });
+                    await mcpServer.loadDefaultActors(apifyToken);
+                }
+
                 // Connect the transport to the MCP server BEFORE handling the request
                 await mcpServer.connect(transport);
 

src/input.ts

@@ -3,7 +3,7 @@
  */
 import log from '@apify/log';
 
-import type { Input } from './types.js';
+import type { Input, ToolCategory } from './types.js';
 
 /**
  * Process input parameters, split Actors string into an array
@@ -30,7 +30,8 @@ export function processInput(originalInput: Partial<Input>): Input {
         input.enableAddingActors = input.enableAddingActors === true || input.enableAddingActors === 'true';
     }
 
-    // If beta present, set input.beta to true
-    input.beta = input.beta !== undefined && (input.beta !== false && input.beta !== 'false');
+    if (input.tools && typeof input.tools === 'string') {
+        input.tools = input.tools.split(',').map((tool: string) => tool.trim()) as ToolCategory[];
+    }
     return input;
 }

src/mcp/server.ts

@@ -23,7 +23,7 @@ import {
     SERVER_NAME,
     SERVER_VERSION,
 } from '../const.js';
-import { addRemoveTools, betaTools, callActorGetDataset, defaultTools, getActorsAsTools } from '../tools/index.js';
+import { addRemoveTools, callActorGetDataset, defaultTools, getActorsAsTools, toolCategories } from '../tools/index.js';
 import { actorNameToToolName, decodeDotPropertyNames } from '../tools/utils.js';
 import type { ActorMcpTool, ActorTool, HelperTool, ToolEntry } from '../types.js';
 import { connectMCPClient } from './client.js';
@@ -33,7 +33,6 @@ import { processParamsGetTools } from './utils.js';
 type ActorsMcpServerOptions = {
     enableAddingActors?: boolean;
     enableDefaultActors?: boolean;
-    enableBeta?: boolean; // Enable beta features
 };
 
 type ToolsChangedHandler = (toolNames: string[]) => void;
@@ -52,7 +51,6 @@ export class ActorsMcpServer {
         this.options = {
             enableAddingActors: options.enableAddingActors ?? true,
             enableDefaultActors: options.enableDefaultActors ?? true, // Default to true for backward compatibility
-            enableBeta: options.enableBeta ?? false, // Disabled by default
         };
         this.server = new Server(
             {
@@ -78,10 +76,6 @@ export class ActorsMcpServer {
             this.enableDynamicActorTools();
         }
 
-        if (this.options.enableBeta) {
-            this.upsertTools(betaTools, false);
-        }
-
         // Initialize automatically for backward compatibility
         this.initialize().catch((error) => {
             log.error('Failed to initialize server:', error);
@@ -172,7 +166,11 @@ export class ActorsMcpServer {
         const loadedTools = this.listAllToolNames();
         const actorsToLoad: string[] = [];
         const toolsToLoad: ToolEntry[] = [];
-        const internalToolMap = new Map([...defaultTools, ...addRemoveTools, ...betaTools].map((tool) => [tool.tool.name, tool]));
+        const internalToolMap = new Map([
+            ...defaultTools,
+            ...addRemoveTools,
+            ...Object.values(toolCategories).flat(),
+        ].map((tool) => [tool.tool.name, tool]));
 
         for (const tool of toolNames) {
             // Skip if the tool is already loaded

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, getActorsAsTools } from '../tools/index.js';
-import type { 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,26 +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);
-    }
-    return tools;
+    return await loadToolsFromInput(input, apifyToken);
 }
 
 export function parseInputParamsFromUrl(url: string): Input {

src/stdio.ts

@@ -22,9 +22,10 @@ import { hideBin } from 'yargs/helpers';
 
 import log from '@apify/log';
 
-import { defaults } from './const.js';
 import { ActorsMcpServer } from './mcp/server.js';
-import { getActorsAsTools } from './tools/index.js';
+import { toolCategories } from './tools/index.js';
+import type { Input, ToolCategory } 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
@@ -36,7 +37,8 @@ interface CliArgs {
     enableAddingActors: boolean;
     /** @deprecated */
     enableActorAutoLoading: boolean;
-    beta: boolean;
+    /** Tool categories to include */
+    tools?: string;
 }
 
 // Configure logging, set to ERROR
@@ -47,24 +49,35 @@ const argv = yargs(hideBin(process.argv))
     .usage('Usage: $0 [options]')
     .option('actors', {
         type: 'string',
-        describe: 'Comma-separated list of Actor full names to add to the server',
+        describe: 'Comma-separated list of Actor full names to add to the server.',
         example: 'apify/google-search-scraper,apify/instagram-scraper',
     })
     .option('enable-adding-actors', {
         type: 'boolean',
         default: true,
-        describe: 'Enable dynamically adding Actors as tools based on user requests',
+        describe: 'Enable dynamically adding Actors as tools based on user requests.',
     })
     .option('enableActorAutoLoading', {
         type: 'boolean',
         default: true,
         hidden: true,
-        describe: 'Deprecated: use enable-adding-actors instead',
+        describe: 'Deprecated: use enable-adding-actors instead.',
     })
-    .option('beta', {
-        type: 'boolean',
-        default: false,
-        describe: 'Enable beta features',
+    .options('tools', {
+        type: 'string',
+        describe: `Comma-separated list of specific tool categories to enable.
+
+Available choices: ${Object.keys(toolCategories).join(', ')}
+
+Tool categories are as follows:
+- docs: Search and fetch Apify documentation tools.
+- runs: Get Actor runs list, run details, and logs from a specific Actor run.
+- storage: Access datasets, key-value stores, and their records.
+- preview: Experimental tools in preview mode.
+
+Note: Tools that enable you to search Actors from the Apify Store and get their details are always enabled by default.
+`,
+        example: 'docs,runs,storage',
     })
     .help('help')
     .alias('h', 'help')
@@ -73,13 +86,14 @@ const argv = yargs(hideBin(process.argv))
         'To connect, set your MCP client server command to `npx @apify/actors-mcp-server`'
         + ' and set the environment variable `APIFY_TOKEN` to your Apify API token.\n',
     )
-    .epilogue('For more information, visit https://github.com/apify/actors-mcp-server')
+    .epilogue('For more information, visit https://mcp.apify.com or https://github.com/apify/actors-mcp-server')
     .parseSync() as CliArgs;
 
 const enableAddingActors = argv.enableAddingActors && argv.enableActorAutoLoading;
 const actors = argv.actors as string || '';
 const actorList = actors ? actors.split(',').map((a: string) => a.trim()) : [];
-const enableBeta = argv.beta;
+// Keys of the tool categories to enable
+const toolCategoryKeys = argv.tools ? argv.tools.split(',').map((t: string) => t.trim()) : [];
 
 // Validate environment
 if (!process.env.APIFY_TOKEN) {
@@ -88,8 +102,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);
+    const mcpServer = new ActorsMcpServer({ enableAddingActors, enableDefaultActors: false });
+
+    // Create an Input object from CLI arguments
+    const input: Input = {
+        actors: actorList.length ? actorList : [],
+        enableAddingActors,
+        tools: toolCategoryKeys as ToolCategory[],
+    };
+
+    // 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/tools/index.ts

@@ -1,45 +1,58 @@
 // Import specific tools that are being used
+import type { ToolCategory } from '../types.js';
 import { callActor, callActorGetDataset, getActorsAsTools } from './actor.js';
+import { getDataset, getDatasetItems } from './dataset.js';
+import { getUserDatasetsList } from './dataset_collection.js';
 import { fetchApifyDocsTool } from './fetch-apify-docs.js';
 import { getActorDetailsTool } from './get-actor-details.js';
-import { addTool, helpTool } from './helpers.js';
+import { addTool } from './helpers.js';
+import { getKeyValueStore, getKeyValueStoreKeys, getKeyValueStoreRecord } from './key_value_store.js';
+import { getUserKeyValueStoresList } from './key_value_store_collection.js';
+import { getActorRun, getActorRunLog } from './run.js';
+import { getUserRunsList } from './run_collection.js';
 import { searchApifyDocsTool } from './search-apify-docs.js';
 import { searchActors } from './store_collection.js';
 
+export const toolCategories = {
+    docs: [
+        searchApifyDocsTool,
+        fetchApifyDocsTool,
+    ],
+    runs: [
+        getActorRun,
+        getUserRunsList,
+        getActorRunLog,
+    ],
+    storage: [
+        getDataset,
+        getDatasetItems,
+        getKeyValueStore,
+        getKeyValueStoreKeys,
+        getKeyValueStoreRecord,
+        getUserDatasetsList,
+        getUserKeyValueStoresList,
+    ],
+    preview: [
+        callActor,
+    ],
+};
+export const toolCategoriesEnabledByDefault: ToolCategory[] = [
+    'docs',
+];
+
 export const defaultTools = [
-    // abortActorRun,
-    // actorDetailsTool,
-    // getActor,
-    // getActorLog,
-    // getActorRun,
-    // getDataset,
-    // getDatasetItems,
-    // getKeyValueStore,
-    // getKeyValueStoreKeys,
-    // getKeyValueStoreRecord,
-    // getUserRunsList,
-    // getUserDatasetsList,
-    // getUserKeyValueStoresList,
     getActorDetailsTool,
-    helpTool,
     searchActors,
-    searchApifyDocsTool,
-    fetchApifyDocsTool,
-];
-
-export const betaTools = [
-    callActor,
+    // Add the tools from the enabled categories
+    ...toolCategoriesEnabledByDefault.map((key) => toolCategories[key]).flat(),
 ];
 
 export const addRemoveTools = [
     addTool,
-    // removeTool,
 ];
 
 // Export only the tools that are being used
 export {
-    addTool,
-    // removeTool,
     getActorsAsTools,
     callActorGetDataset,
 };

src/tools/run.ts

@@ -59,7 +59,7 @@ const GetRunLogArgs = z.object({
  * https://docs.apify.com/api/v2/actor-run-get
  *  /v2/actor-runs/{runId}/log{?token}
  */
-export const getActorLog: ToolEntry = {
+export const getActorRunLog: ToolEntry = {
     type: 'internal',
     tool: {
         name: HelperTools.ACTOR_RUNS_LOG,