Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/server/index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ import {
ListResourcesRequestSchema,
ListToolsRequestSchema,
SetLevelRequestSchema,
ErrorCode,
ErrorCode
} from "../types.js";
import { Transport } from "../shared/transport.js";
import { InMemoryTransport } from "../inMemory.js";
Expand Down
2 changes: 1 addition & 1 deletion src/server/mcp.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ import {
LoggingMessageNotificationSchema,
Notification,
TextContent,
ElicitRequestSchema,
ElicitRequestSchema
} from "../types.js";
import { ResourceTemplate } from "./mcp.js";
import { completable } from "./completable.js";
Expand Down
204 changes: 198 additions & 6 deletions src/server/sse.test.ts
Original file line number Diff line number Diff line change
@@ -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();
Expand Down Expand Up @@ -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`);
});
});
});
9 changes: 5 additions & 4 deletions src/server/sse.ts
Original file line number Diff line number Diff line change
@@ -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";
Expand All @@ -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`.
Expand Down Expand Up @@ -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 {
Expand All @@ -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;
Expand All @@ -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);
Expand Down
65 changes: 65 additions & 0 deletions src/server/streamableHttp.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -216,6 +217,7 @@ describe("StreamableHTTPServerTransport", () => {
const result = await createTestServer();
server = result.server;
transport = result.transport;
mcpServer = result.mcpServer;
baseUrl = result.baseUrl;
});

Expand Down Expand Up @@ -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);

Expand Down
Loading