diff --git a/src/telemetry/telemetry.ts b/src/telemetry/telemetry.ts index c820ec61..339ba419 100644 --- a/src/telemetry/telemetry.ts +++ b/src/telemetry/telemetry.ts @@ -35,14 +35,21 @@ export class Telemetry { userConfig: UserConfig, deviceId: DeviceId, { - commonProperties = { ...MACHINE_METADATA }, + commonProperties = {}, eventCache = EventCache.getInstance(), }: { + commonProperties?: Partial; eventCache?: EventCache; - commonProperties?: CommonProperties; } = {} ): Telemetry { - const instance = new Telemetry(session, userConfig, commonProperties, { eventCache, deviceId }); + const mergedProperties = { + ...MACHINE_METADATA, + ...commonProperties, + }; + const instance = new Telemetry(session, userConfig, mergedProperties, { + eventCache, + deviceId, + }); void instance.setup(); return instance; diff --git a/src/telemetry/types.ts b/src/telemetry/types.ts index f919ab88..f0392344 100644 --- a/src/telemetry/types.ts +++ b/src/telemetry/types.ts @@ -53,11 +53,34 @@ export type ServerEvent = TelemetryEvent; * Interface for static properties, they can be fetched once and reused. */ export type CommonStaticProperties = { + /** + * The version of the MCP server (as read from package.json). + */ mcp_server_version: string; + + /** + * The name of the MCP server (as read from package.json). + */ mcp_server_name: string; + + /** + * The platform/OS the MCP server is running on. + */ platform: string; + + /** + * The architecture of the OS the server is running on. + */ arch: string; + + /** + * Same as platform. + */ os_type: string; + + /** + * The version of the OS the server is running on. + */ os_version?: string; }; @@ -65,12 +88,50 @@ export type CommonStaticProperties = { * Common properties for all events that might change. */ export type CommonProperties = { + /** + * The device id - will be populated with the machine id when it resolves. + */ device_id?: string; + + /** + * A boolean indicating whether the server is running in a container environment. + */ is_container_env?: boolean; + + /** + * The version of the MCP client as reported by the client on session establishment. + */ mcp_client_version?: string; + + /** + * The name of the MCP client as reported by the client on session establishment. + */ mcp_client_name?: string; + + /** + * The transport protocol used by the MCP server. + */ transport?: "stdio" | "http"; + + /** + * A boolean indicating whether Atlas credentials are configured. + */ config_atlas_auth?: TelemetryBoolSet; + + /** + * A boolean indicating whether a connection string is configured. + */ config_connection_string?: TelemetryBoolSet; + + /** + * The randomly generated session id. + */ session_id?: string; + + /** + * The way the MCP server is hosted - e.g. standalone for a server running independently or + * "vscode" if embedded in the VSCode extension. This field should be populated by the hosting + * application to differentiate events coming from an MCP server it's hosting. + */ + hosting_mode?: string; } & CommonStaticProperties; diff --git a/src/transports/base.ts b/src/transports/base.ts index f870514b..9b6abd1e 100644 --- a/src/transports/base.ts +++ b/src/transports/base.ts @@ -13,12 +13,14 @@ import { type ConnectionErrorHandler, connectionErrorHandler as defaultConnectionErrorHandler, } from "../common/connectionErrorHandler.js"; +import type { CommonProperties } from "../telemetry/types.js"; export type TransportRunnerConfig = { userConfig: UserConfig; createConnectionManager?: ConnectionManagerFactoryFn; connectionErrorHandler?: ConnectionErrorHandler; additionalLoggers?: LoggerBase[]; + telemetryProperties?: Partial; }; export abstract class TransportRunnerBase { @@ -27,16 +29,19 @@ export abstract class TransportRunnerBase { protected readonly userConfig: UserConfig; private readonly createConnectionManager: ConnectionManagerFactoryFn; private readonly connectionErrorHandler: ConnectionErrorHandler; + private readonly telemetryProperties: Partial; protected constructor({ userConfig, createConnectionManager = createMCPConnectionManager, connectionErrorHandler = defaultConnectionErrorHandler, additionalLoggers = [], + telemetryProperties = {}, }: TransportRunnerConfig) { this.userConfig = userConfig; this.createConnectionManager = createConnectionManager; this.connectionErrorHandler = connectionErrorHandler; + this.telemetryProperties = telemetryProperties; const loggers: LoggerBase[] = [...additionalLoggers]; if (this.userConfig.loggers.includes("stderr")) { loggers.push(new ConsoleLogger()); @@ -79,7 +84,9 @@ export abstract class TransportRunnerBase { connectionManager, }); - const telemetry = Telemetry.create(session, this.userConfig, this.deviceId); + const telemetry = Telemetry.create(session, this.userConfig, this.deviceId, { + commonProperties: this.telemetryProperties, + }); const result = new Server({ mcpServer, diff --git a/tests/integration/transports/streamableHttp.test.ts b/tests/integration/transports/streamableHttp.test.ts index 6a7b17bf..065f6362 100644 --- a/tests/integration/transports/streamableHttp.test.ts +++ b/tests/integration/transports/streamableHttp.test.ts @@ -1,7 +1,7 @@ import { StreamableHttpRunner } from "../../../src/transports/streamableHttp.js"; import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; -import { describe, expect, it, beforeAll, afterAll, beforeEach } from "vitest"; +import { describe, expect, it, beforeAll, afterAll, beforeEach, afterEach } from "vitest"; import { config } from "../../../src/common/config.js"; import type { LoggerType, LogLevel, LogPayload } from "../../../src/common/logger.js"; import { LoggerBase, LogId } from "../../../src/common/logger.js"; @@ -158,4 +158,26 @@ describe("StreamableHttpRunner", () => { expect(serverStartedMessage?.level).toBe("info"); }); }); + + describe("with telemetry properties", () => { + afterEach(async () => { + await runner.close(); + config.telemetry = oldTelemetry; + config.loggers = oldLoggers; + config.httpHeaders = {}; + }); + + it("merges them with the base properties", async () => { + config.telemetry = "enabled"; + runner = new StreamableHttpRunner({ + userConfig: config, + telemetryProperties: { hosting_mode: "vscode-extension" }, + }); + await runner.start(); + + const server = await runner["setupServer"](); + const properties = server["telemetry"].getCommonProperties(); + expect(properties.hosting_mode).toBe("vscode-extension"); + }); + }); }); diff --git a/tests/unit/telemetry.test.ts b/tests/unit/telemetry.test.ts index 8ce1de2d..51341d56 100644 --- a/tests/unit/telemetry.test.ts +++ b/tests/unit/telemetry.test.ts @@ -1,13 +1,14 @@ import { ApiClient } from "../../src/common/atlas/apiClient.js"; import type { Session } from "../../src/common/session.js"; import { Telemetry } from "../../src/telemetry/telemetry.js"; -import type { BaseEvent, TelemetryResult } from "../../src/telemetry/types.js"; +import type { BaseEvent, CommonProperties, TelemetryEvent, TelemetryResult } from "../../src/telemetry/types.js"; import { EventCache } from "../../src/telemetry/eventCache.js"; import { config } from "../../src/common/config.js"; import { afterEach, beforeEach, describe, it, vi, expect } from "vitest"; import { NullLogger } from "../../src/common/logger.js"; import type { MockedFunction } from "vitest"; import type { DeviceId } from "../../src/helpers/deviceId.js"; +import { expectDefined } from "../integration/helpers.js"; // Mock the ApiClient to avoid real API calls vi.mock("../../src/common/atlas/apiClient.js"); @@ -29,6 +30,7 @@ describe("Telemetry", () => { }; let session: Session; let telemetry: Telemetry; + let mockDeviceId: DeviceId; // Helper function to create properly typed test events function createTestEvent(options?: { @@ -115,7 +117,7 @@ describe("Telemetry", () => { mockEventCache.appendEvents = vi.fn().mockResolvedValue(undefined); MockEventCache.getInstance = vi.fn().mockReturnValue(mockEventCache as unknown as EventCache); - const mockDeviceId = { + mockDeviceId = { get: vi.fn().mockResolvedValue("test-device-id"), } as unknown as DeviceId; @@ -137,183 +139,200 @@ describe("Telemetry", () => { config.telemetry = "enabled"; }); - describe("sending events", () => { - describe("when telemetry is enabled", () => { - it("should send events successfully", async () => { - const testEvent = createTestEvent(); + describe("when telemetry is enabled", () => { + it("should send events successfully", async () => { + const testEvent = createTestEvent(); - await telemetry.setupPromise; + await telemetry.setupPromise; - await telemetry.emitEvents([testEvent]); + await telemetry.emitEvents([testEvent]); - verifyMockCalls({ - sendEventsCalls: 1, - clearEventsCalls: 1, - sendEventsCalledWith: [testEvent], - }); + verifyMockCalls({ + sendEventsCalls: 1, + clearEventsCalls: 1, + sendEventsCalledWith: [testEvent], }); + }); - it("should cache events when sending fails", async () => { - mockApiClient.sendEvents.mockRejectedValueOnce(new Error("API error")); + it("should cache events when sending fails", async () => { + mockApiClient.sendEvents.mockRejectedValueOnce(new Error("API error")); - const testEvent = createTestEvent(); + const testEvent = createTestEvent(); - await telemetry.setupPromise; + await telemetry.setupPromise; - await telemetry.emitEvents([testEvent]); + await telemetry.emitEvents([testEvent]); - verifyMockCalls({ - sendEventsCalls: 1, - appendEventsCalls: 1, - appendEventsCalledWith: [testEvent], - }); + verifyMockCalls({ + sendEventsCalls: 1, + appendEventsCalls: 1, + appendEventsCalledWith: [testEvent], }); + }); - it("should include cached events when sending", async () => { - const cachedEvent = createTestEvent({ - command: "cached-command", - component: "cached-component", - }); + it("should include cached events when sending", async () => { + const cachedEvent = createTestEvent({ + command: "cached-command", + component: "cached-component", + }); - const newEvent = createTestEvent({ - command: "new-command", - component: "new-component", - }); + const newEvent = createTestEvent({ + command: "new-command", + component: "new-component", + }); - // Set up mock to return cached events - mockEventCache.getEvents.mockReturnValueOnce([cachedEvent]); + // Set up mock to return cached events + mockEventCache.getEvents.mockReturnValueOnce([cachedEvent]); - await telemetry.setupPromise; + await telemetry.setupPromise; - await telemetry.emitEvents([newEvent]); + await telemetry.emitEvents([newEvent]); - verifyMockCalls({ - sendEventsCalls: 1, - clearEventsCalls: 1, - sendEventsCalledWith: [cachedEvent, newEvent], - }); + verifyMockCalls({ + sendEventsCalls: 1, + clearEventsCalls: 1, + sendEventsCalledWith: [cachedEvent, newEvent], }); + }); - it("should correctly add common properties to events", async () => { - await telemetry.setupPromise; + it("should correctly add common properties to events", async () => { + await telemetry.setupPromise; + + const commonProps = telemetry.getCommonProperties(); - const commonProps = telemetry.getCommonProperties(); + // Use explicit type assertion + const expectedProps: Record = { + mcp_client_version: "1.0.0", + mcp_client_name: "test-agent", + session_id: "test-session-id", + config_atlas_auth: "true", + config_connection_string: expect.any(String) as unknown as string, + device_id: "test-device-id", + }; - // Use explicit type assertion - const expectedProps: Record = { - mcp_client_version: "1.0.0", - mcp_client_name: "test-agent", - session_id: "test-session-id", - config_atlas_auth: "true", - config_connection_string: expect.any(String) as unknown as string, - device_id: "test-device-id", - }; + expect(commonProps).toMatchObject(expectedProps); + }); - expect(commonProps).toMatchObject(expectedProps); + it("should add hostingMode to events if set", async () => { + telemetry = Telemetry.create(session, config, mockDeviceId, { + eventCache: mockEventCache as unknown as EventCache, + commonProperties: { hosting_mode: "vscode-extension" }, }); + await telemetry.setupPromise; + + const commonProps = telemetry.getCommonProperties(); + expect(commonProps.hosting_mode).toBe("vscode-extension"); + + await telemetry.emitEvents([createTestEvent()]); + + const calls = mockApiClient.sendEvents.mock.calls; + expect(calls).toHaveLength(1); + const event = calls[0]?.[0][0]; + expectDefined(event); + expect((event as TelemetryEvent).properties.hosting_mode).toBe("vscode-extension"); + }); - describe("device ID resolution", () => { - beforeEach(() => { - vi.clearAllMocks(); - }); + describe("device ID resolution", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); - afterEach(() => { - vi.clearAllMocks(); - }); + afterEach(() => { + vi.clearAllMocks(); + }); - it("should successfully resolve the device ID", async () => { - const mockDeviceId = { - get: vi.fn().mockResolvedValue("test-device-id"), - } as unknown as DeviceId; + it("should successfully resolve the device ID", async () => { + const mockDeviceId = { + get: vi.fn().mockResolvedValue("test-device-id"), + } as unknown as DeviceId; - telemetry = Telemetry.create(session, config, mockDeviceId); + telemetry = Telemetry.create(session, config, mockDeviceId); - expect(telemetry["isBufferingEvents"]).toBe(true); - expect(telemetry.getCommonProperties().device_id).toBe(undefined); + expect(telemetry["isBufferingEvents"]).toBe(true); + expect(telemetry.getCommonProperties().device_id).toBe(undefined); - await telemetry.setupPromise; + await telemetry.setupPromise; - expect(telemetry["isBufferingEvents"]).toBe(false); - expect(telemetry.getCommonProperties().device_id).toBe("test-device-id"); - }); + expect(telemetry["isBufferingEvents"]).toBe(false); + expect(telemetry.getCommonProperties().device_id).toBe("test-device-id"); + }); - it("should handle device ID resolution failure gracefully", async () => { - const mockDeviceId = { - get: vi.fn().mockResolvedValue("unknown"), - } as unknown as DeviceId; + it("should handle device ID resolution failure gracefully", async () => { + const mockDeviceId = { + get: vi.fn().mockResolvedValue("unknown"), + } as unknown as DeviceId; - telemetry = Telemetry.create(session, config, mockDeviceId); + telemetry = Telemetry.create(session, config, mockDeviceId); - expect(telemetry["isBufferingEvents"]).toBe(true); - expect(telemetry.getCommonProperties().device_id).toBe(undefined); + expect(telemetry["isBufferingEvents"]).toBe(true); + expect(telemetry.getCommonProperties().device_id).toBe(undefined); - await telemetry.setupPromise; + await telemetry.setupPromise; - expect(telemetry["isBufferingEvents"]).toBe(false); - // Should use "unknown" as fallback when device ID resolution fails - expect(telemetry.getCommonProperties().device_id).toBe("unknown"); - }); + expect(telemetry["isBufferingEvents"]).toBe(false); + // Should use "unknown" as fallback when device ID resolution fails + expect(telemetry.getCommonProperties().device_id).toBe("unknown"); + }); - it("should handle device ID timeout gracefully", async () => { - const mockDeviceId = { - get: vi.fn().mockResolvedValue("unknown"), - } as unknown as DeviceId; + it("should handle device ID timeout gracefully", async () => { + const mockDeviceId = { + get: vi.fn().mockResolvedValue("unknown"), + } as unknown as DeviceId; - telemetry = Telemetry.create(session, config, mockDeviceId); + telemetry = Telemetry.create(session, config, mockDeviceId); - expect(telemetry["isBufferingEvents"]).toBe(true); - expect(telemetry.getCommonProperties().device_id).toBe(undefined); + expect(telemetry["isBufferingEvents"]).toBe(true); + expect(telemetry.getCommonProperties().device_id).toBe(undefined); - await telemetry.setupPromise; + await telemetry.setupPromise; - expect(telemetry["isBufferingEvents"]).toBe(false); - // Should use "unknown" as fallback when device ID times out - expect(telemetry.getCommonProperties().device_id).toBe("unknown"); - }); + expect(telemetry["isBufferingEvents"]).toBe(false); + // Should use "unknown" as fallback when device ID times out + expect(telemetry.getCommonProperties().device_id).toBe("unknown"); }); }); + }); - describe("when telemetry is disabled", () => { - beforeEach(() => { - config.telemetry = "disabled"; - }); + describe("when telemetry is disabled", () => { + beforeEach(() => { + config.telemetry = "disabled"; + }); - afterEach(() => { - config.telemetry = "enabled"; - }); + afterEach(() => { + config.telemetry = "enabled"; + }); - it("should not send events", async () => { - const testEvent = createTestEvent(); + it("should not send events", async () => { + const testEvent = createTestEvent(); - await telemetry.emitEvents([testEvent]); + await telemetry.emitEvents([testEvent]); - verifyMockCalls(); - }); + verifyMockCalls(); }); + }); - describe("when DO_NOT_TRACK environment variable is set", () => { - let originalEnv: string | undefined; + describe("when DO_NOT_TRACK environment variable is set", () => { + let originalEnv: string | undefined; - beforeEach(() => { - originalEnv = process.env.DO_NOT_TRACK; - process.env.DO_NOT_TRACK = "1"; - }); + beforeEach(() => { + originalEnv = process.env.DO_NOT_TRACK; + process.env.DO_NOT_TRACK = "1"; + }); - afterEach(() => { - if (originalEnv) { - process.env.DO_NOT_TRACK = originalEnv; - } else { - delete process.env.DO_NOT_TRACK; - } - }); + afterEach(() => { + if (originalEnv) { + process.env.DO_NOT_TRACK = originalEnv; + } else { + delete process.env.DO_NOT_TRACK; + } + }); - it("should not send events", async () => { - const testEvent = createTestEvent(); + it("should not send events", async () => { + const testEvent = createTestEvent(); - await telemetry.emitEvents([testEvent]); + await telemetry.emitEvents([testEvent]); - verifyMockCalls(); - }); + verifyMockCalls(); }); }); });