renjianbo
876789fac1
- Backend: workspace_id isolation for 14 model tables + safe migration/backfill - Backend: RBAC system with 4 roles and 23 permissions, seeded on startup - Backend: workspace admin endpoints (list/manage all workspaces) - Backend: admin user management API (CRUD, reset password) - Backend: billing API with subscription plans, usage tracking, rate limiting - Backend: fix system_logs.py UNION query and wrong column references - Backend: WebSocket JWT auth and workspace enforcement - Frontend: sidebar navigation replacing top dropdown menu - Frontend: user management page (Users.vue) for admins - Frontend: enhanced Workspaces.vue with admin table view - Frontend: workspace RBAC computed properties in user store - Android: agent marketplace, billing/subscription UI, onboarding wizard - Android: phone login, analytics tracker, crash handler, network diagnostics - Android: splash screen, encrypted token storage, app update enhancements - Docs: multi-tenant RBAC guide with 8 sections and role-permission matrix Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
220 lines
8.0 KiB
Python
220 lines
8.0 KiB
Python
"""
|
||
WebSocket API — 实时推送执行状态（workspace 隔离）
|
||
"""
|
||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
|
||
from app.websocket.manager import websocket_manager
|
||
from app.core.database import SessionLocal
|
||
from app.core.security import decode_access_token
|
||
from app.models.execution import Execution
|
||
from app.models.user import User
|
||
from app.models.workspace import WorkspaceMembership
|
||
from typing import Optional
|
||
import json
|
||
import asyncio
|
||
import logging
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
router = APIRouter()
|
||
|
||
|
||
def _get_progress_from_redis(execution_id: str) -> Optional[dict]:
|
||
"""从 Redis 读取进度数据。"""
|
||
try:
|
||
from app.core.redis_client import get_redis_client
|
||
redis_client = get_redis_client()
|
||
if redis_client:
|
||
raw = redis_client.get(f"workflow:progress:{execution_id}")
|
||
if raw:
|
||
return json.loads(raw)
|
||
except Exception:
|
||
pass
|
||
return None
|
||
|
||
|
||
def _validate_ws_token(token: Optional[str], execution_id: str) -> dict:
|
||
"""验证 WebSocket JWT token，返回 {user_id, workspace_id}。
|
||
|
||
失败时返回包含 error/code 的字典，调用方应据此拒绝连接。
|
||
"""
|
||
if not token:
|
||
return {"error": "缺少认证令牌", "code": 4001}
|
||
|
||
payload = decode_access_token(token)
|
||
if payload is None:
|
||
return {"error": "无效的访问令牌", "code": 4001}
|
||
|
||
user_id = payload.get("sub")
|
||
ws_id = payload.get("ws", "")
|
||
if not user_id:
|
||
return {"error": "无效的令牌载荷", "code": 4001}
|
||
|
||
# 验证用户存在且未被删除
|
||
db = SessionLocal()
|
||
try:
|
||
user = db.query(User).filter(User.id == user_id).first()
|
||
if not user:
|
||
return {"error": "用户不存在", "code": 4001}
|
||
if user.status == "deleted":
|
||
return {"error": "账号已注销", "code": 4003}
|
||
|
||
# 平台管理员绕过 workspace 检查
|
||
if user.role == "admin":
|
||
return {"user_id": user_id, "workspace_id": ws_id, "is_admin": True}
|
||
|
||
# 验证 workspace 成员资格
|
||
if ws_id:
|
||
membership = (
|
||
db.query(WorkspaceMembership)
|
||
.filter(
|
||
WorkspaceMembership.workspace_id == ws_id,
|
||
WorkspaceMembership.user_id == user_id,
|
||
)
|
||
.first()
|
||
)
|
||
if not membership:
|
||
return {"error": "无权访问此工作区", "code": 4003}
|
||
|
||
# 验证 execution 属于该 workspace
|
||
execution = db.query(Execution).filter(Execution.id == execution_id).first()
|
||
if execution and ws_id and execution.workspace_id and execution.workspace_id != ws_id:
|
||
return {"error": "无权查看此执行记录", "code": 4003}
|
||
|
||
return {"user_id": user_id, "workspace_id": ws_id}
|
||
finally:
|
||
db.close()
|
||
|
||
|
||
@router.websocket("/api/v1/ws/executions/{execution_id}")
|
||
async def websocket_execution_status(
|
||
websocket: WebSocket,
|
||
execution_id: str,
|
||
token: Optional[str] = Query(None),
|
||
):
|
||
"""
|
||
WebSocket 实时推送执行状态（workspace 隔离）。
|
||
|
||
需要 JWT token 通过 query 参数传递: ?token=xxx
|
||
连接建立前验证：JWT 有效性 → workspace 成员资格 → execution workspace 归属。
|
||
"""
|
||
# ── 认证 + 鉴权 ──
|
||
auth = _validate_ws_token(token, execution_id)
|
||
if "error" in auth:
|
||
await websocket.accept()
|
||
await websocket.send_json({"type": "error", "message": auth["error"], "code": auth["code"]})
|
||
await websocket.close(code=auth["code"])
|
||
return
|
||
|
||
user_id = auth["user_id"]
|
||
ws_id = auth.get("workspace_id", "")
|
||
|
||
await websocket_manager.connect(websocket, execution_id, workspace_id=ws_id)
|
||
|
||
db = SessionLocal()
|
||
|
||
try:
|
||
# 发送初始状态
|
||
execution = db.query(Execution).filter(Execution.id == execution_id).first()
|
||
if execution:
|
||
# 尝试从 Redis 读取进度
|
||
redis_progress = _get_progress_from_redis(execution_id)
|
||
if redis_progress:
|
||
await websocket_manager.send_personal_message({
|
||
"type": "progress",
|
||
"execution_id": execution_id,
|
||
**redis_progress,
|
||
}, websocket)
|
||
else:
|
||
await websocket_manager.send_personal_message({
|
||
"type": "status",
|
||
"execution_id": execution_id,
|
||
"status": execution.status,
|
||
"progress": 0,
|
||
"message": "连接已建立"
|
||
}, websocket)
|
||
else:
|
||
await websocket_manager.send_personal_message({
|
||
"type": "error",
|
||
"message": f"执行记录 {execution_id} 不存在"
|
||
}, websocket)
|
||
await websocket.close()
|
||
return
|
||
|
||
last_progress = -1
|
||
|
||
# 持续监听并推送状态更新
|
||
while True:
|
||
try:
|
||
# 接收客户端消息（心跳等），超时 1 秒以便轮询进度
|
||
try:
|
||
data = await asyncio.wait_for(websocket.receive_text(), timeout=1.0)
|
||
try:
|
||
message = json.loads(data)
|
||
if message.get("type") == "ping":
|
||
await websocket_manager.send_personal_message({
|
||
"type": "pong"
|
||
}, websocket)
|
||
except Exception:
|
||
pass
|
||
except asyncio.TimeoutError:
|
||
pass # 超时后轮询进度
|
||
|
||
except WebSocketDisconnect:
|
||
break
|
||
|
||
# 优先从 Redis 读取进度（推送模式）
|
||
redis_progress = _get_progress_from_redis(execution_id)
|
||
if redis_progress:
|
||
pct = redis_progress.get("progress", -1)
|
||
if pct != last_progress:
|
||
last_progress = pct
|
||
await websocket_manager.send_personal_message({
|
||
"type": "progress",
|
||
"execution_id": execution_id,
|
||
**redis_progress,
|
||
}, websocket)
|
||
else:
|
||
# Redis 不可用时回退到 DB 轮询
|
||
db.refresh(execution)
|
||
pct = 100 if execution.status in ["completed", "failed"] else (50 if execution.status == "running" else 0)
|
||
if pct != last_progress:
|
||
last_progress = pct
|
||
await websocket_manager.send_personal_message({
|
||
"type": "status",
|
||
"execution_id": execution_id,
|
||
"status": execution.status,
|
||
"progress": pct,
|
||
"message": f"执行中..." if execution.status == "running" else "等待执行"
|
||
}, websocket)
|
||
|
||
# 如果执行完成或失败，发送最终状态并断开
|
||
db.refresh(execution)
|
||
if execution.status in ["completed", "failed"]:
|
||
await websocket_manager.send_personal_message({
|
||
"type": "status",
|
||
"execution_id": execution_id,
|
||
"status": execution.status,
|
||
"progress": 100,
|
||
"result": execution.output_data if execution.status == "completed" else None,
|
||
"error": execution.error_message if execution.status == "failed" else None,
|
||
"execution_time": execution.execution_time
|
||
}, websocket)
|
||
|
||
await asyncio.sleep(1)
|
||
break
|
||
|
||
except WebSocketDisconnect:
|
||
pass
|
||
except Exception as e:
|
||
logger.error("WebSocket 错误: %s", e)
|
||
try:
|
||
await websocket_manager.send_personal_message({
|
||
"type": "error",
|
||
"message": f"发生错误: {str(e)}"
|
||
}, websocket)
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
websocket_manager.disconnect(websocket, execution_id)
|
||
db.close()
|