From ac3cf164f13df403d6dca0a3051365cdd4cea9cc Mon Sep 17 00:00:00 2001
From: Mathew Han <mathewjhan@gmail.com>
Date: Thu, 6 Aug 2026 15:14:43 -0700
Subject: [PATCH] [rollout] feat: add TtlCache utility

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 osmosis_ai/rollout/utils/ttl_cache.py      | 39 +++++++++++++++++++
 tests/unit/rollout/test_utils_ttl_cache.py | 44 ++++++++++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100644 osmosis_ai/rollout/utils/ttl_cache.py
 create mode 100644 tests/unit/rollout/test_utils_ttl_cache.py

diff --git a/osmosis_ai/rollout/utils/ttl_cache.py b/osmosis_ai/rollout/utils/ttl_cache.py
new file mode 100644
index 00000000..dccb84fe
--- /dev/null
+++ b/osmosis_ai/rollout/utils/ttl_cache.py
@@ -0,0 +1,39 @@
+"""Dict with a uniform per-entry TTL and amortized O(1) expiry.
+
+Uniform TTLs expire in insertion order, so a deque of (deadline, key)
+tombstones is pruned from the head on every write. The deadline doubles as
+the tombstone marker: overwriting a key refreshes its deadline, which
+orphans the old tombstone harmlessly.
+"""
+
+from __future__ import annotations
+
+import time
+from collections import deque
+
+
+class TtlCache[K, V]:
+    def __init__(self, ttl_sec: float) -> None:
+        self.ttl_sec = ttl_sec
+        self.entries: dict[K, tuple[float, V]] = {}
+        self.tombstones: deque[tuple[float, K]] = deque()
+
+    def set(self, key: K, value: V) -> None:
+        now = time.monotonic()
+        while self.tombstones and self.tombstones[0][0] <= now:
+            deadline, stale = self.tombstones.popleft()
+            entry = self.entries.get(stale)
+            if entry is not None and entry[0] == deadline:
+                del self.entries[stale]
+        deadline = now + self.ttl_sec
+        self.entries[key] = (deadline, value)
+        self.tombstones.append((deadline, key))
+
+    def get(self, key: K) -> V | None:
+        entry = self.entries.get(key)
+        if entry is None or entry[0] <= time.monotonic():
+            return None
+        return entry[1]
+
+    def __len__(self) -> int:
+        return len(self.entries)
diff --git a/tests/unit/rollout/test_utils_ttl_cache.py b/tests/unit/rollout/test_utils_ttl_cache.py
new file mode 100644
index 00000000..d1436aa9
--- /dev/null
+++ b/tests/unit/rollout/test_utils_ttl_cache.py
@@ -0,0 +1,44 @@
+"""TtlCache: uniform-TTL dict with tombstone-based amortized expiry."""
+
+from unittest.mock import patch
+
+from osmosis_ai.rollout.utils.ttl_cache import TtlCache
+
+
+def test_set_get_and_expiry():
+    clock = [0.0]
+    with patch("osmosis_ai.rollout.utils.ttl_cache.time.monotonic", lambda: clock[0]):
+        cache = TtlCache(ttl_sec=10.0)
+        cache.set("a", 1)
+        assert cache.get("a") == 1
+        clock[0] = 9.9
+        assert cache.get("a") == 1
+        clock[0] = 10.0
+        assert cache.get("a") is None
+
+
+def test_writes_prune_expired_entries():
+    clock = [0.0]
+    with patch("osmosis_ai.rollout.utils.ttl_cache.time.monotonic", lambda: clock[0]):
+        cache = TtlCache(ttl_sec=10.0)
+        cache.set("a", 1)
+        cache.set("b", 2)
+        clock[0] = 11.0
+        cache.set("c", 3)
+        assert len(cache) == 1
+        assert not cache.tombstones or cache.tombstones[0][1] == "c"
+
+
+def test_overwrite_refreshes_and_orphans_old_tombstone():
+    clock = [0.0]
+    with patch("osmosis_ai.rollout.utils.ttl_cache.time.monotonic", lambda: clock[0]):
+        cache = TtlCache(ttl_sec=10.0)
+        cache.set("a", 1)
+        clock[0] = 5.0
+        cache.set("a", 2)
+        # Past the first deadline: the stale tombstone must not evict the refresh.
+        clock[0] = 12.0
+        cache.set("other", 3)
+        assert cache.get("a") == 2
+        clock[0] = 15.0
+        assert cache.get("a") is None