| 1 | import fnmatch |
| 2 | import threading |
| 3 | import time |
| 4 | from dataclasses import dataclass |
| 5 | from typing import Any |
| 6 | |
| 7 | _lock = threading.RLock() |
| 8 | _cache: dict[str, dict[str, "CacheEntry"]] = {} |
| 9 | |
| 10 | _enabled_global: bool = True |
| 11 | _enabled_areas: dict[str, bool] = {} |
| 12 | |
| 13 | |
| 14 | @dataclass(slots=True) |
| 15 | class CacheEntry: |
| 16 | value: Any |
| 17 | timestamp: float |
| 18 | |
| 19 | |
| 20 | def toggle_global(enabled: bool) -> None: |
| 21 | global _enabled_global |
| 22 | _enabled_global = enabled |
| 23 | |
| 24 | |
| 25 | def toggle_area(area: str, enabled: bool) -> None: |
| 26 | _enabled_areas[area] = enabled |
| 27 | |
| 28 | |
| 29 | def has(area: str, key: Any) -> bool: |
| 30 | if not _is_enabled(area): |
| 31 | return False |
| 32 | with _lock: |
| 33 | entry = _cache.get(area, {}).get(key) |
| 34 | if entry is None: |
| 35 | return False |
| 36 | _touch_entry(entry) |
| 37 | return True |
| 38 | |
| 39 | |
| 40 | def add(area: str, key: Any, data: Any) -> None: |
| 41 | if not _is_enabled(area): |
| 42 | return |
| 43 | with _lock: |
| 44 | if area not in _cache: |
| 45 | _cache[area] = {} |
| 46 | _cache[area][key] = _create_entry(data) |
| 47 | |
| 48 | |
| 49 | def get(area: str, key: Any, default: Any = None) -> Any: |
| 50 | if not _is_enabled(area): |
| 51 | return default |
| 52 | with _lock: |
| 53 | entry = _cache.get(area, {}).get(key) |
| 54 | if entry is None: |
| 55 | return default |
| 56 | _touch_entry(entry) |
| 57 | return entry.value |
| 58 | |
| 59 | |
| 60 | def remove(area: str, key: Any) -> None: |
| 61 | if not _is_enabled(area): |
| 62 | return |
| 63 | with _lock: |
| 64 | if area in _cache: |
| 65 | _cache[area].pop(key, None) |
| 66 | |
| 67 | |
| 68 | def clear(area: str) -> None: |
| 69 | with _lock: |
| 70 | if any(ch in area for ch in "*?["): |
| 71 | keys_to_remove = [k for k in _cache.keys() if fnmatch.fnmatch(k, area)] |
| 72 | for k in keys_to_remove: |
| 73 | _cache.pop(k, None) |
| 74 | return |
| 75 | |
| 76 | _cache.pop(area, None) |
| 77 | |
| 78 | |
| 79 | def trim_cache(area: str, seconds: float = 300) -> None: |
| 80 | cutoff = time.time() - seconds |
| 81 | with _lock: |
| 82 | for area_key in _get_matching_areas(area): |
| 83 | area_cache = _cache.get(area_key) |
| 84 | if not area_cache: |
| 85 | continue |
| 86 | |
| 87 | keys_to_remove = [key for key, entry in area_cache.items() if entry.timestamp < cutoff] |
| 88 | for key in keys_to_remove: |
| 89 | area_cache.pop(key, None) |
| 90 | |
| 91 | if not area_cache: |
| 92 | _cache.pop(area_key, None) |
| 93 | |
| 94 | |
| 95 | def clear_all() -> None: |
| 96 | with _lock: |
| 97 | _cache.clear() |
| 98 | |
| 99 | |
| 100 | def _is_enabled(area: str) -> bool: |
| 101 | if not _enabled_global: |
| 102 | return False |
| 103 | return _enabled_areas.get(area, True) |
| 104 | |
| 105 | |
| 106 | def _create_entry(value: Any) -> CacheEntry: |
| 107 | return CacheEntry(value=value, timestamp=time.time()) |
| 108 | |
| 109 | |
| 110 | def _touch_entry(entry: CacheEntry) -> None: |
| 111 | entry.timestamp = time.time() |
| 112 | |
| 113 | |
| 114 | def _get_matching_areas(area: str) -> list[str]: |
| 115 | if any(ch in area for ch in "*?["): |
| 116 | return [k for k in _cache.keys() if fnmatch.fnmatch(k, area)] |
| 117 | return [area] |
| 118 | |
| 119 | |
| 120 | def determine_cache_key(agent, *additional): |
| 121 | if agent: |
| 122 | profile = agent.config.profile or "none" |
| 123 | project = agent.context.get_data("project") or "none" |
| 124 | return (profile, project, *additional) |
| 125 | return ("none", "none", *additional) |