Update app.py

app.py

@@ -114,9 +114,6 @@ mcp_server.tools = tools  # Set tools as attribute for Deep Agent to discover
 # Store write streams for each session ID
 write_streams: Dict[str, anyio.streams.memory.MemoryObjectSendStream] = {}
 
-# Store session IDs for each connection
-session_ids: Dict[str, str] = {}
-
 # Initialize SseServerTransport
 transport = SseServerTransport("/airtable/mcp")
 
@@ -127,8 +124,6 @@ async def handle_sse(request: Request):
     async def sse_writer():
         logger.debug("Starting SSE writer")
         async with sse_stream_writer, write_stream_reader:
-            # Extract session_id dynamically
-            session_id = None
             endpoint_data = "/airtable/mcp?session_id={session_id}"
             await sse_stream_writer.send(
                 {"event": "endpoint", "data": endpoint_data}
@@ -144,25 +139,19 @@ async def handle_sse(request: Request):
                     }
                 )
                 # Extract session_id from the endpoint event
-                if not session_id and session_message.message.method == "endpoint":
-                    try:
-                        message = json.loads(message_data)
-                        endpoint_url = message.get("data", "")
-                        if "session_id=" in endpoint_url:
-                            session_id = endpoint_url.split("session_id=")[1]
-                            session_ids[id(write_stream)] = session_id
-                            write_streams[session_id] = write_stream
-                            logger.debug(f"Extracted session_id: {session_id}")
-                    except Exception as e:
-                        logger.error(f"Error extracting session_id: {str(e)}")
+                message = json.loads(message_data)
+                if message.get("event") == "endpoint":
+                    endpoint_url = message.get("data", "")
+                    if "session_id=" in endpoint_url:
+                        session_id = endpoint_url.split("session_id=")[1]
+                        write_streams[session_id] = write_stream
+                        logger.debug(f"Extracted session_id: {session_id}")
 
     sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream(0)
     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
-            # Use a unique identifier for the stream until session_id is extracted
-            stream_id = id(write_stream)
             logger.debug("Running MCP server with streams")
             await mcp_server.run(read_stream, write_stream, mcp_server.create_initialization_options())
     except Exception as e:
@@ -232,10 +221,15 @@ async def handle_post_message(request: Request):
             )
             await write_stream.send(session_message)
             return Response(status_code=202)
+        # If write_stream is None, log and handle gracefully
+        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)
         logger.debug("POST message handled successfully")
     except Exception as e:
         logger.error(f"Error handling POST message: {str(e)}")
+        return Response(status_code=202)
     return Response(status_code=202)
 
 # Health check endpoint