Merge pull request #37 from OpenWAM-Official/fix/prompt-embed-cache-eviction

fix: prevent prompt embed cache eviction KeyError

Author: wayrise

openwam/deploy/engine.py

@@ -78,7 +78,12 @@ def __setitem__(self, key: Any, value: Any) -> None:
             self.move_to_end(key)
         super().__setitem__(key, value)
         while len(self) > self._maxsize:
-            evicted_key, _ = self.popitem(last=False)
+            # ``OrderedDict.popitem`` retrieves the value through ``self[key]``.
+            # Since ``__getitem__`` updates recency, that lookup runs after the
+            # key has been removed from the ordering map and raises ``KeyError``.
+            # Delete through the base implementation to bypass the override.
+            evicted_key = next(iter(self))
+            super().__delitem__(evicted_key)
             if not self._evict_warned:
                 self._evict_warned = True
                 logger.warning("prompt_embed_cache exceeded maxsize=%d; evicted %r", self._maxsize, evicted_key)

tests/test_deployment_changes.py

@@ -580,6 +580,26 @@ def test_prompt_embed_cache_default_is_bounded_default(self):
         assert isinstance(engine._prompt_embed_cache, _BoundedPromptEmbedCache)
         assert engine._prompt_embed_cache._maxsize == DEFAULT_PROMPT_EMBED_CACHE_MAXSIZE
 
+    def test_prompt_embed_cache_eviction_preserves_lru_order(self, caplog):
+        from openwam.deploy.engine import _BoundedPromptEmbedCache
+
+        caplog.set_level("WARNING", logger="openwam.deploy.engine")
+        cache = _BoundedPromptEmbedCache(maxsize=3)
+        for key in ("probe", "a", "b", "c"):
+            cache[key] = key
+
+        assert list(cache) == ["a", "b", "c"]
+
+        # Reads refresh recency, so "a" survives the next eviction.
+        assert cache["a"] == "a"
+        cache["d"] = "d"
+        assert list(cache) == ["c", "a", "d"]
+
+        cache["e"] = "e"
+        assert list(cache) == ["a", "d", "e"]
+        warnings = [record for record in caplog.records if "prompt_embed_cache exceeded maxsize" in record.getMessage()]
+        assert len(warnings) == 1
+
     def test_prompt_embed_cache_maxsize_from_config(self):
         engine, _ = self._make_engine(prompt_cache_cfg={"enabled": True, "maxsize": 8})
         assert engine._prompt_embed_cache._maxsize == 8