ws_subscribe.py 4.55 KB
"""c_ws_subscribe (03-C §3.14): real-time execution_event push over a raw
WebSocket connection (RFC 6455), hand-rolled with hashlib/base64/socket --
pure stdlib, no `websockets`/`tornado` (05 §5.7).

U-024 scope ("最小连通版"): a poll loop is what actually pushes events, not
a write-side hook -- execution_event rows are inserted from many different
call sites across this codebase (org_node_trigger, work_graph_execute,
work_node_run, task_create, bug_record, ...) with no single write path to
instrument; polling the table every POLL_INTERVAL_SECONDS decouples this
from all of them and easily clears REQ-12's own <=2s target (framed there
as a suggested non-functional goal, not a hard SLA). Reconnect catch-up is
explicitly routed through REST snapshot endpoints per 03-C §3.14's own
"断线重连约束" note, not WS-layer replay -- a fresh connection only sees
events from its own connect time onward, no history replay to build here.
REQ-12's boundary/failure cases are U-024.5's job once more event-generating
Units exist to test against.
"""

import base64
import hashlib
import json

_WS_MAGIC_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
POLL_INTERVAL_SECONDS = 0.5

OPCODE_TEXT = 0x1
OPCODE_CLOSE = 0x8
OPCODE_PING = 0x9
OPCODE_PONG = 0xA


def compute_accept_key(client_key):
    """RFC 6455 §1.3 handshake: SHA1(client_key + magic GUID), base64-encoded."""
    digest = hashlib.sha1((client_key + _WS_MAGIC_GUID).encode("ascii")).digest()
    return base64.b64encode(digest).decode("ascii")


def encode_frame(payload_bytes, opcode=OPCODE_TEXT):
    """Server->client frames are never masked (RFC 6455 §5.1)."""
    length = len(payload_bytes)
    if length < 126:
        header = bytes([0x80 | opcode, length])
    elif length < 65536:
        header = bytes([0x80 | opcode, 126]) + length.to_bytes(2, "big")
    else:
        header = bytes([0x80 | opcode, 127]) + length.to_bytes(8, "big")
    return header + payload_bytes


def encode_text_frame(text):
    return encode_frame(text.encode("utf-8"), opcode=OPCODE_TEXT)


def decode_client_frame(sock):
    """Reads one client->server frame (always masked per RFC 6455 §5.1).
    Returns (opcode, payload_bytes). Raises ConnectionError if the socket
    closes mid-frame; a plain read timeout propagates as socket.timeout,
    left for the caller to catch (that's the normal "nothing arrived this
    poll tick" case, not an error)."""
    first_two = _recv_exact(sock, 2)
    opcode = first_two[0] & 0x0F
    masked = first_two[1] & 0x80
    length = first_two[1] & 0x7F
    if length == 126:
        length = int.from_bytes(_recv_exact(sock, 2), "big")
    elif length == 127:
        length = int.from_bytes(_recv_exact(sock, 8), "big")
    mask_key = _recv_exact(sock, 4) if masked else b""
    payload = _recv_exact(sock, length)
    if masked:
        payload = bytes(b ^ mask_key[i % 4] for i, b in enumerate(payload))
    return opcode, payload


def _recv_exact(sock, n):
    if n == 0:
        return b""
    chunks = []
    remaining = n
    while remaining > 0:
        chunk = sock.recv(remaining)
        if not chunk:
            raise ConnectionError("socket closed while reading a WebSocket frame")
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)


def latest_rowid(conn, project_id):
    """Starting point for a fresh subscription -- 'now', not full history."""
    row = conn.execute(
        "SELECT MAX(rowid) FROM execution_event WHERE project_id = ?", (project_id,)
    ).fetchone()
    return row[0] or 0


def poll_new_events(conn, project_id, after_rowid, task_id=None):
    """Returns (events, new_after_rowid) for execution_event rows with
    rowid > after_rowid, matching project_id (and task_id if given, per
    03-C §3.14's optional server-side filter), oldest first. Each event
    dict matches the output_contract field-for-field."""
    query = (
        "SELECT rowid, event_id, event_type, task_id, run_id, node_exec_id, payload, created_at "
        "FROM execution_event WHERE project_id = ? AND rowid > ?"
    )
    params = [project_id, after_rowid]
    if task_id:
        query += " AND task_id = ?"
        params.append(task_id)
    query += " ORDER BY rowid"
    rows = conn.execute(query, params).fetchall()
    events = [
        {
            "event_id": r[1],
            "event_type": r[2],
            "task_id": r[3],
            "run_id": r[4],
            "node_exec_id": r[5],
            "payload": json.loads(r[6]),
            "created_at": r[7],
        }
        for r in rows
    ]
    new_after_rowid = rows[-1][0] if rows else after_rowid
    return events, new_after_rowid