feat: add tool selection by feature param

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.beta || 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.beta || 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 { FeatureToolKey, Input } from './types.js';
 
 /**
  * Process input parameters, split Actors string into an array
@@ -32,5 +32,9 @@ export function processInput(originalInput: Partial<Input>): Input {
 
     // 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 FeatureToolKey[];
+    }
     return input;
 }

src/mcp/server.ts

@@ -23,9 +23,9 @@ import {
     SERVER_NAME,
     SERVER_VERSION,
 } from '../const.js';
-import { addRemoveTools, betaTools, callActorGetDataset, defaultTools, getActorsAsTools } from '../tools/index.js';
+import { addRemoveTools, betaTools, callActorGetDataset, defaultTools, featureTools, getActorsAsTools } from '../tools/index.js';
 import { actorNameToToolName, decodeDotPropertyNames } from '../tools/utils.js';
-import type { ActorMcpTool, ActorTool, HelperTool, ToolEntry } from '../types.js';
+import type { ActorMcpTool, ActorTool, FeatureToolKey, HelperTool, ToolEntry } from '../types.js';
 import { connectMCPClient } from './client.js';
 import { EXTERNAL_TOOL_CALL_TIMEOUT_MSEC } from './const.js';
 import { processParamsGetTools } from './utils.js';
@@ -173,6 +173,13 @@ export class ActorsMcpServer {
         const actorsToLoad: string[] = [];
         const toolsToLoad: ToolEntry[] = [];
         const internalToolMap = new Map([...defaultTools, ...addRemoveTools, ...betaTools].map((tool) => [tool.tool.name, tool]));
+        // Add all feature tools
+        for (const key of Object.keys(featureTools)) {
+            const tools = featureTools[key as FeatureToolKey];
+            for (const tool of tools) {
+                internalToolMap.set(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 { 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
@@ -37,6 +38,7 @@ interface CliArgs {
     /** @deprecated */
     enableActorAutoLoading: boolean;
     beta: boolean;
+    tools?: string;
 }
 
 // Configure logging, set to ERROR
@@ -47,24 +49,39 @@ 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',
+        describe: 'Enable beta features.',
+    })
+    .options('tools', {
+        type: 'string',
+        describe: `Comma-separated list of specific feature tools to enable.
+
+Available choices: ${Object.keys(featureTools).join(', ')}
+
+Feature tools are categorized 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.
+
+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 +90,15 @@ 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 feature tools to enable
+const featureToolKeys = argv.tools ? argv.tools.split(',').map((t: string) => t.trim()) : [];
 
 // Validate environment
 if (!process.env.APIFY_TOKEN) {
@@ -89,7 +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);
+
+    // 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/tools/index.ts

@@ -1,45 +1,53 @@
 // Import specific tools that are being used
 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 defaultTools = [
-    // abortActorRun,
-    // actorDetailsTool,
-    // getActor,
-    // getActorLog,
-    // getActorRun,
-    // getDataset,
-    // getDatasetItems,
-    // getKeyValueStore,
-    // getKeyValueStoreKeys,
-    // getKeyValueStoreRecord,
-    // getUserRunsList,
-    // getUserDatasetsList,
-    // getUserKeyValueStoresList,
     getActorDetailsTool,
-    helpTool,
     searchActors,
-    searchApifyDocsTool,
-    fetchApifyDocsTool,
 ];
 
+export const featureTools = {
+    docs: [
+        searchApifyDocsTool,
+        fetchApifyDocsTool,
+    ],
+    runs: [
+        getActorRun,
+        getUserRunsList,
+        getActorRunLog,
+    ],
+    storage: [
+        getDataset,
+        getDatasetItems,
+        getKeyValueStore,
+        getKeyValueStoreKeys,
+        getKeyValueStoreRecord,
+        getUserDatasetsList,
+        getUserKeyValueStoresList,
+    ],
+};
+
 export const betaTools = [
     callActor,
 ];
 
 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,

src/types.ts

@@ -6,6 +6,7 @@ import type { ActorDefaultRunOptions, ActorDefinition, ActorStoreList, PricingIn
 
 import type { ACTOR_PRICING_MODEL } from './const.js';
 import type { ActorsMcpServer } from './mcp/server.js';
+import type { featureTools } from './tools/index.js';
 
 export interface ISchemaProperties {
     type: string;
@@ -213,6 +214,8 @@ export interface InternalTool extends ToolBase {
     call: (toolArgs: InternalToolArgs) => Promise<object>;
 }
 
+export type FeatureToolKey = keyof typeof featureTools;
+
 export type Input = {
     actors: string[] | string;
     /**
@@ -225,6 +228,7 @@ export type Input = {
     debugActorInput?: unknown;
     /** Enable beta features flag */
     beta?: boolean | string;
+    tools?: FeatureToolKey[] | string;
 };
 
 // Utility type to get a union of values from an object type