Skip to content

Commit 2fc9b5c

Browse files
authored
Implement plans/01: cache cap + RetryAfter helper + handler refactor + /status command (#22)
* refactor: extract helpers from 5 complex handlers (Phase 3) Behavior-preserving readability refactor with zero test changes: 1. bio_bait.py: extract enforce delete+restrict+notify into _enforce_bio_bait_restriction() — parent raises ApplicationHandlerStop 2. dm.py: extract unrestrict loop into _unrestrict_in_groups() returning success count — parent handles response messages 3. anti_spam.py: extract guard clauses into _should_skip_new_user_spam_check() predicate 4. message.py: extract is_bot guard into _should_skip_profile_check() predicate 5. check.py: extract callback data parsing into _parse_warn_callback_data() returning (user_id, missing_code) or None All 948 tests pass, coverage 99% (22 missed, unchanged). Ruff + mypy clean. * fix(duplicate_spam): cap in-memory recent-messages dict to prevent unbounded growth * add retry-after-helper: send_message_with_retry + restrict_chat_member_with_retry - Two helpers in telegram_utils.py: RetryAfter-aware wrappers that sleep + retry once, return False on second RetryAfter, re-raise other errors - unrestrict_user uses restrict_chat_member_with_retry internally - 4 call sites updated: check.py handle_warn_callback, scheduler.py auto-restrict loop, dm.py unrestrict loop, bio_bait.py monitor alert - 8 unit tests for helpers, 2 integration tests for bio_bait error paths - 958 tests pass, 99% coverage, ruff clean, mypy clean ponytail: scheduler.py RetryAfter-fallback continue (lines 95-98) not covered by tests; add scheduler-level RetryAfter test if it becomes a real problem. * feat(status): add /status DM admin command + job timestamp tracking * fix(retry-after): handle int | timedelta in RetryAfter.retry_after * chore: gitignore .pi-subagents artifacts (consistent with reviews/) * fix(review): inline dead _should_skip_profile_check, gate auto-restrict stamp on real work Applies multi-model review feedback: - Remove premature-abstraction _should_skip_profile_check predicate (M1) - Track processed_any in auto_restrict to skip stamp on no-op cycles (R3) - Add logger.warning for missing message/from_user in status handler (C3) - Drop group_titles read; use str(gid) directly (M2, M6) * chore: gitignore plans/ (consistent with docs/, reviews/) Plan docs are working notes for a session, not source. Untrack the one that snuck into the improvement branch's history. * fix(review): apply additional reviewer feedback - verify.py: catch RuntimeError from unrestrict_user retry exhaustion; wrap clearance notification with send_message_with_retry - bio_bait.py: wrap enforcement path (restrict + notify) with retry helpers (the monitor-only alert path was already wrapped) - scheduler.py: stamp last_auto_restrict unconditionally on every successful run (matches admin_cache behavior) - telegram_utils.py: cap retry sleep at 30s to prevent one bad RetryAfter from stalling per-group loop for a minute-plus - status.py: drop dead title variable (was duplicating group_id) * fix(status): drop duplicate group_id display in /status group line title = str(gid) just re-stringified the same group_id with nothing else feeding it (GroupConfig has no name/title field), so each group line printed the ID twice: `-100123` — -100123. Show it once. * style(telegram_utils): use f-strings for retry warning/error logs
1 parent 339c78d commit 2fc9b5c

22 files changed

Lines changed: 1012 additions & 127 deletions

.gitignore

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,11 @@ data/
1414

1515
# Agent/planning docs
1616
docs/
17+
plans/
1718
.pi/
1819

1920
# Local review artifacts (parallel-reviewer outputs)
2021
reviews/
22+
.pi-subagents/
2123

2224
.tokensave/

src/bot/database/service.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -634,6 +634,17 @@ def get_all_pending_captchas(self) -> list[PendingCaptchaValidation]:
634634
statement = select(PendingCaptchaValidation)
635635
return list(session.exec(statement).all())
636636

637+
def get_all_new_user_probations(self) -> list[NewUserProbation]:
638+
"""
639+
Get all new-user probation records.
640+
641+
Returns:
642+
list[NewUserProbation]: All probation records.
643+
"""
644+
with Session(self._engine) as session:
645+
statement = select(NewUserProbation)
646+
return list(session.exec(statement).all())
647+
637648
def start_new_user_probation(self, user_id: int, group_id: int) -> NewUserProbation:
638649
"""
639650
Start or refresh probation for a new user.

src/bot/handlers/anti_spam.py

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -348,6 +348,18 @@ async def handle_inline_keyboard_spam(
348348
)
349349

350350

351+
def _should_skip_new_user_spam_check(update, context, group_config) -> bool:
352+
"""Check if new user spam handler should skip this message."""
353+
if group_config is None:
354+
return True
355+
user = update.message.from_user
356+
if user.is_bot:
357+
return True
358+
if is_user_admin_or_trusted(context, group_config.group_id, user.id):
359+
return True
360+
return False
361+
362+
351363
async def handle_new_user_spam(
352364
update: Update, context: ContextTypes.DEFAULT_TYPE
353365
) -> None:
@@ -371,15 +383,7 @@ async def handle_new_user_spam(
371383
group_config = get_group_config_for_update(update)
372384
user = update.message.from_user
373385

374-
# Only process messages from monitored groups
375-
if group_config is None:
376-
return
377-
378-
# Ignore bots
379-
if user.is_bot:
380-
return
381-
382-
if is_user_admin_or_trusted(context, group_config.group_id, user.id):
386+
if _should_skip_new_user_spam_check(update, context, group_config):
383387
return
384388

385389
db = get_database()

src/bot/handlers/bio_bait.py

Lines changed: 74 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,13 @@
3737
WHITELISTED_TELEGRAM_PATHS,
3838
)
3939
from bot.group_config import get_group_config_for_update
40-
from bot.services.telegram_utils import get_user_mention, is_user_admin_or_trusted, is_url_whitelisted
40+
from bot.services.telegram_utils import (
41+
get_user_mention,
42+
is_user_admin_or_trusted,
43+
is_url_whitelisted,
44+
restrict_chat_member_with_retry,
45+
send_message_with_retry,
46+
)
4147

4248
# Filter for bio-bait handler registration in main.py.
4349
# Must NOT restrict to TEXT|CAPTION so non-text messages (e.g. photos
@@ -269,7 +275,14 @@ async def send_monitor_alert_to_owner(
269275

270276
try:
271277
for chunk in _chunk_telegram_text(alert_text):
272-
await context.bot.send_message(chat_id=alert_chat_id, text=chunk)
278+
ok = await send_message_with_retry(
279+
context.bot, chat_id=alert_chat_id, text=chunk
280+
)
281+
if not ok:
282+
logger.error(
283+
f"Failed to send bio bait monitor alert chunk: user_id={user_id}, group_id={group_id}"
284+
)
285+
return False
273286
return True
274287
except Exception:
275288
logger.error(f"Failed to send bio bait monitor alert: user_id={user_id}, group_id={group_id}")
@@ -312,6 +325,64 @@ async def get_cached_user_bio(
312325
cache[user_id] = (now + USER_BIO_FAILURE_CACHE_TTL_SECONDS, _BIO_CACHE_FAILURE)
313326
return None
314327

328+
329+
async def _enforce_bio_bait_restriction(
330+
update: Update,
331+
context: ContextTypes.DEFAULT_TYPE,
332+
group_config,
333+
user,
334+
detection_reason: str,
335+
) -> None:
336+
"""Delete, restrict, notify for confirmed bio bait spam. Caller raises ApplicationHandlerStop."""
337+
user_mention = get_user_mention(user)
338+
339+
try:
340+
await update.message.delete()
341+
logger.info(f"Deleted bio bait spam from user_id={user.id}")
342+
except Exception:
343+
logger.error(f"Failed to delete bio bait spam: user_id={user.id}", exc_info=True)
344+
345+
restricted = False
346+
try:
347+
await restrict_chat_member_with_retry(
348+
context.bot,
349+
chat_id=group_config.group_id,
350+
user_id=user.id,
351+
permissions=RESTRICTED_PERMISSIONS,
352+
)
353+
restricted = True
354+
clear_cached_user_bio(context, user.id)
355+
logger.info(f"Restricted user_id={user.id} for bio bait spam")
356+
except Exception:
357+
logger.error(f"Failed to restrict user for bio bait spam: user_id={user.id}", exc_info=True)
358+
359+
try:
360+
if detection_reason == "bio_links":
361+
template = (
362+
BIO_LINK_SPAM_NOTIFICATION if restricted
363+
else BIO_LINK_SPAM_NOTIFICATION_NO_RESTRICT
364+
)
365+
else:
366+
template = (
367+
BIO_BAIT_SPAM_NOTIFICATION if restricted
368+
else BIO_BAIT_SPAM_NOTIFICATION_NO_RESTRICT
369+
)
370+
notification_text = template.format(
371+
user_mention=user_mention,
372+
rules_link=group_config.rules_link,
373+
)
374+
await send_message_with_retry(
375+
context.bot,
376+
chat_id=group_config.group_id,
377+
message_thread_id=group_config.warning_topic_id,
378+
text=notification_text,
379+
parse_mode="Markdown",
380+
)
381+
logger.info(f"Sent bio bait spam notification for user_id={user.id}")
382+
except Exception:
383+
logger.error(f"Failed to send bio bait spam notification: user_id={user.id}", exc_info=True)
384+
385+
315386
async def handle_bio_bait_spam(
316387
update: Update, context: ContextTypes.DEFAULT_TYPE
317388
) -> None:
@@ -392,50 +463,5 @@ async def handle_bio_bait_spam(
392463
)
393464
return
394465

395-
user_mention = get_user_mention(user)
396-
397-
try:
398-
await update.message.delete()
399-
logger.info(f"Deleted bio bait spam from user_id={user.id}")
400-
except Exception:
401-
logger.error(f"Failed to delete bio bait spam: user_id={user.id}", exc_info=True)
402-
403-
restricted = False
404-
try:
405-
await context.bot.restrict_chat_member(
406-
chat_id=group_config.group_id,
407-
user_id=user.id,
408-
permissions=RESTRICTED_PERMISSIONS,
409-
)
410-
restricted = True
411-
clear_cached_user_bio(context, user.id)
412-
logger.info(f"Restricted user_id={user.id} for bio bait spam")
413-
except Exception:
414-
logger.error(f"Failed to restrict user for bio bait spam: user_id={user.id}", exc_info=True)
415-
416-
try:
417-
if detection_reason == "bio_links":
418-
template = (
419-
BIO_LINK_SPAM_NOTIFICATION if restricted
420-
else BIO_LINK_SPAM_NOTIFICATION_NO_RESTRICT
421-
)
422-
else:
423-
template = (
424-
BIO_BAIT_SPAM_NOTIFICATION if restricted
425-
else BIO_BAIT_SPAM_NOTIFICATION_NO_RESTRICT
426-
)
427-
notification_text = template.format(
428-
user_mention=user_mention,
429-
rules_link=group_config.rules_link,
430-
)
431-
await context.bot.send_message(
432-
chat_id=group_config.group_id,
433-
message_thread_id=group_config.warning_topic_id,
434-
text=notification_text,
435-
parse_mode="Markdown",
436-
)
437-
logger.info(f"Sent bio bait spam notification for user_id={user.id}")
438-
except Exception:
439-
logger.error(f"Failed to send bio bait spam notification: user_id={user.id}", exc_info=True)
440-
466+
await _enforce_bio_bait_restriction(update, context, group_config, user, detection_reason)
441467
raise ApplicationHandlerStop

src/bot/handlers/check.py

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
get_user_mention,
3131
get_user_mention_by_id,
3232
require_admin_dm_target,
33+
send_message_with_retry,
3334
)
3435
from bot.services.user_checker import check_user_profile
3536

@@ -205,6 +206,17 @@ async def handle_check_forwarded_message(
205206
logger.error(f"Error checking forwarded user {user_id}: {e}", exc_info=True)
206207

207208

209+
def _parse_warn_callback_data(data: str) -> tuple[int, str] | None:
210+
"""Parse warn callback data (warn:<user_id>:<missing_code>). Returns (user_id, missing_code) or None."""
211+
try:
212+
parts = data.split(":")
213+
user_id = int(parts[1])
214+
missing_code = parts[2] if len(parts) > 2 else ""
215+
return (user_id, missing_code)
216+
except (IndexError, ValueError):
217+
return None
218+
219+
208220
async def handle_warn_callback(
209221
update: Update, context: ContextTypes.DEFAULT_TYPE
210222
) -> None:
@@ -231,14 +243,12 @@ async def handle_warn_callback(
231243
return
232244

233245
# Parse callback data: warn:<user_id>:<missing_code>
234-
try:
235-
parts = query.data.split(":")
236-
target_user_id = int(parts[1])
237-
missing_code = parts[2] if len(parts) > 2 else ""
238-
except (IndexError, ValueError):
246+
parsed = _parse_warn_callback_data(query.data)
247+
if parsed is None:
239248
await query.edit_message_text("❌ Data callback tidak valid.")
240249
logger.error(f"Invalid callback_data format: {query.data}")
241250
return
251+
target_user_id, missing_code = parsed
242252

243253
# Build missing items text
244254
missing_items = []
@@ -264,16 +274,18 @@ async def handle_warn_callback(
264274
rules_link=group_config.rules_link,
265275
)
266276
try:
267-
await context.bot.send_message(
277+
ok = await send_message_with_retry(
278+
context.bot,
268279
chat_id=group_config.group_id,
269280
message_thread_id=group_config.warning_topic_id,
270281
text=warn_message,
271282
parse_mode="Markdown",
272283
)
273-
sent_to_any = True
274-
logger.info(
275-
f"Admin {admin_user_id} sent warning to user {target_user_id} in group {group_config.group_id}"
276-
)
284+
if ok:
285+
sent_to_any = True
286+
logger.info(
287+
f"Admin {admin_user_id} sent warning to user {target_user_id} in group {group_config.group_id}"
288+
)
277289
except Exception as e:
278290
logger.error(f"Failed to send warning to group {group_config.group_id}: {e}")
279291

0 commit comments

Comments
 (0)