aria-code 4.1.3__py3-none-any.whl

aria_relay_client.py

@@ -0,0 +1,182 @@
+#!/usr/bin/env python3
+"""
+aria_relay_client.py — 连接 Aria 中继服务器的 WebSocket 客户端
+================================================================
+本机运行，把中继服务器转发来的飞书消息交给 aria_feishu_bot 处理，
+并把 LLM 回复返回给中继服务器，再由服务器推送到飞书。
+
+启动方式:
+  python3 aria_relay_client.py              # 单次连接（断线自动重连）
+  python3 aria_relay_client.py --once       # 调试：收到第一条消息后退出
+
+所需环境变量（读取 ~/.aria/.env）:
+  ARIA_RELAY_URL        wss://relay.aria.ai（或自建服务器地址）
+  ARIA_RELAY_CLIENT_ID  setup_wizard 生成的 12 位 hex id
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import sys
+import time
+from pathlib import Path
+
+logger = logging.getLogger("aria.relay_client")
+
+# ── 加载 ~/.aria/.env ──────────────────────────────────────────────────────
+
+def _load_env() -> None:
+    env_file = Path.home() / ".aria" / ".env"
+    if env_file.exists():
+        for line in env_file.read_text().splitlines():
+            line = line.strip()
+            if line and not line.startswith("#") and "=" in line:
+                k, _, v = line.partition("=")
+                k = k.strip()
+                if k not in os.environ:
+                    os.environ[k] = v.strip()
+
+
+_load_env()
+
+_RELAY_URL = os.environ.get("ARIA_RELAY_URL", "wss://relay.aria.ai")
+_CLIENT_ID = os.environ.get("ARIA_RELAY_CLIENT_ID", "")
+_RECONNECT_DELAY_MAX = 60   # seconds
+_RECONNECT_DELAY_BASE = 3
+
+
+# ── 本地 aria_feishu_bot import ───────────────────────────────────────────────
+
+def _get_feishu_bot():
+    aria_dir = Path(__file__).parent
+    if str(aria_dir) not in sys.path:
+        sys.path.insert(0, str(aria_dir))
+    try:
+        import aria_feishu_bot
+        return aria_feishu_bot
+    except ImportError as e:
+        logger.warning("aria_feishu_bot not importable: %s", e)
+        return None
+
+
+# ── Message handler ───────────────────────────────────────────────────────────
+
+async def _handle_message(raw_msg: dict, ws) -> None:
+    """
+    Server sends:
+      {"type": "message", "id": "req_xxx", "payload": <feishu_event_dict>}
+
+    We reply:
+      {"type": "response", "id": "req_xxx", "result": <any>}
+    """
+    req_id  = raw_msg.get("id", "")
+    payload = raw_msg.get("payload", {})
+
+    bot = _get_feishu_bot()
+    if bot is None:
+        result = {"error": "aria_feishu_bot unavailable"}
+    else:
+        try:
+            result = await bot.dispatch_event(payload)
+        except Exception as e:
+            logger.exception("dispatch_event error")
+            result = {"error": str(e)[:300]}
+
+    reply = json.dumps({"type": "response", "id": req_id, "result": result})
+    await ws.send(reply)
+
+
+# ── Main loop ─────────────────────────────────────────────────────────────────
+
+async def _connect_and_serve(once: bool = False) -> None:
+    try:
+        import websockets  # type: ignore
+    except ImportError:
+        logger.error("websockets package not installed — run: pip install websockets")
+        sys.exit(1)
+
+    if not _CLIENT_ID:
+        logger.error(
+            "ARIA_RELAY_CLIENT_ID is not set. "
+            "Run setup_wizard.py to generate your client ID."
+        )
+        sys.exit(1)
+
+    delay = _RECONNECT_DELAY_BASE
+    while True:
+        try:
+            logger.info("Connecting to %s (client_id=%s)", _RELAY_URL, _CLIENT_ID)
+            async with websockets.connect(
+                _RELAY_URL,
+                ping_interval=30,
+                ping_timeout=10,
+                open_timeout=15,
+            ) as ws:
+                # Register with server
+                await ws.send(json.dumps({
+                    "type": "register",
+                    "client_id": _CLIENT_ID,
+                }))
+                ack_raw = await asyncio.wait_for(ws.recv(), timeout=10)
+                ack = json.loads(ack_raw)
+                if not ack.get("ok"):
+                    logger.error("Registration rejected: %s", ack.get("reason", "unknown"))
+                    await asyncio.sleep(delay)
+                    continue
+
+                logger.info("Registered. Waiting for messages…")
+                delay = _RECONNECT_DELAY_BASE  # reset on success
+
+                async for raw in ws:
+                    try:
+                        msg = json.loads(raw)
+                    except json.JSONDecodeError:
+                        logger.warning("Invalid JSON from relay: %r", raw[:100])
+                        continue
+
+                    if msg.get("type") == "ping":
+                        await ws.send(json.dumps({"type": "pong"}))
+                        continue
+
+                    if msg.get("type") == "message":
+                        asyncio.create_task(_handle_message(msg, ws))
+
+                    if once:
+                        return
+
+        except (OSError, ConnectionRefusedError) as e:
+            logger.warning("Connection failed: %s — retry in %ds", e, delay)
+        except asyncio.CancelledError:
+            logger.info("Relay client cancelled")
+            return
+        except Exception as e:
+            logger.warning("Relay error: %s — retry in %ds", e, delay)
+
+        await asyncio.sleep(delay)
+        delay = min(delay * 2, _RECONNECT_DELAY_MAX)
+
+
+def main() -> None:
+    import argparse
+    parser = argparse.ArgumentParser(description="Aria 中继客户端")
+    parser.add_argument("--once", action="store_true", help="接收一条消息后退出（调试用）")
+    parser.add_argument("--verbose", "-v", action="store_true")
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=logging.DEBUG if args.verbose else logging.INFO,
+        format="%(asctime)s [aria-relay] %(message)s",
+        datefmt="%H:%M:%S",
+    )
+
+    try:
+        asyncio.run(_connect_and_serve(once=args.once))
+    except KeyboardInterrupt:
+        pass
+
+
+if __name__ == "__main__":
+    main()

aria_relay_server.py

@@ -0,0 +1,405 @@
+#!/usr/bin/env python3
+"""
+aria_relay_server.py — Aria 中继服务器
+=======================================
+你（产品方）部署一次，所有用户共用。
+
+架构:
+  用户飞书消息
+    → 飞书云 → POST /feishu/event
+    → relay_server 查 feishu_user_id → 对应 WebSocket
+    → relay_server 转发 payload 到用户本机
+    → aria_relay_client.dispatch_event()
+    → Aria LLM 回复
+    → relay_server 收到 response
+    → relay_server 调飞书 reply_message API
+    → 飞书消息卡片展示给用户
+
+WebSocket 注册流程:
+  1. 用户本机 aria_relay_client 连接 wss://relay.yourdomain.com
+  2. 发送: {"type": "register", "client_id": "aria-xxxxxxxxxxxx"}
+  3. 用户在飞书向 Aria Bot 发: /bind ARIA-BIND-ARIA-XXXXXXXXXXXX
+  4. 服务器记录: feishu_user_id → client_id
+  5. 后续消息经 WebSocket 透传
+
+依赖:
+  pip install fastapi uvicorn websockets httpx
+
+启动:
+  python3 aria_relay_server.py
+  # 或
+  uvicorn aria_relay_server:app --host 0.0.0.0 --port 8765
+
+所需环境变量:
+  FEISHU_APP_ID         飞书应用 App ID
+  FEISHU_APP_SECRET     飞书应用 App Secret
+  RELAY_SECRET          WebSocket 注册鉴权（可选，留空则不鉴权）
+  DB_PATH               SQLite 路径（默认 ./relay.db）
+  MESSAGE_TIMEOUT       等待用户本机回复的超时秒数（默认 90）
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import logging
+import os
+import sqlite3
+import time
+import uuid
+from contextlib import asynccontextmanager
+from typing import Any, Optional
+
+import httpx
+from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
+from fastapi.responses import JSONResponse
+
+logger = logging.getLogger("aria.relay_server")
+
+# ── Config ────────────────────────────────────────────────────────────────────
+
+_FEISHU_APP_ID     = os.environ.get("FEISHU_APP_ID", "")
+_FEISHU_APP_SECRET = os.environ.get("FEISHU_APP_SECRET", "")
+_RELAY_SECRET      = os.environ.get("RELAY_SECRET", "")
+_DB_PATH           = os.environ.get("DB_PATH", "./relay.db")
+_MSG_TIMEOUT       = int(os.environ.get("MESSAGE_TIMEOUT", "90"))
+_FEISHU_API        = "https://open.feishu.cn/open-apis"
+
+
+# ── SQLite store ──────────────────────────────────────────────────────────────
+
+def _db() -> sqlite3.Connection:
+    conn = sqlite3.connect(_DB_PATH, check_same_thread=False)
+    conn.row_factory = sqlite3.Row
+    conn.execute("""
+        CREATE TABLE IF NOT EXISTS bindings (
+            feishu_user_id TEXT PRIMARY KEY,
+            client_id      TEXT NOT NULL,
+            bound_at       REAL NOT NULL
+        )
+    """)
+    conn.execute("""
+        CREATE TABLE IF NOT EXISTS pending_binds (
+            client_id  TEXT PRIMARY KEY,
+            created_at REAL NOT NULL
+        )
+    """)
+    conn.commit()
+    return conn
+
+
+_db_conn: Optional[sqlite3.Connection] = None
+
+
+def get_db() -> sqlite3.Connection:
+    global _db_conn
+    if _db_conn is None:
+        _db_conn = _db()
+    return _db_conn
+
+
+def _lookup_client(feishu_user_id: str) -> Optional[str]:
+    row = get_db().execute(
+        "SELECT client_id FROM bindings WHERE feishu_user_id = ?",
+        (feishu_user_id,),
+    ).fetchone()
+    return row["client_id"] if row else None
+
+
+def _bind(feishu_user_id: str, client_id: str) -> None:
+    get_db().execute(
+        "INSERT OR REPLACE INTO bindings VALUES (?, ?, ?)",
+        (feishu_user_id, client_id, time.time()),
+    )
+    get_db().execute(
+        "DELETE FROM pending_binds WHERE client_id = ?", (client_id,)
+    )
+    get_db().commit()
+
+
+def _register_pending(client_id: str) -> None:
+    get_db().execute(
+        "INSERT OR REPLACE INTO pending_binds VALUES (?, ?)",
+        (client_id, time.time()),
+    )
+    get_db().commit()
+
+
+def _is_valid_client_id(client_id: str) -> bool:
+    row = get_db().execute(
+        "SELECT client_id FROM pending_binds WHERE client_id = ?", (client_id,)
+    ).fetchone()
+    return row is not None
+
+
+# ── WebSocket connection registry ─────────────────────────────────────────────
+
+_connections: dict[str, WebSocket] = {}   # client_id → WebSocket
+_pending_responses: dict[str, asyncio.Future] = {}   # request_id → Future
+
+
+# ── Feishu API helpers ────────────────────────────────────────────────────────
+
+_feishu_token_cache: dict[str, Any] = {}
+
+
+async def _get_tenant_token() -> str:
+    now = time.time()
+    if _feishu_token_cache.get("expires_at", 0) > now + 60:
+        return _feishu_token_cache["token"]
+
+    async with httpx.AsyncClient() as client:
+        resp = await client.post(
+            f"{_FEISHU_API}/auth/v3/tenant_access_token/internal",
+            json={"app_id": _FEISHU_APP_ID, "app_secret": _FEISHU_APP_SECRET},
+            timeout=10,
+        )
+        data = resp.json()
+
+    token = data.get("tenant_access_token", "")
+    expire = int(data.get("expire", 7200))
+    _feishu_token_cache.update({"token": token, "expires_at": now + expire})
+    return token
+
+
+async def _reply_feishu(message_id: str, content: str, color: str = "blue") -> None:
+    token = await _get_tenant_token()
+    card = {
+        "msg_type": "interactive",
+        "card": {
+            "header": {
+                "title": {"tag": "plain_text", "content": "Aria"},
+                "template": color,
+            },
+            "elements": [{"tag": "div", "text": {"tag": "lark_md", "content": content}}],
+        },
+    }
+    async with httpx.AsyncClient() as client:
+        await client.post(
+            f"{_FEISHU_API}/im/v1/messages/{message_id}/reply",
+            headers={"Authorization": f"Bearer {token}"},
+            json=card,
+            timeout=15,
+        )
+
+
+async def _send_feishu_text(open_id: str, text: str) -> None:
+    token = await _get_tenant_token()
+    async with httpx.AsyncClient() as client:
+        await client.post(
+            f"{_FEISHU_API}/im/v1/messages",
+            headers={"Authorization": f"Bearer {token}"},
+            params={"receive_id_type": "open_id"},
+            json={
+                "receive_id": open_id,
+                "msg_type": "text",
+                "content": json.dumps({"text": text}),
+            },
+            timeout=15,
+        )
+
+
+# ── Route message to local aria instance ─────────────────────────────────────
+
+async def _route_to_local(feishu_user_id: str, payload: dict) -> Optional[Any]:
+    """Forward Feishu event to the user's connected local aria instance."""
+    client_id = _lookup_client(feishu_user_id)
+    if not client_id:
+        return None
+
+    ws = _connections.get(client_id)
+    if not ws:
+        return None
+
+    req_id = f"req_{uuid.uuid4().hex[:10]}"
+    future: asyncio.Future = asyncio.get_event_loop().create_future()
+    _pending_responses[req_id] = future
+
+    try:
+        await ws.send_text(json.dumps({
+            "type": "message",
+            "id": req_id,
+            "payload": payload,
+        }))
+        result = await asyncio.wait_for(future, timeout=_MSG_TIMEOUT)
+        return result
+    except asyncio.TimeoutError:
+        logger.warning("Timeout waiting for response from client_id=%s", client_id)
+        return {"error": "Aria 本机响应超时，请检查 aria_relay_client 是否在线"}
+    finally:
+        _pending_responses.pop(req_id, None)
+
+
+# ── FastAPI app ───────────────────────────────────────────────────────────────
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [relay] %(message)s",
+        datefmt="%H:%M:%S",
+    )
+    logger.info("Aria Relay Server started  db=%s", _DB_PATH)
+    get_db()  # init tables
+    yield
+    logger.info("Relay Server shutting down")
+
+
+app = FastAPI(title="Aria Relay Server", lifespan=lifespan)
+
+
+# ── WebSocket endpoint (users' local machines) ────────────────────────────────
+
+@app.websocket("/ws")
+async def ws_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    client_id: Optional[str] = None
+
+    try:
+        # First message must be register
+        raw = await asyncio.wait_for(websocket.receive_text(), timeout=15)
+        msg = json.loads(raw)
+
+        if msg.get("type") != "register":
+            await websocket.send_text(json.dumps({"ok": False, "reason": "first message must be register"}))
+            return
+
+        client_id = msg.get("client_id", "")
+        if not client_id:
+            await websocket.send_text(json.dumps({"ok": False, "reason": "client_id required"}))
+            return
+
+        # Validate RELAY_SECRET if configured
+        if _RELAY_SECRET and msg.get("secret") != _RELAY_SECRET:
+            await websocket.send_text(json.dumps({"ok": False, "reason": "invalid secret"}))
+            return
+
+        # Register in-memory + mark as pending bind (if first time)
+        _connections[client_id] = websocket
+        if _lookup_client.__module__:  # always true; used as noop to be explicit
+            _register_pending(client_id)
+
+        await websocket.send_text(json.dumps({"ok": True, "client_id": client_id}))
+        logger.info("Client connected: %s", client_id)
+
+        # Message loop
+        async for raw in websocket.iter_text():
+            try:
+                response_msg = json.loads(raw)
+            except json.JSONDecodeError:
+                continue
+
+            if response_msg.get("type") == "pong":
+                continue
+
+            if response_msg.get("type") == "response":
+                req_id = response_msg.get("id", "")
+                future = _pending_responses.get(req_id)
+                if future and not future.done():
+                    future.set_result(response_msg.get("result"))
+
+    except WebSocketDisconnect:
+        pass
+    except asyncio.TimeoutError:
+        logger.warning("Register timeout for new connection")
+    except Exception as e:
+        logger.exception("WebSocket error: %s", e)
+    finally:
+        if client_id:
+            _connections.pop(client_id, None)
+            logger.info("Client disconnected: %s", client_id)
+
+
+# ── Feishu event endpoint ─────────────────────────────────────────────────────
+
+@app.post("/feishu/event")
+async def feishu_event(request: Request):
+    """Feishu Developer Console → Event Subscription → Request URL: /feishu/event"""
+    body = await request.body()
+    try:
+        payload = json.loads(body)
+    except Exception:
+        raise HTTPException(400, "Invalid JSON")
+
+    # URL verification challenge
+    if "challenge" in payload:
+        return {"challenge": payload["challenge"]}
+
+    # Extract sender + message_id
+    event = payload.get("event", {})
+    message = event.get("message", {})
+    sender  = event.get("sender", {})
+    feishu_user_id = sender.get("sender_id", {}).get("open_id", "")
+    message_id     = message.get("message_id", "")
+
+    if not feishu_user_id:
+        return {"code": 0}
+
+    # Handle /bind command
+    msg_type = message.get("message_type", "")
+    if msg_type == "text":
+        try:
+            text_content = json.loads(message.get("content", "{}")).get("text", "").strip()
+        except Exception:
+            text_content = ""
+
+        if text_content.upper().startswith("/BIND ") or text_content.upper().startswith("ARIA-BIND-"):
+            raw_code = text_content.upper().replace("/BIND ", "").strip()
+            # Normalize: "ARIA-BIND-ARIA-XXXX" → "aria-xxxx"
+            client_id_upper = raw_code.replace("ARIA-BIND-", "").replace("ARIA-", "aria-").lower()
+            _bind(feishu_user_id, client_id_upper)
+            await _send_feishu_text(
+                feishu_user_id,
+                f"✅ 绑定成功！你的 Aria 实例已连接。\n"
+                f"Client ID: {client_id_upper}\n"
+                f"现在可以直接发消息与你的 Aria 交互了。"
+            )
+            return {"code": 0}
+
+    # Route to local aria instance
+    result = await _route_to_local(feishu_user_id, payload)
+
+    if result is None:
+        # No binding found or client offline
+        if not _lookup_client(feishu_user_id):
+            await _send_feishu_text(
+                feishu_user_id,
+                "👋 你好！要开始使用 Aria，请：\n"
+                "1. 在你的电脑上安装 Aria Code\n"
+                "2. 运行 `python3 setup_wizard.py` 完成配置\n"
+                "3. 发送绑定码绑定你的账户"
+            )
+        else:
+            await _reply_feishu(
+                message_id,
+                "⚠️ Aria 本机未连接。请确保你的电脑上 `aria_relay_client.py` 正在运行。",
+                color="yellow",
+            )
+
+    return {"code": 0}
+
+
+# ── Status endpoint ───────────────────────────────────────────────────────────
+
+@app.get("/status")
+async def status():
+    return {
+        "connected_clients": len(_connections),
+        "client_ids": list(_connections.keys()),
+        "total_bindings": get_db().execute("SELECT COUNT(*) FROM bindings").fetchone()[0],
+        "feishu_app_configured": bool(_FEISHU_APP_ID),
+    }
+
+
+# ── Entry point ───────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(
+        "aria_relay_server:app",
+        host="0.0.0.0",
+        port=int(os.environ.get("PORT", "8765")),
+        reload=False,
+        log_level="info",
+    )