refactor(memory): refactor TestMemoryStore to use the unified storage interface

strands-py/src/strands/vended_memory_stores/test_memory_store/store.py

@@ -6,9 +6,9 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 import re
-import threading
 import uuid
 from datetime import datetime, timezone
 from pathlib import Path
@@ -17,6 +17,7 @@
 from typing_extensions import Unpack
 
 from ...memory.types import MemoryEntry, MemoryStore, Metadata, SearchOptions
+from ...storage import InMemoryStorage, LocalFileStorage, Storage
 from .types import TestMemoryAddResult, TestMemoryStoreConfig
 
 DEFAULT_MAX_SEARCH_RESULTS = 10
@@ -69,7 +70,7 @@ def _token_overlap_score(query_tokens: set[str], content: str) -> int:
 
 
 class TestMemoryStore(MemoryStore):
-    """A :class:`~strands.memory.types.MemoryStore` backed by an in-memory list and a local JSON file.
+    """A :class:`~strands.memory.types.MemoryStore` backed by a local JSON file.
 
     A zero-infrastructure store for prototyping and testing. It persists to disk by default so memories persist
     across sessions. Set ``persist=False`` for an ephemeral, single-session store.
@@ -81,7 +82,11 @@ class TestMemoryStore(MemoryStore):
 
     Each :meth:`add` rewrites the whole file, so this fits modest volumes (hundreds to low thousands
     of entries), not production workloads — use a managed store like ``BedrockKnowledgeBaseStore`` for
-    that. Writes within a process are serialized; concurrent writers across processes are not.
+    that. Writes within one event loop are serialized; concurrent writers across processes are not.
+
+    Persistence is backed by the unified :class:`~strands.storage.Storage` interface: ``persist=True``
+    (the default) uses a :class:`~strands.storage.LocalFileStorage`, ``persist=False`` an ephemeral
+    :class:`~strands.storage.InMemoryStorage`.
 
     The on-disk format is shared with the TypeScript SDK's ``TestMemoryStore``: records use the same
     camelCase keys (``id``, ``content``, ``metadata``, ``createdAt``) and the same timestamp shape, so
@@ -124,20 +129,36 @@ def __init__(self, **store_config: Unpack[TestMemoryStoreConfig]) -> None:
         self.writable = store_config.get("writable", True)
         self.extraction = store_config.get("extraction")
 
-        self._persist = store_config.get("persist", True)
+        persist = store_config.get("persist", True)
         path = store_config.get("path")
         if path is not None and not path.strip():
             raise ValueError("TestMemoryStore: path must not be empty.")
-        if not self._persist:
-            self._path: Path | None = None
+
+        # Persistence runs on the unified Storage interface. Resolve a (backend, key) pair whose
+        # on-disk location matches the pre-Storage behavior exactly:
+        #   persist=False       -> ephemeral in-memory store
+        #   persist=True + path  -> the file at `path` (backend rooted at its parent dir)
+        #   persist=True default -> ~/.strands/memory/<sanitized-name>.json
+        # LocalFileStorage/InMemoryStorage construction touches no filesystem, so building the store
+        # never does I/O.
+        if not persist:
+            self._storage: Storage = InMemoryStorage()
+            self._key = f"{_sanitize_name(self.name)}.json"
         elif path is not None:
-            self._path = Path(path)
+            file = Path(path)
+            self._storage = LocalFileStorage(str(file.parent))
+            self._key = file.name
         else:
-            self._path = Path.home() / ".strands" / "memory" / f"{_sanitize_name(self.name)}.json"
+            self._storage = LocalFileStorage(str(Path.home() / ".strands" / "memory"))
+            self._key = f"{_sanitize_name(self.name)}.json"
 
