Mirrowel commited on
Commit
9bc26b9
·
1 Parent(s): afe6e70

refactor(providers): 🔨 extract cache logic into shared ProviderCache module

Browse files

Extracted the AntigravityCache class into a new shared ProviderCache module to eliminate code duplication and improve maintainability across providers.

- Created src/rotator_library/providers/provider_cache.py with generic, reusable cache implementation
- Removed 266 lines of cache-specific code from antigravity_provider.py
- Updated AntigravityProvider to use ProviderCache for both signature and thinking caches
- Added configurable env_prefix parameter for flexible environment variable namespacing
- Improved cache naming with _cache_name for better logging context
- Added convenience factory function create_provider_cache() for streamlined cache creation
- Removed unused imports (shutil, tempfile) from antigravity_provider.py
- Updated .gitignore to include cache/ directory

The new ProviderCache maintains full backward compatibility with the previous AntigravityCache implementation while providing a more modular, reusable foundation for other providers.

.gitignore CHANGED
@@ -126,3 +126,4 @@ staged_changes.txt
126
  launcher_config.json
127
  cache/antigravity/thought_signatures.json
128
  logs/
 
 
126
  launcher_config.json
127
  cache/antigravity/thought_signatures.json
128
  logs/
129
+ cache/
src/rotator_library/providers/antigravity_provider.py CHANGED
@@ -23,8 +23,6 @@ import json
23
  import logging
24
  import os
25
  import random
26
- import shutil
27
- import tempfile
28
  import time
29
  import uuid
30
  from datetime import datetime
@@ -37,6 +35,7 @@ import litellm
37
 
38
  from .provider_interface import ProviderInterface
39
  from .antigravity_auth_base import AntigravityAuthBase
 
40
  from ..model_definitions import ModelDefinitions
41
 
42
 
@@ -269,272 +268,6 @@ class AntigravityFileLogger:
269
  lib_logger.error(f"Failed to append to {filename}: {e}")
270
 
271
 
