| """A minimal in-memory TTL cache, used to avoid re-fetching the
|
| (effectively static) molecule catalog on every tool call that needs to
|
| resolve a molecule name, while still eventually picking up a real catalog
|
| change (e.g. after a kernel restart with a different build) instead of
|
| caching forever like the pre-Phase-2 code did (a plain dict, reset only by
|
| the test suite, never by production code). Also caches FAILURES for a
|
| short TTL, so a tight loop of tool calls made while the kernel is down
|
| doesn't each wait out their own full connection/timeout error -- see
|
| `set_failure`/`get`'s re-raise behavior below.
|
| """
|
| import time
|
| from typing import Any, Optional
|
|
|
|
|
| class TTLCache:
|
| def __init__(self, ttl_seconds: float = 300.0, failure_ttl_seconds: float = 10.0):
|
| self._ttl = ttl_seconds
|
| self._failure_ttl = failure_ttl_seconds
|
|
|
| self._store: dict[str, tuple[float, Any, bool]] = {}
|
|
|
| def get(self, key: str):
|
| """Returns the cached value, or None if missing/expired. Re-raises
|
| a cached failure (so callers see the same error shape they would
|
| from a fresh call) instead of returning it as if it were a value."""
|
| entry = self._store.get(key)
|
| if entry is None:
|
| return None
|
| expires_at, value, is_failure = entry
|
| if time.monotonic() >= expires_at:
|
| del self._store[key]
|
| return None
|
| if is_failure:
|
| raise value
|
| return value
|
|
|
| def set(self, key: str, value: Any) -> None:
|
| self._store[key] = (time.monotonic() + self._ttl, value, False)
|
|
|
| def set_failure(self, key: str, exc: Exception) -> None:
|
| self._store[key] = (time.monotonic() + self._failure_ttl, exc, True)
|
|
|
| def invalidate(self, key: Optional[str] = None) -> None:
|
| """Clears one key, or the whole cache if key is None."""
|
| if key is None:
|
| self._store.clear()
|
| else:
|
| self._store.pop(key, None)
|
|
|
