Coverage for src/local_deep_research/notifications/manager.py: 99%
228 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1"""
2High-level notification manager with database integration.
3"""
5from dataclasses import dataclass
6from datetime import datetime, timezone, timedelta
7from enum import Enum
8from typing import Dict, Optional, Any
9from collections import deque
10import threading
12from loguru import logger
14from ..security.notification_validator import parse_notification_url_list
15from ..security.ssrf_validator import redact_url_for_log
16from ..settings.env_registry import get_env_setting
17from .service import NotificationService
18from .templates import EventType
19from .exceptions import (
20 RateLimitError,
21 SendError,
22 ServiceError,
23 SecurityBlockError,
24)
25from ..utilities.type_utils import unwrap_setting
26from ..constants import DEFAULT_SEARCH_TOOL
29class NotificationReason(str, Enum):
30 """Precise reason a notification was (not) sent.
32 Surfaced in :class:`NotificationResult` so callers and log lines can
33 distinguish, for example, a per-user toggle being off from the
34 server-level master switch being unset — see issue #4877.
36 Note: rate-limit refusal is signalled via :class:`RateLimitError`
37 (preserved API for callers that catch it directly) rather than as a
38 ``NotificationResult`` value.
39 """
41 SENT = "sent"
42 SERVER_DISABLED = "server_disabled"
43 EVENT_DISABLED = "event_disabled"
44 UNCONFIGURED = "unconfigured"
45 EGRESS_DENIED = "egress_denied"
46 INVALID_URL = "invalid_url"
47 WEBHOOK_FAILED = "webhook_failed"
48 EXCEPTION = "exception"
51@dataclass(frozen=True)
52class NotificationResult:
53 """Structured outcome of :meth:`NotificationManager.send_notification`.
55 ``sent`` keeps the legacy bool contract (``if result:`` keeps working —
56 see ``__bool__``) while ``reason`` names the precise cause so the queue
57 helpers and error reporter can log a useful, non-misleading line.
58 ``detail`` is a short operator-facing explanation complementing reason.
59 """
61 sent: bool
62 reason: NotificationReason
63 detail: str = ""
65 def __bool__(self) -> bool:
66 return self.sent
69class NotificationManager:
70 """
71 High-level notification manager that uses settings snapshots for
72 thread-safe access to user settings.
74 This manager is designed to be used from background threads (e.g., queue
75 processors) by passing a settings_snapshot dictionary captured from the
76 main thread.
78 **Per-User Rate Limiting:**
79 The rate limiter is shared across ALL NotificationManager instances as a
80 singleton, but supports per-user rate limit configuration. Each user has
81 their own rate limits based on their settings, which are configured when
82 the NotificationManager is initialized with the required user_id parameter.
84 **How It Works:**
85 - The first NotificationManager instance creates the shared RateLimiter
86 with default limits
87 - Each instance configures user-specific limits by passing user_id to __init__
88 - The rate limiter maintains separate counters and limits for each user
89 - Users are completely isolated - one user's limit doesn't affect others
91 Example:
92 >>> # User A with conservative limits
93 >>> snapshot_a = {"notifications.rate_limit_per_hour": 5}
94 >>> manager_a = NotificationManager(snapshot_a, user_id="user_a")
95 >>> # ✅ user_a gets 5/hour
96 >>>
97 >>> # User B with generous limits (doesn't affect User A!)
98 >>> snapshot_b = {"notifications.rate_limit_per_hour": 20}
99 >>> manager_b = NotificationManager(snapshot_b, user_id="user_b")
100 >>> # ✅ user_b gets 20/hour, user_a still has 5/hour
101 """
103 # Shared rate limiter instance across all NotificationManager instances
104 # This ensures rate limits are enforced correctly even when multiple
105 # NotificationManager instances are created
106 _shared_rate_limiter: Optional["RateLimiter"] = None
107 _rate_limiter_lock = threading.Lock()
109 def __init__(self, settings_snapshot: Dict[str, Any], user_id: str):
110 """
111 Initialize the notification manager.
113 Args:
114 settings_snapshot: Dictionary of settings key-value pairs captured
115 from SettingsManager.get_settings_snapshot().
116 This allows thread-safe access to user settings
117 from background threads.
118 user_id: User identifier for per-user rate limiting. The rate limits
119 from settings_snapshot will be configured for this user.
121 Example:
122 >>> # In main thread with database session
123 >>> settings_manager = SettingsManager(session)
124 >>> snapshot = settings_manager.get_settings_snapshot()
125 >>>
126 >>> # In background thread (thread-safe)
127 >>> notification_manager = NotificationManager(
128 ... settings_snapshot=snapshot,
129 ... user_id="user123"
130 ... )
131 >>> notification_manager.send_notification(...)
132 """
133 # Store settings snapshot for thread-safe access
134 self._settings_snapshot = settings_snapshot
135 self._user_id = user_id
137 # Security: read from server-side environment variable only — never from
138 # user-writable DB settings. Previously this was read via
139 # _get_setting("notifications.allow_private_ips"), which allowed any
140 # user to bypass SSRF protection through the settings API.
141 # Registered as env-only in settings/env_definitions/security.py
142 # so SettingsManager will always read LDR_NOTIFICATIONS_ALLOW_PRIVATE_IPS
143 # from the environment, never from the database.
144 allow_private_ips = get_env_setting(
145 "notifications.allow_private_ips", False
146 )
148 # Server-level master switch (env-only). Distinct from the per-user
149 # notifications.enabled toggle in the UI: this gate is set by the
150 # deployment operator and cannot be flipped via the user-writable
151 # settings API. Disabled by default; enabling it accepts the
152 # documented DNS-rebinding TOCTOU residual risk
153 # (see SECURITY.md "Notification Webhook SSRF").
154 self._outbound_allowed = bool(
155 get_env_setting("notifications.allow_outbound", False)
156 )
158 self.service = NotificationService(
159 allow_private_ips=allow_private_ips,
160 outbound_allowed=self._outbound_allowed,
161 )
163 # Initialize shared rate limiter on first use
164 # The shared rate limiter now supports per-user limits, so each user's
165 # settings are respected regardless of initialization order.
166 with NotificationManager._rate_limiter_lock:
167 if NotificationManager._shared_rate_limiter is None:
168 # Create shared rate limiter with default limits
169 # (individual users can have different limits)
170 default_max_per_hour = self._get_setting(
171 "notifications.rate_limit_per_hour", default=10
172 )
173 default_max_per_day = self._get_setting(
174 "notifications.rate_limit_per_day", default=50
175 )
177 logger.info(
178 f"Initializing shared rate limiter with defaults: "
179 f"{default_max_per_hour}/hour, {default_max_per_day}/day"
180 )
182 NotificationManager._shared_rate_limiter = RateLimiter(
183 max_per_hour=default_max_per_hour,
184 max_per_day=default_max_per_day,
185 )
187 # Use the shared instance
188 self._rate_limiter = NotificationManager._shared_rate_limiter
190 # Configure per-user rate limits
191 max_per_hour = self._get_setting(
192 "notifications.rate_limit_per_hour", default=10
193 )
194 max_per_day = self._get_setting(
195 "notifications.rate_limit_per_day", default=50
196 )
198 self._rate_limiter.set_user_limits(
199 user_id, max_per_hour, max_per_day
200 )
201 logger.debug(
202 f"Configured rate limits for user {user_id}: "
203 f"{max_per_hour}/hour, {max_per_day}/day"
204 )
206 def _get_setting(self, key: str, default: Any = None) -> Any:
207 """
208 Get a setting value from snapshot.
210 Args:
211 key: Setting key
212 default: Default value if not found
214 Returns:
215 Setting value or default
216 """
217 return self._settings_snapshot.get(key, default)
219 def send_notification(
220 self,
221 event_type: EventType,
222 context: Dict[str, Any],
223 force: bool = False,
224 ) -> NotificationResult:
225 """
226 Send a notification for an event.
228 Uses the user_id that was provided during initialization for
229 rate limiting and user preferences.
231 Args:
232 event_type: Type of event
233 context: Context data for the notification
234 force: If True, bypass rate limiting
236 Returns:
237 :class:`NotificationResult` describing the outcome. The
238 dataclass is truthy iff the notification was actually sent,
239 so legacy ``if manager.send_notification(...):`` callers
240 continue to work without source changes.
242 Raises:
243 RateLimitError: If rate limit is exceeded and force=False
244 """
245 logger.debug(
246 f"Sending notification: event_type={event_type.value}, "
247 f"user_id={self._user_id}, force={force}"
248 )
249 logger.debug(f"Context keys: {list(context.keys())}")
251 # Server-level master switch (env-only). Disabled by default;
252 # flipping it on is the operator's acknowledgement of the
253 # documented SSRF rebinding residual risk. force=True does NOT
254 # bypass this — it bypasses the per-user event-type and
255 # rate-limit toggles only. WARNING level so an operator wondering
256 # "why aren't notifications firing?" sees the actionable signal
257 # at default log level.
258 if not self._outbound_allowed:
259 logger.warning(
260 "Notification refused: outbound notifications are disabled "
261 "at the server level. Set "
262 "LDR_NOTIFICATIONS_ALLOW_OUTBOUND=true to enable. See "
263 "SECURITY.md 'Notification Webhook SSRF' for the rationale "
264 "and residual risk. (event={}, user={})",
265 event_type.value,
266 self._user_id,
267 )
268 return NotificationResult(
269 sent=False,
270 reason=NotificationReason.SERVER_DISABLED,
271 detail=(
272 "outbound notifications are disabled at the server "
273 "level (LDR_NOTIFICATIONS_ALLOW_OUTBOUND)"
274 ),
275 )
277 # Check if notifications are enabled for this event type
278 should_notify = self._should_notify(event_type)
279 logger.debug(
280 f"Notification enabled check for {event_type.value}: "
281 f"{should_notify}"
282 )
284 if not force and not should_notify:
285 logger.debug(
286 f"Notifications disabled for event type: "
287 f"{event_type.value} (user: {self._user_id})"
288 )
289 return NotificationResult(
290 sent=False,
291 reason=NotificationReason.EVENT_DISABLED,
292 detail=f"notifications.on_{event_type.value} is off",
293 )
295 # Check rate limit using the manager's user_id
296 rate_limit_ok = self._rate_limiter.is_allowed(self._user_id)
297 logger.debug(f"Rate limit check for {self._user_id}: {rate_limit_ok}")
299 if not force and not rate_limit_ok:
300 logger.warning(f"Rate limit exceeded for user {self._user_id}")
301 raise RateLimitError(
302 "Notification rate limit exceeded. "
303 "Please wait before sending more notifications."
304 )
306 try:
307 # Get service URLs from settings (snapshot or database)
308 service_urls = self._get_setting(
309 "notifications.service_url", default=""
310 )
312 if not service_urls or not service_urls.strip():
313 logger.debug(
314 f"No notification service URLs configured for user "
315 f"{self._user_id}"
316 )
317 return NotificationResult(
318 sent=False,
319 reason=NotificationReason.UNCONFIGURED,
320 detail="notifications.service_url is not configured",
321 )
323 # Egress policy: classify raw HTTP(S) URLs with evaluate_url.
324 # PRIVATE_ONLY refuses non-HTTP Apprise schemes because their
325 # effective destinations cannot be classified reliably; other
326 # scopes pass plugin schemes to the service's URL validator.
327 allowed_urls = self._filter_urls_by_egress_policy(service_urls)
328 if not allowed_urls:
329 # Distinguish "the policy was evaluated and refused every
330 # URL" (empty string) from "the policy itself could not be
331 # evaluated and we failed closed" (None) so the operator
332 # fixes the right thing — see issue #5110.
333 if allowed_urls is None:
334 detail = "egress policy could not be evaluated"
335 else:
336 detail = "all configured URLs refused by egress policy"
337 logger.bind(policy_audit=True).warning(
338 "notification dropped: {}",
339 detail,
340 user=self._user_id,
341 event=event_type.value,
342 )
343 return NotificationResult(
344 sent=False,
345 reason=NotificationReason.EGRESS_DENIED,
346 detail=detail,
347 )
349 # Send notification with the allowed subset.
350 logger.debug(f"Calling service.send_event for {event_type.value}")
351 sent = self.service.send_event(
352 event_type, context, service_urls=allowed_urls
353 )
355 # Log to database if enabled
356 if sent:
357 self._log_notification(event_type, context)
358 logger.info(
359 f"Notification sent: {event_type.value} to user "
360 f"{self._user_id}"
361 )
362 return NotificationResult(
363 sent=True,
364 reason=NotificationReason.SENT,
365 detail="",
366 )
367 # A falsy return from send_event means Apprise accepted none
368 # of the URLs (temp_apprise.add() failed — unparseable or
369 # unsupported scheme). Actual delivery failures raise
370 # SendError and are handled below — see issue #5110.
371 logger.warning(
372 f"Notification failed: {event_type.value} to user "
373 f"{self._user_id} — Apprise accepted none of the "
374 f"configured service URLs"
375 )
376 return NotificationResult(
377 sent=False,
378 reason=NotificationReason.INVALID_URL,
379 detail="Apprise accepted none of the configured service URLs",
380 )
382 except SendError:
383 # The delivery layer raises SendError ONLY from
384 # _send_with_retry's own terminal raise once all retries are
385 # exhausted (dead webhook, HTTP 4xx/5xx, network down) — a
386 # transient, worth-retrying failure. NotificationService.send()
387 # re-raises SendError unchanged (no re-wrapping) and does NOT
388 # launder incidental/unexpected dispatch-time exceptions into
389 # SendError — those fall through to the generic
390 # ``except Exception`` branch below and get the EXCEPTION
391 # reason instead, so this branch stays reserved for genuine
392 # delivery failures. It also specifically does NOT mean a
393 # confirmed SSRF/DNS-rebind block: NotificationService.send()
394 # routes that non-retryable case through ServiceError instead
395 # (see the except clause below) precisely so it is never
396 # mislabeled as a retryable webhook failure here. Keep the
397 # detail static — SendError's message wraps the underlying
398 # error, which may contain webhook URLs with embedded tokens.
399 logger.exception(
400 f"Webhook delivery failed for {event_type.value} to user "
401 f"{self._user_id}"
402 )
403 return NotificationResult(
404 sent=False,
405 reason=NotificationReason.WEBHOOK_FAILED,
406 detail="webhook delivery failed after retries",
407 )
409 except SecurityBlockError:
410 # A ServiceError subclass NotificationService.send() raises
411 # specifically for a confirmed send-time SSRF/DNS-rebind block
412 # (destination resolved to an internal/private/metadata
413 # address after dispatch started) — as opposed to the plain
414 # ServiceError below, raised for a pre-dispatch validation
415 # reject. Both are permanent, invalid-destination outcomes
416 # that map to the same non-transient INVALID_URL reason, but
417 # this one gets a distinct detail: the generic "check your
418 # URL/settings" text below would misleadingly suggest a
419 # config mistake when the true cause is an active security
420 # block. Static detail — no URL or resolved IP — the
421 # exception message can echo parts of the URL.
422 logger.exception(
423 f"Notification destination blocked by SSRF/egress "
424 f"protection for {event_type.value} to user "
425 f"{self._user_id}"
426 )
427 return NotificationResult(
428 sent=False,
429 reason=NotificationReason.INVALID_URL,
430 detail=(
431 "notification destination blocked by egress/SSRF protection"
432 ),
433 )
435 except ServiceError:
436 # The service raises ServiceError when the URL security
437 # validator rejects a URL before any dispatch is attempted
438 # (structurally invalid URLs, or valid URLs whose destination
439 # is blocked, e.g. a LAN webhook without
440 # LDR_NOTIFICATIONS_ALLOW_PRIVATE_IPS). A confirmed send-time
441 # SSRF/DNS-rebind block is caught above as SecurityBlockError
442 # (a ServiceError subclass) before reaching this branch. The
443 # detail must not claim the URL itself is malformed. Static
444 # detail: the exception message can echo parts of the URL.
445 logger.exception(
446 f"Service URL rejected for {event_type.value} to user "
447 f"{self._user_id}"
448 )
449 return NotificationResult(
450 sent=False,
451 reason=NotificationReason.INVALID_URL,
452 detail=(
453 "service URL rejected by URL security validation "
454 "(for a private/LAN webhook, check "
455 "LDR_NOTIFICATIONS_ALLOW_PRIVATE_IPS)"
456 ),
457 )
459 except Exception as exc:
460 # Reached for unexpected failures anywhere in the block above
461 # (settings read, egress filtering, or dispatch). The exception
462 # string may contain webhook URLs with embedded tokens, SMTP
463 # credentials, or payload content. Log the raw exception for
464 # operators via logger.exception, but only expose the exception
465 # class name (and not the message) in the structured result so
466 # downstream callers can't echo it into logs or UI surfaces
467 # verbatim.
468 logger.exception(
469 f"Error sending notification for {event_type.value} to user "
470 f"{self._user_id}"
471 )
472 return NotificationResult(
473 sent=False,
474 reason=NotificationReason.EXCEPTION,
475 detail=f"unexpected {type(exc).__name__} while sending",
476 )
478 def test_service(self, url: str) -> Dict[str, Any]:
479 """
480 Test a notification service.
482 Args:
483 url: Service URL to test
485 Returns:
486 Dict with test results
487 """
488 # Egress policy precheck before forwarding to Apprise.
489 # /api/notifications/test-url is a user-supplied URL endpoint.
490 allowed = self._filter_urls_by_egress_policy(url)
491 if allowed is None:
492 # Fail-closed: the policy could not be evaluated at all, so
493 # don't tell the user their scope refused the URL.
494 return {
495 "status": "error",
496 "message": (
497 "Egress policy could not be evaluated; check the "
498 "egress policy settings (policy.egress_scope)."
499 ),
500 }
501 if not allowed:
502 return {
503 "status": "error",
504 "message": (
505 "URL refused by egress policy. Set Egress Scope to "
506 "'Unprotected' (or pick a local webhook) to test this URL."
507 ),
508 }
509 return self.service.test_service(allowed)
511 def _filter_urls_by_egress_policy(self, service_urls: str) -> Optional[str]:
512 """Filter an Apprise URL string by the user's egress policy.
514 Commas inside one service URL are preserved; commas or whitespace
515 followed by another scheme delimit URLs. Returns the joined string
516 of allowed URLs (may be empty — every URL was refused), or ``None``
517 when the policy itself could not be evaluated and we failed closed.
518 Both are falsy; callers that care why nothing was allowed check for
519 ``None``. When no snapshot / context is available, returns the
520 input unchanged for backwards compatibility.
521 """
522 snapshot = getattr(self, "_settings_snapshot", None)
523 if not snapshot:
524 return service_urls
526 url_entries, invalid_fragment = parse_notification_url_list(
527 service_urls
528 )
529 if invalid_fragment is not None:
530 logger.bind(policy_audit=True).warning(
531 "all notification URLs refused: scheme-less URL fragment",
532 fragment=redact_url_for_log(invalid_fragment),
533 )
534 return ""
536 try:
537 from ..security.egress.policy import (
538 EgressScope,
539 PolicyDeniedError,
540 context_from_snapshot,
541 evaluate_url,
542 )
543 except ImportError:
544 logger.debug(
545 "egress_policy unavailable in notifications manager; "
546 "URLs will not be scope-gated"
547 )
548 return service_urls
550 from ..search_system import username_from_snapshot
552 primary_raw = unwrap_setting(
553 snapshot.get("search.tool", DEFAULT_SEARCH_TOOL)
554 )
555 try:
556 ctx = context_from_snapshot(
557 snapshot,
558 primary_raw or DEFAULT_SEARCH_TOOL,
559 username=username_from_snapshot(snapshot),
560 )
561 except (PolicyDeniedError, ValueError) as exc:
562 # Snapshot present but policy cannot be evaluated. The
563 # previous bare-except dropped to "return service_urls"
564 # here, which dispatched every URL unfiltered — fail-open
565 # on a misconfigured policy. Now refuse all URLs instead
566 # so an Apprise dispatch with no scope check is impossible.
567 logger.bind(policy_audit=True).warning(
568 "all notification URLs refused: egress policy could "
569 "not be evaluated",
570 reason=str(exc),
571 )
572 return None
574 # Apprise accepts space- or comma-separated URLs and has non-HTTP
575 # schemes (discord://, slack://, tgram://, mailto://, msteams://,
576 # ntfy://, ...) whose effective destination may be a fixed provider,
577 # a user-supplied authority, or a provider mapping. evaluate_url only
578 # understands HTTP(S), so it cannot classify those modes reliably.
579 #
580 # Under PRIVATE_ONLY ("nothing leaves the box"), a non-HTTP scheme is
581 # therefore refused — fail closed. Services that expose a raw HTTP(S) webhook
582 # can use that URL and be classified by evaluate_url; Apprise SMTP and
583 # other plugin URLs remain unavailable under PRIVATE_ONLY. Under the
584 # other scopes the prior pass-through stands (the modeled threat there
585 # is internal-http SSRF via a raw webhook, not vendor APIs).
586 parts = []
587 for url_entry in url_entries:
588 scheme = (
589 url_entry.split(":", 1)[0].lower() if ":" in url_entry else ""
590 )
591 if scheme not in ("http", "https"):
592 if ctx.scope == EgressScope.PRIVATE_ONLY:
593 logger.bind(policy_audit=True).warning(
594 "notification refused under PRIVATE_ONLY egress "
595 "scope: non-http vendor scheme cannot be verified "
596 "local",
597 scheme=scheme or "(none)",
598 )
599 continue
600 parts.append(url_entry)
601 continue
602 decision = evaluate_url(url_entry, ctx)
603 if decision.allowed:
604 parts.append(url_entry)
605 else:
606 logger.bind(policy_audit=True).warning(
607 "notification URL refused by egress policy",
608 url=url_entry,
609 scope=ctx.scope.value,
610 reason=decision.reason,
611 )
612 return " ".join(parts)
614 def _should_notify(self, event_type: EventType) -> bool:
615 """
616 Check if notifications should be sent for this event type.
618 Uses the manager's settings snapshot to determine if the event type
619 is enabled for the user.
621 Args:
622 event_type: Event type to check
624 Returns:
625 True if notifications should be sent
626 """
627 try:
628 # Check event-specific setting (from snapshot or database)
629 setting_key = f"notifications.on_{event_type.value}"
630 enabled = self._get_setting(setting_key, default=False)
632 return bool(enabled)
634 except Exception:
635 logger.warning("Error checking notification preferences")
636 # Default to disabled on error to avoid infinite loops during login
637 return False
639 def _log_notification(
640 self, event_type: EventType, context: Dict[str, Any]
641 ) -> None:
642 """
643 Log a sent notification (simplified logging to application logs only).
645 Uses the manager's user_id for logging.
647 Args:
648 event_type: Event type
649 context: Notification context
650 """
651 try:
652 title = (
653 context.get("query")
654 or context.get("subscription_name")
655 or "Unknown"
656 )
657 logger.info(
658 f"Notification sent: {event_type.value} - {title} "
659 f"(user: {self._user_id})"
660 )
661 except Exception as e:
662 logger.debug(f"Failed to log notification: {e}")
665class RateLimiter:
666 """
667 Simple in-memory rate limiter for notifications with per-user limit support.
669 This rate limiter tracks notification counts per user and enforces
670 configurable rate limits. Each user can have their own rate limits,
671 which are stored separately from the notification counts.
673 **Per-User Limits:**
674 Rate limits can be configured per-user using `set_user_limits()`.
675 If no user-specific limits are set, the default limits (passed to
676 __init__) are used.
678 **Memory Storage:**
679 This implementation stores rate limits in memory only, which means
680 limits are reset when the server restarts. This is acceptable for normal
681 users since they cannot restart the server. If an admin restarts the server,
682 rate limits reset which is reasonable behavior.
684 **Thread Safety:**
685 This implementation is thread-safe using threading.Lock() for concurrent
686 requests from the same user.
688 **Multi-Worker Limitation:**
689 In multi-worker deployments, each worker process maintains its own rate
690 limit counters. Users could potentially bypass rate limits by distributing
691 requests across different workers, getting up to N × max_per_hour
692 notifications (where N = number of workers). For single-worker deployments
693 (the default for LDR), this is not a concern. For production multi-worker
694 deployments, consider implementing Redis-based rate limiting.
696 Example:
697 >>> limiter = RateLimiter(max_per_hour=10, max_per_day=50)
698 >>> # Set custom limits for specific user
699 >>> limiter.set_user_limits("user_a", max_per_hour=5, max_per_day=25)
700 >>> limiter.set_user_limits("user_b", max_per_hour=20, max_per_day=100)
701 >>> # Users get their configured limits
702 >>> limiter.is_allowed("user_a") # Limited to 5/hour
703 >>> limiter.is_allowed("user_b") # Limited to 20/hour
704 >>> limiter.is_allowed("user_c") # Uses defaults: 10/hour
705 """
707 def __init__(
708 self,
709 max_per_hour: int = 10,
710 max_per_day: int = 50,
711 cleanup_interval_hours: int = 24,
712 ):
713 """
714 Initialize rate limiter with default limits.
716 Args:
717 max_per_hour: Default maximum notifications per hour per user
718 max_per_day: Default maximum notifications per day per user
719 cleanup_interval_hours: How often to run cleanup of inactive users (hours)
720 """
721 # Default limits used when no user-specific limits are set
722 self.max_per_hour = max_per_hour
723 self.max_per_day = max_per_day
724 self.cleanup_interval_hours = cleanup_interval_hours
726 # Per-user rate limit configuration (user_id -> (max_per_hour, max_per_day))
727 self._user_limits: Dict[str, tuple[int, int]] = {}
729 # Per-user notification counts
730 self._hourly_counts: Dict[str, deque] = {}
731 self._daily_counts: Dict[str, deque] = {}
733 self._last_cleanup = datetime.now(timezone.utc)
734 self._lock = threading.Lock() # Thread safety for all operations
736 def set_user_limits(
737 self, user_id: str, max_per_hour: int, max_per_day: int
738 ) -> None:
739 """
740 Set rate limits for a specific user.
742 This allows each user to have their own rate limit configuration.
743 If not set, the user will use the default limits passed to __init__.
745 Args:
746 user_id: User identifier
747 max_per_hour: Maximum notifications per hour for this user
748 max_per_day: Maximum notifications per day for this user
750 Example:
751 >>> limiter = RateLimiter(max_per_hour=10, max_per_day=50)
752 >>> limiter.set_user_limits("power_user", max_per_hour=20, max_per_day=100)
753 >>> limiter.set_user_limits("limited_user", max_per_hour=5, max_per_day=25)
754 """
755 with self._lock:
756 self._user_limits[user_id] = (max_per_hour, max_per_day)
757 logger.debug(
758 f"Set rate limits for user {user_id}: "
759 f"{max_per_hour}/hour, {max_per_day}/day"
760 )
762 def get_user_limits(self, user_id: str) -> tuple[int, int]:
763 """
764 Get the effective rate limits for a user.
766 Returns the user-specific limits if set, otherwise returns defaults.
768 Args:
769 user_id: User identifier
771 Returns:
772 Tuple of (max_per_hour, max_per_day)
773 """
774 with self._lock:
775 return self._user_limits.get(
776 user_id, (self.max_per_hour, self.max_per_day)
777 )
779 def is_allowed(self, user_id: str) -> bool:
780 """
781 Check if a notification is allowed for a user.
783 Uses per-user rate limits if configured via set_user_limits(),
784 otherwise uses the default limits from __init__.
786 Args:
787 user_id: User identifier
789 Returns:
790 True if notification is allowed, False if rate limit exceeded
791 """
792 with self._lock:
793 now = datetime.now(timezone.utc)
795 # Periodic cleanup of inactive users
796 self._cleanup_inactive_users_if_needed(now)
798 # Initialize queues for user if needed
799 if user_id not in self._hourly_counts:
800 self._hourly_counts[user_id] = deque()
801 self._daily_counts[user_id] = deque()
803 # Clean old entries
804 self._clean_old_entries(user_id, now)
806 # Get user-specific limits or defaults
807 max_per_hour, max_per_day = self._user_limits.get(
808 user_id, (self.max_per_hour, self.max_per_day)
809 )
811 # Check limits
812 hourly_count = len(self._hourly_counts[user_id])
813 daily_count = len(self._daily_counts[user_id])
815 if hourly_count >= max_per_hour:
816 logger.warning(
817 f"Hourly rate limit exceeded for user {user_id}: "
818 f"{hourly_count}/{max_per_hour}"
819 )
820 return False
822 if daily_count >= max_per_day:
823 logger.warning(
824 f"Daily rate limit exceeded for user {user_id}: "
825 f"{daily_count}/{max_per_day}"
826 )
827 return False
829 # Record this notification
830 self._hourly_counts[user_id].append(now)
831 self._daily_counts[user_id].append(now)
833 return True
835 def _clean_old_entries(self, user_id: str, now: datetime) -> None:
836 """
837 Remove old entries from rate limit counters.
839 Args:
840 user_id: User identifier
841 now: Current time
842 """
843 hour_ago = now - timedelta(hours=1)
844 day_ago = now - timedelta(days=1)
846 # Clean hourly queue
847 while (
848 self._hourly_counts[user_id]
849 and self._hourly_counts[user_id][0] < hour_ago
850 ):
851 self._hourly_counts[user_id].popleft()
853 # Clean daily queue
854 while (
855 self._daily_counts[user_id]
856 and self._daily_counts[user_id][0] < day_ago
857 ):
858 self._daily_counts[user_id].popleft()
860 def reset(self, user_id: Optional[str] = None) -> None:
861 """
862 Reset rate limits for a user or all users.
864 Args:
865 user_id: User to reset, or None for all users
866 """
867 with self._lock:
868 if user_id:
869 self._hourly_counts.pop(user_id, None)
870 self._daily_counts.pop(user_id, None)
871 else:
872 self._hourly_counts.clear()
873 self._daily_counts.clear()
875 def _cleanup_inactive_users_if_needed(self, now: datetime) -> None:
876 """
877 Periodically clean up data for inactive users to prevent memory leaks.
879 Args:
880 now: Current timestamp
881 """
882 # Check if cleanup is needed
883 if now - self._last_cleanup < timedelta(
884 hours=self.cleanup_interval_hours
885 ):
886 return
888 logger.debug("Running periodic cleanup of inactive notification users")
890 # Define inactive threshold (users with no activity for 7 days)
891 inactive_threshold = now - timedelta(days=7)
893 inactive_users = []
895 # Find users with no recent activity
896 # Convert to list to avoid "dictionary changed size during iteration" error
897 for user_id in list(self._hourly_counts.keys()):
898 # Check if user has any recent entries
899 hourly_entries: list = list(self._hourly_counts.get(user_id, []))
900 daily_entries: list = list(self._daily_counts.get(user_id, []))
902 # If no entries or all entries are old, mark as inactive
903 has_recent_activity = False
904 for entry in hourly_entries + daily_entries:
905 if entry > inactive_threshold:
906 has_recent_activity = True
907 break
909 if not has_recent_activity:
910 inactive_users.append(user_id)
912 # Remove inactive users
913 for user_id in inactive_users:
914 self._hourly_counts.pop(user_id, None)
915 self._daily_counts.pop(user_id, None)
916 logger.debug(
917 f"Cleaned up inactive user {user_id} from rate limiter"
918 )
920 if inactive_users: 920 ↛ 925line 920 didn't jump to line 925 because the condition on line 920 was always true
921 logger.info(
922 f"Cleaned up {len(inactive_users)} inactive users from rate limiter"
923 )
925 self._last_cleanup = now