Update app.py

app.py

@@ -121,7 +121,7 @@ transport = SseServerTransport("/airtable/mcp")
 @app.get("/airtable/mcp")
 async def handle_sse(request: Request):
     logger.debug("Handling SSE connection request")
-    session_id = None  # We'll extract this from transport
+    session_id = None  # We'll extract this from the endpoint event
     async def sse_writer():
         nonlocal session_id
         logger.debug("Starting SSE writer")
@@ -131,35 +131,24 @@ async def handle_sse(request: Request):
                 {"event": "endpoint", "data": endpoint_data}
             )
             logger.debug(f"Sent endpoint event: {endpoint_data}")
-            async for session_message in write_stream_reader:
-                logger.debug(f"Sending message via SSE: {session_message}")
-                message_data = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
-                message = json.loads(message_data)
+            async for event in sse_stream_writer.stream:
+                # Log the raw SSE event to understand its format
+                logger.debug(f"Raw SSE event: {event}")
                 # Extract session_id from the endpoint event
-                if not session_id and message.get("event") == "endpoint":
-                    endpoint_url = message.get("data", "")
+                if event.get("event") == "endpoint":
+                    endpoint_url = event.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}")
-                await sse_stream_writer.send(
-                    {
-                        "event": "message",
-                        "data": message_data
-                    }
-                )
+                # Forward the event to the client
+                await sse_stream_writer.send(event)
 
     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
-            # Directly extract session_id from transport
-            session_id = None
-            # Access transport's internal session_id (this is a simplification; we need to match the session_id)
-            # SseServerTransport sets the session_id during connect_sse
-            # We can get it from the endpoint event or transport's internal state
-            # For now, we'll rely on the sse_writer to extract it
             logger.debug("Running MCP server with streams")
             await mcp_server.run(read_stream, write_stream, mcp_server.create_initialization_options())
     except Exception as e:
@@ -206,20 +195,18 @@ async def handle_post_message(request: Request):
                 }
             }
             logger.debug(f"Manual initialize response: {response}")
-            session_message = mcp_types.SessionMessage(
-                message=mcp_types.JSONRPCResponse(**response),
-                metadata=mcp_types.ServerMessageMetadata(request_context=request)
-            )
-            await write_stream.send(session_message)
+            # Send the response directly as an SSE event
+            response_data = json.dumps(response)
+            await write_stream.send({
+                "event": "message",
+                "data": response_data
+            })
             return Response(status_code=202)
         if message.get("method") == "tools/list":
             logger.debug("Handling tools/list request manually")
             # If write_stream is not found, try to find it by iterating over write_streams
             if not write_stream and session_id:
-                # Since we're not extracting session_id correctly, we'll bypass for now
-                # In a real scenario, we'd need to ensure session_id is set
                 logger.warning(f"Session ID {session_id} not found in write_streams, attempting to proceed")
-                # Send the response to all active write_streams (temporary workaround)
                 response = {
                     "jsonrpc": "2.0",
                     "id": message.get("id"),
@@ -229,14 +216,14 @@ async def handle_post_message(request: Request):
                     }
                 }
                 logger.debug(f"Manual tools/list response: {response}")
-                session_message = mcp_types.SessionMessage(
-                    message=mcp_types.JSONRPCResponse(**response),
-                    metadata=mcp_types.ServerMessageMetadata(request_context=request)
-                )
+                response_data = json.dumps(response)
                 # Send to all active write_streams (temporary workaround)
                 for sid, ws in list(write_streams.items()):
                     try:
-                        await ws.send(session_message)
+                        await ws.send({
+                            "event": "message",
+                            "data": response_data
+                        })
                         logger.debug(f"Sent tools/list response to session {sid}")
                     except Exception as e:
                         logger.error(f"Error sending to session {sid}: {str(e)}")