use threads so that SSE works with shared memory

Author: Chenglong-MS

local_server.sh

@@ -5,4 +5,6 @@
 # export http_proxy=http://127.0.0.1:7890
 # export https_proxy=http://127.0.0.1:7890
 
-env FLASK_APP=py-src/data_formulator/app.py FLASK_RUN_PORT=5000 FLASK_RUN_HOST=0.0.0.0 flask run
+#env FLASK_APP=py-src/data_formulator/app.py FLASK_RUN_PORT=5000 FLASK_RUN_HOST=0.0.0.0 flask run
+export FLASK_RUN_PORT=5000
+python -m py-src.data_formulator.app --port ${FLASK_RUN_PORT} --dev

py-src/data_formulator/agent_routes.py

@@ -425,8 +425,6 @@ def request_code_expl():
     if request.is_json:
         logger.info("# request data: ")
         content = request.get_json()        
-        token = content["token"]
-
         client = get_client(content['model'])
 
         # each table is a dict with {"name": xxx, "rows": [...]}

py-src/data_formulator/app.py

@@ -256,6 +256,8 @@ def parse_args() -> argparse.Namespace:
         help="Whether to execute python in subprocess, it makes the app more secure (reducing the chance for the model to access the local machine), but increases the time of response")
     parser.add_argument("-d", "--disable-display-keys", action='store_true', default=False,
         help="Whether disable displaying keys in the frontend UI, recommended to turn on if you host the app not just for yourself.")
+    parser.add_argument("--dev", action='store_true', default=False,
+        help="Launch the app in development mode (prevents the app from opening the browser automatically)")
     return parser.parse_args()
 
 
@@ -268,11 +270,12 @@ def run_app():
         'disable_display_keys': args.disable_display_keys
     }
 
-    url = "http://localhost:{0}".format(args.port)
-    threading.Timer(2, lambda: webbrowser.open(url, new=2)).start()
+    if not args.dev:
+        url = "http://localhost:{0}".format(args.port)
+        threading.Timer(2, lambda: webbrowser.open(url, new=2)).start()
 
     app.run(host='0.0.0.0', port=args.port, threaded=True)
-    
+
 if __name__ == '__main__':
     #app.run(debug=True, host='127.0.0.1', port=5000)
     #use 0.0.0.0 for public

py-src/data_formulator/sse_routes.py

@@ -8,34 +8,47 @@
 from typing import Dict, Any, Generator
 from flask import Blueprint, Response, session, request, jsonify, current_app
 import threading
+import uuid
+from pprint import pprint
 
 # Get logger for this module
 logger = logging.getLogger(__name__)
 
 # Create blueprint for SSE routes
 sse_bp = Blueprint('sse', __name__, url_prefix='/api/sse')
 
-# Global dictionary to store SSE connections
-sse_connections: Dict[str, queue.Queue] = {}
+# Add a lock for thread safety
+sse_connections_lock = threading.RLock()
+sse_connections: Dict[str, Dict[str, Any]] = {}
 
 @sse_bp.route('/connect')
 def sse_connect():
     """
     SSE endpoint for clients to establish a connection and receive real-time messages
     """
     session_id = session.get('session_id')
+    connection_id = f"conn_{uuid.uuid4().hex[:8]}"
+    
     if not session_id:
         return Response("No session ID found", status=401)
 
-    logger.info(f"SSE connection established for session: {session_id}")
+    logger.info(f"[SSE Connect] Thread {threading.current_thread().name} accessing sse_connections")
+    logger.info(f"[SSE Connect] sse_connections id: {id(sse_connections)}")
 
-    # Create a queue for this session if it doesn't exist
-    if session_id not in sse_connections:
-        sse_connections[session_id] = queue.Queue()
+    # Thread-safe connection creation
+    with sse_connections_lock:
+        if session_id not in sse_connections:
+            sse_connections[session_id] = {
+                'queue': queue.Queue(),
+                'connected_clients': []
+            }
+        sse_connections[session_id]['connected_clients'].append(connection_id)
+        logger.info(f"[SSE Connect] sse_connections after creation: {sse_connections}")
 
