refactor: Move all lookupDeploymentId calls to resolveTelemetryMetadata

Author: Luke-Sanderson

src/tools/atlasLocal/atlasLocalTool.ts

@@ -3,10 +3,11 @@ import type { TelemetryToolMetadata, ToolArgs, ToolCategory } from "../tool.js";
 import { ToolBase } from "../tool.js";
 import type { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
 import type { Client } from "@mongodb-js-preview/atlas-local";
+import { LogId } from "../../common/logger.js";
+import { z } from "zod";
 
 export abstract class AtlasLocalToolBase extends ToolBase {
     public category: ToolCategory = "atlas-local";
-    protected deploymentId?: string;
 
     protected verifyAllowed(): boolean {
         return this.session.atlasLocalClient !== undefined && super.verifyAllowed();
@@ -38,16 +39,10 @@ please log a ticket here: https://github.com/mongodb-js/mongodb-mcp-server/issue
         return this.executeWithAtlasLocalClient(client, ...args);
     }
 
-    protected async lookupDeploymentId(client: Client, containerId: string): Promise<void> {
-        // Don't run if telemetry is disabled
-        if (this.telemetry.isTelemetryEnabled()) {
-            return;
-        }
-
-        // Lookup the deployment id and save it to the deploymentId property.
-        // This property will be added to the telemetry metadata when resolveTelemetryMetadata is called.
+    protected async lookupDeploymentId(client: Client, containerId: string): Promise<string> {
+        // Lookup and return the deployment id for telemetry metadata.
         const deploymentId = await client.getDeploymentId(containerId);
-        this.deploymentId = deploymentId;
+        return deploymentId;
     }
 
     protected abstract executeWithAtlasLocalClient(
@@ -79,9 +74,37 @@ please log a ticket here: https://github.com/mongodb-js/mongodb-mcp-server/issue
         return super.handleError(error, args);
     }
 
-    protected resolveTelemetryMetadata(): TelemetryToolMetadata {
-        return {
-            atlasLocaldeploymentId: this.deploymentId,
-        };
+    protected async resolveTelemetryMetadata(
+        ...args: Parameters<ToolCallback<typeof this.argsShape>>
+    ): Promise<TelemetryToolMetadata> {
+        const toolMetadata: TelemetryToolMetadata = {};
+
+        const client = this.session.atlasLocalClient;
+        if (!args.length || !client) {
+            return toolMetadata;
+        }
+
+        // Create a typed parser for the exact shape we expect
+        const argsShape = z.object(this.argsShape);
+        const parsedResult = argsShape.safeParse(args[0]);
+
+        if (!parsedResult.success) {
+            this.session.logger.debug({
+                id: LogId.telemetryMetadataError,
+                context: "tool",
+                message: `Error parsing tool arguments: ${parsedResult.error.message}`,
+            });
+            return toolMetadata;
+        }
+
+        const data = parsedResult.data;
+
+        // Extract deploymentName using type guard and lookup deployment ID
+        if ("deploymentName" in data && typeof data.deploymentName === "string" && data.deploymentName.trim() !== "") {
+            const deploymentId = await this.lookupDeploymentId(client, data.deploymentName);
+            toolMetadata.atlasLocaldeploymentId = deploymentId;
+        }
+
+        return toolMetadata;
     }
 }

src/tools/atlasLocal/connect/connectDeployment.ts

@@ -22,9 +22,6 @@ export class ConnectDeploymentTool extends AtlasLocalToolBase {
         // Connect to the deployment
         await this.session.connectToMongoDB({ connectionString });
 
-        // Lookup the deployment id and add it to the telemetry metadata
-        await this.lookupDeploymentId(client, deploymentName);
-
         return {
             content: [
                 {

src/tools/atlasLocal/create/createDeployment.ts

@@ -1,8 +1,9 @@
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { AtlasLocalToolBase } from "../atlasLocalTool.js";
-import type { OperationType, ToolArgs } from "../../tool.js";
+import type { OperationType, ToolArgs, TelemetryToolMetadata } from "../../tool.js";
 import type { Client, CreateDeploymentOptions, CreationSourceType } from "@mongodb-js-preview/atlas-local";
 import { CommonArgs } from "../../args.js";
+import type { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
 
 export class CreateDeploymentTool extends AtlasLocalToolBase {
     public name = "atlas-local-create-deployment";
@@ -12,6 +13,8 @@ export class CreateDeploymentTool extends AtlasLocalToolBase {
         deploymentName: CommonArgs.string().describe("Name of the deployment to create").optional(),
     };
 
+    private createdDeploymentId?: string;
+
     protected async executeWithAtlasLocalClient(
         client: Client,
         { deploymentName }: ToolArgs<typeof this.argsShape>
@@ -27,8 +30,8 @@ export class CreateDeploymentTool extends AtlasLocalToolBase {
         // Create the deployment
         const deployment = await client.createDeployment(deploymentOptions);
 
-        // Lookup the deployment id and add it to the telemetry metadata
-        await this.lookupDeploymentId(client, deployment.containerId);
+        // Capture deployment ID for telemetry
+        this.createdDeploymentId = await this.lookupDeploymentId(client, deployment.containerId);
 
         return {
             content: [
@@ -39,4 +42,15 @@ export class CreateDeploymentTool extends AtlasLocalToolBase {
             ],
         };
     }
+
+    // Create tool needs to override resolveTelemetryMetadata because it doesn't
+    // have the deployment name in the arguments, but rather in the response.
+    protected resolveTelemetryMetadata(
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        ...args: Parameters<ToolCallback<typeof this.argsShape>>
+    ): Promise<TelemetryToolMetadata> {
+        return Promise.resolve({
+            atlasLocaldeploymentId: this.createdDeploymentId,
+        });
+    }
 }

src/tools/atlasLocal/delete/deleteDeployment.ts

@@ -16,9 +16,6 @@ export class DeleteDeploymentTool extends AtlasLocalToolBase {
         client: Client,
         { deploymentName }: ToolArgs<typeof this.argsShape>
     ): Promise<CallToolResult> {
-        // Lookup the deployment id and add it to the telemetry metadata
-        await this.lookupDeploymentId(client, deploymentName);
-
         // Delete the deployment
         await client.deleteDeployment(deploymentName);
 

src/tools/tool.ts

@@ -158,7 +158,7 @@ export abstract class ToolBase {
                 });
 
                 const result = await this.execute(...args);
-                this.emitToolEvent(startTime, result, ...args);
+                await this.emitToolEvent(startTime, result, ...args);
 
                 this.session.logger.debug({
                     id: LogId.toolExecute,
@@ -174,7 +174,7 @@ export abstract class ToolBase {
                     message: `Error executing ${this.name}: ${error as string}`,
                 });
                 const toolResult = await this.handleError(error, args[0] as ToolArgs<typeof this.argsShape>);
-                this.emitToolEvent(startTime, toolResult, ...args);
+                await this.emitToolEvent(startTime, toolResult, ...args);
                 return toolResult;
             }
         };
@@ -275,24 +275,24 @@ export abstract class ToolBase {
 
     protected abstract resolveTelemetryMetadata(
         ...args: Parameters<ToolCallback<typeof this.argsShape>>
-    ): TelemetryToolMetadata;
+    ): TelemetryToolMetadata | Promise<TelemetryToolMetadata>;
 
     /**
      * Creates and emits a tool telemetry event
      * @param startTime - Start time in milliseconds
      * @param result - Whether the command succeeded or failed
      * @param args - The arguments passed to the tool
      */
-    private emitToolEvent(
+    private async emitToolEvent(
         startTime: number,
         result: CallToolResult,
         ...args: Parameters<ToolCallback<typeof this.argsShape>>
-    ): void {
+    ): Promise<void> {
         if (!this.telemetry.isTelemetryEnabled()) {
             return;
         }
         const duration = Date.now() - startTime;
-        const metadata = this.resolveTelemetryMetadata(...args);
+        const metadata = await this.resolveTelemetryMetadata(...args);
         const event: ToolEvent = {
             timestamp: new Date().toISOString(),
             source: "mdbmcp",