Update app.py

app.py

@@ -111,9 +111,12 @@ 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 streams for each session ID
+# Store write streams for each session ID (for SseServerTransport messages)
 write_streams: Dict[str, anyio.streams.memory.MemoryObjectSendStream] = {}
 
+# Store SSE stream writers for each session ID (for manual messages)
+sse_stream_writers: Dict[str, anyio.streams.memory.MemoryObjectSendStream] = {}
+
 # Initialize SseServerTransport
 transport = SseServerTransport("/airtable/mcp")
 
@@ -133,15 +136,14 @@ async def handle_sse(request: Request):
             )
             logger.debug(f"Sent endpoint event: {endpoint_data}")
             async for session_message in write_stream_reader:
-                # Since SseServerTransport expects a SessionMessage, we'll handle messages manually
-                # and send raw SSE events to avoid the error in SseServerTransport's sse_writer
+                # Handle messages from SseServerTransport
                 if hasattr(session_message, 'message'):
                     message_data = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
                     event_data = json.loads(message_data)
-                    logger.debug(f"Received SessionMessage: {event_data}")
+                    logger.debug(f"Received SessionMessage from SseServerTransport: {event_data}")
                 else:
                     event_data = session_message
-                    logger.debug(f"Received dict event: {event_data}")
+                    logger.debug(f"Received dict event from SseServerTransport: {event_data}")
                 # Extract session_id from the endpoint event
                 if not session_id and event_data.get("event") == "endpoint":
                     endpoint_url = event_data.get("data", "")
@@ -150,8 +152,9 @@ 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_stream_writers[session_id] = sse_stream_writer
                             logger.debug(f"Updated placeholder {placeholder_id} to session_id {session_id}")
-                # Send the event as a raw SSE event
+                # Forward the event to the client
                 await sse_stream_writer.send({
                     "event": event_data.get("event", "message"),
                     "data": event_data.get("data", json.dumps(event_data))
@@ -170,11 +173,12 @@ async def handle_sse(request: Request):
             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)}")
-        # Clean up write_streams on error
+        # Clean up write_streams and sse_stream_writers on error
         placeholder_id = f"placeholder_{id(write_stream)}"
         write_streams.pop(placeholder_id, None)
         if session_id:
             write_streams.pop(session_id, None)
+            sse_stream_writers.pop(session_id, None)
         raise
     return EventSourceResponse(sse_stream_reader, data_sender_callable=sse_writer)
 
@@ -187,8 +191,10 @@ async def handle_post_message(request: Request):
     try:
         message = json.loads(body.decode())
         session_id = request.query_params.get("session_id")
+        # Use sse_stream_writers to send manual responses directly
+        sse_writer = sse_stream_writers.get(session_id) if session_id else None
         write_stream = write_streams.get(session_id) if session_id else None
-        if message.get("method") == "initialize" and write_stream:
+        if message.get("method") == "initialize" and sse_writer:
             logger.debug("Handling initialize request manually")
             response = {
                 "jsonrpc": "2.0",
@@ -218,10 +224,11 @@ async def handle_post_message(request: Request):
             }
             logger.debug(f"Manual initialize response: {response}")
             response_data = json.dumps(response)
-            await write_stream.send({
+            await sse_writer.send({
                 "event": "message",
                 "data": response_data
             })
+            logger.debug(f"Sent initialize response directly via SSE for session {session_id}")
             return Response(status_code=202)
         if message.get("method") == "tools/list":
             logger.debug("Handling tools/list request manually")
@@ -236,27 +243,27 @@ async def handle_post_message(request: Request):
             logger.debug(f"Manual tools/list response: {response}")
             response_data = json.dumps(response)
             sent = False
-            # First, try the session_id directly
-            if session_id in write_streams:
-                write_stream = write_streams[session_id]
+            # First, try sending directly via sse_writer
+            if sse_writer:
                 try:
-                    await write_stream.send({
+                    await sse_writer.send({
                         "event": "message",
                         "data": response_data
                     })
-                    logger.debug(f"Sent tools/list response to session {session_id}")
+                    logger.debug(f"Sent tools/list response directly via SSE for session {session_id}")
                     sent = True
                 except Exception as e:
-                    logger.error(f"Error sending to session {session_id}: {str(e)}")
-                    write_streams.pop(session_id, None)
-            # If not found, look for a placeholder ID and update it
-            if not sent:
+                    logger.error(f"Error sending to session {session_id} via sse_writer: {str(e)}")
+                    sse_stream_writers.pop(session_id, None)
+            # If not found or failed, look for a placeholder ID and update it
+            if not sent and write_stream:
                 for sid, ws in list(write_streams.items()):
                     if sid.startswith("placeholder_"):
                         try:
                             write_streams[session_id] = ws
+                            sse_stream_writers[session_id] = sse_writer
                             write_streams.pop(sid, None)
-                            await ws.send({
+                            await sse_writer.send({
                                 "event": "message",
                                 "data": response_data
                             })
@@ -266,12 +273,13 @@ async def handle_post_message(request: Request):
                         except Exception as e:
                             logger.error(f"Error sending to placeholder {sid}: {str(e)}")
                             write_streams.pop(sid, None)
+                            sse_stream_writers.pop(session_id, None)
             if not sent:
-                logger.warning(f"Failed to send tools/list response: no active write_streams found")
+                logger.warning(f"Failed to send tools/list response: no active write_streams or sse_writer found")
             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}")
+        # If neither sse_writer nor write_stream is available, log and handle gracefully
+        if not sse_writer and not write_stream:
+            logger.error(f"No sse_writer or 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")