feat: Read, Save docs and chat fixes (#5865)

jyejare · web-flow · commit 2081b55de32b · 2026-01-16T07:08:16.000-05:00
Signed-off-by: jyejare &lt;jyejare@redhat.com&gt;
diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py
@@ -542,7 +542,7 @@ async def chat(request: ChatRequest):
         # For now, just return dummy text
         return {"response": "This is a dummy response from the Feast feature server."}
 
-    @app.post("/read-document")
+    @app.post("/read-document", dependencies=[Depends(inject_user_details)])
     async def read_document_endpoint(request: ReadDocumentRequest):
         try:
             import os
@@ -557,7 +557,7 @@ async def read_document_endpoint(request: ReadDocumentRequest):
         except Exception as e:
             return {"error": str(e)}
 
-    @app.post("/save-document")
+    @app.post("/save-document", dependencies=[Depends(inject_user_details)])
     async def save_document_endpoint(request: SaveDocumentRequest):
         try:
             import json
@@ -668,12 +668,41 @@ async def send_message(self, message: str, websocket: WebSocket):
 
     manager = ConnectionManager()
 
+    MAX_WS_CONNECTIONS = 5
+    MAX_MESSAGE_SIZE = 4096
+    MAX_MESSAGES_PER_MINUTE = 60
+    WS_READ_TIMEOUT_SEC = 60
+
     @app.websocket("/ws/chat")
     async def websocket_endpoint(websocket: WebSocket):
+        if len(manager.active_connections) >= MAX_WS_CONNECTIONS:
+            await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+            return
+
         await manager.connect(websocket)
+        message_timestamps: List[float] = []
         try:
             while True:
-                message = await websocket.receive_text()
+                try:
+                    message = await asyncio.wait_for(
+                        websocket.receive_text(), timeout=WS_READ_TIMEOUT_SEC
+                    )
+                except asyncio.TimeoutError:
+                    await websocket.close(code=status.WS_1001_GOING_AWAY)
+                    return
+
+                if len(message) > MAX_MESSAGE_SIZE:
+                    await websocket.close(code=status.WS_1009_MESSAGE_TOO_BIG)
+                    return
+
+                now = time.time()
+                cutoff = now - 60
+                message_timestamps = [ts for ts in message_timestamps if ts >= cutoff]
+                if len(message_timestamps) >= MAX_MESSAGES_PER_MINUTE:
+                    await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+                    return
+                message_timestamps.append(now)
+
                 # Process the received message (currently unused but kept for future implementation)
                 # For now, just return dummy text
                 response = f"You sent: '{message}'. This is a dummy response from the Feast feature server."