+
     def event_stream() -> Generator[str, None, None]:
         """Generator function that yields SSE formatted messages"""
-        client_queue = sse_connections[session_id]
+        client_queue = sse_connections[session_id]['queue']
 
         # Send initial connection confirmation
         yield format_sse_message({
@@ -45,27 +58,39 @@ def event_stream() -> Generator[str, None, None]:
         })
 
         try:
+            logger.info(f"Starting event stream for connection {connection_id} for session {session_id}")
+            last_heartbeat_time = time.time()
             while True:
                 try:
-                    # Wait for messages with a timeout to allow periodic heartbeat
-                    message = client_queue.get(timeout=30)  # 30 second timeout
+                    message = client_queue.get(timeout=1)  # 1 second timeout
                     yield format_sse_message(message)
                 except queue.Empty:
                     # Send heartbeat to keep connection alive
-                    yield format_sse_message({
-                        "type": "notification",
-                        "text": "Heartbeat",
-                        "timestamp": time.time()
-                    })
+                    if time.time() - last_heartbeat_time > 30:
+                        last_heartbeat_time = time.time()
+                        yield format_sse_message({
+                            "type": "notification",
+                            "text": "Heartbeat",
+                            "timestamp": time.time()
+                        })
+                    else:
+                       # lightweight heartbeat to keep connection alive (no data)
+                       yield ": heartbeat\n\n"
                 except Exception as e:
                     logger.error(f"Error in SSE stream for session {session_id}: {e}")
                     break
         finally:
