642 lines
24 KiB
Python
642 lines
24 KiB
Python
"""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)
|