fix: report registered MCP tool count (#504)

Co-authored-by: CharlieHelps <charlie@charlielabs.ai>

Author: charliecreates[bot]

.changeset/bright-tools-count.md

@@ -0,0 +1,5 @@
+---
+"hevy-mcp": patch
+---
+
+Correct the MCP registration span tool count to reflect successful tool registrations.

src/index.test.ts

@@ -1,3 +1,4 @@
+import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import * as stdioModule from "@modelcontextprotocol/sdk/server/stdio.js";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import createServer, {
@@ -22,6 +23,9 @@ const testDoubles = vi.hoisted(() => ({
 		end: vi.fn(),
 	},
 	connect: vi.fn().mockResolvedValue(undefined),
+	tool: vi.fn(),
+	registerTool: vi.fn(),
+	directRegisterToolCalls: 0,
 	sentry: {
 		init: vi.fn(() => ({})),
 		setUser: vi.fn(),
@@ -68,6 +72,13 @@ vi.mock("./utils/metrics.js", () => ({
 	serverStartups: { add: vi.fn() },
 }));
 
+vi.mock("./tools/user.js", () => ({
+	registerUserTools: vi.fn((server: McpServer) => {
+		testDoubles.directRegisterToolCalls += 1;
+		server.registerTool("get-user-info", {}, vi.fn());
+	}),
+}));
+
 vi.mock("@opentelemetry/api", () => ({
 	SpanStatusCode: { OK: 1, ERROR: 2 },
 	trace: { getTracer: vi.fn(() => ({ startActiveSpan: vi.fn() })) },
@@ -82,8 +93,8 @@ vi.mock("@opentelemetry/api", () => ({
 vi.mock("@modelcontextprotocol/sdk/server/mcp.js", () => {
 	class MockMcpServer {
 		connect = testDoubles.connect;
-		registerTool = vi.fn();
-		tool = vi.fn();
+		tool = testDoubles.tool;
+		registerTool = testDoubles.registerTool;
 		registerResource = vi.fn();
 	}
 
@@ -111,6 +122,12 @@ describe("Server entry", () => {
 		process.env = { ...originalEnv };
 		process.argv = [...originalArgv];
 		vi.clearAllMocks();
+		testDoubles.directRegisterToolCalls = 0;
+		testDoubles.tool.mockImplementation(
+			function (this: { registerTool: () => void }) {
+				this.registerTool();
+			},
+		);
 		const anyStdioModule = stdioModule as { __transports?: unknown[] };
 		if (anyStdioModule.__transports) {
 			anyStdioModule.__transports.length = 0;
@@ -142,6 +159,17 @@ describe("Server entry", () => {
 		);
 	});
 
+	it("reports the number of tool registration calls on the registration span", () => {
+		createServer({ config: { apiKey: "test-key" } });
+
+		const registrationCount = testDoubles.registerTool.mock.calls.length;
+		expect(registrationCount).toBeGreaterThan(0);
+		expect(testDoubles.span.setAttribute).toHaveBeenCalledWith(
+			"mcp.tools.count",
+			registrationCount,
+		);
+	});
+
 	it("exports createServer as both default and named exports", () => {
 		expect(namedCreateServer).toBe(createServer);
 		const server = namedCreateServer({ config: { apiKey: "named-key" } });

src/index.ts

@@ -83,6 +83,42 @@ const serverConfigSchema = z.object({
 export const configSchema = serverConfigSchema;
 type ServerConfig = z.infer<typeof serverConfigSchema>;
 
+function createToolCountingServer(server: McpServer) {
+	let count = 0;
+
+	const countingServer = new Proxy(server, {
+		get(target, property, receiver) {
+			if (property === "tool") {
+				return (...args: Parameters<McpServer["tool"]>) => {
+					const registeredTool = target.tool(...args);
+					count += 1;
+					return registeredTool;
+				};
+			}
+
+			if (property === "registerTool") {
+				const registerTool: McpServer["registerTool"] = (
+					name,
+					config,
+					callback,
+				) => {
+					const registeredTool = target.registerTool(name, config, callback);
+					count += 1;
+					return registeredTool;
+				};
+				return registerTool;
+			}
+
+			return Reflect.get(target, property, receiver);
+		},
+	});
+
+	return {
+		server: countingServer,
+		getCount: () => count,
+	};
+}
+
 function buildServer(apiKey: string) {
 	const userId = fingerprintApiKey(apiKey);
 
@@ -119,26 +155,20 @@ function buildServer(apiKey: string) {
 				);
 				console.error("Hevy client initialized with API key");
 
-				tracer.startActiveSpan(
-					"mcp.tools.register",
-					{
-						attributes: {
-							"mcp.tools.count": 6,
-						},
-					},
-					(toolsSpan) => {
-						try {
-							registerWorkoutTools(server, hevyClient);
-							registerRoutineTools(server, hevyClient);
-							registerTemplateTools(server, hevyClient);
-							registerFolderTools(server, hevyClient);
-							registerBodyMeasurementTools(server, hevyClient);
-							registerUserTools(server, hevyClient);
-						} finally {
-							toolsSpan.end();
-						}
-					},
-				);
+				tracer.startActiveSpan("mcp.tools.register", (toolsSpan) => {
+					try {
+						const counting = createToolCountingServer(server);
+						registerWorkoutTools(counting.server, hevyClient);
+						registerRoutineTools(counting.server, hevyClient);
+						registerTemplateTools(counting.server, hevyClient);
+						registerFolderTools(counting.server, hevyClient);
+						registerBodyMeasurementTools(counting.server, hevyClient);
+						registerUserTools(counting.server, hevyClient);
+						toolsSpan.setAttribute("mcp.tools.count", counting.getCount());
+					} finally {
+						toolsSpan.end();
+					}
+				});
 
 				tracer.startActiveSpan(
 					"mcp.resources.register",