Update server.py

server.py

@@ -18,8 +18,9 @@ app = FastAPI()
 server = Server(name="airtable-mcp")
 transport = SseServerTransport("/airtable/mcp")
 
-# Store write streams for each session ID
+# Store write streams and SSE writers for each session ID
 write_streams = {}
+sse_writers = {}
 
 # Configure environment variables (for logging purposes)
 token = os.getenv("AIRTABLE_API_TOKEN")
@@ -37,7 +38,6 @@ async def handle_sse(request: Request):
                 endpoint_data = f"/airtable/mcp?session_id={{session_id}}"
                 await sse_stream_writer.send({"event": "endpoint", "data": endpoint_data})
                 async for session_message in write_stream_reader:
-                    # Handle both dict and SessionMessage objects
                     if isinstance(session_message, dict):
                         event = session_message.get("event", "message")
                         data = session_message.get("data", json.dumps(session_message))
@@ -52,6 +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
                                 logger.info(f"Updated placeholder {placeholder_id} to session_id {session_id}")
                     await sse_stream_writer.send({"event": event, "data": data})
 
@@ -78,16 +79,18 @@ async def handle_post(request: Request):
 
         # Try to find the write_stream, including checking for placeholder IDs
         write_stream = write_streams.get(session_id)
+        sse_writer = sse_writers.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
+                    sse_writers[session_id] = sse_writers.get(sid)
                     logger.info(f"Associated placeholder {sid} with session_id {session_id}")
                     break
 
-        if message.get("method") == "tools/list" and write_stream:
+        if message.get("method") == "tools/list" and write_stream and sse_writer:
             response = {
                 "jsonrpc": "2.0",
                 "id": message.get("id"),
@@ -100,12 +103,12 @@ 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}")
+            await sse_writer.send({"event": "message", "data": response_data})
+            logger.info(f"Sent tools/list response for session {session_id} via SSE")
             return Response(status_code=202)
 
-        if not write_stream:
-            logger.error(f"No write_stream found for session_id: {session_id}")
+        if not write_stream or not sse_writer:
+            logger.error(f"No write_stream or SSE writer found for session_id: {session_id}")
             return Response(status_code=202)
 
         await transport.handle_post_message(request.scope, request.receive, request._send)