Spaces:
Paused
refactor(providers): 🔨 extract cache logic into shared ProviderCache module
Browse filesExtracted 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.
|
@@ -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/
|
|
@@ -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 =
|
| 576 |
-
GEMINI3_SIGNATURE_CACHE_FILE, memory_ttl, disk_ttl
|
|
|
|
| 577 |
)
|
| 578 |
-
self._thinking_cache =
|
| 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
|
|
@@ -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 |
+
)
|