instrumentation for @modelcontextprotocol/sdk

packages/ai-semantic-conventions/src/SemanticAttributes.ts

@@ -56,6 +56,13 @@ export const SpanAttributes = {
   TRACELOOP_ASSOCIATION_PROPERTIES: "traceloop.association.properties",
   TRACELOOP_ENTITY_INPUT: "traceloop.entity.input",
   TRACELOOP_ENTITY_OUTPUT: "traceloop.entity.output",
+
+  // MCP
+  MCP_METHOD_NAME: "mcp.method.name",
+  MCP_REQUEST_ARGUMENT: "mcp.request.argument",
+  MCP_REQUEST_ID: "mcp.request.id",
+  MCP_SESSION_INIT_OPTIONS: "mcp.session.init_options",
+  MCP_RESPONSE_VALUE: "mcp.response.value",
 };
 
 export const Events = {

packages/instrumentation-mcp/package.json

@@ -0,0 +1,60 @@
+{
+    "name": "@traceloop/instrumentation-mcp",
+    "version": "0.13.0",
+    "description": "OpenTelemetry instrumentation for MCP",
+    "main": "dist/index.js",
+    "module": "dist/index.mjs",
+    "types": "dist/index.d.ts",
+    "repository": "traceloop/openllmetry-js",
+    "scripts": {
+        "build": "rollup -c",
+        "lint": "eslint .",
+        "lint:fix": "eslint . --fix",
+        "test": "ts-mocha -p tsconfig.json 'test/**/*.test.ts' --timeout 20000"
+    },
+    "keywords": [
+        "opentelemetry",
+        "nodejs",
+        "tracing",
+        "attributes",
+        "mcp"
+    ],
+    "author": "Traceloop",
+    "license": "Apache-2.0",
+    "engines": {
+        "node": ">=14"
+    },
+    "files": [
+        "dist/**/*.js",
+        "dist/**/*.mjs",
+        "dist/**/*.js.map",
+        "dist/**/*.d.ts",
+        "doc",
+        "LICENSE",
+        "README.md",
+        "package.json"
+    ],
+    "publishConfig": {
+        "access": "public"
+    },
+    "dependencies": {
+        "@modelcontextprotocol/sdk": "^1.12.0",
+        "@opentelemetry/core": "^1.29.0",
+        "@opentelemetry/instrumentation": "^0.56.0",
+        "@opentelemetry/semantic-conventions": "^1.28.0",
+        "@traceloop/ai-semantic-conventions": "^0.13.0",
+        "tslib": "^2.8.1"
+    },
+    "devDependencies": {
+        "@modelcontextprotocol/sdk": "^1.12.0",
+        "@pollyjs/adapter-node-http": "^6.0.6",
+        "@pollyjs/core": "^6.0.6",
+        "@pollyjs/persister-fs": "^6.0.6",
+        "@types/express": "^5.0.2",
+        "@types/mocha": "^10.0.10",
+        "express": "^5.1.0",
+        "mocha": "^10.8.2",
+        "ts-mocha": "^10.0.0"
+    },
+    "homepage": "https://github.com/traceloop/openllmetry-js/tree/main/packages/instrumentation-mcp"
+}

packages/instrumentation-mcp/rollup.config.js

@@ -0,0 +1,36 @@
+const dts = require("rollup-plugin-dts");
+const typescript = require("@rollup/plugin-typescript");
+const json = require("@rollup/plugin-json");
+
+const name = require("./package.json").main.replace(/\.js$/, "");
+
+const bundle = (config) => ({
+  ...config,
+  input: "src/index.ts",
+  external: (id) => !/^[./]/.test(id),
+});
+
+exports.default = [
+  bundle({
+    plugins: [typescript.default(), json.default()],
+    output: [
+      {
+        file: `${name}.js`,
+        format: "cjs",
+        sourcemap: true,
+      },
+      {
+        file: `${name}.mjs`,
+        format: "es",
+        sourcemap: true,
+      },
+    ],
+  }),
+  bundle({
+    plugins: [dts.default()],
+    output: {
+      file: `${name}.d.ts`,
+      format: "es",
+    },
+  }),
+];

packages/instrumentation-mcp/src/index.ts

@@ -0,0 +1 @@
+export * from "./instrumentation";

packages/instrumentation-mcp/src/instrumentation.ts

@@ -0,0 +1,200 @@
+import {
+  InstrumentationBase,
+  safeExecuteInTheMiddle,
+} from "@opentelemetry/instrumentation";
+import {
+  context,
+  Span,
+  SpanKind,
+  SpanStatusCode,
+  trace,
+} from "@opentelemetry/api";
+import type * as sse from "@modelcontextprotocol/sdk/client/sse";
+import type * as stdio from "@modelcontextprotocol/sdk/client/stdio.js";
+
+import { McpInstrumentationConfig } from "./types";
+import { version } from "../package.json";
+
+export class McpInstrumentation extends InstrumentationBase {
+  declare protected _config: McpInstrumentationConfig;
+
+  constructor(config: McpInstrumentationConfig = {}) {
+    super("@modelcontextprotocol/sdk", version, config);
+  }
+
+  public override setConfig(config: McpInstrumentationConfig = {}) {
+    super.setConfig(config);
+  }
+
+  // auto-instrumentation doesn't work for the modules we're instrumenting
+  protected init() {}
+
+  public manuallyInstrument({
+    sseModule,
+    stdioModule,
+  }: {
+    sseModule?: typeof sse;
+    stdioModule?: typeof stdio;
+  }) {
+    if (sseModule) {
+      this.wrapSSEClient(sseModule);
+    }
+    if (stdioModule) {
+      this.wrapStdioClient(stdioModule);
+    }
+  }
+
+  private wrapSSEClient(module: typeof sse) {
+    this._diag.debug(`Patching @modelcontextprotocol/sdk/client/sse`);
+
+    // Store reference to the wrapper method
+    const onMessageWrapper = this.onMessageWrapper();
+
+    Object.defineProperty(module.SSEClientTransport.prototype, "onmessage", {
+      get(this: any) {
+        return this._wrappedOnMessage || this._originalOnMessage;
+      },
+      set(this: any, newHandler: Function) {
+        // Store the original handler
+        this._originalOnMessage = newHandler;
+
+        // Wrap the new handler if it's a function
+        if (typeof newHandler === "function") {
+          this._wrappedOnMessage = onMessageWrapper(newHandler);
+        } else {
+          this._wrappedOnMessage = newHandler;
+        }
+      },
+      configurable: true,
+      enumerable: true,
+    });
+
+    this._wrap(
+      module.SSEClientTransport.prototype,
+      "send",
+      this.onSendWrapper(),
+    );
+  }
+
+  private wrapStdioClient(module: typeof stdio) {
+    this._diag.debug(`Patching @modelcontextprotocol/sdk/client/stdio`);
+
+    // Store reference to the wrapper method
+    const onMessageWrapper = this.onMessageWrapper();
+
+    Object.defineProperty(module.StdioClientTransport.prototype, "onmessage", {
+      get(this: any) {
+        return this._wrappedOnMessage || this._originalOnMessage;
+      },
+      set(this: any, newHandler: Function) {
+        // Store the original handler
+        this._originalOnMessage = newHandler;
+
+        // Wrap the new handler if it's a function
+        if (typeof newHandler === "function") {
+          this._wrappedOnMessage = onMessageWrapper(newHandler);
+        } else {
+          this._wrappedOnMessage = newHandler;
+        }
+      },
+      configurable: true,
+      enumerable: true,
+    });
+
+    this._wrap(
+      module.StdioClientTransport.prototype,
+      "send",
+      this.onSendWrapper(),
+    );
+  }
+
+  private onMessageWrapper() {
+    const plugin = this;
+    return (original: Function) => {
+      return function method(this: any, ...args: any) {
+        const span = plugin._startSpan(args[0], "response");
+        const execContext = trace.setSpan(context.active(), span);
+        const execPromise = safeExecuteInTheMiddle(
+          (): Promise<void> => {
+            return context.with(execContext, () => {
+              return original.apply(this, args);
+            });
+          },
+          (e) => {
+            if (e) {
+              plugin._diag.error(`Error in mcp instrumentation`, e);
+            }
+          },
+        );
+        const wrappedPromise = plugin._wrapPromise(span, execPromise);
+        return context.bind(execContext, wrappedPromise);
+      };
+    };
+  }
+
+  private onSendWrapper() {
+    const plugin = this;
+    return (original: Function) => {
+      return function method(this: any, ...args: any) {
+        const span = plugin._startSpan(args[0], "request");
+        const execContext = trace.setSpan(context.active(), span);
+        const execPromise = safeExecuteInTheMiddle(
+          (): Promise<void> => {
+            return context.with(execContext, () => {
+              return original.apply(this, args);
+            });
+          },
+          (e) => {
+            if (e) {
+              plugin._diag.error(`Error in mcp instrumentation`, e);
+            }
+          },
+        );
+        const wrappedPromise = plugin._wrapPromise(span, execPromise);
+        return context.bind(execContext, wrappedPromise);
+      };
+    };
+  }
+
+  private _startSpan(params: any, type: "request" | "response") {
+    let attributes;
+    if (type === "response") {
+      attributes = {
+        "mcp.response.id": params.params?.id,
+        "mcp.response.value": JSON.stringify(
+          params.result || params.params?.content || params.params?.data,
+        ),
+        "mcp.method.name": params.params?.method,
+      };
+    } else {
+      attributes = {
+        "mcp.request.id": params.id || params.params?.requestId,
+        "mcp.request.method": params.method,
+        "mcp.request.params": JSON.stringify(params.params),
+      };
+    }
+    return this.tracer.startSpan("mcp.client.response", {
+      kind: SpanKind.CLIENT,
+      attributes,
+    });
+  }
+
+  private async _wrapPromise<T>(span: Span, promise: Promise<T>): Promise<T> {
+    try {
+      const result = await promise;
+      span.end();
+      return result;
+    } catch (error) {
+      span.setStatus({
+        code: SpanStatusCode.ERROR,
+        message: error.message,
+      });
+      span.setAttributes({
+        "mcp.error.message": error.message,
+      });
+      span.recordException(error);
+      span.end();
+      throw error;
+    }
+  }
+}

packages/instrumentation-mcp/src/types.ts

@@ -0,0 +1,8 @@
+import { InstrumentationConfig } from "@opentelemetry/instrumentation";
+
+export interface McpInstrumentationConfig extends InstrumentationConfig {
+  /**
+   * A custom logger to log any exceptions that happen during span creation.
+   */
+  exceptionLogger?: (e: Error) => void;
+}