feat(W-18738193): dynamic tools

README.md

@@ -174,6 +174,12 @@ This example shows how to enable the `data`, `orgs`, and `metadata` toolsets whe
      }
 ```
 
+#### Dynamic Tools
+
+The `--dynamic-tools` flag enables dynamic tool discovery and loading. When this flag is set, the MCP server starts with a minimal set of core tools and will load new tools as the need arises. This is useful for reducing initial context size and improving LLM performance.
+
+**NOTE:** This feature works in VSCode and Cline but may not work in other environments.
+
 #### Core Toolset (always enabled)
 
 Includes this tool:

src/index.ts

@@ -24,11 +24,11 @@ import * as data from './tools/data/index.js';
 import * as users from './tools/users/index.js';
 import * as testing from './tools/testing/index.js';
 import * as metadata from './tools/metadata/index.js';
+import * as dynamic from './tools/dynamic/index.js';
 import Cache from './shared/cache.js';
 import { Telemetry } from './telemetry.js';
 import { SfMcpServer } from './sf-mcp-server.js';
-
-const TOOLSETS = ['all', 'testing', 'orgs', 'data', 'users', 'metadata', 'experimental'] as const;
+import { determineToolsetsToEnable, TOOLSETS } from './shared/tools.js';
 
 /**
  * Sanitizes an array of org usernames by replacing specific orgs with a placeholder.
@@ -89,12 +89,12 @@ You can also use special values to control access to orgs:
       },
     }),
     toolsets: Flags.option({
-      options: TOOLSETS,
+      options: ['all', ...TOOLSETS] as const,
       char: 't',
       summary: 'Toolset to enable',
       multiple: true,
       delimiter: ',',
-      default: ['all'],
+      exclusive: ['dynamic-toolsets'],
     })(),
     version: Flags.version(),
     'no-telemetry': Flags.boolean({
@@ -103,6 +103,11 @@ You can also use special values to control access to orgs:
     debug: Flags.boolean({
       summary: 'Enable debug logging',
     }),
+    'dynamic-tools': Flags.boolean({
+      summary: 'Enable dynamic toolsets',
+      char: 'd',
+      exclusive: ['toolsets'],
+    }),
   };
 
   public static examples = [
@@ -127,7 +132,7 @@ You can also use special values to control access to orgs:
 
     if (!flags['no-telemetry']) {
       this.telemetry = new Telemetry(this.config, {
-        toolsets: flags.toolsets.join(', '),
+        toolsets: (flags.toolsets ?? []).join(', '),
         orgs: sanitizeOrgInput(flags.orgs),
       });
 
@@ -139,7 +144,7 @@ You can also use special values to control access to orgs:
       });
     }
 
-    Cache.getInstance().set('allowedOrgs', new Set(flags.orgs));
+    await Cache.safeSet('allowedOrgs', new Set(flags.orgs));
     this.logToStderr(`Allowed orgs:\n${flags.orgs.map((org) => `- ${org}`).join('\n')}`);
     const server = new SfMcpServer(
       {
@@ -150,24 +155,39 @@ You can also use special values to control access to orgs:
           tools: {},
         },
       },
-      { telemetry: this.telemetry }
+      {
+        telemetry: this.telemetry,
+        dynamicTools: flags['dynamic-tools'] ?? false,
+      }
     );
 
-    const enabledToolsets = new Set(flags.toolsets);
-    const all = enabledToolsets.has('all');
+    const toolsetsToEnable = determineToolsetsToEnable(flags.toolsets ?? ['all'], flags['dynamic-tools'] ?? false);
 
     // ************************
     // CORE TOOLS (always on)
+    // If you're adding a new tool to the core toolset, you MUST add it to the `CORE_TOOLS` array in shared/tools.ts
+    // otherwise SfMcpServer will not register it.
+    //
+    // Long term, we will want to consider a more elegant solution for registering core tools.
     // ************************
     this.logToStderr('Registering core tools');
     // get username
     core.registerToolGetUsername(server);
     core.registerToolResume(server);
 
+    // DYNAMIC TOOLSETS
+    // ************************
+    if (toolsetsToEnable.dynamic) {
+      this.logToStderr('Registering dynamic tools');
+      // Individual tool management
+      dynamic.registerToolEnableTool(server);
+      dynamic.registerToolListTools(server);
+    }
+
     // ************************
     // ORG TOOLS
     // ************************
-    if (all || enabledToolsets.has('orgs')) {
+    if (toolsetsToEnable.orgs) {
       this.logToStderr('Registering org tools');
       // list all orgs
       orgs.registerToolListAllOrgs(server);
@@ -176,7 +196,7 @@ You can also use special values to control access to orgs:
     // ************************
     // DATA TOOLS
     // ************************
-    if (all || enabledToolsets.has('data')) {
+    if (toolsetsToEnable.data) {
       this.logToStderr('Registering data tools');
       // query org
       data.registerToolQueryOrg(server);
@@ -185,7 +205,7 @@ You can also use special values to control access to orgs:
     // ************************
     // USER TOOLS
     // ************************
-    if (all || enabledToolsets.has('users')) {
+    if (toolsetsToEnable.users) {
       this.logToStderr('Registering user tools');
       // assign permission set
       users.registerToolAssignPermissionSet(server);
@@ -194,7 +214,7 @@ You can also use special values to control access to orgs:
     // ************************
     // testing TOOLS
     // ************************
-    if (all || enabledToolsets.has('testing')) {
+    if (toolsetsToEnable.testing) {
       this.logToStderr('Registering testing tools');
       testing.registerToolRunApexTest(server);
       testing.registerToolRunAgentTest(server);
@@ -203,7 +223,7 @@ You can also use special values to control access to orgs:
     // ************************
     // METADATA TOOLS
     // ************************
-    if (all || enabledToolsets.has('metadata')) {
+    if (toolsetsToEnable.metadata) {
       this.logToStderr('Registering metadata tools');
       // deploy metadata
       metadata.registerToolDeployMetadata(server);
@@ -217,7 +237,7 @@ You can also use special values to control access to orgs:
     // This toolset needs to be explicitly enabled ('all' will not include it)
     // Tools don't need to be in an 'experimental' directory, only registered here
     // ************************
-    if (enabledToolsets.has('experimental')) {
+    if (toolsetsToEnable.experimental) {
       this.logToStderr('Registering experimental tools');
       // Add any experimental tools here
     }

src/sf-mcp-server.ts

@@ -21,12 +21,18 @@ import { RequestHandlerExtra } from '@modelcontextprotocol/sdk/shared/protocol.j
 import { Logger } from '@salesforce/core';
 import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
 import { Telemetry } from './telemetry.js';
+import { addTool, CORE_TOOLS } from './shared/tools.js';
 
 type ToolMethodSignatures = {
   tool: McpServer['tool'];
   connect: McpServer['connect'];
 };
 
+type Options = ServerOptions & {
+  telemetry?: Telemetry;
+  dynamicTools?: boolean;
+};
+
 /**
  * A server implementation that extends the base MCP server with telemetry capabilities.
  *
@@ -41,15 +47,18 @@ export class SfMcpServer extends McpServer implements ToolMethodSignatures {
   /** Optional telemetry instance for tracking server events */
   private telemetry?: Telemetry;
 
+  private dynamicTools: boolean;
+
   /**
    * Creates a new SfMcpServer instance
    *
    * @param {Implementation} serverInfo - The server implementation details
-   * @param {ServerOptions & { telemetry?: Telemetry }} [options] - Optional server configuration including telemetry
+   * @param {Options} [options] - Optional server configuration including telemetry and dynamic tools support
    */
-  public constructor(serverInfo: Implementation, options?: ServerOptions & { telemetry?: Telemetry }) {
+  public constructor(serverInfo: Implementation, options?: Options) {
     super(serverInfo, options);
     this.telemetry = options?.telemetry;
+    this.dynamicTools = options?.dynamicTools ?? false;
     this.server.oninitialized = (): void => {
       const clientInfo = this.server.getClientVersion();
       if (clientInfo) {
@@ -101,6 +110,14 @@ export class SfMcpServer extends McpServer implements ToolMethodSignatures {
     };
 
     // @ts-expect-error because we no longer know what the type of rest is
-    return super.tool(name, ...rest.slice(0, -1), wrappedCb);
+    const tool = super.tool(name, ...rest.slice(0, -1), wrappedCb);
+
+    if (this.dynamicTools && !CORE_TOOLS.includes(name)) {
+      tool.disable();
+      addTool(tool, name).catch((error) => {
+        this.logger.error(`Failed to add tool ${name}:`, error);
+      });
+    }
+    return tool;
   };
 }

src/shared/auth.ts

@@ -118,7 +118,7 @@ export async function getAllAllowedOrgs(): Promise<SanitizedOrgAuthorization[]>
   const url = new URL(import.meta.url);
   const params = url.searchParams.get('orgs');
   const paramOrg = params ? params : undefined;
-  const orgAllowList = paramOrg ? new Set([paramOrg]) : Cache.getInstance().get('allowedOrgs') ?? new Set<string>();
+  const orgAllowList = paramOrg ? new Set([paramOrg]) : (await Cache.safeGet('allowedOrgs')) ?? new Set<string>();
   // Get all orgs on the user's machine
   const allOrgs = await AuthInfo.listAllAuthorizations();
 

src/shared/cache.ts

@@ -14,32 +14,113 @@
  * limitations under the License.
  */
 
+import { ToolInfo } from './types.js';
+
 type CacheContents = {
   allowedOrgs: Set<string>;
+  tools: ToolInfo[];
 };
 
 type ValueOf<T> = T[keyof T];
 
 /**
- * A simple cache for storing values that need to be accessed globally.
+ * Simple mutex implementation using promises
+ */
+class Mutex {
+  private mutex = Promise.resolve();
+
+  public async lock<T>(fn: () => Promise<T> | T): Promise<T> {
+    const unlock = await this.acquire();
+    try {
+      return await fn();
+    } finally {
+      unlock();
+    }
+  }
+
+  private async acquire(): Promise<() => void> {
+    let release: () => void;
+    const promise = new Promise<void>((resolve) => {
+      release = resolve;
+    });
+
+    const currentMutex = this.mutex;
+    this.mutex = this.mutex.then(() => promise);
+
+    await currentMutex;
+    return release!;
+  }
+}
+
+/**
+ * A thread-safe cache providing generic Map operations with mutex protection.
+ * Offers atomic read, write, and update operations for concurrent access.
  */
 export default class Cache extends Map<keyof CacheContents, ValueOf<CacheContents>> {
   private static instance: Cache;
 
-  public constructor() {
+  // Mutex for thread-safe cache operations
+  private static mutex = new Mutex();
+
+  private constructor() {
     super();
-    this.set('allowedOrgs', new Set<string>());
+    this.initialize();
   }
 
+  /**
+   * Get the singleton instance of the Cache
+   * Creates a new instance if one doesn't exist
+   *
+   * @returns The singleton Cache instance
+   */
   public static getInstance(): Cache {
-    if (!Cache.instance) {
-      Cache.instance = new Cache();
-    }
-    return Cache.instance;
+    return (Cache.instance ??= new Cache());
+  }
+
+  /**
+   * Thread-safe atomic update operation
+   * Allows safe read-modify-write operations with mutex protection
+   */
+  public static async safeUpdate<K extends keyof CacheContents>(
+    key: K,
+    updateFn: (currentValue: CacheContents[K]) => CacheContents[K]
+  ): Promise<CacheContents[K]> {
+    const cache = Cache.getInstance();
+
+    return Cache.mutex.lock(() => {
+      const currentValue = cache.get(key);
+      const newValue = updateFn(currentValue);
+      cache.set(key, newValue);
+      return newValue;
+    });
+  }
+
+  /**
+   * Thread-safe atomic read operation
+   */
+  public static async safeGet<K extends keyof CacheContents>(key: K): Promise<CacheContents[K]> {
+    const cache = Cache.getInstance();
+
+    return Cache.mutex.lock(() => cache.get(key));
   }
 
-  public get(_key: 'allowedOrgs'): Set<string>;
-  public get(key: keyof CacheContents): ValueOf<CacheContents> {
-    return super.get(key) as ValueOf<CacheContents>;
+  /**
+   * Thread-safe atomic write operation
+   */
+  public static async safeSet<K extends keyof CacheContents>(key: K, value: CacheContents[K]): Promise<void> {
+    const cache = Cache.getInstance();
+
+    return Cache.mutex.lock(() => {
+      cache.set(key, value);
+    });
+  }
+
+  public get<K extends keyof CacheContents>(key: K): CacheContents[K] {
+    return super.get(key) as CacheContents[K];
+  }
+
+  private initialize(): void {
+    this.set('allowedOrgs', new Set<string>());
+    this.set('tools', []);
   }
 }