-            # Clean up connection when client disconnects
-            if session_id in sse_connections:
-                del sse_connections[session_id]
-                logger.info(f"SSE connection closed for session: {session_id}")
-    
+            # Safe cleanup with reference counting
+            with sse_connections_lock:
+                logger.info(f"[SSE Connect] cleaning up connection {connection_id} for session {session_id}")
+                logger.info(f"[SSE Connect] sse_connections before cleanup: {sse_connections}")
+                if session_id in sse_connections:
+                    sse_connections[session_id]['connected_clients'].remove(connection_id)
+                    if len(sse_connections[session_id]['connected_clients']) == 0:
+                        del sse_connections[session_id]
+                        logger.info(f"Last SSE connection ({connection_id}) closed for session {session_id}")
+                logger.info(f"[SSE Connect] sse_connections after cleanup: {sse_connections}")
+                
     return Response(
         event_stream(),
         mimetype='text/event-stream',
@@ -83,20 +108,27 @@ def send_message_to_session():
     Endpoint to send a message to a specific session via SSE
     Expected JSON payload: {
         "type": "action" | "notification" | "heartbeat",
-        "message": {...},
-        "data": {...}
+        "text": "....",
+        "data": {...} (optional)
     }
     """
-    data = request.get_json()
-    if not data:
+    content = request.get_json()
+    if not content:
         return jsonify({"error": "No JSON data provided"}), 400
 
-    target_session_id = data.get('session_id')
-    message = data.get('message', {})
+    target_session_id = content.get('session_id')
+    text = content.get('text')
+    data = content.get('data', {})
 
     if not target_session_id:
         return jsonify({"error": "session_id is required"}), 400
 
+    message = {
+        "type": "action",
+        "text": text,
+        "data": data
+    }
+    
     success = send_sse_message(target_session_id, message)
 
     if success:
@@ -127,52 +159,29 @@ def broadcast_message():
 def get_sse_status():
     """Get the current status of SSE connections"""
 
+    connection_info = {}
+    logger.info(f"[SSE Status] Thread {threading.current_thread().name} accessing sse_connections")
+    logger.info(f"[SSE Status] sse_connections id: {id(sse_connections)}")
+    logger.info(f"[SSE Status] sse_connections: {sse_connections}")
+    
+    with sse_connections_lock:
+        for session_id, connection in sse_connections.items():
+            connection_info[session_id] = {
+                'connected_clients': connection['connected_clients'],
+                'queue_size': connection['queue'].qsize()
+            }
+    
     return jsonify({
-        "connected_sessions": list(sse_connections.keys()),
-        "total_connections": len(sse_connections)
+        "connected_sessions": connection_info
     })
 
-@sse_bp.route('/trigger_notification', methods=['POST'])
-def trigger_notification():
-    """
-    Endpoint to trigger a notification to a specific session
-    Expected JSON payload: {
-        "type": "notification",
-        "text": "Notification message",
-        "data": {...} (optional)
-    }
-    """
-    data = request.get_json()
-    if not data:
-        return jsonify({"error": "No JSON data provided"}), 400
-    
-    session_id = data.get('session_id')
-    text = data.get('text')
-    data = data.get('data', {})
-    
-    # Validate required fields
-    if not session_id:
-        return jsonify({"error": "session_id is required"}), 400
-    if not text:
-        return jsonify({"error": "text is required"}), 400
-    
-    # Extract any additional data
-    additional_data = data.get('data', {})
-    
-    # Send the notification
-    success = send_notification(
-        session_id=session_id,
-        text=text,
-        data=data
-    )
-    
-    if success:
-        return jsonify({
-            "status": "Notification sent successfully",
-            "session_id": session_id,
-        })
+@sse_bp.route('/sse-connection-check', methods=['POST'])
+def sse_connection_check():
+    session_id = request.json.get('session_id')
+    if session_id in sse_connections:
+        return jsonify({"status": "connected"})
     else:
-        return jsonify({"error": "Session not found or not connected"}), 404
+        return jsonify({"status": "disconnected"}), 404
 
 # Utility functions
 
@@ -196,18 +205,18 @@ def send_sse_message(session_id: str, message: Dict[str, Any]) -> bool:
     Returns:
         bool: True if message was sent successfully, False otherwise
     """
-    
-    if session_id not in sse_connections:
-        logger.warning(f"Attempted to send message to non-existent session: {session_id}")
-        return False
-    
-    try:
-        sse_connections[session_id].put(message)
-        logger.info(f"Message sent to session {session_id}: {message.get('type', 'unknown')}")
-        return True
-    except Exception as e:
-        logger.error(f"Failed to send message to session {session_id}: {e}")
-        return False
+    with sse_connections_lock:
+        if session_id not in sse_connections:
+            logger.warning(f"Attempted to send message to non-existent session: {session_id}")
+            return False
+        
+        try:
+            sse_connections[session_id]['queue'].put(message)
+            logger.info(f"Message sent to session {session_id}: {message.get('type', 'unknown')}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to send message to session {session_id}: {e}")
+            return False
 
 def broadcast_sse_message(message: Dict[str, Any]) -> int:
     """
@@ -221,27 +230,13 @@ def broadcast_sse_message(message: Dict[str, Any]) -> int:
     """
     sent_count = 0
 
-    for session_id in list(sse_connections.keys()):
+    # Get a snapshot of session IDs to avoid iteration issues
+    with sse_connections_lock:
+        session_ids = list(sse_connections.keys())
+    
+    for session_id in session_ids:
         if send_sse_message(session_id, message):
             sent_count += 1
 
     logger.info(f"Broadcasted message to {sent_count} clients: {message.get('type', 'unknown')}")
-    return sent_count
-
-def send_notification(session_id: str, text: str, data: Dict[str, Any]):
-    """
-    Send a notification message to a specific session
-    
-    Args:
-        session_id: Target session ID
-        text: Notification text
-        data: Additional data to include in the notification
-    """
-    message = {
-        "type": "notification",
-        "text": text,
-        "data": data
-    }
-    
-    return send_sse_message(session_id, message)
-
+    return sent_count

py-src/data_formulator/tables_routes.py

@@ -59,7 +59,6 @@ def list_tables():
                 if database_name in ['system', 'temp']:
                     continue
 
-                
                 print(f"table_metadata: {table_metadata}")
 
                 try:

pyproject.toml

@@ -35,7 +35,7 @@ dependencies = [
     "python-dotenv",  
     "vega_datasets",
     "litellm",
-    "duckdb"
+    "duckdb",
 ]
 
 [project.urls]