Merge pull request #627 from KKonstantinov/request-propagation-in-tools

Raw request propagation in tools - passed to callbacks via RequestHandlerExtra<ServerRequest, ServerNotification>

Author: ihrpr

src/server/index.test.ts

@@ -15,7 +15,7 @@ import {
   ListResourcesRequestSchema,
   ListToolsRequestSchema,
   SetLevelRequestSchema,
-  ErrorCode,
+  ErrorCode
 } from "../types.js";
 import { Transport } from "../shared/transport.js";
 import { InMemoryTransport } from "../inMemory.js";

src/server/mcp.test.ts

@@ -14,7 +14,7 @@ import {
   LoggingMessageNotificationSchema,
   Notification,
   TextContent,
-  ElicitRequestSchema,
+  ElicitRequestSchema
 } from "../types.js";
 import { ResourceTemplate } from "./mcp.js";
 import { completable } from "./completable.js";

src/server/sse.test.ts

@@ -1,20 +1,146 @@
 import http from 'http'; 
 import { jest } from '@jest/globals';
 import { SSEServerTransport } from './sse.js'; 
+import { McpServer } from './mcp.js';
+import { createServer, type Server } from "node:http";
+import { AddressInfo } from "node:net";
+import { z } from 'zod';
+import { CallToolResult, JSONRPCMessage } from 'src/types.js';
 
 const createMockResponse = () => {
   const res = {
-    writeHead: jest.fn<http.ServerResponse['writeHead']>(),
-    write: jest.fn<http.ServerResponse['write']>().mockReturnValue(true),
-    on: jest.fn<http.ServerResponse['on']>(),
+    writeHead: jest.fn<http.ServerResponse['writeHead']>().mockReturnThis(),
+    write: jest.fn<http.ServerResponse['write']>().mockReturnThis(),
+    on: jest.fn<http.ServerResponse['on']>().mockReturnThis(),
+    end: jest.fn<http.ServerResponse['end']>().mockReturnThis(),
   };
-  res.writeHead.mockReturnThis();
-  res.on.mockReturnThis();
 
-  return res as unknown as http.ServerResponse;
+  return res as unknown as jest.Mocked<http.ServerResponse>;
 };
 