272
- # =============================================================================
273
- # SIGNATURE CACHE
274
- # =============================================================================
275
-
276
- class AntigravityCache:
277
- """
278
- Server-side cache for Antigravity conversation state preservation.
279
-
280
- Supports two types of cached data:
281
- - Gemini 3: thoughtSignatures (tool_call_id → encrypted signature)
282
- - Claude: Thinking content (composite_key → thinking text + signature)
283
-
284
- Features:
285
- - Dual-TTL system: 1hr memory, 24hr disk
286
- - Async disk persistence with batched writes
287
- - Background cleanup task for expired entries
288
- """
289
-
290
- def __init__(
291
- self,
292
- cache_file: Path,
293
- memory_ttl_seconds: int = 3600,
294
- disk_ttl_seconds: int = 86400
295
- ):
296
- # In-memory cache: {cache_key: (data, timestamp)}
297
- self._cache: Dict[str, Tuple[str, float]] = {}
298
- self._memory_ttl = memory_ttl_seconds
299
- self._disk_ttl = disk_ttl_seconds
300
- self._lock = asyncio.Lock()
301
- self._disk_lock = asyncio.Lock()
302
-
303
- # Disk persistence
304
- self._cache_file = cache_file
305
- self._enable_disk = _env_bool("ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", True)
306
- self._dirty = False
307
- self._write_interval = _env_int("ANTIGRAVITY_CACHE_WRITE_INTERVAL", 60)
308
- self._cleanup_interval = _env_int("ANTIGRAVITY_CACHE_CLEANUP_INTERVAL", 1800)
309
-
310
- # Background tasks
311
- self._writer_task: Optional[asyncio.Task] = None
312
- self._cleanup_task: Optional[asyncio.Task] = None
313
- self._running = False
314
-
315
- # Statistics
316
- self._stats = {"memory_hits": 0, "disk_hits": 0, "misses": 0, "writes": 0}
317
-
318
- if self._enable_disk:
319
- lib_logger.debug(
320
- f"AntigravityCache: Disk persistence enabled "
321
- f"(memory_ttl={memory_ttl_seconds}s, disk_ttl={disk_ttl_seconds}s)"
322
- )
323
- asyncio.create_task(self._async_init())
324
- else:
325
- lib_logger.debug("AntigravityCache: Memory-only mode")
326
-
327
- async def _async_init(self) -> None:
328
- """Async initialization: load from disk and start background tasks."""
329
- try:
330
- await self._load_from_disk()
331
- await self._start_background_tasks()
332
- except Exception as e:
333
- lib_logger.error(f"Cache async init failed: {e}")
334
-
335
- async def _load_from_disk(self) -> None:
336
- """Load cache from disk file with TTL validation."""
337
- if not self._enable_disk or not self._cache_file.exists():
338
- return
339
-
340
- try:
341
- async with self._disk_lock:
342
- with open(self._cache_file, 'r', encoding='utf-8') as f:
343
- data = json.load(f)
344
-
345
- if data.get("version") != "1.0":
346
- lib_logger.warning("Cache version mismatch, starting fresh")
347
- return
348
-
349
- now = time.time()
350
- entries = data.get("entries", {})
351
- loaded = expired = 0
352
-
353
- for call_id, entry in entries.items():
354
- age = now - entry.get("timestamp", 0)
355
- if age <= self._disk_ttl:
356
- sig = entry.get("signature", "")
357
- if sig:
358
- self._cache[call_id] = (sig, entry["timestamp"])
359
- loaded += 1
360
- else:
361
- expired += 1
362
-
363
- lib_logger.debug(f"Loaded {loaded} entries from disk ({expired} expired)")
364
- except json.JSONDecodeError as e:
365
- lib_logger.warning(f"Cache file corrupted: {e}")
366
- except Exception as e:
367
- lib_logger.error(f"Failed to load cache: {e}")
368
-
369
- async def _save_to_disk(self) -> None:
370
- """Persist cache to disk using atomic write."""
371
- if not self._enable_disk:
372
- return
373
-
374
- try:
375
- async with self._disk_lock:
376
- self._cache_file.parent.mkdir(parents=True, exist_ok=True)
377
-
378
- cache_data = {
379
- "version": "1.0",
380
- "memory_ttl_seconds": self._memory_ttl,
381
- "disk_ttl_seconds": self._disk_ttl,
382
- "entries": {
383
- cid: {"signature": sig, "timestamp": ts}
384
- for cid, (sig, ts) in self._cache.items()
385
- },
386
- "statistics": {
387
- "total_entries": len(self._cache),
388
- "last_write": time.time(),
389
- **self._stats
390
- }
391
- }
392
-
393
- # Atomic write
394
- parent_dir = self._cache_file.parent
395
- tmp_fd, tmp_path = tempfile.mkstemp(dir=parent_dir, prefix='.tmp_', suffix='.json')
396
-
397
- try:
398
- with os.fdopen(tmp_fd, 'w', encoding='utf-8') as f:
399
- json.dump(cache_data, f, indent=2)
400
-
401
- try:
402
- os.chmod(tmp_path, 0o600)
403
- except (OSError, AttributeError):
404
- pass
405
-
406
- shutil.move(tmp_path, self._cache_file)
407
- self._stats["writes"] += 1
408
- lib_logger.debug(f"Saved {len(self._cache)} entries to disk")
409
- except Exception:
410
- if tmp_path and os.path.exists(tmp_path):
411
- os.unlink(tmp_path)
412
- raise
413
- except Exception as e:
414
- lib_logger.error(f"Disk save failed: {e}")
415
-
416
- async def _start_background_tasks(self) -> None:
417
- """Start background writer and cleanup tasks."""
418
- if not self._enable_disk or self._running:
419
- return
420
-
421
- self._running = True
422
- self._writer_task = asyncio.create_task(self._writer_loop())
423
- self._cleanup_task = asyncio.create_task(self._cleanup_loop())
424
- lib_logger.debug("Started background cache tasks")
425
-
426
- async def _writer_loop(self) -> None:
427
- """Background task: periodically flush dirty cache to disk."""
428
- try:
429
- while self._running:
430
- await asyncio.sleep(self._write_interval)
431
- if self._dirty:
432
- try:
433
- await self._save_to_disk()
434
- self._dirty = False
435
- except Exception as e:
436
- lib_logger.error(f"Background writer error: {e}")
437
- except asyncio.CancelledError:
438
- pass
439
-
440
- async def _cleanup_loop(self) -> None:
441
- """Background task: periodically clean up expired entries."""
442
- try:
443
- while self._running:
444
- await asyncio.sleep(self._cleanup_interval)
445
- await self._cleanup_expired()
446
- except asyncio.CancelledError:
447
- pass
448
-
449
- async def _cleanup_expired(self) -> None:
450
- """Remove expired entries from memory cache."""
451
- async with self._lock:
452
- now = time.time()
453
- expired = [k for k, (_, ts) in self._cache.items() if now - ts > self._memory_ttl]
454
- for k in expired:
455
- del self._cache[k]
456
- if expired:
457
- self._dirty = True
458
- lib_logger.debug(f"Cleaned up {len(expired)} expired entries")
459
-
460
- def store(self, key: str, value: str) -> None:
461
- """Store a value (sync wrapper for async storage)."""
462
- asyncio.create_task(self._async_store(key, value))
463
-
464
- async def _async_store(self, key: str, value: str) -> None:
465
- """Async implementation of store."""
466
- async with self._lock:
467
- self._cache[key] = (value, time.time())
468
- self._dirty = True
469
-
470
- def retrieve(self, key: str) -> Optional[str]:
471
- """Retrieve a value by key (sync method)."""
472
- if key in self._cache:
473
- value, timestamp = self._cache[key]
474
- if time.time() - timestamp <= self._memory_ttl:
475
- self._stats["memory_hits"] += 1
476
- return value
477
- else:
478
- del self._cache[key]
479
- self._dirty = True
480
-
481
- self._stats["misses"] += 1
482
- if self._enable_disk:
483
- asyncio.create_task(self._check_disk_fallback(key))
484
- return None
485
-
486
- async def _check_disk_fallback(self, key: str) -> None:
487
- """Check disk for key and load into memory if found."""
488
- try:
489
- if not self._cache_file.exists():
490
- return
491
-
492
- async with self._disk_lock:
493
- with open(self._cache_file, 'r', encoding='utf-8') as f:
494
- data = json.load(f)
495
-
496
- entries = data.get("entries", {})
497
- if key in entries:
498
- entry = entries[key]
499
- ts = entry.get("timestamp", 0)
500
- if time.time() - ts <= self._disk_ttl:
501
- sig = entry.get("signature", "")
502
- if sig:
503
- async with self._lock:
504
- self._cache[key] = (sig, ts)
505
- self._stats["disk_hits"] += 1
506
- lib_logger.debug(f"Loaded {key} from disk")
507
- except Exception as e:
508
- lib_logger.debug(f"Disk fallback failed: {e}")
509
-
510
- async def clear(self) -> None:
511
- """Clear all cached data."""
512
- async with self._lock:
513
- self._cache.clear()
514
- self._dirty = True
515
- if self._enable_disk:
516
- await self._save_to_disk()
517
-
518
- async def shutdown(self) -> None:
519
- """Graceful shutdown: flush pending writes and stop background tasks."""
520
- lib_logger.info("AntigravityCache shutting down...")
521
- self._running = False
522
-
523
- for task in (self._writer_task, self._cleanup_task):
524
- if task:
525
- task.cancel()
526
- try:
527
- await task
528
- except asyncio.CancelledError:
529
- pass
530
-
531
- if self._dirty and self._enable_disk:
532
- await self._save_to_disk()
533
-
534
- lib_logger.info(
535
- f"Cache shutdown complete (stats: mem_hits={self._stats['memory_hits']}, "
536
- f"disk_hits={self._stats['disk_hits']}, misses={self._stats['misses']})"
537
- )
538
 
