Update server.py

server.py

@@ -18,9 +18,8 @@ app = FastAPI()
 server = Server(name="airtable-mcp")
 transport = SseServerTransport("/airtable/mcp")
 
-# Store write streams and SSE writers for each session ID
+# Store write streams for each session ID
 write_streams = {}
-sse_writers = {}
 
 # Configure environment variables (for logging purposes)
 token = os.getenv("AIRTABLE_API_TOKEN")
@@ -52,7 +51,6 @@ 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})
 
@@ -61,7 +59,6 @@ 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
             logger.info("Starting MCP server with streams")
             await server.run(read_stream, write_stream, server.create_initialization_options())
             logger.info("MCP server running")
@@ -78,21 +75,18 @@ 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 and SSE writer
+        # Try to find the write_stream
         write_stream = write_streams.get(session_id)
-        sse_writer = sse_writers.get(session_id)
-        if not write_stream or not sse_writer:
+        if not write_stream:
             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 = write_streams[session_id]
-                    sse_writer = sse_writers[session_id]
                     logger.info(f"Associated placeholder {sid} with session_id {session_id}")
                     break
 
-        if message.get("method") == "tools/list" and write_stream and sse_writer:
+        if message.get("method") == "tools/list" and write_stream:
             response = {
                 "jsonrpc": "2.0",
                 "id": message.get("id"),
@@ -106,26 +100,15 @@ async def handle_post(request: Request):
             }
             response_data = json.dumps(response)
             try:
-                # Use a timeout to prevent hanging
-                async with anyio.fail_after(5):  # 5-second timeout
-                    await sse_writer.send({"event": "message", "data": response_data})
-                logger.info(f"Successfully sent tools/list response for session {session_id} via SSE")
-            except anyio.get_cancelled_exc_class():
-                logger.error(f"Timeout while sending SSE response for session {session_id}")
-                return Response(status_code=202)
+                await write_stream.send({"event": "message", "data": response_data})
+                logger.info(f"Successfully sent tools/list response for session {session_id} via write_stream")
             except Exception as e:
-                logger.error(f"Failed to send SSE response for session {session_id}: {str(e)}")
-                # Fallback: Try sending directly through write_stream
-                try:
-                    await write_stream.send({"event": "message", "data": response_data})
-                    logger.info(f"Fallback: Successfully sent tools/list response for session {session_id} via write_stream")
-                except Exception as fallback_e:
-                    logger.error(f"Fallback failed for session {session_id}: {str(fallback_e)}")
-                    return Response(status_code=202)
+                logger.error(f"Failed to send response for session {session_id} via write_stream: {str(e)}")
+                return Response(status_code=202)
             return Response(status_code=202)
 
-        if not write_stream or not sse_writer:
-            logger.error(f"No write_stream or SSE writer found for session_id: {session_id}")
+        if not write_stream:
+            logger.error(f"No write_stream found for session_id: {session_id}")
             return Response(status_code=202)
 
         try: