249 lines
11 KiB
Python
249 lines
11 KiB
Python
"""Telegram admission before every PTB handler group, including native plugins.
|
|
|
|
Claims belong to the receiving adapter across Application rebuilds. Completed
|
|
history is bounded and has no TTL in memory; dispatch and its PTB tasks pin active
|
|
claims. Completed IDs are also written to a per-bot receipt file under the adapter's
|
|
Hermes home, so a replacement adapter or a restarted gateway still drops updates that
|
|
Telegram redelivers because their getUpdates/webhook acknowledgement never landed.
|
|
No cross-process coordination or exactly-once effects are promised.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
from contextvars import ContextVar
|
|
from dataclasses import dataclass
|
|
from functools import wraps
|
|
|
|
from telegram import Update
|
|
from telegram.ext import Application, ApplicationHandlerStop, ConversationHandler
|
|
|
|
from gateway.platforms.helpers import bounded_put
|
|
from utils import atomic_json_write
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_DEFAULT_BLOCK = object()
|
|
_SEEN_CAP = 4096
|
|
# The Bot API keeps an unconfirmed update for at most 24 hours (getUpdates), so an older
|
|
# receipt can never match a redelivery. That is also well inside the week of silence after
|
|
# which Telegram may restart update IDs at a random value, so a set lookup stays safe
|
|
# where a persisted numeric high-watermark would not.
|
|
RECEIPT_TTL_SECONDS = 24 * 60 * 60
|
|
|
|
|
|
def _receipt_path(adapter, bot_id):
|
|
return adapter._update_receipt_dir / f"telegram_update_receipts_{bot_id}.json"
|
|
|
|
|
|
def _load_receipts(adapter, bot_id) -> None:
|
|
"""Seed completed history from disk once per bot, before its first admission check."""
|
|
if bot_id in adapter._update_receipts_loaded:
|
|
return
|
|
adapter._update_receipts_loaded.add(bot_id)
|
|
path = _receipt_path(adapter, bot_id)
|
|
try:
|
|
payload = json.loads(path.read_text(encoding="utf-8-sig"))
|
|
except FileNotFoundError:
|
|
return
|
|
except (OSError, ValueError):
|
|
logger.warning("[Telegram] Ignoring unreadable update receipts at %s", path, exc_info=True)
|
|
return
|
|
ids = payload.get("update_ids") if isinstance(payload, dict) else None
|
|
if not isinstance(ids, dict):
|
|
return
|
|
cutoff = time.time() - RECEIPT_TTL_SECONDS
|
|
fresh = sorted((float(ts), uid) for uid, ts in ids.items()
|
|
if isinstance(ts, (int, float)) and ts > cutoff and str(uid).lstrip("-").isdigit())
|
|
seen = adapter._seen_update_ids
|
|
for ts, uid in fresh:
|
|
key = f"{bot_id}:{uid}"
|
|
if key not in seen:
|
|
bounded_put(seen, key, ts, _SEEN_CAP)
|
|
|
|
|
|
def _record_receipt(adapter, key: str) -> None:
|
|
bounded_put(adapter._seen_update_ids, key, time.time(), _SEEN_CAP)
|
|
adapter._update_receipts_dirty.add(key.split(":", 1)[0])
|
|
task = adapter._update_receipt_flush
|
|
if task is None or task.done():
|
|
# One writer per adapter coalesces bursts; the fsync stays off the event loop.
|
|
adapter._update_receipt_flush = asyncio.get_running_loop().create_task(
|
|
_flush_receipts(adapter), name="telegram-update-receipts")
|
|
|
|
|
|
async def _flush_receipts(adapter) -> None:
|
|
from hermes_constants import mkdir_under_hermes_home
|
|
|
|
while adapter._update_receipts_dirty:
|
|
bot_id = adapter._update_receipts_dirty.pop()
|
|
prefix, cutoff = f"{bot_id}:", time.time() - RECEIPT_TTL_SECONDS
|
|
# Snapshot on the loop: admission mutates this dict only from loop callbacks.
|
|
ids = {key[len(prefix):]: ts for key, ts in adapter._seen_update_ids.items()
|
|
if key.startswith(prefix) and isinstance(ts, float) and ts > cutoff}
|
|
path = _receipt_path(adapter, bot_id)
|
|
try:
|
|
mkdir_under_hermes_home(path.parent)
|
|
await asyncio.to_thread(atomic_json_write, path, {"update_ids": ids}, indent=None)
|
|
except Exception:
|
|
logger.warning("[Telegram] Failed to persist update receipts to %s", path, exc_info=True)
|
|
|
|
|
|
@dataclass
|
|
class _Claim:
|
|
key: str
|
|
dispatch_task: asyncio.Task | None
|
|
accepted: bool = False
|
|
completed: bool = False
|
|
failed: bool = False
|
|
owners: int = 1
|
|
|
|
|
|
class _ErrorCallback:
|
|
"""Admit callback entry while keeping PTB's original registration/removal key."""
|
|
|
|
def __init__(self, callback):
|
|
self.__wrapped__ = callback
|
|
|
|
def __hash__(self):
|
|
return hash(self.__wrapped__)
|
|
|
|
def __eq__(self, other):
|
|
if isinstance(other, _ErrorCallback):
|
|
other = other.__wrapped__
|
|
return self.__wrapped__ == other
|
|
|
|
async def __call__(self, update, context):
|
|
# Registration is a native extension boundary, unlike PTB's internal logger.
|
|
# Resolve the receiving app so a shared callback never accepts another bot's claim.
|
|
admission = getattr(context.application, "_current_claim", None)
|
|
claim = admission.get() if admission is not None else None
|
|
if claim is not None and claim.owners:
|
|
claim.accepted = True
|
|
return await self.__wrapped__(update, context)
|
|
|
|
|
|
class TelegramApplication(Application):
|
|
__slots__ = ("adapter", "_current_claim")
|
|
|
|
def __init__(self, *, adapter, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.adapter = adapter
|
|
if adapter._update_admission is None:
|
|
adapter._update_admission = ContextVar[_Claim | None]("telegram_update_claim", default=None)
|
|
self._current_claim: ContextVar[_Claim | None] = adapter._update_admission
|
|
|
|
def add_handler(self, handler, group=0):
|
|
super().add_handler(handler, group)
|
|
self._admit_callbacks(handler)
|
|
|
|
def _admit_callbacks(self, handler):
|
|
# ConversationHandler has no callback. Keep its identity (including persistence,
|
|
# state transitions and block resolution) and instrument its selected children.
|
|
if isinstance(handler, ConversationHandler):
|
|
for child in [*handler.entry_points, *handler.fallbacks,
|
|
*(child for handlers in handler.states.values() for child in handlers)]:
|
|
self._admit_callbacks(child)
|
|
return
|
|
callback = getattr(handler, "callback", None)
|
|
if callback is None or getattr(callback, "_telegram_admitted", False):
|
|
return
|
|
|
|
@wraps(callback)
|
|
async def admitted_callback(update, context):
|
|
# Native plugins own arbitrary effects. Entry is the handoff; retrying a partly
|
|
# completed plugin would be unsafe. Core preparation has explicit handoff markers.
|
|
# A plugin may reuse the same handler across applications/reconnects. Resolve
|
|
# the receiving claim at invocation, rather than closing over its first owner.
|
|
application = context.application
|
|
admission = getattr(application, "_current_claim", None)
|
|
claim = admission.get() if admission is not None else None
|
|
if claim is not None and getattr(callback, "__self__", None) is not application.adapter:
|
|
claim.accepted = True
|
|
result = await callback(update, context)
|
|
if claim is not None:
|
|
if claim.failed and not claim.accepted:
|
|
# Stop subsequent groups only in blocking dispatch. PTB cannot stop
|
|
# groups already scheduled via block=False; their outcomes settle below.
|
|
if asyncio.current_task() is claim.dispatch_task:
|
|
raise ApplicationHandlerStop
|
|
else:
|
|
claim.completed = True
|
|
return result
|
|
|
|
admitted_callback._telegram_admitted = True
|
|
handler.callback = admitted_callback
|
|
|
|
def add_error_handler(self, callback, block=_DEFAULT_BLOCK):
|
|
if not isinstance(callback, _ErrorCallback):
|
|
callback = _ErrorCallback(callback)
|
|
# Omitted block must remain PTB's default, not explicit True (bot defaults differ).
|
|
if block is _DEFAULT_BLOCK:
|
|
super().add_error_handler(callback)
|
|
else:
|
|
super().add_error_handler(callback, block=block)
|
|
|
|
def _Application__create_task(self, coroutine, update=None, is_error_handler=False, name=None):
|
|
# PTB's process_error bypasses public create_task. Delegate at the shared task seam
|
|
# without copying dispatch or losing its is_error_handler recursion protection.
|
|
task = super()._Application__create_task(
|
|
coroutine, update=update, is_error_handler=is_error_handler, name=name)
|
|
claim = self._current_claim.get()
|
|
if claim is not None and claim.owners:
|
|
# Pin scheduled work before callback entry, including cancellation before entry.
|
|
claim.owners += 1
|
|
|
|
def finished(done):
|
|
# PTB's task wrapper may never start, leaving its callback unawaited.
|
|
if done.cancelled() and asyncio.iscoroutine(coroutine):
|
|
coroutine.close()
|
|
self._release_claim(claim, done)
|
|
|
|
task.add_done_callback(finished)
|
|
return task
|
|
|
|
def _release_claim(self, claim, task=None):
|
|
if task is not None and task.cancelled():
|
|
claim.failed = True
|
|
claim.owners -= 1
|
|
if claim.owners:
|
|
return
|
|
del self.adapter._inflight_update_ids[claim.key]
|
|
if claim.accepted or (claim.completed and not claim.failed):
|
|
_record_receipt(self.adapter, claim.key)
|
|
|
|
async def process_error(self, update, error, job=None, coroutine=None):
|
|
claim = self._current_claim.get()
|
|
if claim is not None:
|
|
claim.failed = True
|
|
stopped = await super().process_error(update, error, job, coroutine)
|
|
# PTB catches callback errors. Do not let a later observer accept a failed preparation.
|
|
return stopped or (claim is not None and not claim.accepted)
|
|
|
|
async def process_update(self, update):
|
|
if not isinstance(update, Update):
|
|
return await super().process_update(update)
|
|
bot_id = self.bot.id
|
|
key = f"{bot_id}:{update.update_id}"
|
|
_load_receipts(self.adapter, bot_id)
|
|
# Dispatch happened even when preparation fails before the group-99 observer.
|
|
self.adapter._updates_dispatched_total += 1
|
|
seen = self.adapter._seen_update_ids
|
|
pending = self.adapter._inflight_update_ids
|
|
if key in seen or key in pending:
|
|
return
|
|
claim = _Claim(key, asyncio.current_task())
|
|
# Atomic on PTB's event loop: no await between lookup and claim. Dispatch and its
|
|
# PTB tasks share ownership; completed-history pressure cannot evict active work.
|
|
pending[key] = claim
|
|
token = self._current_claim.set(claim)
|
|
try:
|
|
await super().process_update(update)
|
|
except BaseException:
|
|
claim.failed = True
|
|
raise
|
|
finally:
|
|
self._current_claim.reset(token)
|
|
self._release_claim(claim)
|