"""Бот-соавтор: правит редактируемые ячейки листа по WebSocket с заданной частотой. CLI: python -m perfkit.coauthor --form-id N --sheet S [--direction D] --rate 1|5 --duration SEC [--user rf_coauthor] [--scale 1] [--token T] Протокол (api/src/api/v1/websocket.py, web/.../RealtimeContext.js): ws://HOST/api/v1/ws/form/{id}/sheet/{sheet}?direction=D -> {"event":"user_login","data":{"token":...}} -> {"event":"cell_edit_start","data":{"line_id","column"}} (эхо отправителю, блокировка ячейки) -> {"event":"cell_updated","data":{"line_id","line_id_code","column","value"}} (эхо отправителю с "is_self": true; при отказе то же событие с полем "error") -> {"event":"cell_edit_end","data":{"line_id","column"}} Редактируемые колонки берутся из строки META листа (GET sheet: data.editable, роль EXECUTOR_RF), правятся числовые колонки у строк типа INPUT. Коды выхода: 0 - были подтверждённые эхом правки; 2 - нет эха первой правки за --ack-timeout (в т.ч. неверный токен или закрытое соединение); 3 - ошибка подготовки (вход по HTTP, нет редактируемых ячеек). Печатает "READY" после первого эха и в конце последней строкой JSON {"sent","echoed","errors"}. SIGTERM/SIGINT завершают работу штатно (итог печатается). """ from __future__ import annotations import argparse import asyncio import json import os import random import signal import sys import time from urllib.parse import urlencode import httpx import websockets from . import RESULTS_DIR from .manifest import load as load_manifest API = os.environ.get("PERF_API", "http://localhost:8000") TEXT_HINTS = ("comment", "justification", "date", "note") # нечисловые колонки не трогаем def log(msg: str) -> None: print(msg, file=sys.stderr, flush=True) def pick_cells(rows: list[dict]) -> list[tuple[int, str]]: """(line_id, column) для INPUT-строк в числовых редактируемых колонках META.""" meta = next((r for r in rows if r.get("row_type") == "META"), None) if not meta: return [] cols = [e["column"] for e in meta["data"].get("editable", []) if not any(h in e["column"] for h in TEXT_HINTS)] cells = [] for r in rows: if r.get("row_type") != "INPUT": continue d = r.get("data", {}) lid = d.get("line_id") or d.get("id") if lid is None: continue cells.extend((lid, c) for c in cols) return cells async def fetch_cells(args) -> list[tuple[int, str]]: m = load_manifest(RESULTS_DIR / f"manifest-s{args.scale}.json") u = m.users[args.user] async with httpx.AsyncClient(base_url=API + "/api/v1", timeout=120) as c: r = await c.post("/auth/login", json={"username": u.username, "password": u.password}) r.raise_for_status() tok = r.json()["access_token"] params = {"direction": args.direction} if args.direction else None r = await c.get(f"/form/{args.form_id}/sheet/{args.sheet}", params=params, headers={"Authorization": f"Bearer {tok}"}) r.raise_for_status() args._http_token = tok return pick_cells(r.json()["result"]) def ws_url(args) -> str: base = API.replace("http://", "ws://").replace("https://", "wss://") q = "?" + urlencode({"direction": args.direction}) if args.direction else "" return f"{base}/api/v1/ws/form/{args.form_id}/sheet/{args.sheet}{q}" async def run(args) -> int: t0 = time.monotonic() try: cells = await fetch_cells(args) except Exception as e: # noqa: BLE001 log(f"coauthor: подготовка не удалась: {e!r}") return 3 if not cells: log("coauthor: нет редактируемых ячеек для роли пользователя") return 3 st = {"sent": 0, "echoed": 0, "errors": 0} first_echo = asyncio.Event() stop = asyncio.Event() loop = asyncio.get_running_loop() for sig in (signal.SIGTERM, signal.SIGINT): loop.add_signal_handler(sig, stop.set) token = args.token if args.token is not None else args._http_token try: ws = await asyncio.wait_for(websockets.connect(ws_url(args), max_size=None), 5) except Exception as e: # noqa: BLE001 log(f"coauthor: WS недоступен: {e!r}") return 2 async def reader(): try: async for raw in ws: msg = json.loads(raw) if msg.get("event") != "cell_updated" or not msg.get("is_self"): continue if msg.get("error"): st["errors"] += 1 if st["errors"] <= 3: log(f"coauthor: отказ сервера: {msg['error']}") continue st["echoed"] += 1 if not first_echo.is_set(): first_echo.set() print("READY", flush=True) except websockets.ConnectionClosed as e: log(f"coauthor: соединение закрыто: {e.rcvd}") finally: stop.set() rd = asyncio.create_task(reader()) await ws.send(json.dumps({"event": "user_login", "data": {"token": token}})) async def edit(): lid, col = random.choice(cells) key = {"line_id": lid, "column": col} await ws.send(json.dumps({"event": "cell_edit_start", "data": key})) await ws.send(json.dumps({"event": "cell_updated", "data": {**key, "line_id_code": lid, "value": random.randint(1000, 999_999)}})) st["sent"] += 1 await ws.send(json.dumps({"event": "cell_edit_end", "data": key})) rc = 0 try: # фаза 1: до первого эха, не дольше ack_timeout с момента старта процесса while not first_echo.is_set(): if stop.is_set() or time.monotonic() - t0 > args.ack_timeout: log("coauthor: нет эха первой правки (неверный токен, нет прав или окно закрыто)") rc = 2 break await edit() try: await asyncio.wait_for(first_echo.wait(), 0.7) except asyncio.TimeoutError: pass if rc == 0: # фаза 2: ровная частота rate правок/с до конца duration period = 1.0 / args.rate end = time.monotonic() + args.duration nxt = time.monotonic() + period while not stop.is_set() and time.monotonic() < end: await asyncio.sleep(max(0.0, min(nxt, end) - time.monotonic())) if time.monotonic() >= end or stop.is_set(): break await edit() nxt += period await asyncio.sleep(0.5) # дождаться эха последних правок if st["echoed"] == 0: rc = 2 except websockets.ConnectionClosed: rc = 2 if st["echoed"] == 0 else 0 finally: try: await ws.close() except Exception: # noqa: BLE001 pass rd.cancel() print(json.dumps(st), flush=True) return rc def main(argv=None) -> int: ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) ap.add_argument("--form-id", type=int, required=True) ap.add_argument("--sheet", required=True) ap.add_argument("--direction", default=None) ap.add_argument("--rate", type=float, default=1.0, help="правок в секунду (1 или 5)") ap.add_argument("--duration", type=float, default=10.0, help="секунд работы; SIGTERM завершает раньше") ap.add_argument("--user", default="rf_coauthor", help="ключ пользователя в манифесте") ap.add_argument("--scale", type=int, default=int(os.environ.get("PERF_SCALE", 1))) ap.add_argument("--token", default=None, help="токен для WS вместо полученного входом (проверка неверного токена)") ap.add_argument("--ack-timeout", type=float, default=5.0, help="сколько ждать эха первой правки") return asyncio.run(run(ap.parse_args(argv))) if __name__ == "__main__": sys.exit(main())