diff --git a/packages/@n8n/nodes-langchain/nodes/mcp/shared/__test__/utils.test.ts b/packages/@n8n/nodes-langchain/nodes/mcp/shared/__test__/utils.test.ts index 7f4060f10f8..75941d978e1 100644 --- a/packages/@n8n/nodes-langchain/nodes/mcp/shared/__test__/utils.test.ts +++ b/packages/@n8n/nodes-langchain/nodes/mcp/shared/__test__/utils.test.ts @@ -35,6 +35,9 @@ vi.mock('@n8n/ai-utilities', async () => { }); const MockedClient = Client as MockedClass; +const MockedStreamableHTTPClientTransport = StreamableHTTPClientTransport as MockedClass< + typeof StreamableHTTPClientTransport +>; const createTestEgressFilter = (): NodeEgressFilter => ({ validateUrl: vi.fn().mockResolvedValue(createResultOk(undefined)), @@ -505,10 +508,9 @@ describe('utils', () => { // Trigger the abort; the listener will call client.close (the wrapper) expect(() => abort.abort()).not.toThrow(); - await Promise.resolve(); // The original close function should have been called exactly once - expect(closeSpy).toHaveBeenCalledTimes(1); + await vi.waitFor(() => expect(closeSpy).toHaveBeenCalledTimes(1)); }); it('should remove the abort listener on normal close, preventing double-close on later abort', async () => { @@ -981,6 +983,96 @@ describe('utils', () => { expect(mockedProxyFetch).toHaveBeenCalledTimes(1); }); }); + + describe('httpStreamable session cleanup', () => { + const connectHttpClient = async (signal?: AbortSignal) => { + const result = await connectMcpClient({ + serverTransport: 'httpStreamable', + secureEgressFilter: createTestEgressFilter(), + endpointUrl: 'https://example.com', + name: 'test-client', + version: 1, + signal, + }); + + expect(result.ok).toBe(true); + if (!result.ok) throw result.error.error; + + const transport = MockedStreamableHTTPClientTransport.mock.instances.at(-1); + if (!transport) throw new Error('Expected Streamable HTTP transport'); + + return { client: result.result, transport }; + }; + + it('should terminate the session before closing the client', async () => { + const originalClose = mockClient.close; + const { client, transport } = await connectHttpClient(); + const terminateSession = vi.mocked(transport.terminateSession); + terminateSession.mockResolvedValue(undefined); + + await client.close(); + + expect(terminateSession).toHaveBeenCalledTimes(1); + expect(originalClose).toHaveBeenCalledTimes(1); + expect(terminateSession.mock.invocationCallOrder[0]).toBeLessThan( + originalClose.mock.invocationCallOrder[0], + ); + }); + + it('should close the client when session termination fails', async () => { + const originalClose = mockClient.close; + const { client, transport } = await connectHttpClient(); + vi.mocked(transport.terminateSession).mockRejectedValue(new Error('DELETE failed')); + + await expect(client.close()).resolves.toBeUndefined(); + expect(originalClose).toHaveBeenCalledTimes(1); + }); + + it('should close the client when session termination times out', async () => { + vi.useFakeTimers(); + try { + const originalClose = mockClient.close; + const { client, transport } = await connectHttpClient(); + vi.mocked(transport.terminateSession).mockReturnValue(new Promise(() => {})); + + const closePromise = client.close(); + await vi.runAllTimersAsync(); + await closePromise; + + expect(originalClose).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it('should terminate the session when execution is aborted', async () => { + const abort = new AbortController(); + const originalClose = mockClient.close; + const { transport } = await connectHttpClient(abort.signal); + const terminateSession = vi.mocked(transport.terminateSession); + terminateSession.mockResolvedValue(undefined); + + abort.abort(); + + await vi.waitFor(() => expect(originalClose).toHaveBeenCalledTimes(1)); + expect(terminateSession).toHaveBeenCalledTimes(1); + expect(terminateSession.mock.invocationCallOrder[0]).toBeLessThan( + originalClose.mock.invocationCallOrder[0], + ); + }); + + it('should terminate and close only once', async () => { + const originalClose = mockClient.close; + const { client, transport } = await connectHttpClient(); + const terminateSession = vi.mocked(transport.terminateSession); + terminateSession.mockResolvedValue(undefined); + + await Promise.all([client.close(), client.close()]); + + expect(terminateSession).toHaveBeenCalledTimes(1); + expect(originalClose).toHaveBeenCalledTimes(1); + }); + }); }); describe('mapToNodeOperationError', () => { diff --git a/packages/@n8n/nodes-langchain/nodes/mcp/shared/runtime.test.ts b/packages/@n8n/nodes-langchain/nodes/mcp/shared/runtime.test.ts index 9610c41c545..e2bb9d3aa21 100644 --- a/packages/@n8n/nodes-langchain/nodes/mcp/shared/runtime.test.ts +++ b/packages/@n8n/nodes-langchain/nodes/mcp/shared/runtime.test.ts @@ -606,8 +606,8 @@ describe('runtime', () => { expect(manager.size).toBe(1); fire(); - expect(close).toHaveBeenCalledTimes(1); expect(manager.size).toBe(0); + await vi.waitFor(() => expect(close).toHaveBeenCalledTimes(1)); }, ); }); diff --git a/packages/@n8n/nodes-langchain/nodes/mcp/shared/utils.ts b/packages/@n8n/nodes-langchain/nodes/mcp/shared/utils.ts index 1fcaac3cc53..850152e01d5 100644 --- a/packages/@n8n/nodes-langchain/nodes/mcp/shared/utils.ts +++ b/packages/@n8n/nodes-langchain/nodes/mcp/shared/utils.ts @@ -31,6 +31,8 @@ import { type McpTool, } from './types'; +const MCP_SESSION_TERMINATION_TIMEOUT_MS = 1_000; + export async function getAllTools(client: Client, cursor?: string): Promise { const { tools, nextCursor } = await client.listTools({ cursor }); @@ -214,6 +216,26 @@ export async function connectMcpClient({ fetch: authFetch, ...(signal ? { requestInit: { signal } } : {}), }); + const originalClose = client.close.bind(client); + let closePromise: Promise | undefined; + client.close = async () => { + closePromise ??= (async () => { + let timeout: ReturnType | undefined; + try { + await Promise.race([ + Promise.resolve(transport.terminateSession()).catch(() => {}), + new Promise((resolve) => { + timeout = setTimeout(resolve, MCP_SESSION_TERMINATION_TIMEOUT_MS); + }), + ]); + } finally { + if (timeout !== undefined) clearTimeout(timeout); + await originalClose(); + } + })(); + + await closePromise; + }; await client.connect(transport); return createResultOk(client); } catch (error) {