Update server.py

server.py

@@ -1,7 +1,9 @@
 import os
+import json
 import logging
 import anyio
 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
@@ -16,6 +18,9 @@ app = FastAPI()
 server = Server(name="airtable-mcp")
 transport = SseServerTransport("/airtable/mcp")
 
+# Store write streams for each session ID
+write_streams = {}
+
 # Configure environment variables (for logging purposes)
 token = os.getenv("AIRTABLE_API_TOKEN")
 base_id = os.getenv("AIRTABLE_BASE_ID")
@@ -25,16 +30,29 @@ logger.info(f"Using Airtable base ID: {base_id}")
 @app.get("/airtable/mcp")
 async def handle_sse(request: Request):
     try:
+        session_id = None
         async def sse_writer():
-            async with sse_stream_writer:
-                await sse_stream_writer.send({"event": "endpoint", "data": "/airtable/mcp"})
+            nonlocal session_id
+            async with sse_stream_writer, write_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:
                     message_data = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
+                    message = json.loads(message_data)
+                    if not session_id and message.get("event") == "endpoint":
+                        endpoint_url = message.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)
                     await sse_stream_writer.send({"event": "message", "data": message_data})
 
         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")
@@ -43,6 +61,37 @@ async def handle_sse(request: Request):
         logger.error(f"Error in handle_sse: {str(e)}")
         raise
 
+@app.post("/airtable/mcp")
+async def handle_post(request: Request):
+    try:
+        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
+        if message.get("method") == "tools/list" and write_stream:
+            response = {
+                "jsonrpc": "2.0",
+                "id": message.get("id"),
+                "result": {
+                    "tools": [
+                        {"name": "list_airtable_records", "description": "Lists all records in the specified Airtable table", "inputSchema": {}},
+                        {"name": "create_airtable_record", "description": "Creates a new record in the specified Airtable table", "inputSchema": {"record_data": {"type": "object"}}}
+                    ],
+                    "nextCursor": None
+                }
+            }
+            response_data = json.dumps(response)
+            await write_stream.send({"event": "message", "data": response_data})
+            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:
+        logger.error(f"Error handling POST message: {str(e)}")
+        return Response(status_code=202)
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)