semantic update

backend/src/app.py

@@ -1,8 +1,11 @@
 # [DEF:AppModule:Module]
+# @TIER: CRITICAL
 # @SEMANTICS: app, main, entrypoint, fastapi
 # @PURPOSE: The main entry point for the FastAPI application. It initializes the app, configures CORS, sets up dependencies, includes API routers, and defines the WebSocket endpoint for log streaming.
 # @LAYER: UI (API)
 # @RELATION: Depends on the dependency module and API route modules.
+# @INVARIANT: Only one FastAPI app instance exists per process.
+# @INVARIANT: All WebSocket connections must be properly cleaned up on disconnect.
 import sys
 from pathlib import Path
 
@@ -123,26 +126,65 @@ app.include_router(llm.router)
 app.include_router(storage.router, prefix="/api/storage", tags=["Storage"])
 
 # [DEF:websocket_endpoint:Function]
-# @PURPOSE: Provides a WebSocket endpoint for real-time log streaming of a task.
+# @PURPOSE: Provides a WebSocket endpoint for real-time log streaming of a task with server-side filtering.
 # @PRE: task_id must be a valid task ID.
 # @POST: WebSocket connection is managed and logs are streamed until disconnect.
+# @TIER: CRITICAL
+# @UX_STATE: Connecting -> Streaming -> (Disconnected)
 @app.websocket("/ws/logs/{task_id}")
-async def websocket_endpoint(websocket: WebSocket, task_id: str):
+async def websocket_endpoint(
+    websocket: WebSocket,
+    task_id: str,
+    source: str = None,
+    level: str = None
+):
+    """
+    WebSocket endpoint for real-time log streaming with optional server-side filtering.
+    
+    Query Parameters:
+        source: Filter logs by source component (e.g., "plugin", "superset_api")
+        level: Filter logs by minimum level (DEBUG, INFO, WARNING, ERROR)
+    """
     with belief_scope("websocket_endpoint", f"task_id={task_id}"):
         await websocket.accept()
-    logger.info(f"WebSocket connection accepted for task {task_id}")
+    
+    # Normalize filter parameters
+    source_filter = source.lower() if source else None
+    level_filter = level.upper() if level else None
+    
+    # Level hierarchy for filtering
+    level_hierarchy = {"DEBUG": 0, "INFO": 1, "WARNING": 2, "ERROR": 3}
+    min_level = level_hierarchy.get(level_filter, 0) if level_filter else 0
+    
+    logger.info(f"WebSocket connection accepted for task {task_id} (source={source_filter}, level={level_filter})")
     task_manager = get_task_manager()
     queue = await task_manager.subscribe_logs(task_id)
+    
+    def matches_filters(log_entry) -> bool:
+        """Check if log entry matches the filter criteria."""
+        # Check source filter
+        if source_filter and log_entry.source.lower() != source_filter:
+            return False
+        
+        # Check level filter
+        if level_filter:
+            log_level = level_hierarchy.get(log_entry.level.upper(), 0)
+            if log_level < min_level:
+                return False
+        
+        return True
+    
     try:
         # Stream new logs
         logger.info(f"Starting log stream for task {task_id}")
         
-        # Send initial logs first to build context
+        # Send initial logs first to build context (apply filters)
         initial_logs = task_manager.get_task_logs(task_id)
         for log_entry in initial_logs:
-            log_dict = log_entry.dict()
-            log_dict['timestamp'] = log_dict['timestamp'].isoformat()
-            await websocket.send_json(log_dict)
+            if matches_filters(log_entry):
+                log_dict = log_entry.dict()
+                log_dict['timestamp'] = log_dict['timestamp'].isoformat()
+                await websocket.send_json(log_dict)
 
         # Force a check for AWAITING_INPUT status immediately upon connection
         # This ensures that if the task is already waiting when the user connects, they get the prompt.
@@ -160,6 +202,11 @@ async def websocket_endpoint(websocket: WebSocket, task_id: str):
 
         while True:
             log_entry = await queue.get()
+            
+            # Apply server-side filtering
+            if not matches_filters(log_entry):
+                continue
+            
             log_dict = log_entry.dict()
             log_dict['timestamp'] = log_dict['timestamp'].isoformat()
             await websocket.send_json(log_dict)