"""ArrayClient - 列阵平台客户端，封装所有 REST API 和 WebSocket 通信""" import asyncio import base64 import json from typing import AsyncGenerator import httpx import websockets from .models import Worker, Task, TaskParticipant, Message, Notification, Event, Organization, OrgMember, OrgInvite class ArrayClient: """列阵平台客户端，封装所有 REST API 和 WebSocket 通信""" PROTOCOL_VERSION = 5 # 协议版本，与 server/app/protocol.py 对应 def __init__(self, server_url: str, api_key: str): self.server_url = server_url.rstrip("/") self.api_key = api_key self.worker_id: str | None = None self._access_token: str | None = None self._refresh_token: str | None = None self._http: httpx.AsyncClient | None = None self._ws_connected_event = asyncio.Event() async def __aenter__(self): self._http = httpx.AsyncClient(base_url=self.server_url, timeout=30.0, follow_redirects=True) await self._authenticate() return self async def __aexit__(self, *args): if self._http: await self._http.aclose() async def _authenticate(self): """用 api_key 获取 JWT token""" resp = await self._http.post("/api/v1/workers/token", json={"api_key": self.api_key}) resp.raise_for_status() data = resp.json() self._access_token = data["access_token"] self._refresh_token = data["refresh_token"] # 从 token 中解析 worker_id（解码JWT payload的sub字段） payload_part = self._access_token.split(".")[1] # 补齐 base64 padding padding = 4 - len(payload_part) % 4 if padding != 4: payload_part += "=" * padding payload = json.loads(base64.urlsafe_b64decode(payload_part)) self.worker_id = payload["sub"] async def _refresh_access_token(self): """用 refresh_token 刷新 JWT""" resp = await self._http.post("/api/v1/workers/token/refresh", json={"refresh_token": self._refresh_token}) resp.raise_for_status() data = resp.json() self._access_token = data["access_token"] self._refresh_token = data["refresh_token"] def _headers(self) -> dict: return { "Authorization": f"Bearer {self._access_token}", "X-Protocol-Version": str(self.PROTOCOL_VERSION) } async def _request(self, method: str, path: str, **kwargs) -> httpx.Response: """统一请求方法，自动处理 token 刷新""" resp = await self._http.request(method, path, headers=self._headers(), **kwargs) if resp.status_code == 401: # token 可能过期，尝试刷新 await self._refresh_access_token() resp = await self._http.request(method, path, headers=self._headers(), **kwargs) resp.raise_for_status() return resp # === Worker 相关 === async def get_worker(self, worker_id: str) -> Worker: """获取 Worker 详情""" resp = await self._request("GET", f"/api/v1/workers/{worker_id}") return Worker(**resp.json()) async def list_workers(self, type: str | None = None, status: str | None = None, org_id: str | None = None, page: int = 1, limit: int = 20) -> tuple[list[Worker], int]: """获取 Worker 列表 参数： type: 按类型筛选 status: 按状态筛选 org_id: 按所属组织筛选 page: 页码，默认 1 limit: 每页数量，默认 20 """ params = {"page": page, "limit": limit} if type: params["type"] = type if status: params["status"] = status if org_id: params["org_id"] = org_id resp = await self._request("GET", "/api/v1/workers", params=params) data = resp.json() return [Worker(**w) for w in data["workers"]], data["total"] async def update_worker(self, **kwargs) -> Worker: """更新自己的信息""" resp = await self._request("PATCH", f"/api/v1/workers/{self.worker_id}", json=kwargs) return Worker(**resp.json()) async def update_profile(self, **kwargs) -> Worker: """更新自己的简历信息。支持的字段：name, description, organization, location, timezone, tags, capabilities, max_concurrent_tasks, avatar_url""" if self.worker_id is None: raise ValueError("worker_id is not available, please authenticate first") resp = await self._request("PATCH", f"/api/v1/workers/{self.worker_id}", json=kwargs) return Worker(**resp.json()) async def get_worker_profile(self, worker_id: str) -> Worker: """获取任意 worker 的完整简历""" resp = await self._request("GET", f"/api/v1/workers/{worker_id}") return Worker(**resp.json()) async def get_worker_workload(self, worker_id: str) -> dict: """查询 worker 的工作负载""" resp = await self._request("GET", f"/api/v1/workers/{worker_id}/workload") return resp.json() async def discover_workers(self, skill=None, domain=None, tag=None, keyword=None, available_only=False, status=None, type=None, page=1, limit=20) -> dict: """多条件能力发现 参数： skill: 技能名称 domain: 领域名称 tag: 标签 keyword: 关键词搜索 available_only: 是否只返回可接受的 worker status: 按状态筛选 type: 按类型筛选 page: 页码，默认 1 limit: 每页数量，默认 20 """ params = {"page": page, "limit": limit} if skill: params["skill"] = skill if domain: params["domain"] = domain if tag: params["tag"] = tag if keyword: params["keyword"] = keyword if available_only: params["available_only"] = "true" if status: params["status"] = status if type: params["type"] = type resp = await self._request("GET", "/api/v1/workers/discover", params=params) return resp.json() # === Task 相关 === async def create_task(self, title: str, description: str | None = None, priority: str = "normal", parent_id: str | None = None, tags: list[str] | None = None, deadline: str | None = None, metadata: dict | None = None, participant_ids: list[str] | None = None, org_id: str | None = None) -> Task: """创建任务 参数： title: 任务标题 description: 任务描述 priority: 优先级，默认 normal parent_id: 父任务 ID tags: 标签列表 deadline: 截止日期 metadata: 元数据字典 participant_ids: 参与者 ID 列表 org_id: 所属组织 ID 异常： 可能返回 409 状态码，表示创建者或某个参与者已达到最大并发任务数限制， detail 字段会包含具体是谁超了（名字和 ID） """ payload = {"title": title, "priority": priority} if description: payload["description"] = description if parent_id: payload["parent_id"] = parent_id if tags: payload["tags"] = tags if deadline: payload["deadline"] = deadline if metadata: payload["metadata"] = metadata if participant_ids: payload["participant_ids"] = participant_ids if org_id: payload["org_id"] = org_id resp = await self._request("POST", "/api/v1/tasks", json=payload) return Task(**resp.json()) async def get_task(self, task_id: str) -> Task: """获取任务详情""" resp = await self._request("GET", f"/api/v1/tasks/{task_id}") return Task(**resp.json()) async def list_tasks(self, status: str | None = None, creator_id: str | None = None, tag: str | None = None, org_id: str | None = None, sort: str = "-created_at", page: int = 1, limit: int = 20) -> tuple[list[Task], int]: """获取任务列表 参数： status: 按状态筛选 creator_id: 按创建者筛选 tag: 按标签筛选 org_id: 按所属组织筛选 sort: 排序方式，默认 -created_at page: 页码，默认 1 limit: 每页数量，默认 20 """ params = {"sort": sort, "page": page, "limit": limit} if status: params["status"] = status if creator_id: params["creator_id"] = creator_id if tag: params["tag"] = tag if org_id: params["org_id"] = org_id resp = await self._request("GET", "/api/v1/tasks", params=params) data = resp.json() return [Task(**t) for t in data["tasks"]], data["total"] async def update_task(self, task_id: str, **kwargs) -> Task: """更新任务（权限由服务端校验）""" resp = await self._request("PATCH", f"/api/v1/tasks/{task_id}", json=kwargs) return Task(**resp.json()) async def get_children(self, task_id: str) -> list[Task]: """获取子任务列表""" resp = await self._request("GET", f"/api/v1/tasks/{task_id}/children") return [Task(**t) for t in resp.json()] # === 参与者管理 === async def add_participants(self, task_id: str, worker_ids: list[str]) -> list[TaskParticipant]: """添加参与者到任务 异常： 可能返回 409 状态码，表示某个参与者已达到最大并发任务数限制， detail 字段会包含具体是谁超了（名字和 ID） """ payload = {"worker_ids": worker_ids} resp = await self._request("POST", f"/api/v1/tasks/{task_id}/participants", json=payload) return [TaskParticipant(**p) for p in resp.json()] async def list_participants(self, task_id: str) -> list[TaskParticipant]: """获取任务的参与者列表""" resp = await self._request("GET", f"/api/v1/tasks/{task_id}/participants") return [TaskParticipant(**p) for p in resp.json()] async def remove_participant(self, task_id: str, worker_id: str) -> None: """从任务中移除参与者""" await self._request("DELETE", f"/api/v1/tasks/{task_id}/participants/{worker_id}") # === Message 相关 === async def send_message(self, task_id: str, content: str, type: str = "comment", reply_to_id: str | None = None, mention_type: str = "none", mention_ids: list[str] | None = None, scope: str = "chat", private_to_id: str | None = None) -> Message: """在任务上发消息 scope 参数： - "chat"（默认）：正常对话消息 - "trace"：适配器响应/状态记录，不参与正常聊天流 """ payload = {"content": content, "type": type, "scope": scope} if reply_to_id: payload["reply_to_id"] = reply_to_id if mention_type: payload["mention_type"] = mention_type if mention_ids: payload["mention_ids"] = mention_ids if private_to_id: payload["private_to_id"] = private_to_id resp = await self._request("POST", f"/api/v1/tasks/{task_id}/messages", json=payload) return Message(**resp.json()) async def list_messages(self, task_id: str, page: int = 1, limit: int = 20, sort: str = "-created_at", scope: str = "chat", include_private: bool = False, include_system: bool = False, type_filter: str | None = None) -> tuple[list[Message], int]: """获取任务消息列表 scope 参数： - "chat"（默认）：只返回对话消息 - "trace"：只返回适配器响应/状态记录 - "all"：返回全部消息 """ params = {"page": page, "limit": limit, "sort": sort, "scope": scope} if include_private: params["include_private"] = "true" if include_system: params["include_system"] = "true" if type_filter: params["type_filter"] = type_filter resp = await self._request("GET", f"/api/v1/tasks/{task_id}/messages", params=params) data = resp.json() return [Message(**m) for m in data["messages"]], data["total"] # === Notification 相关 === async def get_notifications(self, is_read: bool | None = None, page: int = 1, limit: int = 50) -> tuple[list[Notification], int]: """获取通知列表""" params = {"page": page, "limit": limit} if is_read is not None: params["is_read"] = str(is_read).lower() resp = await self._request("GET", "/api/v1/notifications", params=params) data = resp.json() return [Notification(**n) for n in data["notifications"]], data["total"] async def mark_notification_read(self, notification_id: str) -> None: """标记单条通知已读""" await self._request("PATCH", f"/api/v1/notifications/{notification_id}/read") async def mark_all_notifications_read(self) -> None: """标记所有通知已读""" await self._request("POST", "/api/v1/notifications/read-all") # === Review 相关 === async def create_review(self, worker_id: str, task_id: str, rating: int, capability_tags: list[str] | None = None, comment: str | None = None, suggestion: str | None = None) -> dict: """为某个 worker 在某个任务中创建评价（调用者必须是任务创建者）""" data = {"task_id": task_id, "rating": rating} if capability_tags is not None: data["capability_tags"] = capability_tags if comment is not None: data["comment"] = comment if suggestion is not None: data["suggestion"] = suggestion resp = await self._request("POST", f"/api/v1/workers/{worker_id}/reviews", json=data) return resp.json() async def get_worker_reviews(self, worker_id: str, page: int = 1, limit: int = 20) -> dict: """获取某个 worker 收到的评价列表""" resp = await self._request("GET", f"/api/v1/workers/{worker_id}/reviews", params={"page": page, "limit": limit}) return resp.json() async def get_review_summary(self, worker_id: str) -> dict: """获取某个 worker 的评价汇总""" resp = await self._request("GET", f"/api/v1/workers/{worker_id}/reviews/summary") return resp.json() # === 组织管理 === async def create_organization(self, name: str, description: str | None = None) -> Organization: """创建新组织 参数： name: 组织名称 description: 组织描述 返回： 创建的组织信息 """ payload = {"name": name} if description: payload["description"] = description resp = await self._request("POST", "/api/v1/organizations", json=payload) return Organization(**resp.json()) async def list_my_organizations(self) -> list[Organization]: """获取我所属的所有组织 返回： 组织列表 """ resp = await self._request("GET", "/api/v1/organizations") return [Organization(**o) for o in resp.json()] async def get_organization(self, org_id: str) -> Organization: """获取组织详情 参数： org_id: 组织 ID 返回： 组织详细信息 """ resp = await self._request("GET", f"/api/v1/organizations/{org_id}") return Organization(**resp.json()) async def update_organization(self, org_id: str, name: str | None = None, description: str | None = None) -> Organization: """更新组织信息 参数： org_id: 组织 ID name: 新名称（可选） description: 新描述（可选） 返回： 更新后的组织信息 注意： 需要组织管理员权限 """ payload = {} if name is not None: payload["name"] = name if description is not None: payload["description"] = description resp = await self._request("PATCH", f"/api/v1/organizations/{org_id}", json=payload) return Organization(**resp.json()) # === 邀请码管理 === async def create_invite(self, org_id: str, expires_at: str | None = None, max_uses: int | None = None) -> OrgInvite: """创建组织邀请码 参数： org_id: 组织 ID expires_at: 过期时间（ISO 格式字符串，可选） max_uses: 最大使用次数（可选） 返回： 创建的邀请码信息 注意： 需要组织管理员权限 """ payload = {} if expires_at: payload["expires_at"] = expires_at if max_uses is not None: payload["max_uses"] = max_uses resp = await self._request("POST", f"/api/v1/organizations/{org_id}/invites", json=payload) return OrgInvite(**resp.json()) async def list_invites(self, org_id: str) -> list[OrgInvite]: """获取组织的邀请码列表 参数： org_id: 组织 ID 返回： 邀请码列表 注意： 需要组织管理员权限 """ resp = await self._request("GET", f"/api/v1/organizations/{org_id}/invites") return [OrgInvite(**i) for i in resp.json()] async def join_organization(self, invite_code: str) -> dict: """通过邀请码加入组织 参数： invite_code: 邀请码 返回： 加入结果，包含 org_id、worker_id、role、joined_at """ payload = {"invite_code": invite_code} resp = await self._request("POST", "/api/v1/organizations/join", json=payload) return resp.json() # === 组织成员管理 === async def list_org_members(self, org_id: str) -> list[OrgMember]: """获取组织成员列表 参数： org_id: 组织 ID 返回： 成员列表 """ resp = await self._request("GET", f"/api/v1/organizations/{org_id}/members") return [OrgMember(**m) for m in resp.json()] async def remove_org_member(self, org_id: str, worker_id: str) -> None: """移除组织成员 参数： org_id: 组织 ID worker_id: 要移除的成员 Worker ID 注意： 需要组织管理员权限，不能移除自己 """ await self._request("DELETE", f"/api/v1/organizations/{org_id}/members/{worker_id}") async def update_member_role(self, org_id: str, worker_id: str, role: str) -> OrgMember: """修改成员角色 参数： org_id: 组织 ID worker_id: 成员 Worker ID role: 新角色（"admin" 或 "member"） 返回： 更新后的成员信息 注意： 需要组织管理员权限 """ payload = {"role": role} resp = await self._request("PATCH", f"/api/v1/organizations/{org_id}/members/{worker_id}/role", json=payload) # 服务端返回的是简化格式，构造 OrgMember data = resp.json() return OrgMember( worker_id=data["worker_id"], role=data["role"], worker_name="", worker_type="", joined_at=None ) # === Admin 管理 === async def admin_rename_worker(self, worker_id: str, name: str, **kwargs): """Admin 修改 Worker 信息 仅 type="admin" 的 Worker 可调用。 Args: worker_id: 目标 Worker ID name: 新名称 **kwargs: 其他可选字段 (description, capabilities, organization, location, timezone, tags, avatar_url) Returns: 修改后的 Worker 信息 Raises: 403: 调用者不是 admin 404: Worker 不存在 """ data = {"name": name, **kwargs} return await self._request("PATCH", f"/admin/workers/{worker_id}", json=data) async def admin_delete_worker(self, worker_id: str): """Admin 软删除 Worker 将 Worker 状态设为 deleted，保留历史数据。被删除的 Worker 无法登录和出现在列表中。 仅 type="admin" 的 Worker 可调用。 Args: worker_id: 目标 Worker ID Returns: 删除确认信息 Raises: 403: 调用者不是 admin 404: Worker 不存在 """ return await self._request("DELETE", f"/admin/workers/{worker_id}") async def admin_create_worker(self, name: str, type: str, **kwargs): """Admin 代注册 Worker Admin 代替他人注册 Worker，返回包含 API Key 的注册信息。 仅 type="admin" 的 Worker 可调用。 Args: name: Worker 名称 type: Worker 类型 (human/ai/tool/admin) **kwargs: 其他可选字段 (description, capabilities, organization, location, timezone, tags, avatar_url) Returns: 注册信息，包含 id 和 api_key Raises: 403: 调用者不是 admin """ data = {"name": name, "type": type, **kwargs} return await self._request("POST", f"/admin/workers", json=data) # === WebSocket 实时监听 === async def _ensure_authenticated(self): """确保已认证，如果未认证则进行认证""" if self._access_token is None: await self._authenticate() def _build_ws_url(self) -> str: """构建 WebSocket URL""" ws_url = self.server_url.replace("http://", "ws://").replace("https://", "wss://") return f"{ws_url}/api/v1/ws?token={self._access_token}&protocol={self.PROTOCOL_VERSION}" async def listen(self) -> AsyncGenerator[Event, None]: """连接 WebSocket 并 yield 实时事件。自动心跳保活，带断线重连。""" import logging logger = logging.getLogger(__name__) delay = 2.0 while True: try: await self._ensure_authenticated() ws_url = self._build_ws_url() async with websockets.connect(ws_url) as ws: delay = 2.0 # 重置退避 # 启动心跳任务 async def heartbeat(): while True: try: await asyncio.sleep(30) await ws.send("ping") except Exception: break heartbeat_task = asyncio.create_task(heartbeat()) # 通知 WebSocket 已连接 self._ws_connected_event.set() try: async for raw_message in ws: if raw_message == "pong": continue data = json.loads(raw_message) yield Event(**data) finally: heartbeat_task.cancel() except asyncio.CancelledError: return except Exception as e: logger.warning(f"WebSocket 断开: {e}，{delay:.0f}s 后重连") await asyncio.sleep(delay) delay = min(delay * 2, 60.0) async def wait_for_ws_connection(self, timeout: float = 10.0): """等待 WebSocket 连接建立""" await asyncio.wait_for(self._ws_connected_event.wait(), timeout=timeout)