ws_subscribe.py
4.55 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
"""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