address github-actions review comments

Author: rohanshrma222

packages/letta-moss/README.md

@@ -74,7 +74,7 @@ Then register the server as an MCP tool source on your Letta agent, pointing at
 
 | Symbol | Kind | Purpose |
 |---|---|---|
-| `MossLettaMemory(*, project_id=None, project_key=None, index_name, top_k=5, alpha=0.8)` | class | Core adapter; `async load_index()`, `async insert_memory(content, tags=, metadata=) -> str`, `async search_memory(query, top_k=, tags=) -> list[ArchivalMemoryItem]`, `async delete_memory(memory_id)`, `async get_memory(memory_id) -> ArchivalMemoryItem \| None`, `async list_memories(limit=) -> list[ArchivalMemoryItem]`. Falls back to `MOSS_PROJECT_ID`/`MOSS_PROJECT_KEY` env vars when credentials are omitted. |
+| `MossLettaMemory(*, project_id=None, project_key=None, index_name, top_k=5, alpha=0.8, auto_refresh=True, refresh_interval_seconds=60)` | class | Core adapter; `async load_index()`, `async insert_memory(content, tags=, metadata=) -> str`, `async search_memory(query, top_k=, tags=) -> list[ArchivalMemoryItem]`, `async delete_memory(memory_id)`, `async get_memory(memory_id) -> ArchivalMemoryItem \| None`, `async list_memories(limit=) -> list[ArchivalMemoryItem]`. Falls back to `MOSS_PROJECT_ID`/`MOSS_PROJECT_KEY` env vars when credentials are omitted. `auto_refresh`/`refresh_interval_seconds` are forwarded to Moss's `load_index()` so this process's in-memory snapshot picks up writes made by *other* processes against the same index, not just its own. |
 | `moss_memory_insert(content, tags=None) -> str` | async function | Custom-tool wrapper; insert a memory, return its id. |
 | `moss_memory_search(query, top_k=5, tags=None) -> list[dict]` | async function | Custom-tool wrapper; search memories, return dicts. |
 | `moss_memory_delete(memory_id) -> None` | async function | Custom-tool wrapper; delete a memory by id. |
@@ -91,6 +91,8 @@ Then register the server as an MCP tool source on your Letta agent, pointing at
 
 **Env var fallback:** `MossLettaMemory` falls back to `MOSS_PROJECT_ID`/`MOSS_PROJECT_KEY` when constructor args are omitted, unlike some other Moss integrations in this repo (e.g. `agora-moss`) that require explicit args. This is deliberate: these tools typically run inside Letta's sandboxed tool-execution environment, where env vars passed via `tool_exec_environment_variables` are the natural way to configure credentials without hardcoding them in tool source.
 
+**Sandbox process reuse:** the Option A custom tools cache a single `MossLettaMemory` instance at module scope per worker process (so `load_index()` only runs once). That cache is keyed on `MOSS_PROJECT_ID`/`MOSS_PROJECT_KEY`/`MOSS_INDEX_NAME` — if Letta reuses the same worker process for a different agent whose `tool_exec_environment_variables` set different values, the cached instance is discarded and rebuilt against the new values rather than leaking reads/writes across agents.
+
 ## Dependencies
 
 - `moss>=1.1.1`

packages/letta-moss/src/letta_moss/memory.py

@@ -126,8 +126,20 @@ def __init__(
         index_name: str,
         top_k: int = 5,
         alpha: float = 0.8,
+        auto_refresh: bool = True,
+        refresh_interval_seconds: int = 60,
     ) -> None:
