Update server.py

server.py

@@ -46,6 +46,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)
+                                logger.info(f"Updated placeholder {placeholder_id} to session_id {session_id}")
                     await sse_stream_writer.send({"event": "message", "data": message_data})
 
         sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream(0)
@@ -67,7 +68,19 @@ async def handle_post(request: Request):
         body = await request.body()
         message = json.loads(body.decode())
         session_id = request.query_params.get("session_id")
-        write_stream = write_streams.get(session_id) if session_id else None
+        logger.info(f"Received POST with session_id: {session_id}, message: {message}")
+
+        # Try to find the write_stream, including checking for placeholder IDs
+        write_stream = write_streams.get(session_id)
+        if not write_stream:
+            for sid, ws in list(write_streams.items()):
+                if sid.startswith("placeholder_"):
+                    write_streams[session_id] = ws
+                    write_streams.pop(sid)
+                    write_stream = ws
+                    logger.info(f"Associated placeholder {sid} with session_id {session_id}")
+                    break
+
         if message.get("method") == "tools/list" and write_stream:
             response = {
                 "jsonrpc": "2.0",
@@ -82,10 +95,13 @@ async def handle_post(request: Request):
             }
             response_data = json.dumps(response)
             await write_stream.send({"event": "message", "data": response_data})
+            logger.info(f"Sent tools/list response for session {session_id}")
             return Response(status_code=202)
+
         if not write_stream:
             logger.error(f"No write_stream found for session_id: {session_id}")
             return Response(status_code=202)
+
         await transport.handle_post_message(request.scope, request.receive, request._send)
         return Response(status_code=202)
     except Exception as e: