From 823affc67b2bf5e56d312d1b6c87699f092c49d6 Mon Sep 17 00:00:00 2001 From: kalenkevich Date: Thu, 29 Jan 2026 16:37:32 -0800 Subject: [PATCH] feat: Add headers option for MCP Session manager and deprecate the header option. --- core/src/tools/mcp/mcp_session_manager.ts | 28 +++- .../tools/mcp/mcp_session_manager_test.ts | 156 ++++++++++++++++++ 2 files changed, 176 insertions(+), 8 deletions(-) create mode 100644 core/test/tools/mcp/mcp_session_manager_test.ts diff --git a/core/src/tools/mcp/mcp_session_manager.ts b/core/src/tools/mcp/mcp_session_manager.ts index 168f403..2fcb543 100644 --- a/core/src/tools/mcp/mcp_session_manager.ts +++ b/core/src/tools/mcp/mcp_session_manager.ts @@ -9,7 +9,10 @@ import { StdioClientTransport, StdioServerParameters, } from '@modelcontextprotocol/sdk/client/stdio.js'; -import {StreamableHTTPClientTransport} from '@modelcontextprotocol/sdk/client/streamableHttp.js'; +import { + StreamableHTTPClientTransport, + StreamableHTTPClientTransportOptions, +} from '@modelcontextprotocol/sdk/client/streamableHttp.js'; /** * Defines the parameters for establishing a connection to an MCP server using @@ -35,10 +38,16 @@ export interface StdioConnectionParams { export interface StreamableHTTPConnectionParams { type: 'StreamableHTTPConnectionParams'; url: string; + /** + * @deprecated + * Use transportOptions.requestInit.headers instead. + * This field will be ignored if transportOptions is provided even if no headers are specified in transportOptions. + */ header?: Record; timeout?: number; sseReadTimeout?: number; terminateOnClose?: boolean; + transportOptions?: StreamableHTTPClientTransportOptions; } /** @@ -78,18 +87,21 @@ export class MCPSessionManager { ); break; case 'StreamableHTTPConnectionParams': { - let transportOptions; - if (this.connectionParams.header) { - transportOptions = { - requestInit: { - headers: this.connectionParams.header as Record, - }, + const options = this.connectionParams.transportOptions ?? {}; + + if ( + !options.requestInit && + this.connectionParams.header !== undefined + ) { + options.requestInit = { + headers: this.connectionParams.header as Record, }; } + await client.connect( new StreamableHTTPClientTransport( new URL(this.connectionParams.url), - transportOptions, + options, ), ); break; diff --git a/core/test/tools/mcp/mcp_session_manager_test.ts b/core/test/tools/mcp/mcp_session_manager_test.ts new file mode 100644 index 0000000..2d0b50d --- /dev/null +++ b/core/test/tools/mcp/mcp_session_manager_test.ts @@ -0,0 +1,156 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Client} from '@modelcontextprotocol/sdk/client/index.js'; +import {StdioClientTransport} from '@modelcontextprotocol/sdk/client/stdio.js'; +import {StreamableHTTPClientTransport} from '@modelcontextprotocol/sdk/client/streamableHttp.js'; +import {describe, expect, it, vi} from 'vitest'; +import {MCPSessionManager} from '../../../src/tools/mcp/mcp_session_manager.js'; + +vi.mock('@modelcontextprotocol/sdk/client/index.js', () => { + return { + Client: vi.fn().mockImplementation(() => ({ + connect: vi.fn().mockResolvedValue(undefined), + })), + }; +}); + +vi.mock('@modelcontextprotocol/sdk/client/stdio.js', () => { + return { + StdioClientTransport: vi.fn(), + }; +}); + +vi.mock('@modelcontextprotocol/sdk/client/streamableHttp.js', () => { + return { + StreamableHTTPClientTransport: vi.fn(), + }; +}); + +describe('MCPSessionManager', () => { + it('creates an stdio client', async () => { + const manager = new MCPSessionManager({ + type: 'StdioConnectionParams', + serverParams: { + command: 'test-command', + args: ['arg1', 'arg2'], + }, + }); + + const client = await manager.createSession(); + + expect(Client).toHaveBeenCalledWith({ + name: 'MCPClient', + version: '1.0.0', + }); + expect(StdioClientTransport).toHaveBeenCalledWith({ + command: 'test-command', + args: ['arg1', 'arg2'], + }); + expect(client.connect).toHaveBeenCalled(); + }); + + it('creates an http client with transport options headers', async () => { + const manager = new MCPSessionManager({ + type: 'StreamableHTTPConnectionParams', + url: 'http://test-url', + transportOptions: { + requestInit: { + headers: { + 'x-test-header': 'test-value', + }, + }, + }, + }); + + const client = await manager.createSession(); + + expect(Client).toHaveBeenCalledWith({ + name: 'MCPClient', + version: '1.0.0', + }); + expect(StreamableHTTPClientTransport).toHaveBeenCalledWith( + new URL('http://test-url'), + { + requestInit: { + headers: {'x-test-header': 'test-value'}, + }, + }, + ); + expect(client.connect).toHaveBeenCalled(); + }); + + it('creates an http client with deprecated header param', async () => { + const manager = new MCPSessionManager({ + type: 'StreamableHTTPConnectionParams', + url: 'http://test-url', + header: { + 'x-test-header': 'test-value', + }, + }); + + await manager.createSession(); + + expect(StreamableHTTPClientTransport).toHaveBeenCalledWith( + new URL('http://test-url'), + { + requestInit: { + headers: {'x-test-header': 'test-value'}, + }, + }, + ); + }); + + it('prioritizes transportOptions headers over header', async () => { + const manager = new MCPSessionManager({ + type: 'StreamableHTTPConnectionParams', + url: 'http://test-url', + transportOptions: { + requestInit: { + headers: { + 'x-priority': 'headers', + }, + }, + }, + header: { + 'x-priority': 'header', + }, + }); + + await manager.createSession(); + + expect(StreamableHTTPClientTransport).toHaveBeenCalledWith( + expect.any(URL), + { + requestInit: { + headers: {'x-priority': 'headers'}, + }, + }, + ); + }); + + it('prioritizes transportOptions over header', async () => { + const manager = new MCPSessionManager({ + type: 'StreamableHTTPConnectionParams', + url: 'http://test-url', + transportOptions: { + requestInit: {}, + }, + header: { + 'x-priority': 'header', + }, + }); + + await manager.createSession(); + + expect(StreamableHTTPClientTransport).toHaveBeenCalledWith( + expect.any(URL), + { + requestInit: {}, + }, + ); + }); +});