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,58 @@
+{
+    "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": {
+        "@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",
+        "@modelcontextprotocol/sdk": "^1.12.0"
+    },
+    "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/mocha": "^10.0.10",
+        "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,148 @@
+import { InstrumentationBase, InstrumentationModuleDefinition, 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"
+
+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);
+    }
+
+    protected init(): InstrumentationModuleDefinition[]{
+
+        return []
+    }
+
+    public manuallyInstrument(sseModule: typeof sse, stdioModule: typeof stdio) {
+        this.wrapSSEClient(sseModule);
+        this.wrapStdioClient(stdioModule);
+    }
+
+    private wrapSSEClient(module: typeof sse) {
+        this._diag.debug(`Patching @modelcontextprotocol/sdk/client/sse`);
+
+        // Store reference to the wrapper method
+        const wrapperMethod = this.wrapperMethod();
+
+        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 = wrapperMethod(newHandler);
+                } else {
+                    this._wrappedOnMessage = newHandler;
+                }
+            },
+            configurable: true,
+            enumerable: true
+        });
+
+        // this._wrap(module.SSEClientTransport.prototype, "send", this.wrapperMethod());
+
+    }
+
+    private wrapStdioClient(module: typeof stdio) {
+        this._diag.debug(`Patching @modelcontextprotocol/sdk/client/stdio`);
+
+        // Store reference to the wrapper method
+        const wrapperMethod = this.wrapperMethod();
+
+        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 = wrapperMethod(newHandler);
+                } else {
+                    this._wrappedOnMessage = newHandler;
+                }
+            },
+            configurable: true,
+            enumerable: true
+        });
+
+        // this._wrap(module.SSEClientTransport.prototype, "send", this.wrapperMethod());
+
+    }
+
+
+
+    private wrapperMethod() {
+        const plugin = this;
+        return (original: Function) => {
+          return function method(this: any, ...args: any) {
+            const span = plugin._startSpan({
+              params: args[0],
+            });
+            const execContext = trace.setSpan(context.active(), span);
+            const execPromise = safeExecuteInTheMiddle(
+              () => {
+                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) {
+        const attributes = {
+            "mcp.request.id": params.params.id,
+            "mcp.response.value": JSON.stringify(params.result || params.params?.content || params.params?.data),
+            "mcp.method.name": params.params.method,
+        }
+        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;
+}

packages/instrumentation-mcp/tsconfig.json

@@ -0,0 +1,10 @@
+{
+  "extends": "../../tsconfig.base.json",
+  "compilerOptions": {
+    "outDir": "dist",
+    "rootDir": "."
+  },
+  "files": [],
+  "include": ["src/**/*.ts", "test/**/*.ts", "package.json"],
+  "references": []
+}

packages/traceloop-sdk/package.json

@@ -50,6 +50,7 @@
     "@traceloop/instrumentation-qdrant": "^0.13.0",
     "@traceloop/instrumentation-together": "^0.13.0",
     "@traceloop/instrumentation-vertexai": "^0.13.0",
+    "@traceloop/instrumentation-mcp": "^0.13.0",
     "@types/nunjucks": "^3.2.6",
     "cross-fetch": "^4.0.0",
     "fetch-retry": "^5.0.6",
@@ -82,6 +83,7 @@
     "langchain": "^0.3.7",
     "llamaindex": "^0.8.28",
     "openai": "^4.76.2",
-    "ts-mocha": "^10.0.0"
+    "ts-mocha": "^10.0.0",
+    "@modelcontextprotocol/sdk": "^1.12.0"
   }
 }

packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts

@@ -17,6 +17,7 @@ import type * as VectorStoreModule from "@langchain/core/vectorstores";
 import type * as llamaindex from "llamaindex";
 import type * as chromadb from "chromadb";
 import type * as qdrant from "@qdrant/js-client-rest";
+import type * as sse from "@modelcontextprotocol/sdk/client/sse.js";
 
 /**
  * Options for initializing the Traceloop SDK.
@@ -111,6 +112,9 @@ export interface InitializeOptions {
     llamaIndex?: typeof llamaindex;
     chromadb?: typeof chromadb;
     qdrant?: typeof qdrant;
+    mcp?: {
+      sseModule?: typeof sse;
+    };
   };
 
   /**

packages/traceloop-sdk/src/lib/tracing/index.ts

@@ -25,6 +25,7 @@ import { LangChainInstrumentation } from "@traceloop/instrumentation-langchain";
 import { ChromaDBInstrumentation } from "@traceloop/instrumentation-chromadb";
 import { QdrantInstrumentation } from "@traceloop/instrumentation-qdrant";
 import { TogetherInstrumentation } from "@traceloop/instrumentation-together";
+import { McpInstrumentation } from "@traceloop/instrumentation-mcp";
 import {
   ALL_INSTRUMENTATION_LIBRARIES,
   createSpanProcessor,
@@ -45,6 +46,7 @@ let pineconeInstrumentation: PineconeInstrumentation | undefined;
 let chromadbInstrumentation: ChromaDBInstrumentation | undefined;
 let qdrantInstrumentation: QdrantInstrumentation | undefined;
 let togetherInstrumentation: TogetherInstrumentation | undefined;
+let mcpInstrumentation: McpInstrumentation | undefined;
 
 const instrumentations: Instrumentation[] = [];
 
@@ -209,6 +211,13 @@ export const manuallyInitInstrumentations = (
     instrumentations.push(togetherInstrumentation);
     togetherInstrumentation.manuallyInstrument(instrumentModules.together);
   }
+
+  if (instrumentModules?.mcp) {
+    mcpInstrumentation = new McpInstrumentation({ exceptionLogger });
+    instrumentations.push(mcpInstrumentation);
+    // @ts-ignore
+    mcpInstrumentation.manuallyInstrument(instrumentModules.mcp);
+  }
 };
 
 /**
@@ -250,6 +259,9 @@ export const startTracing = (options: InitializeOptions) => {
     togetherInstrumentation?.setConfig({
       traceContent: false,
     });
+    mcpInstrumentation?.setConfig({
+      traceContent: false,
+    });
   }
 
   const headers =