Pine
152 lines
5.1 KiB
Python
152 lines
5.1 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""IM 对外桥接路由（server-core 作为唯一对外入口）。
|
||
|
||
- ``/im/ws``：WebSocket 桥接到 im-service（透传 JWT）
|
||
- ``/im/internal/task-sync|park-sync``：内部同步（供上层按需触发）
|
||
- ``/im/{path:path}``：REST 转发到 im-service ``/api/v1/{path}``（透传 JWT）
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
|
||
from fastapi import APIRouter, HTTPException, Request, WebSocket, WebSocketDisconnect
|
||
|
||
from . import client
|
||
|
||
log = logging.getLogger("im")
|
||
|
||
router = APIRouter(prefix="/im", tags=["im"])
|
||
|
||
|
||
def _auth(request: Request) -> str:
|
||
return request.headers.get("Authorization", "")
|
||
|
||
|
||
# ── 内部同步（显式优先于 catch-all） ──────────────────────────
|
||
|
||
@router.get("/health")
|
||
async def im_health():
|
||
try:
|
||
data = await client._request(method="GET", path="/health")
|
||
data["ok"] = True
|
||
return data
|
||
except client.IMError as exc:
|
||
return {"ok": False, "detail": str(exc)}
|
||
|
||
|
||
@router.post("/internal/task-sync")
|
||
async def internal_task_sync(request: Request):
|
||
body = await _json_body(request)
|
||
task_id = (body or {}).get("task_id", "")
|
||
if not task_id:
|
||
raise HTTPException(status_code=400, detail="task_id 必填")
|
||
ok = await client.sync_task_group(task_id, (body or {}).get("task_title", ""))
|
||
if not ok:
|
||
raise HTTPException(status_code=502, detail="IM 服务不可用")
|
||
return {"ok": True}
|
||
|
||
|
||
@router.post("/internal/park-sync")
|
||
async def internal_park_sync(request: Request):
|
||
body = await _json_body(request)
|
||
park_id = (body or {}).get("park_id", "")
|
||
if not park_id:
|
||
raise HTTPException(status_code=400, detail="park_id 必填")
|
||
ok = await client.sync_park_group(park_id, (body or {}).get("park_name", ""))
|
||
if not ok:
|
||
raise HTTPException(status_code=502, detail="IM 服务不可用")
|
||
return {"ok": True}
|
||
|
||
|
||
# ── REST 转发 ────────────────────────────────────────────────
|
||
|
||
@router.api_route("/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE"])
|
||
async def im_forward(path: str, request: Request):
|
||
if path.startswith("internal/"):
|
||
raise HTTPException(status_code=404, detail="unknown internal endpoint")
|
||
body = None
|
||
if request.method in ("POST", "PUT", "PATCH"):
|
||
try:
|
||
body = await request.json()
|
||
except Exception: # noqa: BLE001
|
||
body = None
|
||
# 透传 query 参数（搜索/分页等）
|
||
params = dict(request.query_params) or None
|
||
try:
|
||
return await client.forward(
|
||
"/" + path,
|
||
auth_header=_auth(request),
|
||
method=request.method,
|
||
body=body if body is not None else None,
|
||
params=params,
|
||
)
|
||
except client.IMError as exc:
|
||
raise HTTPException(status_code=502, detail=str(exc)) from exc
|
||
|
||
|
||
# ── WebSocket 桥接 ───────────────────────────────────────────
|
||
|
||
@router.websocket("/ws")
|
||
async def im_ws_bridge(websocket: WebSocket) -> None:
|
||
import websockets
|
||
|
||
token = websocket.query_params.get("token", "")
|
||
await websocket.accept()
|
||
im_ws: websockets.WebSocketClientProtocol | None = None
|
||
try:
|
||
url = f"{client._base()}/ws?token={token}"
|
||
im_ws = await websockets.connect(url, max_size=16 * 1024 * 1024)
|
||
except Exception as exc: # noqa: BLE001
|
||
log.warning("IM WS 桥接上游连接失败: %s", exc)
|
||
await websocket.close(code=1011, reason="IM service unavailable")
|
||
return
|
||
|
||
async def client_to_im() -> None:
|
||
try:
|
||
while True:
|
||
raw = await websocket.receive_text()
|
||
await im_ws.send(raw)
|
||
except (WebSocketDisconnect, Exception): # noqa: BLE001
|
||
pass
|
||
|
||
async def im_to_client() -> None:
|
||
import json as _json
|
||
from .client import resolve_avatar_fields
|
||
try:
|
||
while True:
|
||
raw = await im_ws.recv()
|
||
try:
|
||
parsed = _json.loads(raw)
|
||
if isinstance(parsed, (dict, list)):
|
||
raw = _json.dumps(resolve_avatar_fields(parsed), ensure_ascii=False)
|
||
except (ValueError, TypeError):
|
||
pass # 非 JSON 帧（如 ping）原样透传
|
||
await websocket.send_text(raw)
|
||
except (websockets.ConnectionClosed, Exception): # noqa: BLE001
|
||
pass
|
||
|
||
tasks = [
|
||
asyncio.create_task(client_to_im()),
|
||
asyncio.create_task(im_to_client()),
|
||
]
|
||
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
|
||
for t in pending:
|
||
t.cancel()
|
||
try:
|
||
await websocket.close()
|
||
except Exception: # noqa: BLE001
|
||
pass
|
||
if im_ws is not None:
|
||
try:
|
||
await im_ws.close()
|
||
except Exception: # noqa: BLE001
|
||
pass
|
||
|
||
|
||
async def _json_body(request: Request) -> dict | None:
|
||
try:
|
||
return await request.json()
|
||
except Exception: # noqa: BLE001
|
||
return None
|