Update app.py

app.py

@@ -126,31 +126,35 @@ async def handle_sse(request: Request):
         nonlocal session_id
         logger.debug("Starting SSE writer")
         async with sse_stream_writer, write_stream_reader:
+            # Send the initial endpoint event manually to capture the session_id
             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 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)
+                # Since session_message might be a dict or a SessionMessage, handle both cases
+                if isinstance(session_message, dict):
+                    event_data = session_message
+                    logger.debug(f"Received SSE event (dict): {event_data}")
+                else:
+                    message_data = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
+                    event_data = json.loads(message_data)
+                    logger.debug(f"Received SSE event (SessionMessage): {event_data}")
                 # Extract session_id from the endpoint event
-                if not session_id and message.get("event") == "endpoint":
-                    endpoint_url = message.get("data", "")
+                if not session_id and event_data.get("event") == "endpoint":
+                    endpoint_url = event_data.get("data", "")
                     if "session_id=" in endpoint_url:
                         session_id = endpoint_url.split("session_id=")[1]
-                        # 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",
-                        "data": message_data
-                    }
-                )
+                # Send the event as a raw SSE event
+                await sse_stream_writer.send({
+                    "event": event_data.get("event", "message"),
+                    "data": event_data.get("data", json.dumps(event_data))
+                })
 
     sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream(0)
     try: