wehub-resource-sync
fed8b2eed7
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
261 lines
10 KiB
Python
261 lines
10 KiB
Python
"""Shared snapshot/replay primitives for chat-stream reconnect.
|
||
|
||
The reconnect reader itself is the native-async generator in
|
||
``async_event_replay.build_message_event_stream_async``; this module holds
|
||
the pieces both it and the producer's journal depend on: the SSE wire
|
||
format (``format_sse_event``), the ``message_events`` snapshot read
|
||
(``read_snapshot_lines``), the producer-liveness watchdog probe
|
||
(``_check_producer_liveness``), and the pub/sub envelope encode/decode.
|
||
Keeping them here lets the async reader and the sync journal agree on the
|
||
exact wire shape and dedup/terminal rules. See
|
||
``docs/runbooks/sse-notifications.md``.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import re
|
||
from typing import Optional
|
||
|
||
from sqlalchemy import text as sql_text
|
||
|
||
from application.storage.db.repositories.message_events import (
|
||
MessageEventsRepository,
|
||
)
|
||
from application.storage.db.session import db_readonly
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
DEFAULT_KEEPALIVE_SECONDS = 15.0
|
||
DEFAULT_POLL_TIMEOUT_SECONDS = 1.0
|
||
# When the live tail has no events and no terminal in snapshot, fall
|
||
# back to checking ``conversation_messages`` directly. If the row has
|
||
# already gone terminal (worker journaled ``end``/``error`` to the DB
|
||
# but the matching pub/sub publish was lost, or the row was finalized
|
||
# without a journal write at all) we surface a terminal event so the
|
||
# client doesn't hang on keepalives. If the row is still non-terminal
|
||
# but the producer heartbeat is older than ``PRODUCER_IDLE_SECONDS``
|
||
# the producer is presumed dead (worker crash / recycle between chunks
|
||
# and finalize) and we emit a terminal ``error`` so the UI can recover.
|
||
DEFAULT_WATCHDOG_INTERVAL_SECONDS = 5.0
|
||
# 1.5× the route's 60s heartbeat interval — long enough that a normal
|
||
# heartbeat skew doesn't false-positive, short enough that a stuck
|
||
# stream surfaces before the 5-minute reconciler sweep escalates.
|
||
DEFAULT_PRODUCER_IDLE_SECONDS = 90.0
|
||
|
||
# WHATWG SSE accepts CRLF, CR, LF — split on any of them so a stray CR
|
||
# can't smuggle a record boundary into the wire format.
|
||
_SSE_LINE_SPLIT_PATTERN = re.compile(r"\r\n|\r|\n")
|
||
|
||
# Event types that mark the end of a chat answer. After delivering one
|
||
# we close the reconnect stream — keeping the connection open past a
|
||
# terminal event would leak both the client's reconnect promise and
|
||
# the server's WSGI thread waiting on keepalives that the user no
|
||
# longer cares about. The agent loop emits ``end`` for normal /
|
||
# tool-paused completion and ``error`` for the catch-all failure path
|
||
# (which doesn't get a trailing ``end``).
|
||
_TERMINAL_EVENT_TYPES = frozenset({"end", "error"})
|
||
|
||
|
||
def _payload_is_terminal(
|
||
payload: object, event_type: Optional[str] = None
|
||
) -> bool:
|
||
"""True if ``payload['type']`` or ``event_type`` is a terminal sentinel."""
|
||
if isinstance(payload, dict) and payload.get("type") in _TERMINAL_EVENT_TYPES:
|
||
return True
|
||
return event_type in _TERMINAL_EVENT_TYPES
|
||
|
||
|
||
def format_sse_event(payload: dict, sequence_no: int) -> str:
|
||
"""Encode a journal event as one ``id:``/``data:`` SSE record.
|
||
|
||
The body is the payload's JSON serialisation. ``complete_stream``
|
||
payloads are flat JSON dicts with no embedded newlines, so a
|
||
single ``data:`` line is sufficient — but we still split on any
|
||
line terminator in case a future caller passes a multi-line string.
|
||
"""
|
||
body = json.dumps(payload)
|
||
lines = [f"id: {sequence_no}"]
|
||
for line in _SSE_LINE_SPLIT_PATTERN.split(body):
|
||
lines.append(f"data: {line}")
|
||
return "\n".join(lines) + "\n\n"
|
||
|
||
|
||
def _check_producer_liveness(
|
||
message_id: str, user_id: Optional[str], idle_seconds: float
|
||
) -> Optional[dict]:
|
||
"""Inspect ``conversation_messages`` and return a terminal SSE
|
||
payload when the producer is no longer alive, else ``None``.
|
||
|
||
When ``user_id`` is given the lookup is scoped to ``AND user_id = :u``
|
||
(defence in depth: this long-lived re-read re-asserts the ownership the
|
||
route gated on, so a stream cannot keep tailing a row it no longer
|
||
owns). A non-matching row reads as missing → a terminal ``error``.
|
||
|
||
Three terminal cases collapse into a single DB round-trip:
|
||
|
||
- ``status='complete'`` — the live finalize ran but its journal
|
||
terminal write didn't reach us (or never happened). Synthesise
|
||
``end`` so the client closes cleanly on the row's user-visible
|
||
state.
|
||
- ``status='failed'`` — same, but for the failure path. Carry the
|
||
stashed ``error`` from ``message_metadata`` so the UI shows the
|
||
real reason.
|
||
- non-terminal status and ``last_heartbeat_at`` (or ``timestamp``)
|
||
older than ``idle_seconds`` — the producing worker is gone.
|
||
Synthesise ``error`` so the client doesn't hang on keepalives
|
||
until the proxy idle-timeout kicks in.
|
||
"""
|
||
owner_clause = " AND user_id = :u" if user_id is not None else ""
|
||
params = {"id": message_id, "idle_secs": float(idle_seconds)}
|
||
if user_id is not None:
|
||
params["u"] = user_id
|
||
try:
|
||
with db_readonly() as conn:
|
||
row = conn.execute(
|
||
sql_text(
|
||
# ``owner_clause`` is a fixed literal (no user input in the
|
||
# SQL string); ``user_id`` is bound via ``:u``.
|
||
f"""
|
||
SELECT
|
||
status,
|
||
message_metadata->>'error' AS err,
|
||
GREATEST(
|
||
timestamp,
|
||
COALESCE(
|
||
(message_metadata->>'last_heartbeat_at')
|
||
::timestamptz,
|
||
timestamp
|
||
)
|
||
) < now() - make_interval(secs => :idle_secs)
|
||
AS is_stale
|
||
FROM conversation_messages
|
||
WHERE id = CAST(:id AS uuid){owner_clause}
|
||
"""
|
||
),
|
||
params,
|
||
).first()
|
||
except Exception:
|
||
logger.exception(
|
||
"Watchdog liveness check failed for message_id=%s", message_id
|
||
)
|
||
return None
|
||
|
||
if row is None:
|
||
# Row deleted out from under us — treat as terminal so the
|
||
# client doesn't keep tailing a message that no longer exists.
|
||
return {
|
||
"type": "error",
|
||
"error": "Message no longer exists; please refresh.",
|
||
"code": "message_missing",
|
||
"message_id": message_id,
|
||
}
|
||
|
||
status, err, is_stale = row[0], row[1], bool(row[2])
|
||
if status == "complete":
|
||
return {"type": "end"}
|
||
if status == "failed":
|
||
return {
|
||
"type": "error",
|
||
"error": err or "Stream failed; please try again.",
|
||
"code": "producer_failed",
|
||
"message_id": message_id,
|
||
}
|
||
if is_stale:
|
||
return {
|
||
"type": "error",
|
||
"error": (
|
||
"Stream producer is no longer responding; please try again."
|
||
),
|
||
"code": "producer_stale",
|
||
"message_id": message_id,
|
||
}
|
||
return None
|
||
|
||
|
||
def read_snapshot_lines(
|
||
message_id: str, last_event_id: Optional[int], user_id: Optional[str] = None
|
||
) -> tuple[list[str], Optional[int], bool]:
|
||
"""Read journal rows after ``last_event_id`` as SSE-formatted lines.
|
||
|
||
Returns ``(lines, max_sequence_no, terminal)``: ``max_sequence_no`` is
|
||
seeded with ``last_event_id`` and advanced past every row read,
|
||
``terminal`` is True if any row carried a terminal ``end``/``error``.
|
||
Raises on DB error so the caller can drive its replay-failed path.
|
||
|
||
Used by ``async_event_replay.build_message_event_stream_async`` (the
|
||
reconnect reader); it shares ``format_sse_event`` / ``_payload_is_terminal``
|
||
with the producer's journal writer so reader and writer never drift on
|
||
wire shape or terminal semantics.
|
||
"""
|
||
lines: list[str] = []
|
||
max_seq = last_event_id
|
||
terminal = False
|
||
with db_readonly() as conn:
|
||
rows = MessageEventsRepository(conn).read_after(
|
||
message_id, last_sequence_no=last_event_id, user_id=user_id
|
||
)
|
||
for row in rows:
|
||
seq = int(row["sequence_no"])
|
||
payload = row.get("payload")
|
||
if not isinstance(payload, dict):
|
||
# ``record_event`` rejects non-dict payloads at the write gate,
|
||
# so this is a legacy/direct-SQL row — drop it rather than ship
|
||
# a malformed envelope that would poison a reconnect.
|
||
logger.warning(
|
||
"Skipping non-dict payload from message_events: "
|
||
"message_id=%s seq=%s type=%s",
|
||
message_id,
|
||
seq,
|
||
row.get("event_type"),
|
||
)
|
||
continue
|
||
lines.append(format_sse_event(payload, seq))
|
||
if max_seq is None or seq > max_seq:
|
||
max_seq = seq
|
||
if _payload_is_terminal(payload, row.get("event_type")):
|
||
terminal = True
|
||
return lines, max_seq, terminal
|
||
|
||
|
||
def _decode_pubsub_message(raw) -> Optional[dict]:
|
||
"""Parse a ``Topic.publish`` payload to ``{sequence_no, payload, ...}``.
|
||
|
||
Returns ``None`` for malformed messages (drop silently — the
|
||
journal is still authoritative on reconnect).
|
||
"""
|
||
try:
|
||
if isinstance(raw, (bytes, bytearray)):
|
||
text_value = raw.decode("utf-8")
|
||
else:
|
||
text_value = str(raw)
|
||
envelope = json.loads(text_value)
|
||
except Exception:
|
||
return None
|
||
if not isinstance(envelope, dict):
|
||
return None
|
||
return envelope
|
||
|
||
|
||
def encode_pubsub_message(
|
||
message_id: str,
|
||
sequence_no: int,
|
||
event_type: str,
|
||
payload: dict,
|
||
) -> str:
|
||
"""Build the JSON envelope used for ``channel:{message_id}`` publishes.
|
||
|
||
Kept here (not in ``message_journal.py``) so the encode/decode pair
|
||
stays in one file — replay's ``_decode_pubsub_message`` and the
|
||
journal's publish must agree on the shape exactly.
|
||
"""
|
||
return json.dumps(
|
||
{
|
||
"message_id": str(message_id),
|
||
"sequence_no": int(sequence_no),
|
||
"event_type": event_type,
|
||
"payload": payload,
|
||
}
|
||
)
|