Fix test suite for refactored shttp and redisTransport

- Update tests to handle async session initialization
- Fix session ID retrieval from SSE/JSON responses
- Update DELETE request tests to match new behavior
- Remove references to deprecated startServerListeningToRedis
- Add proper waits for async operations

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

src/services/redisTransport.integration.test.ts

@@ -2,9 +2,10 @@ import { jest } from '@jest/globals';
 import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
 import { MockRedisClient, setRedisClient } from '../redis.js';
 import { 
-  startServerListeningToRedis,
+  ServerRedisTransport,
   redisRelayToMcpServer,
-  shutdownSession
+  shutdownSession,
+  setSessionOwner
 } from './redisTransport.js';
 import { createMcpServer } from './mcp.js';
 import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
@@ -27,8 +28,10 @@ describe('Redis Transport Integration', () => {
 
     it('should relay initialization request from client to server through Redis', async () => {
       // 1. Start the server listening to Redis
-      const serverWrapper = createMcpServer();
-      const serverTransport = await startServerListeningToRedis(serverWrapper, sessionId);
+      const { server, cleanup: serverCleanup } = createMcpServer();
+      const serverTransport = new ServerRedisTransport(sessionId);
+      serverTransport.onclose = serverCleanup;
+      await server.connect(serverTransport);
 
       // 2. Create a mock client transport (simulating the streamable HTTP client side)
       const mockClientTransport: Transport = {
@@ -121,13 +124,20 @@ describe('Redis Transport Integration', () => {
       // Cleanup
       await cleanup();
       await shutdownSession(sessionId);
+      serverCleanup(); // Clean up MCP server intervals
+      
+      // Ensure server transport is closed
+      await serverTransport.close();
+      
       await new Promise(resolve => setTimeout(resolve, 10));
     });
 
     it('should handle tools/list request through Redis relay', async () => {
       // Set up server and mock client
-      const serverWrapper = createMcpServer();
-      const serverTransport = await startServerListeningToRedis(serverWrapper, sessionId);
+      const { server, cleanup: serverCleanup } = createMcpServer();
+      const serverTransport = new ServerRedisTransport(sessionId);
+      serverTransport.onclose = serverCleanup;
+      await server.connect(serverTransport);
       
       const mockClientTransport: Transport = {
         onmessage: undefined,
@@ -168,13 +178,20 @@ describe('Redis Transport Integration', () => {
       // Cleanup
       await cleanup();
       await shutdownSession(sessionId);
+      serverCleanup(); // Clean up MCP server intervals
+      
+      // Ensure server transport is closed
+      await serverTransport.close();
+      
       await new Promise(resolve => setTimeout(resolve, 10));
     });
 
     it('should handle notifications through Redis relay', async () => {
       // Set up server and mock client
-      const serverWrapper = createMcpServer();
-      const serverTransport = await startServerListeningToRedis(serverWrapper, sessionId);
+      const { server, cleanup: serverCleanup } = createMcpServer();
+      const serverTransport = new ServerRedisTransport(sessionId);
+      serverTransport.onclose = serverCleanup;
+      await server.connect(serverTransport);
       
       const mockClientTransport: Transport = {
         onmessage: undefined,
@@ -187,6 +204,17 @@ describe('Redis Transport Integration', () => {
 
       const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport);
 
+      // Set up notification subscription manually since notifications don't have an id
+      const notificationChannel = `mcp:shttp:toclient:${sessionId}:__GET_stream`;
+      const notificationCleanup = await mockRedis.createSubscription(notificationChannel, async (redisMessageJson) => {
+        const redisMessage = JSON.parse(redisMessageJson);
+        if (redisMessage.type === 'mcp') {
+          await mockClientTransport.send(redisMessage.message, redisMessage.options);
+        }
+      }, (error) => {
+        mockClientTransport.onerror?.(error);
+      });
+
       // Send a notification from server (notifications don't have an id)
       const notification: JSONRPCMessage = {
         jsonrpc: '2.0',
@@ -216,9 +244,17 @@ describe('Redis Transport Integration', () => {
         undefined
       );
 
+      // Cleanup notification subscription
+      await notificationCleanup();
+
       // Cleanup
       await cleanup();
       await shutdownSession(sessionId);
+      serverCleanup(); // Clean up MCP server intervals
+      
+      // Ensure server transport is closed
+      await serverTransport.close();
+      
       await new Promise(resolve => setTimeout(resolve, 10));
     });
 

src/services/redisTransport.test.ts

@@ -5,8 +5,11 @@ import {
   ServerRedisTransport, 
   redisRelayToMcpServer,
   isLive,
-  startServerListeningToRedis,
-  shutdownSession
+  shutdownSession,
+  setSessionOwner,
+  getSessionOwner,
+  validateSessionOwnership,
+  isSessionOwnedBy
 } from './redisTransport.js';
 import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
 import { createMcpServer } from './mcp.js';
@@ -278,36 +281,45 @@ describe('Redis Transport', () => {
     });
   });
 
-  describe('startServerListeningToRedis', () => {
-    const sessionId = 'test-session-server';
+  describe('Session Ownership', () => {
+    const sessionId = 'test-session-ownership';
+    const userId = 'test-user-123';
 
-    it('should connect server with ServerRedisTransport', async () => {
-      const serverWrapper = createMcpServer();
-      const connectSpy = jest.spyOn(serverWrapper.server, 'connect');
-
-      await startServerListeningToRedis(serverWrapper, sessionId);
-
-      expect(connectSpy).toHaveBeenCalledWith(
-        expect.any(ServerRedisTransport)
-      );
-
-      // Clean up properly
-      await shutdownSession(sessionId);
-      await new Promise(resolve => setTimeout(resolve, 10));
+    it('should set and get session owner', async () => {
+      await setSessionOwner(sessionId, userId);
+      const owner = await getSessionOwner(sessionId);
+      expect(owner).toBe(userId);
     });
 
-    it('should create transport that can send responses and shutdown properly', async () => {
-      const serverWrapper = createMcpServer();
+    it('should validate session ownership correctly', async () => {
+      await setSessionOwner(sessionId, userId);
       
-      const transport = await startServerListeningToRedis(serverWrapper, sessionId);
+      expect(await validateSessionOwnership(sessionId, userId)).toBe(true);
+      expect(await validateSessionOwnership(sessionId, 'different-user')).toBe(false);
+    });
 
-      // The server should now be connected and able to handle requests via Redis
-      expect(serverWrapper.server).toBeDefined();
-      expect(transport).toBeInstanceOf(ServerRedisTransport);
-
-      // Test proper shutdown
-      await shutdownSession(sessionId);
-      await new Promise(resolve => setTimeout(resolve, 10));
+    it('should check if session is owned by user including liveness', async () => {
+      // Session not live yet
+      expect(await isSessionOwnedBy(sessionId, userId)).toBe(false);
+      
+      // Make session live
+      await mockRedis.createSubscription(
+        `mcp:shttp:toserver:${sessionId}`,
+        jest.fn(),
+        jest.fn()
+      );
+      
+      // Still false because no owner set
+      expect(await isSessionOwnedBy(sessionId, userId)).toBe(false);
+      
+      // Set owner
+      await setSessionOwner(sessionId, userId);
+      
+      // Now should be true
+      expect(await isSessionOwnedBy(sessionId, userId)).toBe(true);
+      
+      // False for different user
+      expect(await isSessionOwnedBy(sessionId, 'different-user')).toBe(false);
     });
   });