diff --git a/dev-packages/e2e-tests/test-applications/node-express/package.json b/dev-packages/e2e-tests/test-applications/node-express/package.json index c1a608833bdb..d036d2621fa2 100644 --- a/dev-packages/e2e-tests/test-applications/node-express/package.json +++ b/dev-packages/e2e-tests/test-applications/node-express/package.json @@ -11,6 +11,7 @@ "test:assert": "pnpm test" }, "dependencies": { + "@modelcontextprotocol/sdk": "^1.10.2", "@sentry/core": "latest || *", "@sentry/node": "latest || *", "@trpc/server": "10.45.2", @@ -19,7 +20,7 @@ "@types/node": "^18.19.1", "express": "^4.21.2", "typescript": "~5.0.0", - "zod": "~3.22.4" + "zod": "~3.24.3" }, "devDependencies": { "@playwright/test": "~1.50.0", diff --git a/dev-packages/e2e-tests/test-applications/node-express/src/app.ts b/dev-packages/e2e-tests/test-applications/node-express/src/app.ts index d756c0e08372..6b320e26eb8a 100644 --- a/dev-packages/e2e-tests/test-applications/node-express/src/app.ts +++ b/dev-packages/e2e-tests/test-applications/node-express/src/app.ts @@ -19,10 +19,13 @@ import { TRPCError, initTRPC } from '@trpc/server'; import * as trpcExpress from '@trpc/server/adapters/express'; import express from 'express'; import { z } from 'zod'; +import { mcpRouter } from './mcp'; const app = express(); const port = 3030; +app.use(mcpRouter); + app.get('/test-success', function (req, res) { res.send({ version: 'v1' }); }); diff --git a/dev-packages/e2e-tests/test-applications/node-express/src/mcp.ts b/dev-packages/e2e-tests/test-applications/node-express/src/mcp.ts new file mode 100644 index 000000000000..7565e08f7c85 --- /dev/null +++ b/dev-packages/e2e-tests/test-applications/node-express/src/mcp.ts @@ -0,0 +1,64 @@ +import express from 'express'; +import { McpServer, ResourceTemplate } from '@modelcontextprotocol/sdk/server/mcp.js'; +import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js'; +import { z } from 'zod'; +import { wrapMcpServerWithSentry } from '@sentry/core'; + +const mcpRouter = express.Router(); + +const server = wrapMcpServerWithSentry( + new McpServer({ + name: 'Echo', + version: '1.0.0', + }), +); + +server.resource('echo', new ResourceTemplate('echo://{message}', { list: undefined }), async (uri, { message }) => ({ + contents: [ + { + uri: uri.href, + text: `Resource echo: ${message}`, + }, + ], +})); + +server.tool('echo', { message: z.string() }, async ({ message }, rest) => { + return { + content: [{ type: 'text', text: `Tool echo: ${message}` }], + }; +}); + +server.prompt('echo', { message: z.string() }, ({ message }, extra) => ({ + messages: [ + { + role: 'user', + content: { + type: 'text', + text: `Please process this message: ${message}`, + }, + }, + ], +})); + +const transports: Record = {}; + +mcpRouter.get('/sse', async (_, res) => { + const transport = new SSEServerTransport('/messages', res); + transports[transport.sessionId] = transport; + res.on('close', () => { + delete transports[transport.sessionId]; + }); + await server.connect(transport); +}); + +mcpRouter.post('/messages', async (req, res) => { + const sessionId = req.query.sessionId; + const transport = transports[sessionId as string]; + if (transport) { + await transport.handlePostMessage(req, res); + } else { + res.status(400).send('No transport found for sessionId'); + } +}); + +export { mcpRouter }; diff --git a/dev-packages/e2e-tests/test-applications/node-express/tests/mcp.test.ts b/dev-packages/e2e-tests/test-applications/node-express/tests/mcp.test.ts new file mode 100644 index 000000000000..a36e635685b6 --- /dev/null +++ b/dev-packages/e2e-tests/test-applications/node-express/tests/mcp.test.ts @@ -0,0 +1,109 @@ +import { expect, test } from '@playwright/test'; +import { waitForTransaction } from '@sentry-internal/test-utils'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'; + +test('Should record transactions for mcp handlers', async ({ baseURL }) => { + const transport = new SSEClientTransport(new URL(`${baseURL}/sse`)); + + const client = new Client({ + name: 'test-client', + version: '1.0.0', + }); + + await client.connect(transport); + + await test.step('tool handler', async () => { + const postTransactionPromise = waitForTransaction('node-express', transactionEvent => { + return transactionEvent.transaction === 'POST /messages'; + }); + const toolTransactionPromise = waitForTransaction('node-express', transactionEvent => { + return transactionEvent.transaction === 'mcp-server/tool:echo'; + }); + + const toolResult = await client.callTool({ + name: 'echo', + arguments: { + message: 'foobar', + }, + }); + + expect(toolResult).toMatchObject({ + content: [ + { + text: 'Tool echo: foobar', + type: 'text', + }, + ], + }); + + const postTransaction = await postTransactionPromise; + expect(postTransaction).toBeDefined(); + + const toolTransaction = await toolTransactionPromise; + expect(toolTransaction).toBeDefined(); + + // TODO: When https://github.com/modelcontextprotocol/typescript-sdk/pull/358 is released check for trace id equality between the post transaction and the handler transaction + }); + + await test.step('resource handler', async () => { + const postTransactionPromise = waitForTransaction('node-express', transactionEvent => { + return transactionEvent.transaction === 'POST /messages'; + }); + const resourceTransactionPromise = waitForTransaction('node-express', transactionEvent => { + return transactionEvent.transaction === 'mcp-server/resource:echo'; + }); + + const resourceResult = await client.readResource({ + uri: 'echo://foobar', + }); + + expect(resourceResult).toMatchObject({ + contents: [{ text: 'Resource echo: foobar', uri: 'echo://foobar' }], + }); + + const postTransaction = await postTransactionPromise; + expect(postTransaction).toBeDefined(); + + const resourceTransaction = await resourceTransactionPromise; + expect(resourceTransaction).toBeDefined(); + + // TODO: When https://github.com/modelcontextprotocol/typescript-sdk/pull/358 is released check for trace id equality between the post transaction and the handler transaction + }); + + await test.step('prompt handler', async () => { + const postTransactionPromise = waitForTransaction('node-express', transactionEvent => { + return transactionEvent.transaction === 'POST /messages'; + }); + const promptTransactionPromise = waitForTransaction('node-express', transactionEvent => { + return transactionEvent.transaction === 'mcp-server/prompt:echo'; + }); + + const promptResult = await client.getPrompt({ + name: 'echo', + arguments: { + message: 'foobar', + }, + }); + + expect(promptResult).toMatchObject({ + messages: [ + { + content: { + text: 'Please process this message: foobar', + type: 'text', + }, + role: 'user', + }, + ], + }); + + const postTransaction = await postTransactionPromise; + expect(postTransaction).toBeDefined(); + + const promptTransaction = await promptTransactionPromise; + expect(promptTransaction).toBeDefined(); + + // TODO: When https://github.com/modelcontextprotocol/typescript-sdk/pull/358 is released check for trace id equality between the post transaction and the handler transaction + }); +}); diff --git a/dev-packages/e2e-tests/test-applications/node-express/tsconfig.json b/dev-packages/e2e-tests/test-applications/node-express/tsconfig.json index ce4fafb745ad..0060abd94682 100644 --- a/dev-packages/e2e-tests/test-applications/node-express/tsconfig.json +++ b/dev-packages/e2e-tests/test-applications/node-express/tsconfig.json @@ -4,7 +4,8 @@ "esModuleInterop": true, "lib": ["es2020"], "strict": true, - "outDir": "dist" + "outDir": "dist", + "skipLibCheck": true }, "include": ["src/**/*.ts"] } diff --git a/packages/core/src/mcp-server.ts b/packages/core/src/mcp-server.ts index 3290b5b674b1..56c53010b58a 100644 --- a/packages/core/src/mcp-server.ts +++ b/packages/core/src/mcp-server.ts @@ -4,8 +4,17 @@ import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, SEMANTIC_ATTRIBUTE_SENTRY_SOURCE, } from './semanticAttributes'; -import { startSpan } from './tracing'; +import { startSpan, withActiveSpan } from './tracing'; +import type { Span } from './types-hoist/span'; import { logger } from './utils-hoist/logger'; +import { getActiveSpan } from './utils/spanUtils'; + +interface MCPTransport { + // The first argument is a JSON RPC message + onmessage?: (...args: unknown[]) => void; + onclose?: (...args: unknown[]) => void; + sessionId?: string; +} interface MCPServerInstance { // The first arg is always a name, the last arg should always be a callback function (ie a handler). @@ -15,6 +24,7 @@ interface MCPServerInstance { tool: (name: string, ...args: unknown[]) => void; // The first arg is always a name, the last arg should always be a callback function (ie a handler). prompt: (name: string, ...args: unknown[]) => void; + connect(transport: MCPTransport): Promise; } const wrappedMcpServerInstances = new WeakSet(); @@ -35,6 +45,59 @@ export function wrapMcpServerWithSentry(mcpServerInstance: S): return mcpServerInstance; } + // eslint-disable-next-line @typescript-eslint/unbound-method + mcpServerInstance.connect = new Proxy(mcpServerInstance.connect, { + apply(target, thisArg, argArray) { + const [transport, ...restArgs] = argArray as [MCPTransport, ...unknown[]]; + + if (!transport.onclose) { + transport.onclose = () => { + if (transport.sessionId) { + handleTransportOnClose(transport.sessionId); + } + }; + } + + if (!transport.onmessage) { + transport.onmessage = jsonRpcMessage => { + if (transport.sessionId && isJsonRPCMessageWithRequestId(jsonRpcMessage)) { + handleTransportOnMessage(transport.sessionId, jsonRpcMessage.id); + } + }; + } + + const patchedTransport = new Proxy(transport, { + set(target, key, value) { + if (key === 'onmessage') { + target[key] = new Proxy(value, { + apply(onMessageTarget, onMessageThisArg, onMessageArgArray) { + const [jsonRpcMessage] = onMessageArgArray; + if (transport.sessionId && isJsonRPCMessageWithRequestId(jsonRpcMessage)) { + handleTransportOnMessage(transport.sessionId, jsonRpcMessage.id); + } + return Reflect.apply(onMessageTarget, onMessageThisArg, onMessageArgArray); + }, + }); + } else if (key === 'onclose') { + target[key] = new Proxy(value, { + apply(onCloseTarget, onCloseThisArg, onCloseArgArray) { + if (transport.sessionId) { + handleTransportOnClose(transport.sessionId); + } + return Reflect.apply(onCloseTarget, onCloseThisArg, onCloseArgArray); + }, + }); + } else { + target[key as keyof MCPTransport] = value; + } + return true; + }, + }); + + return Reflect.apply(target, thisArg, [patchedTransport, ...restArgs]); + }, + }); + mcpServerInstance.resource = new Proxy(mcpServerInstance.resource, { apply(target, thisArg, argArray) { const resourceName: unknown = argArray[0]; @@ -44,19 +107,28 @@ export function wrapMcpServerWithSentry(mcpServerInstance: S): return target.apply(thisArg, argArray); } - return startSpan( - { - name: `mcp-server/resource:${resourceName}`, - forceTransaction: true, - attributes: { - [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server', - [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server', - [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route', - 'mcp_server.resource': resourceName, - }, + const wrappedResourceHandler = new Proxy(resourceHandler, { + apply(resourceHandlerTarget, resourceHandlerThisArg, resourceHandlerArgArray) { + const extraHandlerDataWithRequestId = resourceHandlerArgArray.find(isExtraHandlerDataWithRequestId); + return associateContextWithRequestSpan(extraHandlerDataWithRequestId, () => { + return startSpan( + { + name: `mcp-server/resource:${resourceName}`, + forceTransaction: true, + attributes: { + [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server', + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server', + [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route', + 'mcp_server.resource': resourceName, + }, + }, + () => resourceHandlerTarget.apply(resourceHandlerThisArg, resourceHandlerArgArray), + ); + }); }, - () => target.apply(thisArg, argArray), - ); + }); + + return Reflect.apply(target, thisArg, [...argArray.slice(0, -1), wrappedResourceHandler]); }, }); @@ -69,19 +141,28 @@ export function wrapMcpServerWithSentry(mcpServerInstance: S): return target.apply(thisArg, argArray); } - return startSpan( - { - name: `mcp-server/tool:${toolName}`, - forceTransaction: true, - attributes: { - [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server', - [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server', - [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route', - 'mcp_server.tool': toolName, - }, + const wrappedToolHandler = new Proxy(toolHandler, { + apply(toolHandlerTarget, toolHandlerThisArg, toolHandlerArgArray) { + const extraHandlerDataWithRequestId = toolHandlerArgArray.find(isExtraHandlerDataWithRequestId); + return associateContextWithRequestSpan(extraHandlerDataWithRequestId, () => { + return startSpan( + { + name: `mcp-server/tool:${toolName}`, + forceTransaction: true, + attributes: { + [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server', + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server', + [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route', + 'mcp_server.tool': toolName, + }, + }, + () => toolHandlerTarget.apply(toolHandlerThisArg, toolHandlerArgArray), + ); + }); }, - () => target.apply(thisArg, argArray), - ); + }); + + return Reflect.apply(target, thisArg, [...argArray.slice(0, -1), wrappedToolHandler]); }, }); @@ -94,19 +175,28 @@ export function wrapMcpServerWithSentry(mcpServerInstance: S): return target.apply(thisArg, argArray); } - return startSpan( - { - name: `mcp-server/resource:${promptName}`, - forceTransaction: true, - attributes: { - [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server', - [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server', - [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route', - 'mcp_server.prompt': promptName, - }, + const wrappedPromptHandler = new Proxy(promptHandler, { + apply(promptHandlerTarget, promptHandlerThisArg, promptHandlerArgArray) { + const extraHandlerDataWithRequestId = promptHandlerArgArray.find(isExtraHandlerDataWithRequestId); + return associateContextWithRequestSpan(extraHandlerDataWithRequestId, () => { + return startSpan( + { + name: `mcp-server/prompt:${promptName}`, + forceTransaction: true, + attributes: { + [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server', + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server', + [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route', + 'mcp_server.prompt': promptName, + }, + }, + () => promptHandlerTarget.apply(promptHandlerThisArg, promptHandlerArgArray), + ); + }); }, - () => target.apply(thisArg, argArray), - ); + }); + + return Reflect.apply(target, thisArg, [...argArray.slice(0, -1), wrappedPromptHandler]); }, }); @@ -124,6 +214,79 @@ function isMcpServerInstance(mcpServerInstance: unknown): mcpServerInstance is M 'tool' in mcpServerInstance && typeof mcpServerInstance.tool === 'function' && 'prompt' in mcpServerInstance && - typeof mcpServerInstance.prompt === 'function' + typeof mcpServerInstance.prompt === 'function' && + 'connect' in mcpServerInstance && + typeof mcpServerInstance.connect === 'function' + ); +} + +function isJsonRPCMessageWithRequestId(target: unknown): target is { id: RequestId } { + return ( + typeof target === 'object' && + target !== null && + 'id' in target && + (typeof target.id === 'number' || typeof target.id === 'string') ); } + +interface ExtraHandlerDataWithRequestId { + sessionId: SessionId; + requestId: RequestId; +} + +// Note that not all versions of the MCP library have `requestId` as a field on the extra data. +function isExtraHandlerDataWithRequestId(target: unknown): target is ExtraHandlerDataWithRequestId { + return ( + typeof target === 'object' && + target !== null && + 'sessionId' in target && + typeof target.sessionId === 'string' && + 'requestId' in target && + (typeof target.requestId === 'number' || typeof target.requestId === 'string') + ); +} + +type SessionId = string; +type RequestId = string | number; + +const sessionAndRequestToRequestParentSpanMap = new Map>(); + +function handleTransportOnClose(sessionId: SessionId): void { + sessionAndRequestToRequestParentSpanMap.delete(sessionId); +} + +function handleTransportOnMessage(sessionId: SessionId, requestId: RequestId): void { + const activeSpan = getActiveSpan(); + if (activeSpan) { + const requestIdToSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId) ?? new Map(); + requestIdToSpanMap.set(requestId, activeSpan); + sessionAndRequestToRequestParentSpanMap.set(sessionId, requestIdToSpanMap); + } +} + +function associateContextWithRequestSpan( + extraHandlerData: ExtraHandlerDataWithRequestId | undefined, + cb: () => T, +): T { + if (extraHandlerData) { + const { sessionId, requestId } = extraHandlerData; + const requestIdSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId); + + if (!requestIdSpanMap) { + return cb(); + } + + const span = requestIdSpanMap.get(requestId); + if (!span) { + return cb(); + } + + // remove the span from the map so it can be garbage collected + requestIdSpanMap.delete(requestId); + return withActiveSpan(span, () => { + return cb(); + }); + } + + return cb(); +} diff --git a/packages/core/test/lib/mcp-server.test.ts b/packages/core/test/lib/mcp-server.test.ts index 70904409e06d..12e85f9f370e 100644 --- a/packages/core/test/lib/mcp-server.test.ts +++ b/packages/core/test/lib/mcp-server.test.ts @@ -26,6 +26,7 @@ describe('wrapMcpServerWithSentry', () => { resource: mockResource, tool: mockTool, prompt: mockPrompt, + connect: vi.fn(), }; // Wrap the MCP server @@ -93,11 +94,19 @@ describe('wrapMcpServerWithSentry', () => { resource: vi.fn(), tool: vi.fn(), prompt: vi.fn(), + connect: vi.fn(), }; const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); wrappedMcpServer.resource(resourceName, {}, mockResourceHandler); + // The original registration should use a wrapped handler + expect(mockMcpServer.resource).toHaveBeenCalledWith(resourceName, {}, expect.any(Function)); + + // Invoke the wrapped handler to trigger Sentry span + const wrappedResourceHandler = (mockMcpServer.resource as any).mock.calls[0][2]; + wrappedResourceHandler('test-uri', { foo: 'bar' }); + expect(tracingModule.startSpan).toHaveBeenCalledTimes(1); expect(tracingModule.startSpan).toHaveBeenCalledWith( { @@ -113,8 +122,8 @@ describe('wrapMcpServerWithSentry', () => { expect.any(Function), ); - // Verify the original method was called with all arguments - expect(mockMcpServer.resource).toHaveBeenCalledWith(resourceName, {}, mockResourceHandler); + // Verify the original handler was called within the span + expect(mockResourceHandler).toHaveBeenCalledWith('test-uri', { foo: 'bar' }); }); it('should call the original resource method directly if name or handler is not valid', () => { @@ -122,6 +131,7 @@ describe('wrapMcpServerWithSentry', () => { resource: vi.fn(), tool: vi.fn(), prompt: vi.fn(), + connect: vi.fn(), }; const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); @@ -147,11 +157,19 @@ describe('wrapMcpServerWithSentry', () => { resource: vi.fn(), tool: vi.fn(), prompt: vi.fn(), + connect: vi.fn(), }; const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); wrappedMcpServer.tool(toolName, {}, mockToolHandler); + // The original registration should use a wrapped handler + expect(mockMcpServer.tool).toHaveBeenCalledWith(toolName, {}, expect.any(Function)); + + // Invoke the wrapped handler to trigger Sentry span + const wrappedToolHandler = (mockMcpServer.tool as any).mock.calls[0][2]; + wrappedToolHandler({ arg: 'value' }, { foo: 'baz' }); + expect(tracingModule.startSpan).toHaveBeenCalledTimes(1); expect(tracingModule.startSpan).toHaveBeenCalledWith( { @@ -167,8 +185,8 @@ describe('wrapMcpServerWithSentry', () => { expect.any(Function), ); - // Verify the original method was called with all arguments - expect(mockMcpServer.tool).toHaveBeenCalledWith(toolName, {}, mockToolHandler); + // Verify the original handler was called within the span + expect(mockToolHandler).toHaveBeenCalledWith({ arg: 'value' }, { foo: 'baz' }); }); it('should call the original tool method directly if name or handler is not valid', () => { @@ -176,6 +194,7 @@ describe('wrapMcpServerWithSentry', () => { resource: vi.fn(), tool: vi.fn(), prompt: vi.fn(), + connect: vi.fn(), }; const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); @@ -198,15 +217,23 @@ describe('wrapMcpServerWithSentry', () => { resource: vi.fn(), tool: vi.fn(), prompt: vi.fn(), + connect: vi.fn(), }; const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); wrappedMcpServer.prompt(promptName, {}, mockPromptHandler); + // The original registration should use a wrapped handler + expect(mockMcpServer.prompt).toHaveBeenCalledWith(promptName, {}, expect.any(Function)); + + // Invoke the wrapped handler to trigger Sentry span + const wrappedPromptHandler = (mockMcpServer.prompt as any).mock.calls[0][2]; + wrappedPromptHandler({ msg: 'hello' }, { data: 123 }); + expect(tracingModule.startSpan).toHaveBeenCalledTimes(1); expect(tracingModule.startSpan).toHaveBeenCalledWith( { - name: `mcp-server/resource:${promptName}`, + name: `mcp-server/prompt:${promptName}`, forceTransaction: true, attributes: { [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server', @@ -218,8 +245,8 @@ describe('wrapMcpServerWithSentry', () => { expect.any(Function), ); - // Verify the original method was called with all arguments - expect(mockMcpServer.prompt).toHaveBeenCalledWith(promptName, {}, mockPromptHandler); + // Verify the original handler was called within the span + expect(mockPromptHandler).toHaveBeenCalledWith({ msg: 'hello' }, { data: 123 }); }); it('should call the original prompt method directly if name or handler is not valid', () => { @@ -227,6 +254,7 @@ describe('wrapMcpServerWithSentry', () => { resource: vi.fn(), tool: vi.fn(), prompt: vi.fn(), + connect: vi.fn(), }; const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);