Update server.py

server.py

@@ -7,7 +7,6 @@ from fastapi import FastAPI, Request
 from fastapi.responses import Response
 from sse_starlette import EventSourceResponse
 from mcp.server.lowlevel import Server
-from mcp.server.sse import SseServerTransport
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -17,7 +16,6 @@ app = FastAPI()
 
 # Define the MCP server
 server = Server(name="airtable-mcp")
-transport = SseServerTransport("/airtable/mcp")
 
 # Store write streams for each session ID
 write_streams = {}
@@ -34,39 +32,33 @@ async def handle_sse(request: Request):
         session_id = None
         async def sse_writer():
             nonlocal session_id
-            async with sse_stream_writer, write_stream_reader:
+            async with sse_stream_writer, read_stream_reader:
                 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:
-                    if isinstance(session_message, dict):
-                        event = session_message.get("event", "message")
-                        data = session_message.get("data", json.dumps(session_message))
-                    else:
-                        event = "message"
-                        data = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
-                    message = json.loads(data) if isinstance(data, str) else data
-                    if not session_id and event == "endpoint":
-                        endpoint_url = data
+                async for message in read_stream_reader:
+                    message_data = json.loads(message) if isinstance(message, str) else message
+                    if message_data.get("event") == "endpoint":
+                        endpoint_url = message_data.get("data", "")
                         if "session_id=" in endpoint_url:
                             session_id = endpoint_url.split("session_id=")[1]
                             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": event, "data": data})
+                    await sse_stream_writer.send(message_data)
                 # Keep-alive loop to maintain the SSE connection
                 while True:
                     await sse_stream_writer.send({"event": "ping", "data": "keep-alive"})
                     await asyncio.sleep(15)  # Send keep-alive every 15 seconds
 
         sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream(0)
-        async with transport.connect_sse(request.scope, request.receive, request._send) as (read_stream, write_stream):
-            write_stream_reader = write_stream
-            placeholder_id = f"placeholder_{id(write_stream)}"
-            write_streams[placeholder_id] = write_stream
-            logger.info("Starting MCP server with streams")
-            await server.run(read_stream, write_stream, server.create_initialization_options())
-            logger.info("MCP server running")
+        read_stream, write_stream = anyio.create_memory_object_stream(0)
+        placeholder_id = f"placeholder_{id(write_stream)}"
+        write_streams[placeholder_id] = write_stream
+        read_stream_reader = read_stream
+        logger.info("Starting MCP server with streams")
+        await server.run(read_stream, write_stream, server.create_initialization_options())
+        logger.info("MCP server running")
         return EventSourceResponse(sse_stream_reader, data_sender_callable=sse_writer)
     except Exception as e:
         logger.error(f"Error in handle_sse: {str(e)}")
@@ -141,11 +133,6 @@ async def handle_post(request: Request):
             logger.error(f"No write_stream found for session_id: {session_id}")
             return Response(status_code=202)
 
-        try:
-            await transport.handle_post_message(request.scope, request.receive, request._send)
-            logger.info(f"Handled post message for session {session_id} via transport")
-        except Exception as e:
-            logger.error(f"Failed to handle post message via transport for session {session_id}: {str(e)}")
         return Response(status_code=202)
     except Exception as e:
         logger.error(f"Error handling POST message: {str(e)}")