Update app.py

app.py

@@ -121,22 +121,16 @@ 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 the endpoint event
+    session_id = None  # We'll extract this later
     async def sse_writer():
         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}")
-            # 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)
@@ -146,8 +140,11 @@ async def handle_sse(request: Request):
                     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}")
+                        # Update write_streams with the real session_id
+                        placeholder_id = f"placeholder_{id(write_stream)}"
+                        if placeholder_id in write_streams:
+                            write_streams[session_id] = write_streams.pop(placeholder_id)
+                            logger.debug(f"Updated placeholder {placeholder_id} to session_id {session_id}")
                 await sse_stream_writer.send(
                     {
                         "event": "message",
@@ -160,19 +157,19 @@ async def handle_sse(request: Request):
         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))
+            # Store the write_stream with a placeholder ID
+            placeholder_id = f"placeholder_{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:
         logger.error(f"Error in handle_sse: {str(e)}")
+        # Clean up write_streams on error
+        placeholder_id = f"placeholder_{id(write_stream)}"
+        write_streams.pop(placeholder_id, None)
+        if session_id:
+            write_streams.pop(session_id, None)
         raise
     return EventSourceResponse(sse_stream_reader, data_sender_callable=sse_writer)
 
@@ -233,30 +230,37 @@ async def handle_post_message(request: Request):
             }
             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()):
+            # First, try the session_id directly
+            if session_id in write_streams:
+                write_stream = write_streams[session_id]
                 try:
-                    if sid == session_id:
-                        await ws.send({
-                            "event": "message",
-                            "data": response_data
-                        })
-                        logger.debug(f"Sent tools/list response to session {sid}")
-                        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
+                    await write_stream.send({
+                        "event": "message",
+                        "data": response_data
+                    })
+                    logger.debug(f"Sent tools/list response to session {session_id}")
+                    sent = True
                 except Exception as e:
-                    logger.error(f"Error sending to session {sid}: {str(e)}")
-                    write_streams.pop(sid, None)  # Remove closed streams
+                    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:
+                for sid, ws in list(write_streams.items()):
+                    if sid.startswith("placeholder_"):
+                        try:
+                            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
+                            break
+                        except Exception as e:
+                            logger.error(f"Error sending to placeholder {sid}: {str(e)}")
+                            write_streams.pop(sid, None)
             if not sent:
                 logger.warning(f"Failed to send tools/list response: no active write_streams found")
             return Response(status_code=202)