Update server.py

server.py

@@ -52,7 +52,7 @@ async def handle_sse(request: Request):
                             placeholder_id = f"placeholder_{id(write_stream)}"
                             if placeholder_id in write_streams:
                                 write_streams[session_id] = write_streams.pop(placeholder_id)
-                                sse_writers[session_id] = sse_stream_writer
+                                sse_writers[session_id] = sse_writers.pop(placeholder_id)
                                 logger.info(f"Updated placeholder {placeholder_id} to session_id {session_id}")
                     await sse_stream_writer.send({"event": event, "data": data})
 
@@ -61,6 +61,7 @@ async def handle_sse(request: Request):
             write_stream_reader = write_stream
             placeholder_id = f"placeholder_{id(write_stream)}"
             write_streams[placeholder_id] = write_stream
+            sse_writers[placeholder_id] = sse_stream_writer  # Store the SSE writer with placeholder ID
             logger.info("Starting MCP server with streams")
             await server.run(read_stream, write_stream, server.create_initialization_options())
             logger.info("MCP server running")
@@ -77,16 +78,17 @@ async def handle_post(request: Request):
         session_id = request.query_params.get("session_id")
         logger.info(f"Received POST with session_id: {session_id}, message: {message}")
 
-        # Try to find the write_stream, including checking for placeholder IDs
+        # Try to find the write_stream and SSE writer, including checking for placeholder IDs
         write_stream = write_streams.get(session_id)
         sse_writer = sse_writers.get(session_id)
-        if not write_stream:
+        if not write_stream or not sse_writer:
             for sid, ws in list(write_streams.items()):
                 if sid.startswith("placeholder_"):
                     write_streams[session_id] = ws
+                    sse_writers[session_id] = sse_writers.pop(sid, None)
                     write_streams.pop(sid)
-                    write_stream = ws
-                    sse_writers[session_id] = sse_writers.get(sid)
+                    write_stream = write_streams[session_id]
+                    sse_writer = sse_writers[session_id]
                     logger.info(f"Associated placeholder {sid} with session_id {session_id}")
                     break