539
 
540
  # =============================================================================
@@ -571,12 +304,14 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface):
571
  memory_ttl = _env_int("ANTIGRAVITY_SIGNATURE_CACHE_TTL", 3600)
572
  disk_ttl = _env_int("ANTIGRAVITY_SIGNATURE_DISK_TTL", 86400)
573
 
574
- # Initialize caches
575
- self._signature_cache = AntigravityCache(
576
- GEMINI3_SIGNATURE_CACHE_FILE, memory_ttl, disk_ttl
 
577
  )
578
- self._thinking_cache = AntigravityCache(
579
- CLAUDE_THINKING_CACHE_FILE, memory_ttl, disk_ttl
 
580
  )
581
 
582
  # Feature flags
 
23
  import logging
24
  import os
25
  import random
 
 
26
  import time
27
  import uuid
28
  from datetime import datetime
 
35
 
36
  from .provider_interface import ProviderInterface
37
  from .antigravity_auth_base import AntigravityAuthBase
38
+ from .provider_cache import ProviderCache
39
  from ..model_definitions import ModelDefinitions
40
 
41
 
 
268
  lib_logger.error(f"Failed to append to {filename}: {e}")
269
 
270
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
271
 
272
 
273
  # =============================================================================
 
304
  memory_ttl = _env_int("ANTIGRAVITY_SIGNATURE_CACHE_TTL", 3600)
