Update app/main.py

app/main.py

@@ -1,10 +1,11 @@
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, Header, HTTPException
 from fastapi.responses import HTMLResponse, JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
+from fastapi.encoders import jsonable_encoder
+
 from uuid import uuid4
 from datetime import datetime, timezone, timedelta
-from typing import Optional, Dict, Any, List
+from typing import Optional
 import os
 import asyncio
 
@@ -53,7 +54,7 @@ async def create_post(
 ):
     client_ip = get_client_ip(request)
 
-    # 1. IP block check
+    # 1) IP block check
     if is_ip_blocked(client_ip):
         logger.warning(
             f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: blocked_ip - Endpoint: /api/v1/posts"
@@ -63,7 +64,7 @@ async def create_post(
             content={"error": "Too many failed authentication attempts"},
         )
 
-    # 2. API key validation
+    # 2) API key validation
     if not x_api_key or x_api_key != POSTER_KEY:
         record_auth_failure(client_ip)
         logger.warning(
@@ -71,13 +72,12 @@ async def create_post(
         )
         raise HTTPException(status_code=401, detail="Invalid API key")
 
-    # 3. Validate content length (pydantic already enforces max_length, but do a clear error)
+    # 3) Defensive length check (pydantic also enforces this)
     if len(payload.content) > 1000:
         raise HTTPException(
             status_code=413, detail="Content exceeds 1000 character limit"
         )
 
-    # 4. Create message
     now = datetime.now(timezone.utc)
     msg = Message(
         id=str(uuid4()),
@@ -88,25 +88,25 @@ async def create_post(
         metadata=payload.metadata or {},
     )
 
-    # Insert + prune by age
     async with storage_lock:
         messages.append(msg)
+
+        # Retention: expire >48h (50-message maxlen still takes precedence)
         cutoff = now - timedelta(hours=48)
-        # Age-based pruning (after maxlen already applied)
         tmp = [m for m in messages if m.timestamp >= cutoff]
         messages.clear()
         for m in tmp:
             messages.append(m)
 
-        # Broadcast to readers
-        payload_out = NewPostPayload(type="new_post", message=msg).dict()
-        stale_connections = []
+        # Broadcast to all connected readers (JSON-safe)
+        payload_out = jsonable_encoder(NewPostPayload(type="new_post", message=msg))
+        stale = []
         for ws in connected_readers:
             try:
                 await ws.send_json(payload_out)
             except Exception:
-                stale_connections.append(ws)
-        for ws in stale_connections:
+                stale.append(ws)
+        for ws in stale:
             connected_readers.discard(ws)
 
     logger.info(
@@ -126,10 +126,9 @@ async def create_post(
 
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
-    # We need to manually pull query params, including api_key
     client_ip = websocket.client.host if websocket.client else "unknown"
 
-    # 1. IP block check
+    # 1) IP block check
     if is_ip_blocked(client_ip):
         logger.warning(
             f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: blocked_ip - Endpoint: /ws"
@@ -137,16 +136,16 @@ async def websocket_endpoint(websocket: WebSocket):
         await websocket.close(code=1008)
         return
 
+    # 2) API key validation (query param)
     api_key = websocket.query_params.get("api_key")
     if not api_key or api_key != READER_KEY:
         record_auth_failure(client_ip)
         logger.warning(
             f"{datetime.utcnow().isoformat()} - AUTH_FAILURE - IP: {client_ip} - Reason: invalid_reader_key - Endpoint: /ws"
         )
-        # Send error payload once; then close.
         await websocket.accept()
         await websocket.send_json(
-            ErrorPayload(type="error", error="Invalid API key").dict()
+            jsonable_encoder(ErrorPayload(type="error", error="Invalid API key"))
         )
         await websocket.close(code=1008)
         return
@@ -157,28 +156,21 @@ async def websocket_endpoint(websocket: WebSocket):
     )
     record_auth_success(client_ip)
 
-    # Register connection
+    # Register + send initial history (JSON-safe)
     async with storage_lock:
         connected_readers.add(websocket)
-        # Send history
-        history_payload = HistoryPayload(
-            type="history",
-            messages=list(messages),
-        )
-        await websocket.send_json(history_payload.dict())
+        history_payload = HistoryPayload(type="history", messages=list(messages))
+        await websocket.send_json(jsonable_encoder(history_payload))
 
     try:
         while True:
-            # Keep the connection alive and support ping/pong at transport layer
-            # We don't expect client messages; just read and ignore or timeout.
+            # Keepalive: wait for any client message; if none, send ping
             try:
                 await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
             except asyncio.TimeoutError:
-                # Send ping as a small text; client can ignore it.
                 await websocket.send_json({"type": "ping"})
     except WebSocketDisconnect:
         pass
     finally:
         async with storage_lock:
-            if websocket in connected_readers:
-                connected_readers.discard(websocket)
+            connected_readers.discard(websocket)