-        """Initialize the adapter with Moss credentials and index-query defaults."""
+        """Initialize the adapter with Moss credentials and index-query defaults.
+
+        ``auto_refresh``/``refresh_interval_seconds`` are forwarded to the
+        underlying ``MossClient.load_index()`` call: they make the SDK poll
+        for and pull in changes made by *other* processes writing to the same
+        index (e.g. another sandboxed tool worker, or a script inserting
+        memories directly). This instance's own ``insert_memory``/
+        ``delete_memory`` calls are reflected immediately regardless, via the
+        generation-counter invalidation in ``load_index`` below — auto-refresh
+        only covers the cross-process case that local invalidation can't see.
+        """
         project_id = project_id or os.getenv("MOSS_PROJECT_ID")
         project_key = project_key or os.getenv("MOSS_PROJECT_KEY")
         if not project_id or not project_key:
@@ -139,6 +151,8 @@ def __init__(
         self._index_name = index_name
         self._top_k = top_k
         self._alpha = alpha
+        self._auto_refresh = auto_refresh
+        self._refresh_interval_seconds = refresh_interval_seconds
         self._index_loaded = False
         self._index_created = False
         # Bumped by insert_memory/delete_memory. load_index() only marks the
@@ -160,7 +174,11 @@ async def load_index(self) -> None:
             return
         generation_at_start = self._generation
         try:
-            await self._client.load_index(self._index_name)
+            await self._client.load_index(
+                self._index_name,
+                auto_refresh=self._auto_refresh,
+                polling_interval_in_seconds=self._refresh_interval_seconds,
+            )
         except RuntimeError as e:
             if "not found" not in str(e).lower():
                 raise

packages/letta-moss/src/letta_moss/tools.py

@@ -21,28 +21,42 @@
 from .memory import MossLettaMemory
 
 _memory: MossLettaMemory | None = None
+_memory_key: tuple[str | None, str | None, str] | None = None
 _memory_lock = asyncio.Lock()
 
 
+def _current_memory_key() -> tuple[str | None, str | None, str]:
+    """Read the env vars that identify which Moss index/credentials to use."""
+    index_name = os.getenv("MOSS_INDEX_NAME")
+    if not index_name:
+        raise ValueError("MOSS_INDEX_NAME env var is required.")
+    return (os.getenv("MOSS_PROJECT_ID"), os.getenv("MOSS_PROJECT_KEY"), index_name)
+
+
 async def _get_memory() -> MossLettaMemory:
     """Lazily build and load the module-level ``MossLettaMemory`` singleton.
 
     A single instance is reused across tool calls within a sandbox process so
-    ``load_index()`` (and the underlying index download) only happens once.
-    Guarded by a lock so concurrent tool calls can't each construct and load
-    their own instance before the first one is stored.
+    ``load_index()`` (and the underlying index download) only happens once —
+    but only for as long as ``MOSS_PROJECT_ID``/``MOSS_PROJECT_KEY``/
+    ``MOSS_INDEX_NAME`` keep resolving to the same values. If a worker process
+    gets reused for a different agent/config (a different set of
+    ``tool_exec_environment_variables``), the cached instance is discarded and
+    rebuilt against the new env vars instead of silently leaking reads/writes
+    to the previous agent's index. Guarded by a lock so concurrent tool calls
+    can't each construct and load their own instance before the first one is
+    stored.
     """
-    global _memory
-    if _memory is not None:
+    global _memory, _memory_key
+    key = _current_memory_key()
+    if _memory is not None and _memory_key == key:
         return _memory
     async with _memory_lock:
-        if _memory is None:
-            index_name = os.getenv("MOSS_INDEX_NAME")
-            if not index_name:
-                raise ValueError("MOSS_INDEX_NAME env var is required.")
-            memory = MossLettaMemory(index_name=index_name)
+        if _memory is None or _memory_key != key:
+            memory = MossLettaMemory(index_name=key[2])
             await memory.load_index()
             _memory = memory
+            _memory_key = key
     return _memory
 
 

packages/letta-moss/tests/test_memory.py

@@ -29,8 +29,10 @@ def __init__(self, *, project_id=None, project_key=None):
         self._query_return = FakeQueryResult(docs=[])
         self._create_index_error = None
         self._load_index_error = None
+        self.load_index_calls = []
 
-    async def load_index(self, index_name):
+    async def load_index(self, index_name, auto_refresh=False, polling_interval_in_seconds=600):
+        self.load_index_calls.append((index_name, auto_refresh, polling_interval_in_seconds))
         if self._load_index_error is not None:
             raise self._load_index_error
 
@@ -103,6 +105,22 @@ async def test_delegates_and_marks_loaded(self, monkeypatch):
         await m.load_index()
         assert m._index_loaded is True
 
+    async def test_forwards_auto_refresh_settings_to_client(self, monkeypatch):
+        import letta_moss.memory as memory_mod
+
+        monkeypatch.setattr(memory_mod, "MossClient", FakeClient)
+        from letta_moss.memory import MossLettaMemory
+
+        m = MossLettaMemory(
+            project_id="p",
+            project_key="k",
+            index_name="idx",
+            auto_refresh=True,
+            refresh_interval_seconds=45,
+        )
+        await m.load_index()
+        assert m._client.load_index_calls == [("idx", True, 45)]
+
     async def test_marks_index_created_on_successful_load(self, monkeypatch):
         import letta_moss.memory as memory_mod
 
@@ -120,7 +138,9 @@ async def test_is_idempotent(self, monkeypatch):
         load_calls = []
 
         class TrackingClient(FakeClient):
-            async def load_index(self, index_name):
+            async def load_index(
+                self, index_name, auto_refresh=False, polling_interval_in_seconds=600
+            ):
                 load_calls.append(index_name)
 
         monkeypatch.setattr(memory_mod, "MossClient", TrackingClient)
@@ -164,9 +184,11 @@ async def test_mutation_during_load_prevents_stale_loaded_flag(self, monkeypatch
         release_load = asyncio.Event()
 
         class SlowClient(FakeClient):
-            async def load_index(self, index_name):
+            async def load_index(
+                self, index_name, auto_refresh=False, polling_interval_in_seconds=600
+            ):
                 await release_load.wait()
-                await super().load_index(index_name)
+                await super().load_index(index_name, auto_refresh, polling_interval_in_seconds)
 
         monkeypatch.setattr(memory_mod, "MossClient", SlowClient)
         from letta_moss.memory import MossLettaMemory

packages/letta-moss/tests/test_tools.py

@@ -44,9 +44,11 @@ def _reset_singleton(monkeypatch):
     FakeMemory.instances = []
     monkeypatch.setattr(tools_mod, "MossLettaMemory", FakeMemory)
     monkeypatch.setattr(tools_mod, "_memory", None)
+    monkeypatch.setattr(tools_mod, "_memory_key", None)
     monkeypatch.setenv("MOSS_INDEX_NAME", "idx")
     yield
     monkeypatch.setattr(tools_mod, "_memory", None)
+    monkeypatch.setattr(tools_mod, "_memory_key", None)
 
 
 class TestLazySingleton:
@@ -67,6 +69,27 @@ async def test_concurrent_calls_construct_only_one_instance(self):
         assert len(FakeMemory.instances) == 1
         assert all(r is results[0] for r in results)
 
+    async def test_rebuilds_when_index_name_env_var_changes(self, monkeypatch):
+        from letta_moss.tools import _get_memory
+
+        first = await _get_memory()
+        monkeypatch.setenv("MOSS_INDEX_NAME", "other-idx")
+        second = await _get_memory()
+
+        assert first is not second
+        assert len(FakeMemory.instances) == 2
+        assert second.index_name == "other-idx"
+
+    async def test_rebuilds_when_project_credentials_env_vars_change(self, monkeypatch):
+        from letta_moss.tools import _get_memory
+
+        first = await _get_memory()
+        monkeypatch.setenv("MOSS_PROJECT_ID", "new-project")
+        second = await _get_memory()
+
+        assert first is not second
+        assert len(FakeMemory.instances) == 2
+
 
 class TestMossMemoryInsert:
     async def test_delegates_to_memory(self):