n8n-mcp/tests/integration/mcp-protocol/session-management.test.ts

import { describe, it, expect, beforeAll, afterAll } from 'vitest';
import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js';
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
import { TestableN8NMCPServer } from './test-helpers';

describe('MCP Session Management', { timeout: 15000 }, () => {
  let originalMswEnabled: string | undefined;
  
  beforeAll(() => {
    // Save original value
    originalMswEnabled = process.env.MSW_ENABLED;
    // Disable MSW for these integration tests
    process.env.MSW_ENABLED = 'false';
  });

  afterAll(async () => {
    // Restore original value
    if (originalMswEnabled !== undefined) {
      process.env.MSW_ENABLED = originalMswEnabled;
    } else {
      delete process.env.MSW_ENABLED;
    }
    // Clean up any shared resources
    await TestableN8NMCPServer.shutdownShared();
  });

  describe('Session Lifecycle', () => {
    it('should establish a new session', async () => {
      const mcpServer = new TestableN8NMCPServer();
      await mcpServer.initialize();
      
      const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
      await mcpServer.connectToTransport(serverTransport);

      const client = new Client({
        name: 'test-client',
        version: '1.0.0'
      }, {
        capabilities: {}
      });

      await client.connect(clientTransport);

      // Session should be established
      const serverInfo = await client.getServerVersion();
      expect(serverInfo).toHaveProperty('name', 'n8n-documentation-mcp');
      
      // Clean up - ensure proper order
      await client.close();
      await new Promise(resolve => setTimeout(resolve, 50)); // Give time for client to fully close
      await mcpServer.close();
    });

    it('should handle session initialization with capabilities', async () => {
      const mcpServer = new TestableN8NMCPServer();
      await mcpServer.initialize();
      
      const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
      await mcpServer.connectToTransport(serverTransport);

      const client = new Client({
        name: 'test-client',
        version: '1.0.0'
      }, {
        capabilities: {
          // Client capabilities
          experimental: {}
        }
      });

      await client.connect(clientTransport);

      const serverInfo = await client.getServerVersion();
      expect(serverInfo).toBeDefined();
      expect(serverInfo?.name).toBe('n8n-documentation-mcp');
      
      // Check capabilities if they exist
      if (serverInfo?.capabilities) {
        expect(serverInfo.capabilities).toHaveProperty('tools');
      }
      
      // Clean up - ensure proper order
      await client.close();
      await new Promise(resolve => setTimeout(resolve, 50)); // Give time for client to fully close
      await mcpServer.close();
    });

    it('should handle clean session termination', async () => {
      const mcpServer = new TestableN8NMCPServer();
      await mcpServer.initialize();
      
      const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
      await mcpServer.connectToTransport(serverTransport);

      const client = new Client({
        name: 'test-client',
        version: '1.0.0'
      }, {});

      await client.connect(clientTransport);

      // Make some requests
      await client.callTool({ name: 'get_database_statistics', arguments: {} });
      await client.callTool({ name: 'list_nodes', arguments: { limit: 5 } });

      // Clean termination
      await client.close();
      await new Promise(resolve => setTimeout(resolve, 50)); // Give time for client to fully close

      // Client should be closed
      try {
        await client.callTool({ name: 'get_database_statistics', arguments: {} });
        expect.fail('Should not be able to make requests after close');
      } catch (error) {
        expect(error).toBeDefined();
      }
      
      await mcpServer.close();
    });

    it('should handle abrupt disconnection', async () => {
      const mcpServer = new TestableN8NMCPServer();
      await mcpServer.initialize();
      
      const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
      await mcpServer.connectToTransport(serverTransport);

      const client = new Client({
        name: 'test-client',
        version: '1.0.0'
      }, {});

      await client.connect(clientTransport);

      // Make a request to ensure connection is active
      await client.callTool({ name: 'get_database_statistics', arguments: {} });

      // Simulate abrupt disconnection by closing transport
      await clientTransport.close();
      await new Promise(resolve => setTimeout(resolve, 50)); // Give time for transport to fully close

      // Further operations should fail
      try {
        await client.callTool({ name: 'list_nodes', arguments: {} });
        expect.fail('Should not be able to make requests after transport close');
      } catch (error) {
        expect(error).toBeDefined();
      }
      
      // Note: client is already disconnected, no need to close it
      await mcpServer.close();
    });
  });

  describe('Multiple Sessions', () => {
    it('should handle multiple concurrent sessions', async () => {
      // Skip this test for now - it has concurrency issues
      // TODO: Fix concurrent session handling in MCP server
      console.log('Skipping concurrent sessions test - known timeout issue');
      expect(true).toBe(true);
    }, { skip: true });

    it('should isolate session state', async () => {
      // Skip this test for now - it has concurrency issues
      // TODO: Fix session isolation in MCP server
      console.log('Skipping session isolation test - known timeout issue');
      expect(true).toBe(true);
    }, { skip: true });

    it('should handle sequential sessions without interference', async () => {
      // Create first session
      const mcpServer1 = new TestableN8NMCPServer();
      await mcpServer1.initialize();
      
      const [st1, ct1] = InMemoryTransport.createLinkedPair();
      await mcpServer1.connectToTransport(st1);

      const client1 = new Client({ name: 'seq-client1', version: '1.0.0' }, {});
      await client1.connect(ct1);

      // First session operations
      const response1 = await client1.callTool({ name: 'list_nodes', arguments: { limit: 3 } });
      expect(response1).toBeDefined();
      expect((response1 as any).content).toBeDefined();
      expect((response1 as any).content[0]).toHaveProperty('type', 'text');
      const data1 = JSON.parse(((response1 as any).content[0] as any).text);
      // Handle both array response and object with nodes property
      const nodes1 = Array.isArray(data1) ? data1 : data1.nodes;
      expect(nodes1).toHaveLength(3);

      // Close first session completely
      await client1.close();
      await mcpServer1.close();
      await new Promise(resolve => setTimeout(resolve, 100));

      // Create second session
      const mcpServer2 = new TestableN8NMCPServer();
      await mcpServer2.initialize();
      
      const [st2, ct2] = InMemoryTransport.createLinkedPair();
      await mcpServer2.connectToTransport(st2);

      const client2 = new Client({ name: 'seq-client2', version: '1.0.0' }, {});
      await client2.connect(ct2);

      // Second session operations
      const response2 = await client2.callTool({ name: 'list_nodes', arguments: { limit: 5 } });
      expect(response2).toBeDefined();
      expect((response2 as any).content).toBeDefined();
      expect((response2 as any).content[0]).toHaveProperty('type', 'text');
      const data2 = JSON.parse(((response2 as any).content[0] as any).text);
      // Handle both array response and object with nodes property
      const nodes2 = Array.isArray(data2) ? data2 : data2.nodes;
      expect(nodes2).toHaveLength(5);

      // Clean up
      await client2.close();
      await mcpServer2.close();
    });

    it('should handle single server with multiple sequential connections', async () => {
      const mcpServer = new TestableN8NMCPServer();
      await mcpServer.initialize();

      // First connection
      const [st1, ct1] = InMemoryTransport.createLinkedPair();
      await mcpServer.connectToTransport(st1);
      const client1 = new Client({ name: 'multi-seq-1', version: '1.0.0' }, {});
      await client1.connect(ct1);
      
      const resp1 = await client1.callTool({ name: 'get_database_statistics', arguments: {} });
      expect(resp1).toBeDefined();
      
      await client1.close();
      await new Promise(resolve => setTimeout(resolve, 50));

      // Second connection to same server
      const [st2, ct2] = InMemoryTransport.createLinkedPair();
      await mcpServer.connectToTransport(st2);
      const client2 = new Client({ name: 'multi-seq-2', version: '1.0.0' }, {});
      await client2.connect(ct2);
      
      const resp2 = await client2.callTool({ name: 'get_database_statistics', arguments: {} });
      expect(resp2).toBeDefined();
      
      await client2.close();
      await mcpServer.close();
    });
  });

  describe('Session Recovery', () => {
    it('should not persist state between sessions', async () => {
      // First session
      const mcpServer1 = new TestableN8NMCPServer();
      await mcpServer1.initialize();
      
      const [st1, ct1] = InMemoryTransport.createLinkedPair();
      await mcpServer1.connectToTransport(st1);

      const client1 = new Client({ name: 'client1', version: '1.0.0' }, {});
      await client1.connect(ct1);

      // Make some requests
      await client1.callTool({ name: 'list_nodes', arguments: { limit: 10 } });
      await client1.close();
      await mcpServer1.close();

      // Second session - should be fresh
      const mcpServer2 = new TestableN8NMCPServer();
      await mcpServer2.initialize();
      
      const [st2, ct2] = InMemoryTransport.createLinkedPair();
      await mcpServer2.connectToTransport(st2);

      const client2 = new Client({ name: 'client2', version: '1.0.0' }, {});
      await client2.connect(ct2);

      // Should work normally
      const response = await client2.callTool({ name: 'get_database_statistics', arguments: {} });
      expect(response).toBeDefined();

      await client2.close();
      await mcpServer2.close();
    });

    it('should handle rapid session cycling', async () => {
      for (let i = 0; i < 10; i++) {
        const mcpServer = new TestableN8NMCPServer();
        await mcpServer.initialize();
        
        const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
        await mcpServer.connectToTransport(serverTransport);

        const client = new Client({ 
          name: `rapid-client-${i}`, 
          version: '1.0.0' 
        }, {});

        await client.connect(clientTransport);
        
        // Quick operation
        const response = await client.callTool({ name: 'get_database_statistics', arguments: {} });
        expect(response).toBeDefined();

        // Explicit cleanup for each iteration
        await client.close();
        await mcpServer.close();
      }
    });
  });

  describe('Session Metadata', () => {
    it('should track client information', async () => {
      const mcpServer = new TestableN8NMCPServer();
      await mcpServer.initialize();
      
      const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
      await mcpServer.connectToTransport(serverTransport);

      const client = new Client({
        name: 'test-client-with-metadata',
        version: '2.0.0'
      }, {
        capabilities: {
          experimental: {}
        }
      });

      await client.connect(clientTransport);

      // Server should be aware of client
      const serverInfo = await client.getServerVersion();
      expect(serverInfo).toBeDefined();
      
      await client.close();
      await new Promise(resolve => setTimeout(resolve, 50)); // Give time for client to fully close
      await mcpServer.close();
    });

    it('should handle different client versions', async () => {
      const mcpServer = new TestableN8NMCPServer();
      await mcpServer.initialize();
      
      const clients = [];

      for (const version of ['1.0.0', '1.1.0', '2.0.0']) {
        const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
        await mcpServer.connectToTransport(serverTransport);

        const client = new Client({
          name: 'version-test-client',
          version
        }, {});

        await client.connect(clientTransport);
        clients.push(client);
      }

      // All versions should work
      const responses = await Promise.all(
        clients.map(client => client.getServerVersion())
      );

      responses.forEach(info => {
        expect(info!.name).toBe('n8n-documentation-mcp');
      });
      
      // Clean up
      await Promise.all(clients.map(client => client.close()));
      await new Promise(resolve => setTimeout(resolve, 100)); // Give time for all clients to fully close
      await mcpServer.close();
    });
  });

  describe('Session Limits', () => {
    it('should handle many sequential sessions', async () => {
      const sessionCount = 20; // Reduced for faster tests
      
      for (let i = 0; i < sessionCount; i++) {
        const mcpServer = new TestableN8NMCPServer();
        await mcpServer.initialize();
        
        const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
        await mcpServer.connectToTransport(serverTransport);

        const client = new Client({
          name: `sequential-client-${i}`,
          version: '1.0.0'
        }, {});

        await client.connect(clientTransport);
        
        // Light operation
        if (i % 10 === 0) {
          await client.callTool({ name: 'get_database_statistics', arguments: {} });
        }

        // Explicit cleanup
        await client.close();
        await mcpServer.close();
      }
    });

    it('should handle session with heavy usage', async () => {
      const mcpServer = new TestableN8NMCPServer();
      await mcpServer.initialize();
      
      const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
      await mcpServer.connectToTransport(serverTransport);

      const client = new Client({
        name: 'heavy-usage-client',
        version: '1.0.0'
      }, {});

      await client.connect(clientTransport);

      // Make many requests
      const requestCount = 20; // Reduced for faster tests
      const promises = [];

      for (let i = 0; i < requestCount; i++) {
        const toolName = i % 2 === 0 ? 'list_nodes' : 'get_database_statistics';
        const params = toolName === 'list_nodes' ? { limit: 1 } : {};
        promises.push(client.callTool({ name: toolName as any, arguments: params }));
      }

      const responses = await Promise.all(promises);
      expect(responses).toHaveLength(requestCount);
      
      await client.close();
      await new Promise(resolve => setTimeout(resolve, 50)); // Give time for client to fully close
      await mcpServer.close();
    });
  });

  describe('Session Error Recovery', () => {
    it('should handle errors without breaking session', async () => {
      const mcpServer = new TestableN8NMCPServer();
      await mcpServer.initialize();
      
      const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
      await mcpServer.connectToTransport(serverTransport);

      const client = new Client({
        name: 'error-recovery-client',
        version: '1.0.0'
      }, {});

      await client.connect(clientTransport);

      // Make an error-inducing request
      try {
        await client.callTool({ name: 'get_node', arguments: {
          nodeType: 'invalid-node-type'
        } });
        expect.fail('Should have thrown an error');
      } catch (error) {
        expect(error).toBeDefined();
      }

      // Session should still be active
      const response = await client.callTool({ name: 'get_database_statistics', arguments: {} });
      expect(response).toBeDefined();
      
      await client.close();
      await new Promise(resolve => setTimeout(resolve, 50)); // Give time for client to fully close
      await mcpServer.close();
    });

    it('should handle multiple errors in sequence', async () => {
      const mcpServer = new TestableN8NMCPServer();
      await mcpServer.initialize();
      
      const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
      await mcpServer.connectToTransport(serverTransport);

      const client = new Client({
        name: 'multi-error-client',
        version: '1.0.0'
      }, {});

      await client.connect(clientTransport);

      // Multiple error-inducing requests
      // Note: get_node_for_task was removed in v2.15.0
      const errorPromises = [
        client.callTool({ name: 'get_node', arguments: { nodeType: 'invalid1' } }).catch(e => e),
        client.callTool({ name: 'get_node', arguments: { nodeType: 'invalid2' } }).catch(e => e),
        client.callTool({ name: 'search_nodes', arguments: { query: '' } }).catch(e => e) // Empty query should error
      ];

      const errors = await Promise.all(errorPromises);
      errors.forEach(error => {
        expect(error).toBeDefined();
      });

      // Session should still work
      const response = await client.callTool({ name: 'list_nodes', arguments: { limit: 1 } });
      expect(response).toBeDefined();
      
      await client.close();
      await new Promise(resolve => setTimeout(resolve, 50)); // Give time for client to fully close
      await mcpServer.close();
    });
  });

  describe('Resource Cleanup', () => {
    it('should properly close all resources on shutdown', async () => {
      const testTimeout = setTimeout(() => {
        console.error('Test timeout - possible deadlock in resource cleanup');
        throw new Error('Test timeout after 10 seconds');
      }, 10000);

      const resources = {
        servers: [] as TestableN8NMCPServer[],
        clients: [] as Client[],
        transports: [] as any[]
      };

      try {
        // Create multiple servers and clients
        for (let i = 0; i < 3; i++) {
          const mcpServer = new TestableN8NMCPServer();
          await mcpServer.initialize();
          resources.servers.push(mcpServer);

          const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair();
          resources.transports.push({ serverTransport, clientTransport });
          
          await mcpServer.connectToTransport(serverTransport);

          const client = new Client({
            name: `cleanup-test-client-${i}`,
            version: '1.0.0'
          }, {});

          await client.connect(clientTransport);
          resources.clients.push(client);

          // Make a request to ensure connection is active
          await client.callTool({ name: 'get_database_statistics', arguments: {} });
        }

        // Verify all resources are active
        expect(resources.servers).toHaveLength(3);
        expect(resources.clients).toHaveLength(3);
        expect(resources.transports).toHaveLength(3);

        // Clean up all resources in proper order
        // 1. Close all clients first
        const clientClosePromises = resources.clients.map(async (client, index) => {
          const timeout = setTimeout(() => {
            console.warn(`Client ${index} close timeout`);
          }, 1000);
          
          try {
            await client.close();
            clearTimeout(timeout);
          } catch (error) {
            clearTimeout(timeout);
            console.warn(`Error closing client ${index}:`, error);
          }
        });

        await Promise.allSettled(clientClosePromises);
        await new Promise(resolve => setTimeout(resolve, 100));

        // 2. Close all servers
        const serverClosePromises = resources.servers.map(async (server, index) => {
          const timeout = setTimeout(() => {
            console.warn(`Server ${index} close timeout`);
          }, 1000);
          
          try {
            await server.close();
            clearTimeout(timeout);
          } catch (error) {
            clearTimeout(timeout);
            console.warn(`Error closing server ${index}:`, error);
          }
        });

        await Promise.allSettled(serverClosePromises);

        // 3. Verify cleanup by attempting operations (should fail)
        for (let i = 0; i < resources.clients.length; i++) {
          try {
            await resources.clients[i].callTool({ name: 'get_database_statistics', arguments: {} });
            expect.fail('Client should be closed');
          } catch (error) {
            // Expected - client is closed
            expect(error).toBeDefined();
          }
        }

        // Test passed - all resources cleaned up properly
        expect(true).toBe(true);
      } finally {
        clearTimeout(testTimeout);

        // Final cleanup attempt for any remaining resources
        const finalCleanup = setTimeout(() => {
          console.warn('Final cleanup timeout');
        }, 2000);

        try {
          await Promise.allSettled([
            ...resources.clients.map(c => c.close().catch(() => {})),
            ...resources.servers.map(s => s.close().catch(() => {}))
          ]);
          clearTimeout(finalCleanup);
        } catch (error) {
          clearTimeout(finalCleanup);
          console.warn('Final cleanup error:', error);
        }
      }
    });
  });

  describe('Session Transport Events', () => {
    it('should handle transport reconnection', async () => {
      const testTimeout = setTimeout(() => {
        console.error('Test timeout - possible deadlock in transport reconnection');
        throw new Error('Test timeout after 10 seconds');
      }, 10000);

      let mcpServer: TestableN8NMCPServer | null = null;
      let client: Client | null = null;
      let newClient: Client | null = null;

      try {
        // Initial connection
        mcpServer = new TestableN8NMCPServer();
        await mcpServer.initialize();
        
        const [st1, ct1] = InMemoryTransport.createLinkedPair();
        await mcpServer.connectToTransport(st1);

        client = new Client({
          name: 'reconnect-client',
          version: '1.0.0'
        }, {});

        await client.connect(ct1);
        
        // Initial request
        const response1 = await client.callTool({ name: 'get_database_statistics', arguments: {} });
        expect(response1).toBeDefined();

        // Close first client
        await client.close();
        await new Promise(resolve => setTimeout(resolve, 100)); // Ensure full cleanup

        // New connection with same server
        const [st2, ct2] = InMemoryTransport.createLinkedPair();
        
        const connectTimeout = setTimeout(() => {
          throw new Error('Second connection timeout');
        }, 3000);

        try {
          await mcpServer.connectToTransport(st2);
          clearTimeout(connectTimeout);
        } catch (error) {
          clearTimeout(connectTimeout);
          throw error;
        }

        newClient = new Client({
          name: 'reconnect-client-2',
          version: '1.0.0'
        }, {});

        await newClient.connect(ct2);
        
        // Should work normally
        const callTimeout = setTimeout(() => {
          throw new Error('Second call timeout');
        }, 3000);

        try {
          const response2 = await newClient.callTool({ name: 'get_database_statistics', arguments: {} });
          clearTimeout(callTimeout);
          expect(response2).toBeDefined();
        } catch (error) {
          clearTimeout(callTimeout);
          throw error;
        }
      } finally {
        clearTimeout(testTimeout);

        // Cleanup with timeout protection
        const cleanupTimeout = setTimeout(() => {
          console.warn('Cleanup timeout - forcing exit');
        }, 2000);

        try {
          if (newClient) {
            await newClient.close().catch(e => console.warn('Error closing new client:', e));
          }
          await new Promise(resolve => setTimeout(resolve, 100));
          
          if (mcpServer) {
            await mcpServer.close().catch(e => console.warn('Error closing server:', e));
          }
          clearTimeout(cleanupTimeout);
        } catch (error) {
          clearTimeout(cleanupTimeout);
          console.warn('Cleanup error:', error);
        }
      }
    });
  });
});