305
  disk_ttl = _env_int("ANTIGRAVITY_SIGNATURE_DISK_TTL", 86400)
306
 
307
+ # Initialize caches using shared ProviderCache
308
+ self._signature_cache = ProviderCache(
309
+ GEMINI3_SIGNATURE_CACHE_FILE, memory_ttl, disk_ttl,
310
+ env_prefix="ANTIGRAVITY_SIGNATURE"
311
  )
312
+ self._thinking_cache = ProviderCache(
313
+ CLAUDE_THINKING_CACHE_FILE, memory_ttl, disk_ttl,
314
+ env_prefix="ANTIGRAVITY_THINKING"
315
  )
316
 
317
  # Feature flags
src/rotator_library/providers/provider_cache.py ADDED
@@ -0,0 +1,498 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # src/rotator_library/providers/provider_cache.py
2
+ """
3
+ Shared cache utility for providers.
4
+
5
+ A modular, async-capable cache system supporting:
6
+ - Dual-TTL: short-lived memory cache, longer-lived disk persistence
7
+ - Background persistence with batched writes
8
+ - Automatic cleanup of expired entries
9
+ - Generic key-value storage for any provider-specific needs
10
+
11
+ Usage examples:
12
+ - Gemini 3: thoughtSignatures (tool_call_id → encrypted signature)
13
+ - Claude: Thinking content (composite_key → thinking text + signature)
14
+ - General: Any transient data that benefits from persistence across requests
15
+ """
16
+
17
+ from __future__ import annotations
18
+
19
+ import asyncio
20
+ import json
21
+ import logging
22
+ import os
23
+ import shutil
24
+ import tempfile
25
+ import time
26
+ from pathlib import Path
27
+ from typing import Any, Dict, Optional, Tuple
28
+
29
+ lib_logger = logging.getLogger('rotator_library')
30
+
31
+
32
+ # =============================================================================
33
+ # UTILITY FUNCTIONS
34
+ # =============================================================================
35
+
36
+ def _env_bool(key: str, default: bool = False) -> bool:
37
+ """Get boolean from environment variable."""
38
+ return os.getenv(key, str(default).lower()).lower() in ("true", "1", "yes")
39
+
40
+
41
+ def _env_int(key: str, default: int) -> int:
42
+ """Get integer from environment variable."""
43
+ return int(os.getenv(key, str(default)))
44
+
45
+
46
+ # =============================================================================
47
+ # PROVIDER CACHE CLASS
48
+ # =============================================================================
49
+
50
+ class ProviderCache:
51
+ """
52
+ Server-side cache for provider conversation state preservation.
53
+
54
+ A generic, modular cache supporting any key-value data that providers need
55
+ to persist across requests. Features:
56
+
57
+ - Dual-TTL system: configurable memory TTL, longer disk TTL
58
+ - Async disk persistence with batched writes
59
+ - Background cleanup task for expired entries
60
+ - Statistics tracking (hits, misses, writes)
61
+
62
+ Args:
63
+ cache_file: Path to disk cache file
64
+ memory_ttl_seconds: In-memory entry lifetime (default: 1 hour)
65
+ disk_ttl_seconds: Disk entry lifetime (default: 24 hours)
66
+ enable_disk: Whether to enable disk persistence (default: from env or True)
67
+ write_interval: Seconds between background disk writes (default: 60)
68
+ cleanup_interval: Seconds between expired entry cleanup (default: 30 min)
69
+ env_prefix: Environment variable prefix for configuration overrides
70
+
71
+ Environment Variables (with default prefix "PROVIDER_CACHE"):
72
+ {PREFIX}_ENABLE: Enable/disable disk persistence
73
+ {PREFIX}_WRITE_INTERVAL: Background write interval in seconds
74
+ {PREFIX}_CLEANUP_INTERVAL: Cleanup interval in seconds
75
+ """
76
+
77
+ def __init__(
78
+ self,
79
+ cache_file: Path,
80
+ memory_ttl_seconds: int = 3600,
81
+ disk_ttl_seconds: int = 86400,
82
+ enable_disk: Optional[bool] = None,
83
+ write_interval: Optional[int] = None,
84
+ cleanup_interval: Optional[int] = None,
85
+ env_prefix: str = "PROVIDER_CACHE"
86
+ ):
87
+ # In-memory cache: {cache_key: (data, timestamp)}
88
+ self._cache: Dict[str, Tuple[str, float]] = {}
89
+ self._memory_ttl = memory_ttl_seconds
90
+ self._disk_ttl = disk_ttl_seconds
91
+ self._lock = asyncio.Lock()
92
+ self._disk_lock = asyncio.Lock()
93
+
94
+ # Disk persistence configuration
95
+ self._cache_file = cache_file
96
+ self._enable_disk = enable_disk if enable_disk is not None else _env_bool(f"{env_prefix}_ENABLE", True)
97
+ self._dirty = False
98
+ self._write_interval = write_interval or _env_int(f"{env_prefix}_WRITE_INTERVAL", 60)
99
+ self._cleanup_interval = cleanup_interval or _env_int(f"{env_prefix}_CLEANUP_INTERVAL", 1800)
100
+
101
+ # Background tasks
102
+ self._writer_task: Optional[asyncio.Task] = None
103
+ self._cleanup_task: Optional[asyncio.Task] = None
104
+ self._running = False
105
+
106
+ # Statistics
107
+ self._stats = {"memory_hits": 0, "disk_hits": 0, "misses": 0, "writes": 0}
108
+
109
+ # Metadata about this cache instance
110
+ self._cache_name = cache_file.stem if cache_file else "unnamed"
111
+
112
+ if self._enable_disk:
113
+ lib_logger.debug(
114
+ f"ProviderCache[{self._cache_name}]: Disk enabled "
115
+ f"(memory_ttl={memory_ttl_seconds}s, disk_ttl={disk_ttl_seconds}s)"
116
+ )
117
+ asyncio.create_task(self._async_init())
118
+ else:
119
+ lib_logger.debug(f"ProviderCache[{self._cache_name}]: Memory-only mode")
120
+
121
+ # =========================================================================
122
+ # INITIALIZATION
123
+ # =========================================================================
124
+
125
+ async def _async_init(self) -> None:
126
+ """Async initialization: load from disk and start background tasks."""
127
+ try:
128
+ await self._load_from_disk()
129
+ await self._start_background_tasks()
130
+ except Exception as e:
131
+ lib_logger.error(f"ProviderCache[{self._cache_name}] async init failed: {e}")
132
+
133
+ async def _load_from_disk(self) -> None:
134
+ """Load cache from disk file with TTL validation."""
135
+ if not self._enable_disk or not self._cache_file.exists():
136
+ return
137
+
138
+ try:
139
+ async with self._disk_lock:
140
+ with open(self._cache_file, 'r', encoding='utf-8') as f:
141
+ data = json.load(f)
142
+
143
+ if data.get("version") != "1.0":
144
+ lib_logger.warning(f"ProviderCache[{self._cache_name}]: Version mismatch, starting fresh")
145
+ return
146
+
147
+ now = time.time()
148
+ entries = data.get("entries", {})
149
+ loaded = expired = 0
150
+
151
+ for cache_key, entry in entries.items():
152
+ age = now - entry.get("timestamp", 0)
153
+ if age <= self._disk_ttl:
154
+ value = entry.get("value", entry.get("signature", "")) # Support both formats
155
+ if value:
156
+ self._cache[cache_key] = (value, entry["timestamp"])
157
+ loaded += 1
158
+ else:
159
+ expired += 1
160
+
161
+ lib_logger.debug(
162
+ f"ProviderCache[{self._cache_name}]: Loaded {loaded} entries ({expired} expired)"
163
+ )
164
+ except json.JSONDecodeError as e:
165
+ lib_logger.warning(f"ProviderCache[{self._cache_name}]: File corrupted: {e}")
166
+ except Exception as e:
167
+ lib_logger.error(f"ProviderCache[{self._cache_name}]: Load failed: {e}")
168
+
169
+ # =========================================================================
170
+ # DISK PERSISTENCE
171
+ # =========================================================================
172
+
173
+ async def _save_to_disk(self) -> None:
174
+ """Persist cache to disk using atomic write."""
175
+ if not self._enable_disk:
176
+ return
177
+
178
+ try:
179
+ async with self._disk_lock:
180
+ self._cache_file.parent.mkdir(parents=True, exist_ok=True)
181
+
182
+ cache_data = {
183
+ "version": "1.0",
184
+ "memory_ttl_seconds": self._memory_ttl,
185
+ "disk_ttl_seconds": self._disk_ttl,
186
+ "entries": {
187
+ key: {"value": val, "timestamp": ts}
188
+ for key, (val, ts) in self._cache.items()
189
+ },
190
+ "statistics": {
191
+ "total_entries": len(self._cache),
192
+ "last_write": time.time(),
193
+ **self._stats
194
+ }
195
+ }
196
+
197
+ # Atomic write using temp file
198
+ parent_dir = self._cache_file.parent
199
+ tmp_fd, tmp_path = tempfile.mkstemp(dir=parent_dir, prefix='.tmp_', suffix='.json')
200
+
201
+ try:
202
+ with os.fdopen(tmp_fd, 'w', encoding='utf-8') as f:
203
+ json.dump(cache_data, f, indent=2)
204
+
205
+ # Set restrictive permissions (if supported)
206
+ try:
207
+ os.chmod(tmp_path, 0o600)
208
+ except (OSError, AttributeError):
209
+ pass
210
+
211
+ shutil.move(tmp_path, self._cache_file)
212
+ self._stats["writes"] += 1
213
+ lib_logger.debug(
214
+ f"ProviderCache[{self._cache_name}]: Saved {len(self._cache)} entries"
215
+ )
216
+ except Exception:
217
+ if tmp_path and os.path.exists(tmp_path):
218
+ os.unlink(tmp_path)
219
+ raise
220
+ except Exception as e:
221
+ lib_logger.error(f"ProviderCache[{self._cache_name}]: Disk save failed: {e}")
222
+
223
+ # =========================================================================
224
+ # BACKGROUND TASKS
225
+ # =========================================================================
226
+
227
+ async def _start_background_tasks(self) -> None:
228
+ """Start background writer and cleanup tasks."""
229
+ if not self._enable_disk or self._running:
230
+ return
231
+
232
+ self._running = True
233
+ self._writer_task = asyncio.create_task(self._writer_loop())
234
+ self._cleanup_task = asyncio.create_task(self._cleanup_loop())
235
+ lib_logger.debug(f"ProviderCache[{self._cache_name}]: Started background tasks")
236
+
237
+ async def _writer_loop(self) -> None:
238
+ """Background task: periodically flush dirty cache to disk."""
239
+ try:
240
+ while self._running:
241
+ await asyncio.sleep(self._write_interval)
242
+ if self._dirty:
243
+ try:
244
+ await self._save_to_disk()
245
+ self._dirty = False
246
+ except Exception as e:
247
+ lib_logger.error(f"ProviderCache[{self._cache_name}]: Writer error: {e}")
248
+ except asyncio.CancelledError:
249
+ pass
250
+
251
+ async def _cleanup_loop(self) -> None:
252
+ """Background task: periodically clean up expired entries."""
253
+ try:
254
+ while self._running:
255
+ await asyncio.sleep(self._cleanup_interval)
256
+ await self._cleanup_expired()
257
+ except asyncio.CancelledError:
258
+ pass
259
+
260
+ async def _cleanup_expired(self) -> None:
261
+ """Remove expired entries from memory cache."""
262
+ async with self._lock:
263
+ now = time.time()
264
+ expired = [k for k, (_, ts) in self._cache.items() if now - ts > self._memory_ttl]
265
+ for k in expired:
266
+ del self._cache[k]
267
+ if expired:
268
+ self._dirty = True
269
+ lib_logger.debug(
270
+ f"ProviderCache[{self._cache_name}]: Cleaned {len(expired)} expired entries"
271
+ )
272
+
273
+ # =========================================================================
274
+ # CORE OPERATIONS
275
+ # =========================================================================
276
+
277
+ def store(self, key: str, value: str) -> None:
278
+ """
279
+ Store a value synchronously (schedules async storage).
280
+
281
+ Args:
282
+ key: Cache key
283
+ value: Value to store (typically JSON-serialized data)
284
+ """
285
+ asyncio.create_task(self._async_store(key, value))
286
+
287
+ async def _async_store(self, key: str, value: str) -> None:
288
+ """Async implementation of store."""
289
+ async with self._lock:
290
+ self._cache[key] = (value, time.time())
291
+ self._dirty = True
292
+
293
+ async def store_async(self, key: str, value: str) -> None:
294
+ """
295
+ Store a value asynchronously (awaitable).
296
+
297
+ Use this when you need to ensure the value is stored before continuing.
298
+ """
299
+ await self._async_store(key, value)
300
+
301
+ def retrieve(self, key: str) -> Optional[str]:
302
+ """
303
+ Retrieve a value by key (synchronous, with optional async disk fallback).
304
+
305
+ Args:
306
+ key: Cache key
307
+
308
+ Returns:
309
+ Cached value if found and not expired, None otherwise
310
+ """
311
+ if key in self._cache:
312
+ value, timestamp = self._cache[key]
313
+ if time.time() - timestamp <= self._memory_ttl:
314
+ self._stats["memory_hits"] += 1
315
+ return value
316
+ else:
317
+ del self._cache[key]
318
+ self._dirty = True
319
+
320
+ self._stats["misses"] += 1
321
+ if self._enable_disk:
322
+ # Schedule async disk lookup for next time
323
+ asyncio.create_task(self._check_disk_fallback(key))
324
+ return None
325
+
326
+ async def retrieve_async(self, key: str) -> Optional[str]:
327
+ """
328
+ Retrieve a value asynchronously (checks disk if not in memory).
329
+
330
+ Use this when you can await and need guaranteed disk fallback.
331
+ """
332
+ # Check memory first
333
+ if key in self._cache:
334
+ value, timestamp = self._cache[key]
335
+ if time.time() - timestamp <= self._memory_ttl:
336
+ self._stats["memory_hits"] += 1
337
+ return value
338
+ else:
339
+ async with self._lock:
340
+ if key in self._cache:
341
+ del self._cache[key]
342
+ self._dirty = True
343
+
344
+ # Check disk
345
+ if self._enable_disk:
346
+ return await self._disk_retrieve(key)
347
+
348
+ self._stats["misses"] += 1
349
+ return None
350
+
351
+ async def _check_disk_fallback(self, key: str) -> None:
352
+ """Check disk for key and load into memory if found (background)."""
353
+ try:
354
+ if not self._cache_file.exists():
355
+ return
356
+
357
+ async with self._disk_lock:
358
+ with open(self._cache_file, 'r', encoding='utf-8') as f:
359
+ data = json.load(f)
360
+
361
+ entries = data.get("entries", {})
362
+ if key in entries:
363
+ entry = entries[key]
364
+ ts = entry.get("timestamp", 0)
365
+ if time.time() - ts <= self._disk_ttl:
366
+ value = entry.get("value", entry.get("signature", ""))
367
+ if value:
368
+ async with self._lock:
369
+ self._cache[key] = (value, ts)
370
+ self._stats["disk_hits"] += 1
371
+ lib_logger.debug(
372
+ f"ProviderCache[{self._cache_name}]: Loaded {key} from disk"
373
+ )
374
+ except Exception as e:
375
+ lib_logger.debug(f"ProviderCache[{self._cache_name}]: Disk fallback failed: {e}")
376
+
377
+ async def _disk_retrieve(self, key: str) -> Optional[str]:
378
+ """Direct disk retrieval with loading into memory."""
379
+ try:
380
+ if not self._cache_file.exists():
381
+ self._stats["misses"] += 1
382
+ return None
383
+
384
+ async with self._disk_lock:
385
+ with open(self._cache_file, 'r', encoding='utf-8') as f:
386
+ data = json.load(f)
387
+
388
+ entries = data.get("entries", {})
389
+ if key in entries:
390
+ entry = entries[key]
391
+ ts = entry.get("timestamp", 0)
392
+ if time.time() - ts <= self._disk_ttl:
393
+ value = entry.get("value", entry.get("signature", ""))
394
+ if value:
395
+ async with self._lock:
396
+ self._cache[key] = (value, ts)
397
+ self._stats["disk_hits"] += 1
398
+ return value
399
+
400
+ self._stats["misses"] += 1
401
+ return None
402
+ except Exception as e:
403
+ lib_logger.debug(f"ProviderCache[{self._cache_name}]: Disk retrieve failed: {e}")
404
+ self._stats["misses"] += 1
405
+ return None
406
+
407
+ # =========================================================================
408
+ # UTILITY METHODS
409
+ # =========================================================================
410
+
411
+ def contains(self, key: str) -> bool:
412
+ """Check if key exists in memory cache (without updating stats)."""
413
+ if key in self._cache:
414
+ _, timestamp = self._cache[key]
415
+ return time.time() - timestamp <= self._memory_ttl
416
+ return False
417
+
418
+ def get_stats(self) -> Dict[str, Any]:
419
+ """Get cache statistics."""
420
+ return {
421
+ **self._stats,
422
+ "memory_entries": len(self._cache),
423
+ "dirty": self._dirty,
424
+ "disk_enabled": self._enable_disk
425
+ }
426
+
427
+ async def clear(self) -> None:
428
+ """Clear all cached data."""
429
+ async with self._lock:
430
+ self._cache.clear()
431
+ self._dirty = True
432
+ if self._enable_disk:
433
+ await self._save_to_disk()
434
+
435
+ async def shutdown(self) -> None:
436
+ """Graceful shutdown: flush pending writes and stop background tasks."""
437
+ lib_logger.info(f"ProviderCache[{self._cache_name}]: Shutting down...")
438
+ self._running = False
439
+
440
+ # Cancel background tasks
441
+ for task in (self._writer_task, self._cleanup_task):
442
+ if task:
443
+ task.cancel()
444
+ try:
445
+ await task
446
+ except asyncio.CancelledError:
447
+ pass
448
+
449
+ # Final save
450
+ if self._dirty and self._enable_disk:
451
+ await self._save_to_disk()
452
+
453
+ lib_logger.info(
454
+ f"ProviderCache[{self._cache_name}]: Shutdown complete "
455
+ f"(stats: mem_hits={self._stats['memory_hits']}, "
456
+ f"disk_hits={self._stats['disk_hits']}, misses={self._stats['misses']})"
457
+ )
458
+
459
+
460
+ # =============================================================================
461
+ # CONVENIENCE FACTORY
462
+ # =============================================================================
463
+
464
+ def create_provider_cache(
465
+ name: str,
466
+ cache_dir: Optional[Path] = None,
467
+ memory_ttl_seconds: int = 3600,
468
+ disk_ttl_seconds: int = 86400,
469
+ env_prefix: Optional[str] = None
470
+ ) -> ProviderCache:
471
+ """
472
+ Factory function to create a provider cache with sensible defaults.
473
+
474
+ Args:
475
+ name: Cache name (used as filename and for logging)
476
+ cache_dir: Directory for cache file (default: project_root/cache/provider_name)
477
+ memory_ttl_seconds: In-memory TTL
478
+ disk_ttl_seconds: Disk TTL
479
+ env_prefix: Environment variable prefix (default: derived from name)
480
+
481
+ Returns:
482
+ Configured ProviderCache instance
483
+ """
484
+ if cache_dir is None:
485
+ cache_dir = Path(__file__).resolve().parent.parent.parent.parent / "cache"
486
+
487
+ cache_file = cache_dir / f"{name}.json"
488
+
489
+ if env_prefix is None:
490
+ # Convert name to env prefix: "gemini3_signatures" -> "GEMINI3_SIGNATURES_CACHE"
491
+ env_prefix = f"{name.upper().replace('-', '_')}_CACHE"
492
+
493
+ return ProviderCache(
494
+ cache_file=cache_file,
495
+ memory_ttl_seconds=memory_ttl_seconds,
496
+ disk_ttl_seconds=disk_ttl_seconds,
497
+ env_prefix=env_prefix
498
+ )