fix: stop get_cache handing out the stored KV cache by reference

src/memos/memories/activation/kv.py

@@ -80,6 +80,40 @@ def get_cache(self, cache_ids: list[str]) -> DynamicCache | None:
 
         return self._concat_caches(caches_to_merge)
 
+    @staticmethod
+    def _copy_cache(cache: DynamicCache) -> DynamicCache:
+        """Shallow-structure copy of a DynamicCache that shares no mutable state.
+
+        The tensors themselves are not cloned -- ``generate`` appends along the
+        sequence axis rather than writing into the existing rows, so a fresh
+        container with the same tensor objects is enough to protect the stored
+        item, without paying to duplicate the cache.
+        """
+        copy = DynamicCache()
+        if hasattr(cache, "layers"):
+            if not hasattr(copy, "layers"):
+                copy.layers = []
+            layer_cls = type(cache.layers[0]) if cache.layers else None
+            while layer_cls is not None and len(copy.layers) < len(cache.layers):
+                copy.layers.append(layer_cls())
+            for i, src in enumerate(cache.layers):
+                dst = copy.layers[i]
+                if not getattr(dst, "is_initialized", True) and hasattr(
+                    dst, "lazy_initialization"
+                ):
+                    dst.lazy_initialization(src.keys, src.values)
+                dst.keys = src.keys
+                dst.values = src.values
+        elif hasattr(cache, "key_cache"):
+            for i in range(len(cache.key_cache)):
+                copy.key_cache.append(cache.key_cache[i])
+                copy.value_cache.append(cache.value_cache[i])
+        else:
+            raise AttributeError(
+                "DynamicCache object has neither 'layers' nor 'key_cache' attributes"
+            )
+        return copy
+
     def get(self, memory_id: str) -> KVCacheItem | None:
         """Get a memory by its ID.
 
@@ -206,7 +240,13 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
 
         assert caches, "Need at least one cache"
         if len(caches) == 1:
-            return caches[0]
+            # Do NOT hand back the stored object. The caller passes this cache to
+            # ``generate``, which appends to it in place, so returning the stored
+            # item makes every chat turn grow the saved activation memory:
+            # observed stored length 6 -> 19 -> 32 -> 45 over three turns. The
+            # multi-cache path below already builds a new container, so only this
+            # early return leaked the reference.
+            return self._copy_cache(caches[0])
 
         merged = DynamicCache()
 

tests/memories/activation/test_kv.py

@@ -33,11 +33,11 @@ def kv_memory(dummy_config):
         yield KVCacheMemory(dummy_config)
 
 
-def make_filled_cache():
-    # Create a DynamicCache with at least one dummy tensor layer
+def make_filled_cache(seq_len: int = 3, n_layers: int = 1):
+    """Create a DynamicCache with dummy tensors, on any transformers version."""
     cache = DynamicCache()
-    cache.key_cache.append(torch.zeros(1, 2, 3))
-    cache.value_cache.append(torch.zeros(1, 2, 3))
+    for layer_idx in range(n_layers):
+        cache.update(torch.zeros(1, 2, seq_len, 4), torch.zeros(1, 2, seq_len, 4), layer_idx)
     return cache
 
 
@@ -59,8 +59,11 @@ def test_get_cache_merge(kv_memory):
     merged = kv_memory.get_cache([item1.id, item2.id])
     assert isinstance(merged, DynamicCache)
     # Check the number of layers in merged key/value cache
-    assert len(merged.key_cache) == 1
-    assert len(merged.value_cache) == 1
+    if hasattr(merged, "layers"):
+        assert len(merged.layers) == 1
+    else:
+        assert len(merged.key_cache) == 1
+        assert len(merged.value_cache) == 1
 
 
 def test_delete_and_get_all(kv_memory):
@@ -84,3 +87,33 @@ class DummyTextualMemory:
     item = kv_memory.from_textual_memory(DummyTextualMemory())
     assert isinstance(item, KVCacheItem)
     assert item.metadata["bar"] == 1
+
+
+def test_get_cache_does_not_alias_stored_memory(kv_memory):
+    """A cache handed to the caller must not be the stored object itself.
+
+    ``get_cache`` returned ``caches[0]`` directly whenever a single id was
+    requested. The caller passes that cache to ``generate``, which appends to it
+    in place, so the saved activation memory grew on every chat turn -- observed
+    as a stored length of 6 -> 19 -> 32 -> 45 across three turns. Nothing in the
+    API suggests retrieving a memory mutates it.
+
+    Simulating ``generate`` by appending to the returned cache must leave the
+    stored item unchanged. Fails before this change, passes after.
+    """
+    item = KVCacheItem(memory=make_filled_cache(seq_len=5, n_layers=2))
+    kv_memory.add([item])
+    stored_len_before = kv_memory.get(item.id).memory.get_seq_length()
+
+    handed_out = kv_memory.get_cache([item.id])
+    assert handed_out is not item.memory, "get_cache returned the stored object itself"
+
+    # what generate() does: append one step of new keys/values
+    for layer_idx in range(2):
+        handed_out.update(torch.zeros(1, 2, 1, 4), torch.zeros(1, 2, 1, 4), layer_idx)
+
+    stored_len_after = kv_memory.get(item.id).memory.get_seq_length()
+    assert stored_len_after == stored_len_before, (
+        f"stored memory grew {stored_len_before} -> {stored_len_after} because the "
+        f"caller mutated the cache it was handed"
+    )