-        # Records load lazily on first read/write so construction never touches the filesystem.
-        self._records: list[dict[str, Any]] | None = None
-        self._lock = threading.Lock()
+        # Serializes the read-modify-write cycle of add so concurrent adds don't each read the same
+        # snapshot and clobber one another (last-write-wins). The lock is created lazily per running
+        # loop (see _get_lock): an asyncio.Lock binds to the first loop that uses it, so a store
+        # reused across the fresh loops a synchronous Agent creates per invocation would otherwise
+        # raise "bound to a different event loop".
+        self._lock: asyncio.Lock | None = None
+        self._lock_loop: asyncio.AbstractEventLoop | None = None
 
     async def search(self, query: str, options: SearchOptions | None = None) -> list[MemoryEntry]:
         """Search stored entries for those whose content overlaps the query.
@@ -154,7 +175,9 @@ async def search(self, query: str, options: SearchOptions | None = None) -> list
             token-less query returns no results.
 
         Raises:
-            ValueError: If ``options.max_search_results`` is less than 1.
+            ValueError: If ``options.max_search_results`` is less than 1, or the backing file is
+                malformed (invalid JSON, not an array, or a record missing required fields).
+            StorageError: If the backend read fails.
         """
         caller_max = options.get("max_search_results") if options is not None else None
         if caller_max is not None and caller_max < 1:
@@ -165,7 +188,7 @@ async def search(self, query: str, options: SearchOptions | None = None) -> list
         if not query_tokens:
             return []
 
-        records = self._load()
+        records = await self._read()
 
         scored: list[tuple[dict[str, Any], int]] = []
         for record in records:
@@ -198,21 +221,21 @@ async def add(self, content: str, metadata: Metadata | None = None) -> TestMemor
             The id of the stored (or already-present) record.
 
         Raises:
-            ValueError: If the store is not writable or ``content`` is empty/whitespace.
-            OSError: If persisting the entry to disk fails (e.g. the path is unreachable or not
-                writable), with the target path in the message.
+            ValueError: If the store is not writable, ``content`` is empty/whitespace, or the
+                existing backing file is malformed.
+            StorageError: If the backend read or write fails.
         """
         if not self.writable:
             raise ValueError("TestMemoryStore: store is not writable. Set writable=True in config to enable add().")
         if not content.strip():
             raise ValueError("TestMemoryStore: content must not be empty.")
 
-        # The lock serializes the whole load-modify-flush cycle so concurrent adds don't each load
-        # the same snapshot and clobber one another (last-write-wins). Within a single event loop the
-        # synchronous critical section is already atomic; the lock additionally guards a store shared
-        # across OS threads.
-        with self._lock:
-            records = self._load()
+        # The lock serializes the whole read-modify-write cycle so concurrent adds on the same event
+        # loop don't each read the same snapshot and clobber one another. Reading inside the critical
+        # section guarantees add #N sees add #N-1's write. Serialization is per event loop; adds
+        # driven from separate loops/processes against a shared file remain last-write-wins.
+        async with self._get_lock():
+            records = await self._read()
 
             normalized_content = content.strip()
             for record in records:
@@ -223,35 +246,46 @@ async def add(self, content: str, metadata: Metadata | None = None) -> TestMemor
             if metadata is not None:
                 new_record["metadata"] = metadata
 
-            # Flush the candidate list first and only commit it to the in-memory cache once the write
-            # succeeds, so a failed flush never leaves a phantom record that later writes resurrect.
-            next_records = [*records, new_record]
-            self._flush(next_records)
-            self._records = next_records
+            await self._write([*records, new_record])
             return TestMemoryAddResult(id=new_record["id"])
 
-    def _load(self) -> list[dict[str, Any]]:
-        """Load records from disk on first use; ephemeral stores (and a missing file) start empty."""
-        if self._records is not None:
-            return self._records
+    def _get_lock(self) -> asyncio.Lock:
+        """Return the write lock for the running event loop, creating a fresh one when the loop changes.
+
+        An ``asyncio.Lock`` binds to the first loop that uses it, so a lock created once and reused
+        across loops raises ``RuntimeError``. A synchronous ``Agent`` runs each invocation on a fresh
+        loop, so a store reused across invocations must rebind. Rebinding per loop keeps the
+        serialization guarantee within a single loop (the only scope concurrency happens in) while
+        never carrying a lock across loops.
+        """
+        running_loop = asyncio.get_running_loop()
+        if self._lock is None or self._lock_loop is not running_loop:
+            self._lock = asyncio.Lock()
+            self._lock_loop = running_loop
+        return self._lock
+
+    async def _read(self) -> list[dict[str, Any]]:
+        """Read and parse the record file from storage; a missing key (or empty store) starts empty.
 
-        if self._path is None:
-            self._records = []
-            return self._records
+        Reads fresh on every call — there is no in-memory cache, so a search always reflects the
+        latest write.
+
+        Raises:
+            ValueError: If the stored file is not valid JSON, is not an array, or holds a record
+                missing the required string fields.
+            StorageError: If the backend read fails.
+        """
+        data = await self._storage.read(self._key)
+        if data is None:
+            return []
 
         try:
-            with open(self._path, encoding="utf-8") as file:
-                parsed_file = json.load(file)
-        except FileNotFoundError:
-            self._records = []
-            return self._records
+            parsed_file = json.loads(data)
         except json.JSONDecodeError as error:
-            raise ValueError(f"TestMemoryStore: invalid JSON in {self._path}: {error}") from error
-        except OSError as error:
-            raise OSError(f"TestMemoryStore: failed to read {self._path}: {error}") from error
+            raise ValueError(f"TestMemoryStore: invalid JSON in {self._key}: {error}") from error
 
         if not isinstance(parsed_file, list):
-            raise ValueError(f"TestMemoryStore: invalid backing file {self._path}: expected a JSON array of records")
+            raise ValueError(f"TestMemoryStore: invalid backing file {self._key}: expected a JSON array of records")
         for record in parsed_file:
             if (
                 not isinstance(record, dict)
@@ -260,29 +294,22 @@ def _load(self) -> list[dict[str, Any]]:
                 or not isinstance(record.get("createdAt"), str)
             ):
                 raise ValueError(
-                    f"TestMemoryStore: invalid backing file {self._path}: "
+                    f"TestMemoryStore: invalid backing file {self._key}: "
                     "each record must have string 'id', 'content', and 'createdAt' fields"
                 )
-        self._records = parsed_file
-        return self._records
-
-    def _flush(self, records: list[dict[str, Any]]) -> None:
-        """Persist ``records`` with an atomic write (write to a ``.tmp`` file, then replace).
+            metadata = record.get("metadata")
+            if metadata is not None and not isinstance(metadata, dict):
+                raise ValueError(
+                    f"TestMemoryStore: invalid backing file {self._key}: "
+                    "a record's 'metadata', when present, must be a JSON object"
+                )
+        return parsed_file
 
-        A crash mid-write can never leave a partially written file. A no-op for ephemeral stores.
+    async def _write(self, records: list[dict[str, Any]]) -> None:
+        """Persist ``records`` as a single JSON file through the storage backend.
 
-        Raises:
-            OSError: If the backing directory cannot be created or the file cannot be written
-                (e.g. the path is unreachable or not writable), with the target path in the message.
+        Callers serialize invocations via the instance lock; atomicity is the backend's
+        responsibility. A backend I/O failure surfaces as its own ``StorageError``, naming the key.
         """
-        if self._path is None:
-            return
-
-        try:
-            self._path.parent.mkdir(parents=True, exist_ok=True)
-            tmp_path = self._path.with_name(f"{self._path.name}.tmp")
-            with open(tmp_path, "w", encoding="utf-8", newline="\n") as file:
-                json.dump(records, file, indent=2, ensure_ascii=False)
-            tmp_path.replace(self._path)
-        except OSError as error:
-            raise OSError(f"TestMemoryStore: failed to write {self._path}: {error}") from error
+        data = json.dumps(records, indent=2, ensure_ascii=False).encode("utf-8")
+        await self._storage.write(self._key, data)

strands-py/tests/strands/vended_memory_stores/test_memory_store/test_store.py

@@ -24,13 +24,43 @@
 from strands.memory.extraction.triggers import InvocationTrigger
 from strands.memory.extraction.types import ExtractionConfig, ExtractionResult
 from strands.memory.memory_manager import MemoryManager
+from strands.storage import Storage
+from strands.types.exceptions import StorageError
 from strands.vended_memory_stores.test_memory_store import (
     TestMemoryAddResult,
     TestMemoryStore,
     TestMemoryStoreConfig,
 )
 
 
