From e9b3ef032819bde22a8a2e93b6de5776a15f630a Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Fri, 10 Apr 2026 15:53:02 +0900 Subject: [PATCH] feat(skill-mcp-manager): improve connection handling and client implementations Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../connection-env-vars.test.ts | 76 ++++++++---- .../skill-mcp-manager/connection-race.test.ts | 34 +++-- src/features/skill-mcp-manager/connection.ts | 13 +- src/features/skill-mcp-manager/http-client.ts | 41 +++++- .../skill-mcp-manager/manager.test.ts | 117 ++++++++++++------ src/features/skill-mcp-manager/manager.ts | 8 +- .../skill-mcp-manager/stdio-client.ts | 40 +++++- src/features/skill-mcp-manager/types.ts | 24 +++- 8 files changed, 259 insertions(+), 94 deletions(-) diff --git a/src/features/skill-mcp-manager/connection-env-vars.test.ts b/src/features/skill-mcp-manager/connection-env-vars.test.ts index 60cf20ce6..728d3ab81 100644 --- a/src/features/skill-mcp-manager/connection-env-vars.test.ts +++ b/src/features/skill-mcp-manager/connection-env-vars.test.ts @@ -1,6 +1,10 @@ import { afterAll, afterEach, beforeEach, describe, expect, it, mock, test } from "bun:test" +import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js" +import type { StdioServerParameters } from "@modelcontextprotocol/sdk/client/stdio.js" import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types" import type { SkillMcpClientInfo, SkillMcpManagerState } from "./types" +import { setHttpClientDependenciesForTesting } from "./http-client" +import { setStdioClientDependenciesForTesting } from "./stdio-client" const trackedStates: SkillMcpManagerState[] = [] const createdStdioTransports: MockStdioClientTransport[] = [] @@ -8,33 +12,62 @@ const createdHttpTransports: MockStreamableHTTPClientTransport[] = [] class MockClient { readonly close = mock(async () => {}) + readonly listTools = mock(async () => ({ tools: [] })) + readonly listResources = mock(async () => ({ resources: [] })) + readonly listPrompts = mock(async () => ({ prompts: [] })) + readonly callTool = mock(async () => ({ content: [] })) + readonly readResource = mock(async () => ({ contents: [] })) + readonly getPrompt = mock(async () => ({ messages: [] })) constructor( _clientInfo: { name: string; version: string }, _options: { capabilities: Record } ) {} - async connect(_transport: unknown): Promise { + async connect(_transport: Transport): Promise { // Successful connect, env-related assertions happen on transport constructor args } } class MockStdioClientTransport { readonly close = mock(async () => {}) - readonly options: { command: string; args?: string[]; env?: Record; stderr?: string } + readonly start = mock(async () => {}) + readonly send = mock(async () => {}) + readonly options: StdioServerParameters - constructor(options: { command: string; args?: string[]; env?: Record; stderr?: string }) { + constructor(options: StdioServerParameters) { this.options = options createdStdioTransports.push(this) } } interface MockHttpTransportOptions { - requestInit?: { headers?: Record } + requestInit?: RequestInit +} + +function getHeaderValue( + headers: HeadersInit | undefined, + name: string, +): string | undefined { + if (!headers) { + return undefined + } + + if (headers instanceof Headers) { + return headers.get(name) ?? undefined + } + + if (Array.isArray(headers)) { + const entry = headers.find(([headerName]) => headerName.toLowerCase() === name.toLowerCase()) + return entry?.[1] + } + + return headers[name] } class MockStreamableHTTPClientTransport { readonly close = mock(async () => {}) + readonly send = mock(async () => {}) readonly url: URL readonly options?: MockHttpTransportOptions @@ -47,18 +80,6 @@ class MockStreamableHTTPClientTransport { async start() {} } -mock.module("@modelcontextprotocol/sdk/client/index.js", () => ({ - Client: MockClient, -})) - -mock.module("@modelcontextprotocol/sdk/client/stdio.js", () => ({ - StdioClientTransport: MockStdioClientTransport, -})) - -mock.module("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({ - StreamableHTTPClientTransport: MockStreamableHTTPClientTransport, -})) - afterAll(() => { mock.restore() }) @@ -110,6 +131,14 @@ const ORIGINAL_ENV = { ...process.env } beforeEach(() => { createdStdioTransports.length = 0 createdHttpTransports.length = 0 + setStdioClientDependenciesForTesting({ + createClient: (clientInfo, options) => new MockClient(clientInfo, options), + createTransport: (options) => new MockStdioClientTransport(options), + }) + setHttpClientDependenciesForTesting({ + createClient: (clientInfo, options) => new MockClient(clientInfo, options), + createTransport: (url, options) => new MockStreamableHTTPClientTransport(url, options), + }) }) afterEach(async () => { @@ -126,6 +155,9 @@ afterEach(async () => { for (const [key, value] of Object.entries(ORIGINAL_ENV)) { process.env[key] = value } + + setStdioClientDependenciesForTesting() + setHttpClientDependenciesForTesting() }) describe("getOrCreateClient env var expansion", () => { @@ -265,11 +297,11 @@ describe("getOrCreateClient env var expansion", () => { // when await getOrCreateClient({ state, clientKey, info, config }) - // then - expect(createdHttpTransports).toHaveLength(1) - expect(createdHttpTransports[0]?.options?.requestInit?.headers?.Authorization).toBe( - "Bearer xoxp-http-secret" - ) - }) + // then + expect(createdHttpTransports).toHaveLength(1) + expect(getHeaderValue(createdHttpTransports[0]?.options?.requestInit?.headers, "Authorization")).toBe( + "Bearer xoxp-http-secret" + ) + }) }) }) diff --git a/src/features/skill-mcp-manager/connection-race.test.ts b/src/features/skill-mcp-manager/connection-race.test.ts index 652987f67..2d78bd735 100644 --- a/src/features/skill-mcp-manager/connection-race.test.ts +++ b/src/features/skill-mcp-manager/connection-race.test.ts @@ -1,5 +1,8 @@ import { afterEach, beforeEach, describe, expect, it, mock, afterAll } from "bun:test" +import type { StdioServerParameters } from "@modelcontextprotocol/sdk/client/stdio.js" +import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js" import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types" +import { setStdioClientDependenciesForTesting } from "./stdio-client" import type { SkillMcpClientInfo, SkillMcpManagerState } from "./types" type Deferred = { @@ -23,7 +26,14 @@ class MockClient { createdClients.push(this) } - async connect(_transport: MockStdioClientTransport): Promise { + readonly listTools = mock(async () => ({ tools: [] })) + readonly listResources = mock(async () => ({ resources: [] })) + readonly listPrompts = mock(async () => ({ prompts: [] })) + readonly callTool = mock(async () => ({ content: [] })) + readonly readResource = mock(async () => ({ contents: [] })) + readonly getPrompt = mock(async () => ({ messages: [] })) + + async connect(_transport: Transport): Promise { const pendingConnect = pendingConnects.shift() if (pendingConnect) { await pendingConnect.promise @@ -33,20 +43,14 @@ class MockClient { class MockStdioClientTransport { readonly close = mock(async () => {}) + readonly start = mock(async () => {}) + readonly send = mock(async () => {}) - constructor(_options: { command: string; args?: string[]; env?: Record; stderr?: string }) { + constructor(_options: StdioServerParameters) { createdTransports.push(this) } } -mock.module("@modelcontextprotocol/sdk/client/index.js", () => ({ - Client: MockClient, -})) - -mock.module("@modelcontextprotocol/sdk/client/stdio.js", () => ({ - StdioClientTransport: MockStdioClientTransport, -})) - afterAll(() => { mock.restore() }) const { disconnectAll, disconnectSession } = await import("./cleanup") @@ -84,6 +88,11 @@ function createState(): SkillMcpManagerState { shutdownGeneration: 0, inFlightConnections: new Map(), disposed: false, + createOAuthProvider: () => ({ + tokens: () => null, + login: async () => ({ accessToken: "test-token" }), + refresh: async () => ({ accessToken: "test-token" }), + }), } trackedStates.push(state) @@ -111,6 +120,10 @@ beforeEach(() => { pendingConnects.length = 0 createdClients.length = 0 createdTransports.length = 0 + setStdioClientDependenciesForTesting({ + createClient: (clientInfo, options) => new MockClient(clientInfo, options), + createTransport: (options) => new MockStdioClientTransport(options), + }) }) afterEach(async () => { @@ -122,6 +135,7 @@ afterEach(async () => { pendingConnects.length = 0 createdClients.length = 0 createdTransports.length = 0 + setStdioClientDependenciesForTesting() }) describe("getOrCreateClient disconnect race", () => { diff --git a/src/features/skill-mcp-manager/connection.ts b/src/features/skill-mcp-manager/connection.ts index 2fa4dc3a3..e7c32085a 100644 --- a/src/features/skill-mcp-manager/connection.ts +++ b/src/features/skill-mcp-manager/connection.ts @@ -1,13 +1,12 @@ -import type { Client } from "@modelcontextprotocol/sdk/client/index.js" import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types" import { expandEnvVarsInObject } from "../claude-code-mcp-loader/env-expander" import { forceReconnect } from "./cleanup" import { getConnectionType } from "./connection-type" import { createHttpClient } from "./http-client" import { createStdioClient } from "./stdio-client" -import type { SkillMcpClientConnectionParams, SkillMcpClientInfo, SkillMcpManagerState } from "./types" +import type { McpClient, SkillMcpClientConnectionParams, SkillMcpClientInfo, SkillMcpManagerState } from "./types" -function removeClientIfCurrent(state: SkillMcpManagerState, clientKey: string, client: Client): void { +function removeClientIfCurrent(state: SkillMcpManagerState, clientKey: string, client: McpClient): void { const managed = state.clients.get(clientKey) if (managed?.client === client) { state.clients.delete(clientKey) @@ -21,7 +20,7 @@ export async function getOrCreateClient(params: { clientKey: string info: SkillMcpClientInfo config: ClaudeCodeMcpServer -}): Promise { +}): Promise { const { state, clientKey, info, config } = params if (state.disposed) { @@ -42,7 +41,7 @@ export async function getOrCreateClient(params: { const isTrusted = !PROJECT_SCOPES.has(info.scope ?? "") const expandedConfig = expandEnvVarsInObject(config, { trusted: isTrusted }) - let currentConnectionPromise!: Promise + let currentConnectionPromise!: Promise state.inFlightConnections.set(info.sessionID, (state.inFlightConnections.get(info.sessionID) ?? 0) + 1) currentConnectionPromise = (async () => { const disconnectGenAtStart = state.disconnectedSessions.get(info.sessionID) ?? 0 @@ -96,7 +95,7 @@ export async function getOrCreateClientWithRetryImpl(params: { clientKey: string info: SkillMcpClientInfo config: ClaudeCodeMcpServer -}): Promise { +}): Promise { const { state, clientKey } = params try { @@ -115,7 +114,7 @@ async function createClient(params: { clientKey: string info: SkillMcpClientInfo config: ClaudeCodeMcpServer -}): Promise { +}): Promise { const { info, config } = params const connectionType = getConnectionType(config) diff --git a/src/features/skill-mcp-manager/http-client.ts b/src/features/skill-mcp-manager/http-client.ts index 74bc43598..d674c20f4 100644 --- a/src/features/skill-mcp-manager/http-client.ts +++ b/src/features/skill-mcp-manager/http-client.ts @@ -2,7 +2,40 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js" import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js" import { registerProcessCleanup, startCleanupTimer } from "./cleanup" import { buildHttpRequestInit } from "./oauth-handler" -import type { ManagedClient, SkillMcpClientConnectionParams } from "./types" +import type { ManagedClient, McpClient, McpTransport, SkillMcpClientConnectionParams } from "./types" + +type HttpClientFactory = ( + clientInfo: { name: string; version: string }, + options: { capabilities: Record } +) => McpClient + +type HttpTransportFactory = ( + url: URL, + options?: { requestInit?: RequestInit } +) => McpTransport + +interface HttpClientDependencies { + createClient: HttpClientFactory + createTransport: HttpTransportFactory +} + +const defaultHttpClientDependencies: HttpClientDependencies = { + createClient: (clientInfo, options) => new Client(clientInfo, options), + createTransport: (url, options) => new StreamableHTTPClientTransport(url, options), +} + +let httpClientDependencies: HttpClientDependencies = defaultHttpClientDependencies + +export function setHttpClientDependenciesForTesting( + dependencies?: Partial +): void { + httpClientDependencies = dependencies + ? { + ...defaultHttpClientDependencies, + ...dependencies, + } + : defaultHttpClientDependencies +} function redactUrl(urlStr: string): string { try { @@ -22,7 +55,7 @@ function redactUrl(urlStr: string): string { } } -export async function createHttpClient(params: SkillMcpClientConnectionParams): Promise { +export async function createHttpClient(params: SkillMcpClientConnectionParams): Promise { const { state, clientKey, info, config } = params const shutdownGenAtStart = state.shutdownGeneration @@ -43,11 +76,11 @@ export async function createHttpClient(params: SkillMcpClientConnectionParams): registerProcessCleanup(state) const requestInit = await buildHttpRequestInit(config, state.authProviders, state.createOAuthProvider) - const transport = new StreamableHTTPClientTransport(url, { + const transport: McpTransport = httpClientDependencies.createTransport(url, { requestInit, }) - const client = new Client( + const client: McpClient = httpClientDependencies.createClient( { name: `skill-mcp-${info.skillName}-${info.serverName}`, version: "1.0.0" }, { capabilities: {} } ) diff --git a/src/features/skill-mcp-manager/manager.test.ts b/src/features/skill-mcp-manager/manager.test.ts index bdbc316a1..f3ef6f51e 100644 --- a/src/features/skill-mcp-manager/manager.test.ts +++ b/src/features/skill-mcp-manager/manager.test.ts @@ -1,44 +1,84 @@ import { describe, it, expect, beforeEach, afterEach, afterAll, mock, spyOn } from "bun:test" +import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js" import type { SkillMcpClientInfo, SkillMcpServerContext } from "./types" import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types" import type { OAuthTokenData } from "../mcp-oauth/storage" +import { setHttpClientDependenciesForTesting } from "./http-client" +import { setStdioClientDependenciesForTesting } from "./stdio-client" +import { SkillMcpManager } from "./manager" -// Mock the MCP SDK transports to avoid network calls const mockHttpConnect = mock(() => Promise.reject(new Error("Mocked HTTP connection failure"))) const mockHttpClose = mock(() => Promise.resolve()) let lastTransportInstance: { url?: URL; options?: { requestInit?: RequestInit } } = {} +class MockHttpClient { + readonly close = mock(() => Promise.resolve()) + readonly listTools = mock(async () => ({ tools: [] })) + readonly listResources = mock(async () => ({ resources: [] })) + readonly listPrompts = mock(async () => ({ prompts: [] })) + readonly callTool = mock(async () => ({ content: [] })) + readonly readResource = mock(async () => ({ contents: [] })) + readonly getPrompt = mock(async () => ({ messages: [] })) + + constructor( + _clientInfo: { name: string; version: string }, + _options: { capabilities: Record } + ) {} + + async connect(transport: Transport): Promise { + await transport.start() + } +} + +class MockStreamableHTTPClientTransport { + constructor(public url: URL, public options?: { requestInit?: RequestInit }) { + lastTransportInstance = { url, options } + } + + async start(): Promise { + await mockHttpConnect() + } + + async send(): Promise {} + + async close(): Promise { + await mockHttpClose() + } +} + +function getHeaderValue(headers: HeadersInit | undefined, name: string): string | undefined { + if (!headers) { + return undefined + } + + if (headers instanceof Headers) { + return headers.get(name) ?? undefined + } + + if (Array.isArray(headers)) { + const entry = headers.find(([headerName]) => headerName.toLowerCase() === name.toLowerCase()) + return entry?.[1] + } + + return headers[name] +} + const mockTokens = mock(() => null as OAuthTokenData | null) const mockLogin = mock(() => Promise.resolve({ accessToken: "test-token" } satisfies OAuthTokenData)) const mockRefresh = mock((_: string) => Promise.resolve({ accessToken: "refreshed-token" } satisfies OAuthTokenData)) -async function importFreshManagerModule(): Promise { - mock.module("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({ - StreamableHTTPClientTransport: class MockStreamableHTTPClientTransport { - constructor(public url: URL, public options?: { requestInit?: RequestInit }) { - lastTransportInstance = { url, options } - } - async start() { - await mockHttpConnect() - } - async close() { - await mockHttpClose() - } - }, - })) - - const module = await import(`./manager?test=${Date.now()}-${Math.random()}`) - mock.restore() - return module -} - afterAll(() => { mock.restore() }) describe("SkillMcpManager", () => { - let manager: any + let manager: SkillMcpManager + + beforeEach(() => { + setHttpClientDependenciesForTesting({ + createClient: (clientInfo, options) => new MockHttpClient(clientInfo, options), + createTransport: (url, options) => new MockStreamableHTTPClientTransport(url, options), + }) + setStdioClientDependenciesForTesting() - beforeEach(async () => { - const { SkillMcpManager } = await importFreshManagerModule() manager = new SkillMcpManager({ createOAuthProvider: () => ({ tokens: () => mockTokens(), @@ -51,10 +91,13 @@ describe("SkillMcpManager", () => { mockTokens.mockClear() mockLogin.mockClear() mockRefresh.mockClear() + lastTransportInstance = {} }) afterEach(async () => { await manager.disconnectAll() + setHttpClientDependenciesForTesting() + setStdioClientDependenciesForTesting() }) describe("getOrCreateClient", () => { @@ -697,8 +740,8 @@ describe("SkillMcpManager", () => { } catch { /* connection fails in test */ } // then - const headers = lastTransportInstance.options?.requestInit?.headers as Record | undefined - expect(headers?.Authorization).toBe("Bearer stored-access-token") + const headers = lastTransportInstance.options?.requestInit?.headers + expect(getHeaderValue(headers, "Authorization")).toBe("Bearer stored-access-token") }) it("does not inject Authorization header when no stored tokens exist and login fails", async () => { @@ -724,8 +767,8 @@ describe("SkillMcpManager", () => { } catch { /* connection fails in test */ } // then - const headers = lastTransportInstance.options?.requestInit?.headers as Record | undefined - expect(headers?.Authorization).toBeUndefined() + const headers = lastTransportInstance.options?.requestInit?.headers + expect(getHeaderValue(headers, "Authorization")).toBeUndefined() }) it("preserves existing static headers alongside OAuth token", async () => { @@ -753,9 +796,9 @@ describe("SkillMcpManager", () => { } catch { /* connection fails in test */ } // then - const headers = lastTransportInstance.options?.requestInit?.headers as Record | undefined - expect(headers?.["X-Custom"]).toBe("custom-value") - expect(headers?.Authorization).toBe("Bearer oauth-token") + const headers = lastTransportInstance.options?.requestInit?.headers + expect(getHeaderValue(headers, "X-Custom")).toBe("custom-value") + expect(getHeaderValue(headers, "Authorization")).toBe("Bearer oauth-token") }) it("attempts silent refresh for expired stored tokens before login", async () => { @@ -785,8 +828,8 @@ describe("SkillMcpManager", () => { } catch { /* connection fails in test */ } // then - const headers = lastTransportInstance.options?.requestInit?.headers as Record | undefined - expect(headers?.Authorization).toBe("Bearer refreshed-token") + const headers = lastTransportInstance.options?.requestInit?.headers + expect(getHeaderValue(headers, "Authorization")).toBe("Bearer refreshed-token") expect(mockRefresh).toHaveBeenCalledWith("refresh-token") expect(mockLogin).not.toHaveBeenCalled() }) @@ -819,8 +862,8 @@ describe("SkillMcpManager", () => { } catch { /* connection fails in test */ } // then - const headers = lastTransportInstance.options?.requestInit?.headers as Record | undefined - expect(headers?.Authorization).toBe("Bearer login-token") + const headers = lastTransportInstance.options?.requestInit?.headers + expect(getHeaderValue(headers, "Authorization")).toBe("Bearer login-token") expect(mockRefresh).toHaveBeenCalledWith("refresh-token") expect(mockLogin).toHaveBeenCalled() }) @@ -846,8 +889,8 @@ describe("SkillMcpManager", () => { } catch { /* connection fails in test */ } // then - const headers = lastTransportInstance.options?.requestInit?.headers as Record | undefined - expect(headers?.Authorization).toBe("Bearer static-token") + const headers = lastTransportInstance.options?.requestInit?.headers + expect(getHeaderValue(headers, "Authorization")).toBe("Bearer static-token") expect(mockTokens).not.toHaveBeenCalled() }) diff --git a/src/features/skill-mcp-manager/manager.ts b/src/features/skill-mcp-manager/manager.ts index f91524be4..98dbb6d0f 100644 --- a/src/features/skill-mcp-manager/manager.ts +++ b/src/features/skill-mcp-manager/manager.ts @@ -1,4 +1,3 @@ -import type { Client } from "@modelcontextprotocol/sdk/client/index.js" import type { Prompt, Resource, Tool } from "@modelcontextprotocol/sdk/types.js" import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types" import { McpOAuthProvider } from "../mcp-oauth/provider" @@ -6,6 +5,7 @@ import { disconnectAll, disconnectSession, forceReconnect } from "./cleanup" import { getOrCreateClient, getOrCreateClientWithRetryImpl } from "./connection" import { handlePostRequestAuthError, handleStepUpIfNeeded } from "./oauth-handler" import type { + McpClient, OAuthProviderFactory, SkillMcpClientInfo, SkillMcpManagerState, @@ -36,7 +36,7 @@ export class SkillMcpManager { return `${info.sessionID}:${info.skillName}:${info.serverName}` } - async getOrCreateClient(info: SkillMcpClientInfo, config: ClaudeCodeMcpServer): Promise { + async getOrCreateClient(info: SkillMcpClientInfo, config: ClaudeCodeMcpServer): Promise { const clientKey = this.getClientKey(info) return await getOrCreateClient({ state: this.state, @@ -106,7 +106,7 @@ export class SkillMcpManager { private async withOperationRetry( info: SkillMcpClientInfo, config: ClaudeCodeMcpServer, - operation: (client: Client) => Promise + operation: (client: McpClient) => Promise ): Promise { const maxRetries = 3 let lastError: Error | null = null @@ -158,7 +158,7 @@ export class SkillMcpManager { } // NOTE: tests spy on this exact method name via `spyOn(manager as any, 'getOrCreateClientWithRetry')`. - private async getOrCreateClientWithRetry(info: SkillMcpClientInfo, config: ClaudeCodeMcpServer): Promise { + private async getOrCreateClientWithRetry(info: SkillMcpClientInfo, config: ClaudeCodeMcpServer): Promise { const clientKey = this.getClientKey(info) return await getOrCreateClientWithRetryImpl({ state: this.state, diff --git a/src/features/skill-mcp-manager/stdio-client.ts b/src/features/skill-mcp-manager/stdio-client.ts index 3a5c796a4..a7be4c39b 100644 --- a/src/features/skill-mcp-manager/stdio-client.ts +++ b/src/features/skill-mcp-manager/stdio-client.ts @@ -4,7 +4,39 @@ import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types" import { createCleanMcpEnvironment } from "./env-cleaner" import { registerProcessCleanup, startCleanupTimer } from "./cleanup" import { redactSensitiveData } from "./error-redaction" -import type { ManagedClient, SkillMcpClientConnectionParams } from "./types" +import type { ManagedClient, McpClient, McpTransport, SkillMcpClientConnectionParams } from "./types" + +type StdioClientFactory = ( + clientInfo: { name: string; version: string }, + options: { capabilities: Record } +) => McpClient + +type StdioTransportFactory = ( + options: ConstructorParameters[0] +) => McpTransport + +interface StdioClientDependencies { + createClient: StdioClientFactory + createTransport: StdioTransportFactory +} + +const defaultStdioClientDependencies: StdioClientDependencies = { + createClient: (clientInfo, options) => new Client(clientInfo, options), + createTransport: (options) => new StdioClientTransport(options), +} + +let stdioClientDependencies: StdioClientDependencies = defaultStdioClientDependencies + +export function setStdioClientDependenciesForTesting( + dependencies?: Partial +): void { + stdioClientDependencies = dependencies + ? { + ...defaultStdioClientDependencies, + ...dependencies, + } + : defaultStdioClientDependencies +} function getStdioCommand(config: ClaudeCodeMcpServer, serverName: string): string { if (!config.command) { @@ -13,7 +45,7 @@ function getStdioCommand(config: ClaudeCodeMcpServer, serverName: string): strin return config.command } -export async function createStdioClient(params: SkillMcpClientConnectionParams): Promise { +export async function createStdioClient(params: SkillMcpClientConnectionParams): Promise { const { state, clientKey, info, config } = params const shutdownGenAtStart = state.shutdownGeneration @@ -23,14 +55,14 @@ export async function createStdioClient(params: SkillMcpClientConnectionParams): registerProcessCleanup(state) - const transport = new StdioClientTransport({ + const transport: McpTransport = stdioClientDependencies.createTransport({ command, args, env: mergedEnv, stderr: "ignore", }) - const client = new Client( + const client: McpClient = stdioClientDependencies.createClient( { name: `skill-mcp-${info.skillName}-${info.serverName}`, version: "1.0.0" }, { capabilities: {} } ) diff --git a/src/features/skill-mcp-manager/types.ts b/src/features/skill-mcp-manager/types.ts index 75ef396cf..bf7d71d15 100644 --- a/src/features/skill-mcp-manager/types.ts +++ b/src/features/skill-mcp-manager/types.ts @@ -1,12 +1,24 @@ import type { Client } from "@modelcontextprotocol/sdk/client/index.js" -import type { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js" -import type { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js" +import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js" import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types" import type { McpOAuthProvider } from "../mcp-oauth/provider" import type { SkillScope } from "../opencode-skill-loader/types" export type SkillMcpConfig = Record +export type McpTransport = Transport + +export interface McpClient { + connect: Client["connect"] + close: Client["close"] + listTools: Client["listTools"] + listResources: Client["listResources"] + listPrompts: Client["listPrompts"] + callTool: Client["callTool"] + readResource: Client["readResource"] + getPrompt: Client["getPrompt"] +} + export interface SkillMcpClientInfo { serverName: string skillName: string @@ -27,7 +39,7 @@ export interface SkillMcpServerContext { export type ConnectionType = "stdio" | "http" export interface ManagedClientBase { - client: Client + client: McpClient skillName: string lastUsedAt: number connectionType: ConnectionType @@ -35,12 +47,12 @@ export interface ManagedClientBase { export interface ManagedStdioClient extends ManagedClientBase { connectionType: "stdio" - transport: StdioClientTransport + transport: McpTransport } export interface ManagedHttpClient extends ManagedClientBase { connectionType: "http" - transport: StreamableHTTPClientTransport + transport: McpTransport } export type ManagedClient = ManagedStdioClient | ManagedHttpClient @@ -63,7 +75,7 @@ export type OAuthProviderFactory = (options: { export interface SkillMcpManagerState { clients: Map - pendingConnections: Map> + pendingConnections: Map> disconnectedSessions: Map authProviders: Map cleanupRegistered: boolean