aliakbar
249 lines
8.6 KiB
Python
249 lines
8.6 KiB
Python
import json
|
|
import asyncio
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy import update
|
|
from datetime import datetime, timezone
|
|
from app.core.database import get_db
|
|
from app.chat.redis_manager import redis_manager
|
|
from app.chat import service as chat_service
|
|
from app.users.service import get_user_by_username
|
|
from app.users.models import User
|
|
from app.core.security import settings
|
|
from jose import jwt, JWTError
|
|
|
|
router = APIRouter()
|
|
|
|
# Store active WebSocket connections per user
|
|
active_connections = {} # user_id -> list of websockets
|
|
|
|
def get_active_websockets(user_id: int):
|
|
"""Get all active WebSocket connections for a user"""
|
|
return active_connections.get(user_id, [])
|
|
|
|
def add_active_connection(user_id: int, websocket: WebSocket):
|
|
"""Add a WebSocket connection for a user"""
|
|
if user_id not in active_connections:
|
|
active_connections[user_id] = []
|
|
active_connections[user_id].append(websocket)
|
|
|
|
def remove_active_connection(user_id: int, websocket: WebSocket):
|
|
"""Remove a WebSocket connection for a user"""
|
|
if user_id in active_connections:
|
|
if websocket in active_connections[user_id]:
|
|
active_connections[user_id].remove(websocket)
|
|
if not active_connections[user_id]:
|
|
del active_connections[user_id]
|
|
|
|
async def broadcast_user_online_status(user_id: int, lab_id: int, is_online: bool):
|
|
"""Broadcast user online/offline status to lab members"""
|
|
status_update = {
|
|
"type": "user_status",
|
|
"user_id": user_id,
|
|
"lab_id": lab_id,
|
|
"is_online": is_online,
|
|
"timestamp": datetime.now(timezone.utc).isoformat()
|
|
}
|
|
await redis_manager.publish(f"lab_{lab_id}_status", json.dumps(status_update))
|
|
|
|
async def get_user_from_token(token: str, db: AsyncSession):
|
|
try:
|
|
payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
|
|
username: str = payload.get("sub")
|
|
if username is None:
|
|
return None
|
|
except JWTError:
|
|
return None
|
|
|
|
user = await get_user_by_username(db, username)
|
|
return user
|
|
|
|
@router.websocket("/ws/{lab_id}/{chat_id}")
|
|
async def websocket_endpoint(
|
|
websocket: WebSocket,
|
|
lab_id: int,
|
|
chat_id: int,
|
|
token: str,
|
|
db: AsyncSession = Depends(get_db)
|
|
):
|
|
user = await get_user_from_token(token, db)
|
|
if not user:
|
|
await websocket.close(code=1008)
|
|
return
|
|
|
|
await websocket.accept()
|
|
user_id = user.id
|
|
|
|
# Set user as online and add connection
|
|
await db.execute(
|
|
update(User)
|
|
.where(User.id == user_id)
|
|
.values(is_online=True, last_seen=datetime.now(timezone.utc))
|
|
)
|
|
await db.commit()
|
|
add_active_connection(user_id, websocket)
|
|
|
|
# Broadcast user came online
|
|
await broadcast_user_online_status(user_id, lab_id, True)
|
|
|
|
redis_task = asyncio.create_task(
|
|
redis_manager.subscribe(f"chat_{chat_id}", redis_handler)
|
|
)
|
|
|
|
# Subscribe to lab status updates
|
|
status_task = asyncio.create_task(
|
|
redis_manager.subscribe(f"lab_{lab_id}_status", status_handler)
|
|
)
|
|
|
|
def status_handler(message):
|
|
if not websocket.client_connected:
|
|
return
|
|
try:
|
|
import asyncio
|
|
asyncio.run_coroutine_threadsafe(
|
|
websocket.send_text(message),
|
|
asyncio.get_event_loop()
|
|
)
|
|
except:
|
|
pass
|
|
|
|
def redis_handler(message):
|
|
if not websocket.client_connected:
|
|
return
|
|
|
|
try:
|
|
import asyncio
|
|
asyncio.run_coroutine_threadsafe(
|
|
websocket.send_text(message),
|
|
asyncio.get_event_loop()
|
|
)
|
|
except:
|
|
pass
|
|
|
|
# Initialize bot user
|
|
bot_enabled = user.bot_username is not None and user.bot_username != ""
|
|
|
|
# Send connection confirmation
|
|
await websocket.send_text(json.dumps({"type": "connected"}))
|
|
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
message_type = data.get("type")
|
|
|
|
if message_type == "ping":
|
|
await websocket.send_text(json.dumps({"type": "pong"}))
|
|
|
|
elif message_type == "join":
|
|
pass
|
|
|
|
elif message_type == "bot_request":
|
|
pass
|
|
|
|
elif message_type == "message":
|
|
# Check if user is lab member
|
|
is_member = await chat_service.is_lab_member(db, lab_id, user_id)
|
|
if not is_member:
|
|
await websocket.send_text(json.dumps({"type": "error", "message": "Not a lab member"}))
|
|
continue
|
|
|
|
# Parse message
|
|
message_data = data.get("message", {})
|
|
content = message_data.get("content", "")
|
|
reply_to_id = message_data.get("reply_to_id")
|
|
|
|
if reply_to_id is not None:
|
|
try:
|
|
reply_to_id = int(reply_to_id)
|
|
except (ValueError, TypeError):
|
|
reply_to_id = None
|
|
|
|
# Check if bot is enabled and message mentions user's bot
|
|
has_bot_mention = False
|
|
clean_message = message_data["content"]
|
|
|
|
if bot_enabled and user.bot_username:
|
|
bot_mention_pattern = f"@{user.bot_username}"
|
|
has_bot_mention = bot_mention_pattern in message_data["content"]
|
|
if has_bot_mention:
|
|
# Remove bot mention from message
|
|
clean_message = message_data["content"].replace(bot_mention_pattern, "").strip()
|
|
|
|
# Save to DB
|
|
db_message = await chat_service.save_message(
|
|
db, chat_id, lab_id, user.id, message_data["content"],
|
|
is_bot=False, reply_to_id=reply_to_id
|
|
)
|
|
|
|
# Build reply preview if available
|
|
reply_to_preview = None
|
|
if db_message.reply_to is not None:
|
|
reply_sender = db_message.reply_to.user.username if db_message.reply_to.user else "Unknown"
|
|
reply_to_preview = {
|
|
"id": db_message.reply_to.id,
|
|
"sender_name": reply_sender,
|
|
"content": db_message.reply_to.content[:200],
|
|
}
|
|
|
|
# Construct response
|
|
response = {
|
|
"type": "user_message",
|
|
"id": db_message.id,
|
|
"sender": user.username,
|
|
"sender_id": user.id,
|
|
"content": message_data["content"],
|
|
"lab_id": lab_id,
|
|
"chat_id": chat_id,
|
|
"timestamp": db_message.timestamp.isoformat(),
|
|
"reply_to_id": reply_to_id,
|
|
"reply_to": reply_to_preview,
|
|
"sender_avatar_url": user.avatar_url,
|
|
}
|
|
|
|
# Publish to Redis for other users
|
|
await redis_manager.publish(f"chat_{chat_id}", json.dumps(response))
|
|
|
|
# If bot is enabled and mentioned, process bot message
|
|
if has_bot_mention:
|
|
# Send bot request back to the sender for local processing
|
|
bot_request = {
|
|
"type": "bot_request",
|
|
"sender": user.username,
|
|
"sender_id": user.id,
|
|
"bot_username": user.bot_username,
|
|
"original_message": message_data["content"],
|
|
"clean_message": clean_message,
|
|
"lab_id": lab_id,
|
|
"chat_id": chat_id
|
|
}
|
|
await websocket.send_text(json.dumps(bot_request))
|
|
|
|
except WebSocketDisconnect:
|
|
redis_task.cancel()
|
|
status_task.cancel()
|
|
|
|
# Set user as offline and remove connection
|
|
await db.execute(
|
|
update(User)
|
|
.where(User.id == user_id)
|
|
.values(is_online=False, last_seen=datetime.now(timezone.utc))
|
|
)
|
|
await db.commit()
|
|
remove_active_connection(user_id, websocket)
|
|
|
|
# Broadcast user went offline
|
|
await broadcast_user_online_status(user_id, lab_id, False)
|
|
|
|
except Exception as e:
|
|
redis_task.cancel()
|
|
status_task.cancel()
|
|
|
|
# Set user as offline and remove connection
|
|
await db.execute(
|
|
update(User)
|
|
.where(User.id == user_id)
|
|
.values(is_online=False, last_seen=datetime.now(timezone.utc))
|
|
)
|
|
await db.commit()
|
|
remove_active_connection(user_id, websocket)
|