+/**
+ * Helper to create and start test HTTP server with MCP setup
+ */
+async function createTestServerWithSse(args: {
+  mockRes: http.ServerResponse;
+}): Promise<{
+  server: Server;
+  transport: SSEServerTransport;
+  mcpServer: McpServer;
+  baseUrl: URL;
+  sessionId: string
+  serverPort: number;
+}> {
+  const mcpServer = new McpServer(
+    { name: "test-server", version: "1.0.0" },
+    { capabilities: { logging: {} } }
+  );
+
+  mcpServer.tool(
+    "greet",
+    "A simple greeting tool",
+    { name: z.string().describe("Name to greet") },
+    async ({ name }): Promise<CallToolResult> => {
+      return { content: [{ type: "text", text: `Hello, ${name}!` }] };
+    }
+  );
+
+  const endpoint = '/messages';
+
+  const transport = new SSEServerTransport(endpoint, args.mockRes);
+  const sessionId = transport.sessionId;
+
+  await mcpServer.connect(transport);
+
+  const server = createServer(async (req, res) => {
+    try {
+        await transport.handlePostMessage(req, res);
+    } catch (error) {
+      console.error("Error handling request:", error);
+      if (!res.headersSent) res.writeHead(500).end();
+    }
+  });
+
+  const baseUrl = await new Promise<URL>((resolve) => {
+    server.listen(0, "127.0.0.1", () => {
+      const addr = server.address() as AddressInfo;
+      resolve(new URL(`http://127.0.0.1:${addr.port}`));
+    });
+  });
+
+  const port = (server.address() as AddressInfo).port;
+
+  return { server, transport, mcpServer, baseUrl, sessionId, serverPort: port };
+}
+
+async function readAllSSEEvents(response: Response): Promise<string[]> {
+  const reader = response.body?.getReader();
+  if (!reader) throw new Error('No readable stream');
+  
+  const events: string[] = [];
+  const decoder = new TextDecoder();
+  
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      
+      if (value) {
+        events.push(decoder.decode(value));
+      }
+    }
+  } finally {
+    reader.releaseLock();
+  }
+  
+  return events;
+}
+
+/**
+ * Helper to send JSON-RPC request
+ */
+async function sendSsePostRequest(baseUrl: URL, message: JSONRPCMessage | JSONRPCMessage[], sessionId?: string, extraHeaders?: Record<string, string>): Promise<Response> {
+  const headers: Record<string, string> = {
+    "Content-Type": "application/json",
+    Accept: "application/json, text/event-stream",
+    ...extraHeaders
+  };
+
+  if (sessionId) {
+    baseUrl.searchParams.set('sessionId', sessionId);
+  }
+
+  return fetch(baseUrl, {
+    method: "POST",
+    headers,
+    body: JSON.stringify(message),
+  });
+}
+
 describe('SSEServerTransport', () => {
+
+  async function initializeServer(baseUrl: URL): Promise<void> {
+    const response = await sendSsePostRequest(baseUrl, {
+      jsonrpc: "2.0",
+      method: "initialize",
+      params: {
+        clientInfo: { name: "test-client", version: "1.0" },
+        protocolVersion: "2025-03-26",
+        capabilities: {
+        },
+      },
+  
+      id: "init-1",
+    } as JSONRPCMessage);
+
+    expect(response.status).toBe(202);
+
+    const text = await readAllSSEEvents(response);
+
+    expect(text).toHaveLength(1);
+    expect(text[0]).toBe('Accepted');
+  }
+
   describe('start method', () => { 
     it('should correctly append sessionId to a simple relative endpoint', async () => { 
       const mockRes = createMockResponse();
@@ -105,5 +231,71 @@ describe('SSEServerTransport', () => {
         `event: endpoint\ndata: /?sessionId=${expectedSessionId}\n\n`
       );
     });
+
+  /**
+   * Test: Tool With Request Info
+   */
+  it("should pass request info to tool callback", async () => {
+    const mockRes = createMockResponse();
+    const { mcpServer, baseUrl, sessionId, serverPort } = await createTestServerWithSse({ mockRes });
+    await initializeServer(baseUrl);
+
+    mcpServer.tool(
+      "test-request-info",
+      "A simple test tool with request info",
+      { name: z.string().describe("Name to greet") },
+      async ({ name }, { requestInfo }): Promise<CallToolResult> => {
+        return { content: [{ type: "text", text: `Hello, ${name}!` }, { type: "text", text: `${JSON.stringify(requestInfo)}` }] };
+      }
+    );
+   
+    const toolCallMessage: JSONRPCMessage = {
+      jsonrpc: "2.0",
+      method: "tools/call",
+      params: {
+        name: "test-request-info",
+        arguments: {
+          name: "Test User",
+        },
+      },
+      id: "call-1",
+    };
+
+    const response = await sendSsePostRequest(baseUrl, toolCallMessage, sessionId);
+
+    expect(response.status).toBe(202);
+
+    expect(mockRes.write).toHaveBeenCalledWith(`event: endpoint\ndata: /messages?sessionId=${sessionId}\n\n`);
+
+    const expectedMessage = {
+      result: {
+        content: [
+          {
+            type: "text",
+            text: "Hello, Test User!",
+          },
+          {
+            type: "text",
+            text: JSON.stringify({
+              headers: {
+                host: `127.0.0.1:${serverPort}`,
+                connection: 'keep-alive',
+                'content-type': 'application/json',
+                accept: 'application/json, text/event-stream',
+                'accept-language': '*',
+                'sec-fetch-mode': 'cors',
+                'user-agent': 'node',
+                'accept-encoding': 'gzip, deflate',
+                'content-length': '124'
+              },
+            })
+          },
+        ],
+      },
+      jsonrpc: "2.0",
+      id: "call-1",
+    };
+    expect(mockRes.write).toHaveBeenCalledWith(`event: message\ndata: ${JSON.stringify(expectedMessage)}\n\n`);
+  });
   });
 });

src/server/sse.ts

@@ -1,7 +1,7 @@
 import { randomUUID } from "node:crypto";
 import { IncomingMessage, ServerResponse } from "node:http";
 import { Transport } from "../shared/transport.js";
-import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
+import { JSONRPCMessage, JSONRPCMessageSchema, MessageExtraInfo, RequestInfo } from "../types.js";
 import getRawBody from "raw-body";
 import contentType from "content-type";
 import { AuthInfo } from "./auth/types.js";
@@ -19,7 +19,7 @@ export class SSEServerTransport implements Transport {
   private _sessionId: string;
   onclose?: () => void;
   onerror?: (error: Error) => void;
-  onmessage?: (message: JSONRPCMessage, extra?: { authInfo?: AuthInfo }) => void;
+  onmessage?: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void;
 
   /**
    * Creates a new SSE server transport, which will direct the client to POST messages to the relative or absolute URL identified by `_endpoint`.
@@ -86,6 +86,7 @@ export class SSEServerTransport implements Transport {
       throw new Error(message);
     }
     const authInfo: AuthInfo | undefined = req.auth;
+    const requestInfo: RequestInfo = { headers: req.headers };
 
     let body: string | unknown;
     try {
@@ -105,7 +106,7 @@ export class SSEServerTransport implements Transport {
     }
 
     try {
-      await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body, { authInfo });
+      await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body, { requestInfo, authInfo });
     } catch {
       res.writeHead(400).end(`Invalid message: ${body}`);
       return;
@@ -117,7 +118,7 @@ export class SSEServerTransport implements Transport {
   /**
    * Handle a client message, regardless of how it arrived. This can be used to inform the server of messages that arrive via a means different than HTTP POST.
    */
-  async handleMessage(message: unknown, extra?: { authInfo?: AuthInfo }): Promise<void> {
+  async handleMessage(message: unknown, extra?: MessageExtraInfo): Promise<void> {
     let parsedMessage: JSONRPCMessage;
     try {
       parsedMessage = JSONRPCMessageSchema.parse(message);

src/server/streamableHttp.test.ts

@@ -208,6 +208,7 @@ function expectErrorResponse(data: unknown, expectedCode: number, expectedMessag
 
 describe("StreamableHTTPServerTransport", () => {
   let server: Server;
+  let mcpServer: McpServer;
   let transport: StreamableHTTPServerTransport;
   let baseUrl: URL;
   let sessionId: string;
@@ -216,6 +217,7 @@ describe("StreamableHTTPServerTransport", () => {
     const result = await createTestServer();
     server = result.server;
     transport = result.transport;
+    mcpServer = result.mcpServer;
     baseUrl = result.baseUrl;
   });
 
@@ -347,6 +349,69 @@ describe("StreamableHTTPServerTransport", () => {
     });
   });
 
+  /***
+   * Test: Tool With Request Info
+   */
+  it("should pass request info to tool callback", async () => {
+    sessionId = await initializeServer();
+
+    mcpServer.tool(
+      "test-request-info",
+      "A simple test tool with request info",
+      { name: z.string().describe("Name to greet") },
+      async ({ name }, { requestInfo }): Promise<CallToolResult> => {
+        return { content: [{ type: "text", text: `Hello, ${name}!` }, { type: "text", text: `${JSON.stringify(requestInfo)}` }] };
+      }
+    );
+   
+    const toolCallMessage: JSONRPCMessage = {
+      jsonrpc: "2.0",
+      method: "tools/call",
+      params: {
+        name: "test-request-info",
+        arguments: {
+          name: "Test User",
+        },
+      },
+      id: "call-1",
+    };
+
+    const response = await sendPostRequest(baseUrl, toolCallMessage, sessionId);
+    expect(response.status).toBe(200);
+
+    const text = await readSSEEvent(response);
+    const eventLines = text.split("\n");
+    const dataLine = eventLines.find(line => line.startsWith("data:"));
+    expect(dataLine).toBeDefined();
+
+    const eventData = JSON.parse(dataLine!.substring(5));
+
+    expect(eventData).toMatchObject({
+      jsonrpc: "2.0",
+      result: {
+        content: [
+          { type: "text", text: "Hello, Test User!" },
+          { type: "text", text: expect.any(String) }
+        ],
+      },
+      id: "call-1",
+    });
+
+    const requestInfo = JSON.parse(eventData.result.content[1].text);
+    expect(requestInfo).toMatchObject({
+      headers: {
+        'content-type': 'application/json',
+        accept: 'application/json, text/event-stream',
+        connection: 'keep-alive',
+        'mcp-session-id': sessionId,
+        'accept-language': '*',
+        'user-agent': expect.any(String),
+        'accept-encoding': expect.any(String),
+        'content-length': expect.any(String),
+      },
+    });
+  });
+
   it("should reject requests without a valid session ID", async () => {
     const response = await sendPostRequest(baseUrl, TEST_MESSAGES.toolsList);