Update app.py

app.py

@@ -126,29 +126,49 @@ async def handle_sse(request: Request):
         nonlocal session_id
         logger.debug("Starting SSE writer")
         async with sse_stream_writer, write_stream_reader:
+            # Since we can't iterate over sse_stream_writer.stream directly in this context,
+            # we'll rely on the session_id being set by SseServerTransport
             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}")
-            async for event in sse_stream_writer.stream:
-                # Log the raw SSE event to understand its format
-                logger.debug(f"Raw SSE event: {event}")
+            # The session_id is set by SseServerTransport and should be available after connect_sse
+            # We need to extract it from the endpoint_data after substitution
+            # SseServerTransport substitutes {session_id} with the actual session_id
+            # We'll extract it from the logs or directly from transport if possible
+            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)
                 # Extract session_id from the endpoint event
-                if event.get("event") == "endpoint":
-                    endpoint_url = event.get("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]
                         write_streams[session_id] = write_stream
                         logger.debug(f"Extracted session_id: {session_id}")
-                # Forward the event to the client
-                await sse_stream_writer.send(event)
+                await sse_stream_writer.send(
+                    {
+                        "event": "message",
+                        "data": message_data
+                    }
+                )
 
     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
+            # Extract session_id directly from transport's internal state
+            # SseServerTransport creates a session_id during connect_sse
+            # We can access it by inspecting transport's state
+            # Unfortunately, transport's internals are not directly accessible
+            # We'll use a temporary workaround by storing the write_stream with a placeholder
+            # and update it once we extract the session_id
+            placeholder_id = str(id(write_stream))
+            write_streams[placeholder_id] = write_stream
+            logger.debug(f"Stored write_stream with placeholder_id: {placeholder_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:
@@ -195,7 +215,6 @@ async def handle_post_message(request: Request):
                 }
             }
             logger.debug(f"Manual initialize response: {response}")
-            # Send the response directly as an SSE event
             response_data = json.dumps(response)
             await write_stream.send({
                 "event": "message",
@@ -204,31 +223,43 @@ async def handle_post_message(request: Request):
             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:
-                logger.warning(f"Session ID {session_id} not found in write_streams, attempting to proceed")
-                response = {
-                    "jsonrpc": "2.0",
-                    "id": message.get("id"),
-                    "result": {
-                        "tools": [tool.model_dump(by_alias=True, exclude_none=True) for tool in tools],
-                        "nextCursor": None
-                    }
+            response = {
+                "jsonrpc": "2.0",
+                "id": message.get("id"),
+                "result": {
+                    "tools": [tool.model_dump(by_alias=True, exclude_none=True) for tool in tools],
+                    "nextCursor": None
                 }
-                logger.debug(f"Manual tools/list response: {response}")
-                response_data = json.dumps(response)
-                # Send to all active write_streams (temporary workaround)
-                for sid, ws in list(write_streams.items()):
-                    try:
+            }
+            logger.debug(f"Manual tools/list response: {response}")
+            response_data = json.dumps(response)
+            # Send to all active write_streams (temporary workaround)
+            sent = False
+            for sid, ws in list(write_streams.items()):
+                try:
+                    if sid == session_id:
                         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)}")
-                        write_streams.pop(sid, None)  # Remove closed streams
-                return Response(status_code=202)
+                        sent = True
+                    elif sid.startswith("placeholder_"):
+                        # Update the placeholder with the real session_id
+                        write_streams[session_id] = ws
+                        write_streams.pop(sid, None)
+                        await ws.send({
+                            "event": "message",
+                            "data": response_data
+                        })
+                        logger.debug(f"Updated placeholder {sid} to session_id {session_id} and sent tools/list response")
+                        sent = True
+                except Exception as e:
+                    logger.error(f"Error sending to session {sid}: {str(e)}")
+                    write_streams.pop(sid, None)  # Remove closed streams
+            if not sent:
+                logger.warning(f"Failed to send tools/list response: no active write_streams 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}")