Coverage for src/local_deep_research/utilities/metric_write_health.py: 100%

1"""Bounded, per-user diagnostics for best-effort metric writes. 

2 

3Only fixed source/reason categories and counts are retained. This module stays 

4independent of database and metrics imports so optional/programmatic callers 

5can report a failure without initializing either subsystem. 

6""" 

7 

8from collections import Counter, OrderedDict 

9from threading import Lock 

10from typing import Literal, TypedDict 

11 

12MetricWriteSource = Literal["tokens", "search", "rate_limits"] 

13MetricWriteReason = Literal[ 

14 "write_error", 

15 "missing_credentials", 

16 "pending_transaction", 

17 "database_unavailable", 

18] 

19 

20_SOURCES: tuple[MetricWriteSource, ...] = ("tokens", "search", "rate_limits") 

21METRIC_WRITE_REASON_LABELS: dict[MetricWriteReason, str] = { 

22 "write_error": "Database write failures", 

23 "missing_credentials": "Missing database credentials", 

24 "pending_transaction": "Updates skipped while another write was pending", 

25 "database_unavailable": "Metrics storage unavailable", 

26} 

27MAX_TRACKED_USERS = 512 

28 

29 

30class MetricWriteSnapshot(TypedDict): 

31 dropped_writes: int 

32 by_source: dict[MetricWriteSource, int] 

33 by_reason: dict[MetricWriteReason, int] 

34 

35 

36class MetricWriteHealth: 

37 """Count dropped updates for at most ``max_users`` recent writers. 

38 

39 Counts are process-local lower bounds, reset by restart or eviction. Reads 

40 never allocate an entry or expose another user's counts. 

41 """ 

42 

43 def __init__(self, max_users: int = MAX_TRACKED_USERS) -> None: 

44 if max_users < 1: 

45 raise ValueError("max_users must be positive") 

46 self._max_users = max_users 

47 self._lock = Lock() 

48 self._users: OrderedDict[ 

49 str, Counter[tuple[MetricWriteSource, MetricWriteReason]] 

50 ] = OrderedDict() 

51 

52 def record( 

53 self, 

54 username: str | None, 

55 source: MetricWriteSource, 

56 reason: MetricWriteReason, 

57 ) -> None: 

58 if not isinstance(username, str) or not username: 

59 return 

60 if source not in _SOURCES or reason not in METRIC_WRITE_REASON_LABELS: 

61 return 

62 with self._lock: 

63 counts = self._users.setdefault(username, Counter()) 

64 counts[(source, reason)] += 1 

65 self._users.move_to_end(username) 

66 if len(self._users) > self._max_users: 

67 self._users.popitem(last=False) 

68 

69 def snapshot(self, username: str | None) -> MetricWriteSnapshot: 

70 counts: Counter[tuple[MetricWriteSource, MetricWriteReason]] = Counter() 

71 with self._lock: 

72 if isinstance(username, str): 

73 counts = self._users.get(username, counts).copy() 

74 return { 

75 "dropped_writes": sum(counts.values()), 

76 "by_source": { 

77 source: sum(n for (s, _), n in counts.items() if s == source) 

78 for source in _SOURCES 

79 }, 

80 "by_reason": { 

81 reason: sum(n for (_, r), n in counts.items() if r == reason) 

82 for reason in METRIC_WRITE_REASON_LABELS 

83 }, 

84 } 

85 

86 

87_health = MetricWriteHealth() 

88 

89 

90def record_metric_write_drop( 

91 username: str | None, source: MetricWriteSource, reason: MetricWriteReason 

92) -> None: 

93 _health.record(username, source, reason) 

94 

95 

96def get_metric_write_health(username: str) -> MetricWriteSnapshot: 

97 return _health.snapshot(username)