+class _YieldingStorage:
+    """A ``Storage`` wrapper that awaits a real suspension point before delegating each operation.
+
+    Forcing the write lock to actually suspend (rather than complete synchronously) is what binds a
+    lock's internal waiter to the running loop, so a store reused across loops fails loudly on a
+    loop-bound lock instead of passing by luck.
+    """
+
+    def __init__(self, inner: Storage) -> None:
+        self._inner = inner
+
+    async def write(self, key: str, data: bytes) -> None:
+        await asyncio.sleep(0)
+        await self._inner.write(key, data)
+
+    async def read(self, key: str) -> bytes | None:
+        await asyncio.sleep(0)
+        return await self._inner.read(key)
+
+    async def delete(self, key: str) -> None:
+        await asyncio.sleep(0)
+        await self._inner.delete(key)
+
+    async def list(self, query: str = "") -> list[str]:
+        await asyncio.sleep(0)
+        return await self._inner.list(query)
+
+
 @pytest.fixture
 def store_path(tmp_path: Path) -> str:
     """A unique backing-file path under the test's temp dir."""
@@ -274,20 +304,56 @@ async def test_raises_clear_error_on_malformed_record(self, make_store, store_pa
         with pytest.raises(ValueError, match="each record must have string"):
             await store.search("anything")
 
+    @pytest.mark.asyncio
+    async def test_raises_clear_error_on_non_object_metadata(self, make_store, store_path):
+        # A present-but-non-object metadata (e.g. a hand-edited or cross-SDK file) must fail fast
+        # rather than crashing opaquely when search spreads it into the result.
+        record = [{"id": "a", "content": "hi", "createdAt": "2026-01-01T00:00:00.000Z", "metadata": "oops"}]
+        Path(store_path).write_text(json.dumps(record), encoding="utf-8")
+        store = make_store()
+        with pytest.raises(ValueError, match="'metadata', when present, must be a JSON object"):
+            await store.search("hi")
+
+    @pytest.mark.asyncio
+    async def test_accepts_null_metadata_as_absent(self, make_store, store_path):
+        # A record with null metadata is accepted and treated as absent, matching the TS store.
+        record = [{"id": "a", "content": "hi there", "createdAt": "2026-01-01T00:00:00.000Z", "metadata": None}]
+        Path(store_path).write_text(json.dumps(record), encoding="utf-8")
+        store = make_store()
+        results = await store.search("hi")
+        assert len(results) == 1
+        assert results[0].metadata == {"_relevanceScore": 1}
+
     @pytest.mark.asyncio
     async def test_keeps_all_entries_under_concurrent_writes(self, make_store, store_path):
         store = make_store()
         await asyncio.gather(*(store.add(f"fact number {index}") for index in range(10)))
         assert len(json.loads(Path(store_path).read_text(encoding="utf-8"))) == 10
 
+    def test_reused_across_separate_event_loops(self, make_store, store_path):
+        # A synchronous Agent runs each invocation on a fresh event loop, so a store reused across
+        # invocations acquires its write lock on a different loop each time. A loop-bound lock created
+        # once would raise "bound to a different event loop"; the lock must rebind per loop. Wrapping
+        # the backend so each op suspends forces the lock to actually bind a waiter to the loop, so a
+        # regression surfaces rather than passing by luck.
+        store = make_store()
+        store._storage = _YieldingStorage(store._storage)
+
+        async def batch(tag: str) -> None:
+            await asyncio.gather(store.add(f"{tag} one"), store.add(f"{tag} two"))
+
+        asyncio.run(batch("first"))
+        asyncio.run(batch("second"))
+        assert len(json.loads(Path(store_path).read_text(encoding="utf-8"))) == 4
+
     @pytest.mark.asyncio
     async def test_raises_clear_error_when_path_is_unreachable(self, tmp_path):
-        # Point the store under an existing FILE, so the backing path can't be reached — surfacing a
-        # wrapped "failed to read/write" error naming the path rather than a bare OSError.
+        # Point the store under an existing FILE, so the backing path can't be reached — the backend
+        # raises a StorageError naming the key rather than a bare OSError.
         blocker = tmp_path / "blocker"
         blocker.write_text("not a directory", encoding="utf-8")
         store = TestMemoryStore(name="notes", path=str(blocker / "notes.json"))
-        with pytest.raises(OSError, match="failed to"):
+        with pytest.raises(StorageError, match="Failed to write"):
             await store.add("user prefers dark mode")