Fix Redis transport refactoring with control message system

- Fix timing issue by awaiting startServerListeningToRedis
- Add request-id based response channel routing
- Implement control message system for proper server shutdown
- Add comprehensive integration tests
- Ensure response subscription is established before message sending
- Fix MCP server cleanup integration to prevent async hanging

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

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

src/handlers/shttp.ts

@@ -44,7 +44,7 @@ export async function handleStreamableHTTP(req: Request, res: Response) {
       const server = createMcpServer();
       console.log('Created MCP server instance');
       
-      startServerListeningToRedis(server.server, sessionId)
+      await startServerListeningToRedis(server, sessionId)
       console.log('Started server listening to Redis for session:', sessionId);
       
       ({ shttpTransport, cleanup } = await getFirstShttpTransport(sessionId));

src/redis.ts

@@ -104,18 +104,34 @@ export class RedisClientImpl implements RedisClient {
     onMessage: (message: string) => void,
     onError: (error: Error) => void,
   ): Promise<() => Promise<void>> {
+    console.log(`[RedisClient] Creating subscription for channel: ${channel}`);
     const subscriber = this.redis.duplicate();
-    subscriber.on("error", onError);
+    subscriber.on("error", (error) => {
+      console.error(`[RedisClient] Subscriber error for channel ${channel}:`, error);
+      onError(error);
+    });
+    
+    console.log(`[RedisClient] Connecting subscriber for channel: ${channel}`);
     await subscriber.connect();
-    await subscriber.subscribe(channel, onMessage);
+    console.log(`[RedisClient] Connected subscriber for channel: ${channel}`);
+    
+    console.log(`[RedisClient] Subscribing to channel: ${channel}`);
+    await subscriber.subscribe(channel, (message) => {
+      console.log(`[RedisClient] Received message on channel ${channel}:`, message.substring(0, 100));
+      onMessage(message);
+    });
+    console.log(`[RedisClient] Successfully subscribed to channel: ${channel}`);
 
     return async () => {
+      console.log(`[RedisClient] Cleaning up subscription for channel: ${channel}`);
       await subscriber.disconnect();
     };
   }
 
   async publish(channel: string, message: string): Promise<void> {
+    console.log(`[RedisClient] Publishing to channel: ${channel}, message length: ${message.length}`);
     await this.redis.publish(channel, message);
+    console.log(`[RedisClient] Successfully published to channel: ${channel}`);
   }
 
   async exists(key: string): Promise<boolean> {

src/services/auth.test.ts

@@ -153,10 +153,16 @@ describe("auth utils", () => {
       const first = await exchangeToken(authCode);
       expect(first).toBeDefined();
       
+      // Mock console.error to suppress expected error message
+      const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation(() => {});
+      
       // Second exchange throws
       await expect(exchangeToken(authCode)).rejects.toThrow(
         "Duplicate use of authorization code detected"
       );
+      
+      // Restore console.error
+      consoleErrorSpy.mockRestore();
     });
     
     it("returns undefined for non-existent code", async () => {

src/services/auth.ts

@@ -243,7 +243,7 @@ export async function exchangeToken(
 
   const tokenExchange: TokenExchange = JSON.parse(decoded);
   if (tokenExchange.alreadyUsed) {
-    console.error("1Duplicate use of authorization code detected; revoking tokens");
+    console.error("Duplicate use of authorization code detected; revoking tokens");
     await revokeMcpInstallation(tokenExchange.mcpAccessToken);
     throw new Error("Duplicate use of authorization code detected; tokens revoked");
   }

src/services/mcp.ts

@@ -95,7 +95,12 @@ enum PromptName {
   RESOURCE = "resource_prompt",
 }
 
-export const createMcpServer = () => {
+interface McpServerWrapper {
+  server: Server;
+  cleanup: () => void;
+}
+
+export const createMcpServer = (): McpServerWrapper => {
   const server = new Server(
     {
       name: "example-servers/everything",

src/services/redisTransport.integration.test.ts

@@ -0,0 +1,256 @@
+import { jest } from '@jest/globals';
+import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
+import { MockRedisClient, setRedisClient } from '../redis.js';
+import { 
+  startServerListeningToRedis,
+  redisRelayToMcpServer,
+  shutdownSession
+} from './redisTransport.js';
+import { createMcpServer } from './mcp.js';
+import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
+
+describe('Redis Transport Integration', () => {
+  let mockRedis: MockRedisClient;
+
+  beforeEach(() => {
+    mockRedis = new MockRedisClient();
+    setRedisClient(mockRedis);
+    jest.resetAllMocks();
+  });
+
+  afterEach(() => {
+    mockRedis.clear();
+  });
+
+  describe('MCP Initialization Flow', () => {
+    const sessionId = 'test-init-session';
+
+    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);
+
+      // 2. Create a mock client transport (simulating the streamable HTTP client side)
+      const mockClientTransport: Transport = {
+        onmessage: undefined,
+        onclose: undefined,
+        onerror: undefined,
+        send: jest.fn(() => Promise.resolve()),
+        close: jest.fn(() => Promise.resolve()),
+        start: jest.fn(() => Promise.resolve())
+      };
+
+      // 3. Set up the Redis relay (this is what happens in the HTTP handler)
+      const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport);
+
+      // Track messages received by server
+      const serverReceivedMessages: JSONRPCMessage[] = [];
+      const originalServerOnMessage = serverTransport.onmessage;
+      serverTransport.onmessage = (message, extra) => {
+        serverReceivedMessages.push(message);
+        originalServerOnMessage?.(message, extra);
+      };
+
+      // 4. Simulate client sending initialization request
+      const initMessage: JSONRPCMessage = {
+        jsonrpc: '2.0',
+        id: 'init-1',
+        method: 'initialize',
+        params: {
+          protocolVersion: '2024-11-05',
+          capabilities: {},
+          clientInfo: {
+            name: 'test-client',
+            version: '1.0.0'
+          }
+        }
+      };
+
+      // Trigger the client transport onmessage (simulates HTTP request)
+      mockClientTransport.onmessage?.(initMessage);
+
+      // Wait for message to be relayed through Redis
+      await new Promise(resolve => setTimeout(resolve, 50));
+
+      // 5. Verify server received the init message
+      expect(serverReceivedMessages).toHaveLength(1);
+      expect(serverReceivedMessages[0]).toMatchObject({
+        jsonrpc: '2.0',
+        id: 'init-1',
+        method: 'initialize'
+      });
+
+      // 6. Simulate server responding (this should get relayed back to client)
+      const initResponse: JSONRPCMessage = {
+        jsonrpc: '2.0',
+        id: 'init-1',
+        result: {
+          protocolVersion: '2024-11-05',
+          capabilities: {
+            tools: {},
+            prompts: {},
+            resources: {}
+          },
+          serverInfo: {
+            name: 'example-server',
+            version: '1.0.0'
+          }
+        }
+      };
+
+      await serverTransport.send(initResponse, { relatedRequestId: 'init-1' });
+
+      // Wait for response to be relayed back
+      await new Promise(resolve => setTimeout(resolve, 50));
+
+      // 7. Verify client transport received the response
+      expect(mockClientTransport.send).toHaveBeenCalledWith(
+        expect.objectContaining({
+          jsonrpc: '2.0',
+          id: 'init-1',
+          result: expect.objectContaining({
+            protocolVersion: '2024-11-05',
+            serverInfo: expect.objectContaining({
+              name: 'example-server'
+            })
+          })
+        }),
+        { relatedRequestId: 'init-1' }
+      );
+
+      // Cleanup
+      await cleanup();
+      await shutdownSession(sessionId);
+      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 mockClientTransport: Transport = {
+        onmessage: undefined,
+        onclose: undefined,
+        onerror: undefined,
+        send: jest.fn(() => Promise.resolve()),
+        close: jest.fn(() => Promise.resolve()),
+        start: jest.fn(() => Promise.resolve())
+      };
+
+      const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport);
+
+      // Send tools/list request
+      const toolsListMessage: JSONRPCMessage = {
+        jsonrpc: '2.0',
+        id: 'tools-1',
+        method: 'tools/list',
+        params: {}
+      };
+
+      mockClientTransport.onmessage?.(toolsListMessage);
+
+      // Wait for processing and response
+      await new Promise(resolve => setTimeout(resolve, 100));
+
+      // Verify client received a response with tools
+      expect(mockClientTransport.send).toHaveBeenCalledWith(
+        expect.objectContaining({
+          jsonrpc: '2.0',
+          id: 'tools-1',
+          result: expect.objectContaining({
+            tools: expect.any(Array)
+          })
+        }),
+        undefined
+      );
+
+      // Cleanup
+      await cleanup();
+      await shutdownSession(sessionId);
+      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 mockClientTransport: Transport = {
+        onmessage: undefined,
+        onclose: undefined,
+        onerror: undefined,
+        send: jest.fn(() => Promise.resolve()),
+        close: jest.fn(() => Promise.resolve()),
+        start: jest.fn(() => Promise.resolve())
+      };
+
+      const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport);
+
+      // Send a notification from server (notifications don't have an id)
+      const notification: JSONRPCMessage = {
+        jsonrpc: '2.0',
+        method: 'notifications/message',
+        params: {
+          level: 'info',
+          logger: 'test',
+          data: 'Test notification'
+        }
+      };
+
+      await serverTransport.send(notification);
+
+      // Wait for notification to be delivered
+      await new Promise(resolve => setTimeout(resolve, 50));
+
+      // Verify client received the notification
+      expect(mockClientTransport.send).toHaveBeenCalledWith(
+        expect.objectContaining({
+          jsonrpc: '2.0',
+          method: 'notifications/message',
+          params: expect.objectContaining({
+            level: 'info',
+            data: 'Test notification'
+          })
+        }),
+        undefined
+      );
+
+      // Cleanup
+      await cleanup();
+      await shutdownSession(sessionId);
+      await new Promise(resolve => setTimeout(resolve, 10));
+    });
+
+    it('should not create response subscriptions for notifications', async () => {
+      const mockClientTransport: Transport = {
+        onmessage: undefined,
+        onclose: undefined,
+        onerror: undefined,
+        send: jest.fn(() => Promise.resolve()),
+        close: jest.fn(() => Promise.resolve()),
+        start: jest.fn(() => Promise.resolve())
+      };
+
+      const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport);
+
+      // Send a notification (no id field)
+      const notification: JSONRPCMessage = {
+        jsonrpc: '2.0',
+        method: 'notifications/initialized',
+        params: {}
+      };
+
+      mockClientTransport.onmessage?.(notification);
+
+      // Wait for processing
+      await new Promise(resolve => setTimeout(resolve, 50));
+
+      // Should not create any response channel subscriptions for notifications
+      // (we can't easily test this directly, but we can ensure no errors occur)
+      expect(mockClientTransport.send).not.toHaveBeenCalled();
+
+      await cleanup();
+    });
+  });
+});

src/services/redisTransport.test.ts

@@ -5,7 +5,8 @@ import {
   ServerRedisTransport, 
   redisRelayToMcpServer,
   isLive,
-  startServerListeningToRedis
+  startServerListeningToRedis,
+  shutdownSession
 } from './redisTransport.js';
 import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
 import { createMcpServer } from './mcp.js';
@@ -60,6 +61,7 @@ describe('Redis Transport', () => {
 
       expect(mockSubscriber).toHaveBeenCalledWith(
         JSON.stringify({
+          type: 'mcp',
           message: responseMessage,
           options: { relatedRequestId: 123 }
         })
@@ -84,6 +86,7 @@ describe('Redis Transport', () => {
 
       expect(mockSubscriber).toHaveBeenCalledWith(
         JSON.stringify({
+          type: 'mcp',
           message: notificationMessage,
           options: undefined
         })
@@ -98,6 +101,55 @@ describe('Redis Transport', () => {
 
       expect(onCloseMock).toHaveBeenCalled();
     });
+
+    it('should respond to shutdown control messages', async () => {
+      await transport.start();
+      
+      const onCloseMock = jest.fn();
+      transport.onclose = onCloseMock;
+
+      // Send a shutdown control message
+      await shutdownSession(sessionId);
+
+      // Wait for async processing
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      expect(onCloseMock).toHaveBeenCalled();
+    });
+
+    it('should receive MCP messages from clients and call onmessage', async () => {
+      const onMessageMock = jest.fn();
+      transport.onmessage = onMessageMock;
+
+      await transport.start();
+
+      // Simulate client sending a message to server
+      const clientMessage: JSONRPCMessage = {
+        jsonrpc: '2.0',
+        id: 'test-req',
+        method: 'tools/list',
+        params: {}
+      };
+
+      await mockRedis.publish(
+        `mcp:shttp:toserver:${sessionId}`,
+        JSON.stringify({
+          type: 'mcp',
+          message: clientMessage,
+          extra: { authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } }
+        })
+      );
+
+      // Wait for async processing
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      expect(onMessageMock).toHaveBeenCalledWith(
+        clientMessage,
+        { authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } }
+      );
+
+      await transport.close();
+    });
   });
 
 
@@ -110,9 +162,9 @@ describe('Redis Transport', () => {
         onmessage: undefined,
         onclose: undefined,
         onerror: undefined,
-        send: jest.fn(),
+        send: jest.fn(() => Promise.resolve()),
-        close: jest.fn(),
+        close: jest.fn(() => Promise.resolve()),
-        start: jest.fn()
+        start: jest.fn(() => Promise.resolve())
       };
     });
 
@@ -128,7 +180,7 @@ describe('Redis Transport', () => {
       };
 
       // Trigger the onmessage handler
-      mockTransport.onmessage?.(requestMessage, { authInfo: { userId: 'test' } });
+      mockTransport.onmessage?.(requestMessage, { authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } });
 
       // Wait a bit for async processing
       await new Promise(resolve => setTimeout(resolve, 10));
@@ -158,7 +210,7 @@ describe('Redis Transport', () => {
       };
 
       // Trigger the onmessage handler
-      mockTransport.onmessage?.(requestMessage, { authInfo: { userId: 'test' } });
+      mockTransport.onmessage?.(requestMessage, { authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } });
 
       // Wait for subscription setup
       await new Promise(resolve => setTimeout(resolve, 10));
@@ -173,6 +225,7 @@ describe('Redis Transport', () => {
       await mockRedis.publish(
         `mcp:shttp:toclient:${sessionId}:req-456`,
         JSON.stringify({
+          type: 'mcp',
           message: responseMessage,
           options: undefined
         })
@@ -229,24 +282,32 @@ describe('Redis Transport', () => {
     const sessionId = 'test-session-server';
 
     it('should connect server with ServerRedisTransport', async () => {
-      const { server } = createMcpServer();
+      const serverWrapper = createMcpServer();
-      const connectSpy = jest.spyOn(server, 'connect');
+      const connectSpy = jest.spyOn(serverWrapper.server, 'connect');
 
-      await startServerListeningToRedis(server, sessionId);
+      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 create transport that can send responses', async () => {
+    it('should create transport that can send responses and shutdown properly', async () => {
-      const { server } = createMcpServer();
+      const serverWrapper = createMcpServer();
       
-      await startServerListeningToRedis(server, sessionId);
+      const transport = await startServerListeningToRedis(serverWrapper, sessionId);
 
       // The server should now be connected and able to handle requests via Redis
-      // This is tested implicitly by the connection succeeding
+      expect(serverWrapper.server).toBeDefined();
-      expect(server).toBeDefined();
+      expect(transport).toBeInstanceOf(ServerRedisTransport);
+
+      // Test proper shutdown
+      await shutdownSession(sessionId);
+      await new Promise(resolve => setTimeout(resolve, 10));
     });
   });
 
@@ -259,9 +320,9 @@ describe('Redis Transport', () => {
         onmessage: undefined,
         onclose: undefined,
         onerror: undefined,
-        send: jest.fn(),
+        send: jest.fn(() => Promise.resolve()),
-        close: jest.fn(),
+        close: jest.fn(() => Promise.resolve()),
-        start: jest.fn()
+        start: jest.fn(() => Promise.resolve())
       };
 
       const cleanup = redisRelayToMcpServer(sessionId, clientTransport);
@@ -291,6 +352,7 @@ describe('Redis Transport', () => {
       // Verify the message was published to server channel
       expect(serverSubscriber).toHaveBeenCalledWith(
         JSON.stringify({
+          type: 'mcp',
           message: listToolsRequest,
           extra: undefined,
           options: undefined
@@ -307,6 +369,7 @@ describe('Redis Transport', () => {
       await mockRedis.publish(
         `mcp:shttp:toclient:${sessionId}:integration-req-1`,
         JSON.stringify({
+          type: 'mcp',
           message: serverResponse,
           options: undefined
         })
@@ -321,4 +384,42 @@ describe('Redis Transport', () => {
       await cleanup();
     });
   });
+
+  describe('Control Messages', () => {
+    const sessionId = 'test-control-session';
+
+    it('should send shutdown control messages', async () => {
+      const controlSubscriber = jest.fn();
+      await mockRedis.createSubscription(
+        `mcp:control:${sessionId}`,
+        controlSubscriber,
+        jest.fn()
+      );
+
+      await shutdownSession(sessionId);
+
+      const callArgs = controlSubscriber.mock.calls[0][0] as string;
+      const message = JSON.parse(callArgs);
+      
+      expect(message.type).toBe('control');
+      expect(message.action).toBe('SHUTDOWN');
+      expect(typeof message.timestamp).toBe('number');
+    });
+
+    it('should properly shutdown server transport via control message', async () => {
+      const transport = new ServerRedisTransport(sessionId);
+      const onCloseMock = jest.fn();
+      transport.onclose = onCloseMock;
+
+      await transport.start();
+
+      // Send shutdown signal
+      await shutdownSession(sessionId);
+      
+      // Wait for async processing
+      await new Promise(resolve => setTimeout(resolve, 10));
+
+      expect(onCloseMock).toHaveBeenCalled();
+    });
+  });
 });

src/services/redisTransport.ts

@@ -8,16 +8,25 @@ import { Server } from "@modelcontextprotocol/sdk/server/index.js";
 let redisTransportCounter = 0;
 const notificationStreamId = "__GET_stream";
 
-interface RedisMessage {
+// Message types for Redis transport
+type RedisMessage = 
+  | {
+      type: 'mcp';
       message: JSONRPCMessage;
       extra?: MessageExtraInfo;
       options?: TransportSendOptions;
     }
+  | {
+      type: 'control';
+      action: 'SHUTDOWN' | 'PING' | 'STATUS';
+      timestamp?: number;
+    };
 
 function sendToMcpServer(sessionId: string, message: JSONRPCMessage, extra?: { authInfo?: AuthInfo; }, options?: TransportSendOptions): Promise<void> {
   const toServerChannel = getToServerChannel(sessionId);
-  const redisMessage: RedisMessage = { message, extra, options };
+  const redisMessage: RedisMessage = { type: 'mcp', message, extra, options };
-  console.log(`[sendToServerChannel] Publishing to ${toServerChannel}:`, JSON.stringify(redisMessage).substring(0, 100));
+  console.log(`[sendToMcpServer] Publishing to Redis channel: ${toServerChannel}`);
+  console.log(`[sendToMcpServer] Message:`, JSON.stringify(message));
   return redisClient.publish(toServerChannel, JSON.stringify(redisMessage));
 }
 
@@ -29,48 +38,79 @@ function getToClientChannel(sessionId: string, relatedRequestId: string): string
   return `mcp:shttp:toclient:${sessionId}:${relatedRequestId}`;
 }
 
+function getControlChannel(sessionId: string): string {
+  return `mcp:control:${sessionId}`;
+}
+
+function sendControlMessage(sessionId: string, action: 'SHUTDOWN' | 'PING' | 'STATUS'): Promise<void> {
+  const controlChannel = getControlChannel(sessionId);
+  const redisMessage: RedisMessage = {
+    type: 'control',
+    action,
+    timestamp: Date.now()
+  };
+  return redisClient.publish(controlChannel, JSON.stringify(redisMessage));
+}
+
+export async function shutdownSession(sessionId: string): Promise<void> {
+  return sendControlMessage(sessionId, 'SHUTDOWN');
+}
+
 export async function isLive(sessionId: string): Promise<boolean> {
   // Check if the session is live by checking if the key exists in Redis
   const numSubs = await redisClient.numsub(getToServerChannel(sessionId));
-  console.log(`[isLive] Session ${sessionId}: Redis subscribers on ${getToServerChannel(sessionId)} = ${numSubs}`);
   return numSubs > 0;
 }
 
 
 export function redisRelayToMcpServer(sessionId: string, transport: Transport): () => Promise<void> {
+  console.log(`[redisRelayToMcpServer] Setting up relay for session: ${sessionId}`);
   let redisCleanup: (() => Promise<void>) | undefined = undefined;
   const cleanup = async () => {
+    console.log(`[redisRelayToMcpServer] Cleaning up relay for session: ${sessionId}`);
     // TODO: solve race conditions where we call cleanup while the subscription is being created / before it is created
     if (redisCleanup) {
       await redisCleanup();
     }
   }
 
-  new Promise<JSONRPCMessage>((resolve) => {
+  const messagePromise = new Promise<JSONRPCMessage>((resolve) => {
+    console.log(`[redisRelayToMcpServer] Setting up transport.onmessage handler for session: ${sessionId}`);
     transport.onmessage = async (message, extra) => {
-      await sendToMcpServer(sessionId, message, extra);
+      console.log(`[redisRelayToMcpServer] Received message from transport for session ${sessionId}:`, JSON.stringify(message));
-      resolve(message);
-    }
-  }).then(async (message) => {
-    // check for request id in the message
-    if (!("id" in message)) {
-      // if no id, it's a notification, so we return
-      return cleanup;
-    }
-    // otherwise we subscribe to the response channel
-    const toClientChannel = getToClientChannel(sessionId, message.id.toString());
       
-    console.log(`[redisRelayToMcpServer] Subscribing to ${toClientChannel} for response to request ${message.id}`);
+      // First, set up response subscription if needed
+      if ("id" in message) {
+        console.log(`[redisRelayToMcpServer] Setting up response subscription for message ID: ${message.id}`);
+        const toClientChannel = getToClientChannel(sessionId, message.id.toString());
+        console.log(`[redisRelayToMcpServer] Subscribing to response channel: ${toClientChannel}`);
 
         redisCleanup = await redisClient.createSubscription(toClientChannel, async (redisMessageJson) => {
+          console.log(`[redisRelayToMcpServer] Received response from Redis for session ${sessionId}:`, redisMessageJson.substring(0, 200));
           const redisMessage = JSON.parse(redisMessageJson) as RedisMessage;
+          if (redisMessage.type === 'mcp') {
+            console.log(`[redisRelayToMcpServer] Sending response back to transport for session ${sessionId}`);
             await transport.send(redisMessage.message, redisMessage.options);
+            console.log(`[redisRelayToMcpServer] Response sent to transport for session ${sessionId}`);
+          }
         }, (error) => {
-      console.error(`[redisRelayToMcpServer] Error in Redis subscriber for ${toClientChannel}:`, error);
+          console.error(`[redisRelayToMcpServer] Error in response subscription for session ${sessionId}:`, error);
           transport.onerror?.(error);
         });
-  }).catch((error) => {
+        console.log(`[redisRelayToMcpServer] Response subscription established for session ${sessionId}`);
-    console.error(`[redisRelayToMcpServer] Error setting up Redis relay for session ${sessionId}:`, error);
+      } else {
+        console.log(`[redisRelayToMcpServer] Message is notification (no ID) for session ${sessionId}, skipping response subscription`);
+      }
+      
+      // Now send the message to the MCP server
+      await sendToMcpServer(sessionId, message, extra);
+      console.log(`[redisRelayToMcpServer] Sent message to MCP server for session ${sessionId}`);
+      resolve(message);
+    }
+  });
+
+  messagePromise.catch((error) => {
+    console.error(`[redisRelayToMcpServer] Error setting up relay for session ${sessionId}:`, error);
     transport.onerror?.(error);
     cleanup();
   });
@@ -83,6 +123,9 @@ export function redisRelayToMcpServer(sessionId: string, transport: Transport):
 export class ServerRedisTransport implements Transport {
   private counter: number;
   private _sessionId: string;
+  private controlCleanup?: (() => Promise<void>);
+  private serverCleanup?: (() => Promise<void>);
+  private shouldShutdown = false;
 
   onclose?: (() => void) | undefined;
   onerror?: ((error: Error) => void) | undefined;
@@ -94,45 +137,108 @@ export class ServerRedisTransport implements Transport {
   }
 
   async start(): Promise<void> {
-    console.log(`[ServerRedisTransport.${this.counter}.start] Starting server transport for session: ${this._sessionId}`);
+    
-    // Server transport doesn't need to subscribe to anything - it only sends responses
+    // Subscribe to MCP messages from clients
+    const serverChannel = getToServerChannel(this._sessionId);
+    console.log(`[ServerRedisTransport.${this.counter}] Subscribing to server channel: ${serverChannel}`);
+    this.serverCleanup = await redisClient.createSubscription(
+      serverChannel,
+      (messageJson) => {
+        console.log(`[ServerRedisTransport.${this.counter}] Received message from Redis:`, messageJson.substring(0, 200));
+        const redisMessage = JSON.parse(messageJson) as RedisMessage;
+        if (redisMessage.type === 'mcp') {
+          console.log(`[ServerRedisTransport.${this.counter}] Processing MCP message:`, JSON.stringify(redisMessage.message));
+          this.onmessage?.(redisMessage.message, redisMessage.extra);
+          console.log(`[ServerRedisTransport.${this.counter}] MCP message processed`);
+        }
+      },
+      (error) => {
+        console.error(`[ServerRedisTransport.${this.counter}] Server channel error:`, error);
+        this.onerror?.(error);
+      }
+    );
+    
+    // Subscribe to control messages for shutdown
+    const controlChannel = getControlChannel(this._sessionId);
+    this.controlCleanup = await redisClient.createSubscription(
+      controlChannel,
+      (messageJson) => {
+        const redisMessage = JSON.parse(messageJson) as RedisMessage;
+        if (redisMessage.type === 'control') {
+          if (redisMessage.action === 'SHUTDOWN') {
+            this.shouldShutdown = true;
+            this.close();
+          }
+        }
+      },
+      (error) => {
+        this.onerror?.(error);
+      }
+    );
+    
   }
 
   async send(message: JSONRPCMessage, options?: TransportSendOptions): Promise<void> {
-    const relatedRequestId = options?.relatedRequestId?.toString() ?? notificationStreamId;
+    const relatedRequestId = options?.relatedRequestId?.toString() ?? "id" in message ? message.id : notificationStreamId;
     const channel = getToClientChannel(this._sessionId, relatedRequestId)
+    console.log(`[ServerRedisTransport.${this.counter}] Sending message to channel: ${channel}`);
+    console.log(`[ServerRedisTransport.${this.counter}] Message:`, JSON.stringify(message));
+    console.log(`[ServerRedisTransport.${this.counter}] Options:`, JSON.stringify(options));
 
-    const redisMessage: RedisMessage = { message, options };
+    const redisMessage: RedisMessage = { type: 'mcp', message, options };
     const messageStr = JSON.stringify(redisMessage);
-    console.log(`[ServerRedisTransport.${this.counter}.send] Publishing to ${channel}:`, messageStr.substring(0, 100));
     await redisClient.publish(channel, messageStr);
-    console.log(`[ServerRedisTransport.${this.counter}.send] Published successfully to ${channel}`);
+    console.log(`[ServerRedisTransport.${this.counter}] Message published to Redis`);
   }
 
   async close(): Promise<void> {
-    console.log(`[ServerRedisTransport.${this.counter}.close] Closing server transport for session: ${this._sessionId}`);
+    
+    // Clean up server message subscription
+    if (this.serverCleanup) {
+      await this.serverCleanup();
+      this.serverCleanup = undefined;
+    }
+    
+    // Clean up control message subscription
+    if (this.controlCleanup) {
+      await this.controlCleanup();
+      this.controlCleanup = undefined;
+    }
+    
     this.onclose?.();
   }
 }
 
-export async function startServerListeningToRedis(server: Server, sessionId: string) {
+export async function startServerListeningToRedis(serverWrapper: { server: Server; cleanup: () => void }, sessionId: string): Promise<ServerRedisTransport> {
-  console.log(`[startServerListeningToRedis] Starting background server for session ${sessionId}`);
   const serverRedisTransport = new ServerRedisTransport(sessionId);
   
+  // Add cleanup callback to the transport
+  const originalClose = serverRedisTransport.close.bind(serverRedisTransport);
+  serverRedisTransport.close = async () => {
+    serverWrapper.cleanup();
+    await originalClose();
+  };
+  
   // The server.connect() will call start() on the transport
-  await server.connect(serverRedisTransport)
+  await serverWrapper.server.connect(serverRedisTransport)
-  console.log(`[startServerListeningToRedis] Background server connected for session ${sessionId}`);
+  
+  return serverRedisTransport;
 }
 
 export async function getFirstShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, cleanup: () => Promise<void>}> {
+  console.log(`[getFirstShttpTransport] Creating transport for session: ${sessionId}`);
   const shttpTransport = new StreamableHTTPServerTransport({
     sessionIdGenerator: () => sessionId,
     enableJsonResponse: true, // Enable JSON response mode
   });
   
+  console.log(`[getFirstShttpTransport] Created transport, sessionId: ${shttpTransport.sessionId}`);
+  
   // Use the new request-id based relay approach
+  console.log(`[getFirstShttpTransport] Setting up Redis relay for session: ${sessionId}`);
   const cleanup = redisRelayToMcpServer(sessionId, shttpTransport);
   
+  console.log(`[getFirstShttpTransport] Transport setup complete for session: ${sessionId}`);
   return { shttpTransport, cleanup };
 }