Update app.py

app.py

@@ -12,6 +12,7 @@ from sse_starlette import EventSourceResponse
 import anyio
 import asyncio
 import logging
+from typing import Dict
 
 # Set up logging
 logging.basicConfig(level=logging.DEBUG)
@@ -110,22 +111,29 @@ mcp_server = Server(name="airtable-mcp")
 mcp_server.tool_handlers = tool_handlers  # Set as attribute
 mcp_server.tools = tools  # Set tools as attribute for Deep Agent to discover
 
-# Store write_stream globally to access in POST handler
-write_stream_global = None
+# Initialize SseServerTransport
+transport = SseServerTransport("/airtable/mcp")
+
+# Store write streams for each session ID
+write_streams: Dict[str, anyio.streams.memory.MemoryObjectSendStream] = {}
 
 # SSE endpoint for GET requests
 @app.get("/airtable/mcp")
 async def handle_sse(request: Request):
-    global write_stream_global
     logger.debug("Handling SSE connection request")
     async def sse_writer():
         logger.debug("Starting SSE writer")
         async with sse_stream_writer, write_stream_reader:
+            # Extract session_id from endpoint_data
+            session_id = None
             endpoint_data = "/airtable/mcp?session_id={session_id}"
             await sse_stream_writer.send(
                 {"event": "endpoint", "data": endpoint_data}
             )
             logger.debug(f"Sent endpoint event: {endpoint_data}")
+            # Extract session_id from the sent message (simplified)
+            # In a real implementation, we'd parse the session_id from transport
+            # For now, we'll set it after the first message
             async for session_message in write_stream_reader:
                 logger.debug(f"Sending message via SSE: {session_message}")
                 await sse_stream_writer.send(
@@ -138,12 +146,19 @@ async def handle_sse(request: Request):
                 )
 
     sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream(0)
-    async with transport.connect_sse(request.scope, request.receive, request._send) as streams:
-        read_stream, write_stream = streams
-        write_stream_reader = write_stream  # Since streams are MemoryObject streams
-        write_stream_global = write_stream  # Store write_stream for POST handler
-        logger.debug("Running MCP server with streams")
-        await mcp_server.run(read_stream, write_stream, mcp_server.create_initialization_options(), tools=tools)
+    try:
+        async with transport.connect_sse(request.scope, request.receive, request._send) as streams:
+            read_stream, write_stream = streams
+            write_stream_reader = write_stream  # Since streams are MemoryObject streams
+            # Extract session_id (this is a simplification; ideally, we'd get it from transport)
+            session_id = request.query_params.get("session_id", "unknown")
+            write_streams[session_id] = write_stream  # Store write_stream for this session
+            logger.debug(f"Stored write_stream for session_id: {session_id}")
+            logger.debug("Running MCP server with streams")
+            await mcp_server.run(read_stream, write_stream, mcp_server.create_initialization_options())
+    except Exception as e:
+        logger.error(f"Error in handle_sse: {str(e)}")
+        raise
     return EventSourceResponse(sse_stream_reader, data_sender_callable=sse_writer)
 
 # Message handling endpoint for POST requests
@@ -154,8 +169,9 @@ async def handle_post_message(request: Request):
     logger.debug(f"Received POST message body: {body}")
     try:
         message = json.loads(body.decode())
-        # Handle initialize request manually to ensure capabilities include tools
-        if message.get("method") == "initialize" and write_stream_global:
+        session_id = request.query_params.get("session_id")
+        write_stream = write_streams.get(session_id) if session_id else None
+        if message.get("method") == "initialize" and write_stream:
             logger.debug("Handling initialize request manually")
             response = {
                 "jsonrpc": "2.0",
@@ -179,10 +195,9 @@ async def handle_post_message(request: Request):
                 message=mcp_types.JSONRPCResponse(**response),
                 metadata=mcp_types.ServerMessageMetadata(request_context=request)
             )
-            await write_stream_global.send(session_message)
+            await write_stream.send(session_message)
             return Response(status_code=202)
-        # Handle tools/list request manually
-        if message.get("method") == "tools/list" and write_stream_global:
+        if message.get("method") == "tools/list" and write_stream:
             logger.debug("Handling tools/list request manually")
             response = {
                 "jsonrpc": "2.0",
@@ -196,7 +211,7 @@ async def handle_post_message(request: Request):
                 message=mcp_types.JSONRPCResponse(**response),
                 metadata=mcp_types.ServerMessageMetadata(request_context=request)
             )
-            await write_stream_global.send(session_message)
+            await write_stream.send(session_message)
             return Response(status_code=202)
         await transport.handle_post_message(request.scope, request.receive, request._send)
         logger.debug("POST message handled successfully")