# -*- 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