memwal 0.1.5.dev0__tar.gz → 0.1.5.dev2__tar.gz

{memwal-0.1.5.dev0 → memwal-0.1.5.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: memwal
-Version: 0.1.5.dev0
+Version: 0.1.5.dev2
 Summary: Python SDK for Walrus Memory — Privacy-first AI memory with Ed25519 signing
 Project-URL: Homepage, https://memory.walrus.xyz
 Project-URL: Documentation, https://memory.walrus.xyz

{memwal-0.1.5.dev0 → memwal-0.1.5.dev2}/memwal/__init__.py

@@ -116,4 +116,4 @@ __all__ = [
     "RecallManualResult",
 ]
 
-__version__ = "0.1.5.dev0"
+__version__ = "0.1.5.dev2"

{memwal-0.1.5.dev0 → memwal-0.1.5.dev2}/memwal/client.py

@@ -784,22 +784,30 @@ class MemWal:
           do **not** count as either restored or skipped.
         * ``total`` — count of on-chain blobs the relayer saw for
           ``(owner, namespace)`` before the limit was applied.
+        * ``truncated`` — True when this restore is known-incomplete (limit
+          truncation or an upstream on-chain candidate-fetch cap; see
+          :class:`~memwal.types.RestoreResult`). Defaults to ``False`` when
+          talking to a relayer older than WALM-319 that omits the field.
 
         Args:
             namespace: Namespace to restore. Exact match — no prefix or
                 hierarchy semantics (see SKILL.md "Namespace Semantics").
             limit: Max blobs the relayer will *inspect* this call (default:
                 10, matches the server-side default and the TypeScript SDK).
-                The relayer fetches blobs newest-first and caps the work at
-                this number; it does **not** cap ``restored`` separately.
+                The relayer selects up to this many missing blobs from the
+                candidates it sees, in unspecified order; it does **not** cap
+                ``restored`` separately.
 
         Returns:
             :class:`RestoreResult`.
 
         Notes:
-            * **No pagination cursor** — restore is single-shot. To rebuild a
-              large namespace, call repeatedly with a growing ``limit`` or
-              prune the local index first.
+            * **No pagination cursor** — restore is single-shot. A larger
+              ``limit`` may help when ``truncated`` is caused by the request
+              limit, but candidate discovery also has a per-owner source cap
+              shared across namespaces. Hitting that cap requires cursor/
+              pagination support for guaranteed full recovery; increasing
+              ``limit`` or retrying cannot guarantee it.
             * **Performance** scales linearly in ``limit``: up to 10 Walrus
               downloads in parallel, then 3 SEAL decrypts in parallel, then
               embeddings. Expect seconds-per-blob on cold caches.
@@ -814,6 +822,10 @@ class MemWal:
             total=data["total"],
             namespace=data["namespace"],
             owner=data["owner"],
+            # Relayers older than WALM-319 omit `truncated` entirely --
+            # treat "not present" as "not known to be truncated" rather
+            # than require every relayer version to send it.
+            truncated=data.get("truncated", False),
         )
 
     async def health(self) -> HealthResult:

{memwal-0.1.5.dev0 → memwal-0.1.5.dev2}/memwal/middleware.py

@@ -130,14 +130,124 @@ def _format_memories(
     ])
 
 
-def _fire_and_forget(coro: Any) -> None:
-    """Schedule an async coroutine as fire-and-forget.
+class _PendingSaves:
+    """Tracks in-flight fire-and-forget auto-save work (asyncio Tasks and
+    background Threads) so a caller can drain it deterministically via
+    :meth:`flush`/:meth:`flush_sync` instead of losing it silently when the
+    process exits before it completes. Tasks scheduled via
+    ``loop.create_task()`` with no other reference are only weakly held by
+    the event loop and can be cancelled or garbage-collected before they
+    run — this keeps a strong reference until each one is done.
+    """
+
+    def __init__(self) -> None:
+        self._tasks: "set[asyncio.Task[Any]]" = set()
+        self._threads: List[threading.Thread] = []
+        self._lock = threading.Lock()
+
+    def track_task(self, task: "asyncio.Task[Any]") -> None:
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+
+    def spawn_thread(self, target: Callable[[], None]) -> None:
+        """Run `target` in a new daemon thread, tracked until it completes.
+        The thread untracks itself on completion -- a long-lived client
+        that keeps using fire-and-forget saves but never calls flush()
+        would otherwise accumulate one Thread object per save, forever.
+        """
+        def _run_and_untrack() -> None:
+            try:
+                target()
+            finally:
+                with self._lock:
+                    if thread in self._threads:
+                        self._threads.remove(thread)
+
+        thread = threading.Thread(target=_run_and_untrack, daemon=True)
+        with self._lock:
+            self._threads.append(thread)
+        thread.start()
+
+    async def flush(self) -> None:
+        """Await every pending task, then join every pending thread."""
+        if self._tasks:
+            await asyncio.gather(*list(self._tasks), return_exceptions=True)
+        self._join_threads()
+
+    def flush_sync(self) -> None:
+        """Join every pending thread. A pending Task belongs to the loop
+        that created it (e.g. an earlier `await llm.ainvoke(...)` call) and
+        cannot be awaited from a different one -- if any are still pending
+        here, the caller mixed sync and async entry points on the same
+        wrapped client, and this cleanup path can't safely drain them. Log
+        and continue draining threads rather than raise out of what's
+        usually shutdown code."""
+        if self._tasks:
+            async def _drain(tasks: "list[asyncio.Task[Any]]") -> None:
+                await asyncio.gather(*tasks, return_exceptions=True)
+
+            try:
+                asyncio.run(_drain(list(self._tasks)))
+            except RuntimeError:
+                logger.warning(
+                    "Walrus Memory flush_sync() could not await %d pending "
+                    "task(s) bound to a different event loop (mixing sync "
+                    "and async calls on the same wrapped client?) — those "
+                    "saves may be lost.",
+                    len(self._tasks),
+                )
+        self._join_threads()
+
+    def _join_threads(self) -> None:
+        with self._lock:
+            threads, self._threads = self._threads, []
+        for thread in threads:
+            thread.join()
+
+
+def _expose_memwal_controls(obj: Any, memwal: MemWal, pending: _PendingSaves) -> None:
+    """Attach the underlying client and a way to drain pending auto-saves —
+    without this, a short-lived process has no way to avoid silently
+    losing writes still in flight when it exits.
+
+    Uses object.__setattr__ rather than plain attribute assignment: a
+    LangChain BaseChatModel is a Pydantic model, and Pydantic's __setattr__
+    rejects assignment to any name that isn't a declared field (or a
+    leading-underscore private attribute) — memwal_flush/memwal_flush_sync
+    are neither, so plain `obj.memwal_flush = ...` raises ValueError on a
+    real LangChain model (masked by tests using MagicMock, which doesn't
+    enforce this). object.__setattr__ bypasses that check and writes
+    straight into the instance's __dict__, where normal attribute lookup
+    (obj.memwal_flush) still finds it — OpenAI clients aren't Pydantic
+    models and work the same way either way.
+    """
+    object.__setattr__(obj, "_memwal", memwal)
+    object.__setattr__(obj, "memwal_flush", pending.flush)
+    object.__setattr__(obj, "memwal_flush_sync", pending.flush_sync)
+
+
+async def _warn_if_cancelled(coro: Any, label: str) -> None:
+    try:
+        await coro
+    except asyncio.CancelledError:
+        logger.warning(
+            "Walrus Memory %s was cancelled before it completed — the process "
+            "likely exited without draining pending saves. Call flush()/"
+            "flush_sync() before exiting to avoid silent data loss.",
+            label,
+        )
+        raise
+
+
+def _fire_and_forget(coro: Any, pending: _PendingSaves, label: str = "auto-save") -> None:
+    """Schedule an async coroutine as fire-and-forget, tracked in `pending`.
 
     Works whether or not an event loop is already running.
     """
     try:
         loop = asyncio.get_running_loop()
-        loop.create_task(coro)
+        task = loop.create_task(_warn_if_cancelled(coro, label))
+        pending.track_task(task)
     except RuntimeError:
         # No running loop -- run in a background thread
         def _run() -> None:
@@ -146,8 +256,7 @@ def _fire_and_forget(coro: Any) -> None:
             except Exception:
                 logger.debug("Fire-and-forget analyze() failed", exc_info=True)
 
-        thread = threading.Thread(target=_run, daemon=True)
-        thread.start()
+        pending.spawn_thread(_run)
 
 
 def _run_blocking(coro_factory: Callable[[], Any]) -> Any:
@@ -226,6 +335,7 @@ def with_memwal_langchain(
     )
 
     log = logger.debug if not debug else logger.warning
+    pending = _PendingSaves()
 
     original_agenerate = llm._agenerate
     original_generate = llm._generate
@@ -289,7 +399,7 @@ def with_memwal_langchain(
         result = await original_agenerate(enriched, *args, **kwargs)
 
         for msg_list in messages:
-            _fire_and_forget(_post_analyze(msg_list))
+            _fire_and_forget(_post_analyze(msg_list), pending)
 
         return result
 
@@ -315,7 +425,7 @@ def with_memwal_langchain(
         result = original_generate(enriched, *args, **kwargs)
 
         for msg_list in messages:
-            _fire_and_forget(_post_analyze(msg_list))
+            _fire_and_forget(_post_analyze(msg_list), pending)
 
         return result
 
@@ -323,6 +433,8 @@ def with_memwal_langchain(
     llm._agenerate = patched_agenerate  # type: ignore[assignment]
     llm._generate = patched_generate  # type: ignore[assignment]
 
+    _expose_memwal_controls(llm, memwal, pending)
+
     return llm
 
 
@@ -379,13 +491,20 @@ def with_memwal_openai(
     )
 
     log = logger.debug if not debug else logger.warning
+    pending = _PendingSaves()
 
     is_async = hasattr(client, "_async_client") or type(client).__name__ == "AsyncOpenAI"
 
     if is_async:
-        _wrap_async_openai(client, memwal, namespace, max_memories, auto_save, min_relevance, log)
+        _wrap_async_openai(
+            client, memwal, namespace, max_memories, auto_save, min_relevance, log, pending
+        )
     else:
-        _wrap_sync_openai(client, memwal, namespace, max_memories, auto_save, min_relevance, log)
+        _wrap_sync_openai(
+            client, memwal, namespace, max_memories, auto_save, min_relevance, log, pending
+        )
+
+    _expose_memwal_controls(client, memwal, pending)
 
     return client
 
@@ -398,6 +517,7 @@ def _wrap_async_openai(
     auto_save: bool,
     min_relevance: float,
     log: Callable[..., Any],
+    pending: _PendingSaves,
 ) -> None:
     """Wrap an async OpenAI client's chat.completions.create."""
     original_create = client.chat.completions.create
@@ -437,7 +557,7 @@ def _wrap_async_openai(
                 except Exception as e:
                     log(f"[Walrus Memory] Auto-save failed: {e}")
 
-            _fire_and_forget(_analyze())
+            _fire_and_forget(_analyze(), pending)
 
         return result
 
@@ -452,6 +572,7 @@ def _wrap_sync_openai(
     auto_save: bool,
     min_relevance: float,
     log: Callable[..., Any],
+    pending: _PendingSaves,
 ) -> None:
     """Wrap a sync OpenAI client's chat.completions.create."""
     original_create = client.chat.completions.create
@@ -498,8 +619,7 @@ def _wrap_sync_openai(
                 except Exception as e:
                     log(f"[Walrus Memory] Auto-save failed: {e}")
 
-            thread = threading.Thread(target=_analyze, daemon=True)
-            thread.start()
+            pending.spawn_thread(_analyze)
 
         return result
 

{memwal-0.1.5.dev0 → memwal-0.1.5.dev2}/memwal/types.py

@@ -201,6 +201,17 @@ class RestoreResult:
     total: int
     namespace: str
     owner: str
+    #: True when this restore is known-incomplete: either more on-chain
+    #: blobs were missing locally than ``limit`` allowed this call to
+    #: restore, or the sidecar's raw on-chain candidate fetch hit its own
+    #: cap before this namespace's blobs were even filtered out of that
+    #: set (WALM-319) -- this can be ``True`` even when ``total == 0``,
+    #: since a cap hit elsewhere can starve this namespace's fetch
+    #: entirely. Raising ``limit`` only helps with the first case.
+    #:
+    #: Relayers older than WALM-319 don't send this field at all; the SDK
+    #: defaults it to ``False`` in that case rather than requiring it.
+    truncated: bool = False
 
 
 @dataclass

{memwal-0.1.5.dev0 → memwal-0.1.5.dev2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "memwal"
-version = "0.1.5.dev0"
+version = "0.1.5.dev2"
 description = "Python SDK for Walrus Memory — Privacy-first AI memory with Ed25519 signing"
 readme = "README.md"
 license = "MIT"

{memwal-0.1.5.dev0 → memwal-0.1.5.dev2}/tests/test_client.py

@@ -630,6 +630,7 @@ class TestRestore:
                     "total": 7,
                     "namespace": "my-app",
                     "owner": "0xowner",
+                    "truncated": False,
                 },
             )
         )
@@ -641,6 +642,57 @@ class TestRestore:
         assert body["limit"] == 100
         assert result.restored == 5
         assert result.skipped == 2
+        assert result.truncated is False
+
+    @respx.mock
+    async def test_restore_preserves_truncated_true(
+        self, memwal_client: MemWal
+    ) -> None:
+        """truncated=True from the relayer must survive into RestoreResult
+        instead of being dropped when the dataclass is reconstructed."""
+        mock_seal_session_prereqs()
+        respx.post(f"{_TEST_SERVER}/api/restore").mock(
+            return_value=httpx.Response(
+                200,
+                json={
+                    "restored": 5,
+                    "skipped": 2,
+                    "total": 7,
+                    "namespace": "my-app",
+                    "owner": "0xowner",
+                    "truncated": True,
+                },
+            )
+        )
+
+        result = await memwal_client.restore("my-app", limit=100)
+
+        assert result.truncated is True
+
+    @respx.mock
+    async def test_restore_truncated_defaults_false_when_omitted(
+        self, memwal_client: MemWal
+    ) -> None:
+        """Relayers older than WALM-319 omit `truncated` entirely — the
+        SDK must default it to False rather than requiring the field."""
+        mock_seal_session_prereqs()
+        respx.post(f"{_TEST_SERVER}/api/restore").mock(
+            return_value=httpx.Response(
+                200,
+                json={
+                    "restored": 5,
+                    "skipped": 2,
+                    "total": 7,
+                    "namespace": "my-app",
+                    "owner": "0xowner",
+                    # no "truncated" key — simulates an older relayer
+                },
+            )
+        )
+
+        result = await memwal_client.restore("my-app", limit=100)
+
+        assert result.truncated is False
 
 
 class TestHealth:

{memwal-0.1.5.dev0 → memwal-0.1.5.dev2}/tests/test_middleware.py

@@ -17,16 +17,19 @@ from __future__ import annotations
 
 import asyncio
 import json
-from unittest.mock import AsyncMock, MagicMock
+import threading
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import nacl.signing
 import respx
 
+from memwal.client import MemWal
 from memwal.middleware import (
     _find_last_user_message,
     _format_memories,
     _inject_openai_memory,
+    _PendingSaves,
     with_memwal_langchain,
     with_memwal_openai,
 )
@@ -444,6 +447,24 @@ class TestWithMemWalLangChain:
         await smart_llm._agenerate([[SystemMessage("only system")]])
         assert not recall_route.called
 
+    def test_wraps_a_real_pydantic_backed_chat_model(self) -> None:
+        """with_memwal_langchain must work on a real BaseChatModel, not just
+        a MagicMock. LangChain chat models are Pydantic models that reject
+        assignment of undeclared fields under normal attribute assignment
+        -- MagicMock doesn't enforce that, so it silently hides this."""
+        from langchain_core.language_models.fake_chat_models import FakeListChatModel
+
+        llm = FakeListChatModel(responses=["canned response"])
+
+        smart_llm = with_memwal_langchain(
+            llm, key=_KEY_HEX, account_id=_ACCOUNT_ID, server_url=_SERVER, auto_save=False
+        )
+
+        assert smart_llm is llm
+        assert smart_llm._memwal is not None
+        assert callable(smart_llm.memwal_flush)
+        assert callable(smart_llm.memwal_flush_sync)
+
 
 # ============================================================
 # OpenAI middleware tests
@@ -587,6 +608,41 @@ class TestWithMemWalOpenAI:
         await asyncio.sleep(0.05)
         assert analyze_route.called
 
+    @respx.mock
+    async def test_memwal_flush_awaits_pending_autosave(self) -> None:
+        """memwal_flush() deterministically waits for the fire-and-forget
+        analyze() call, instead of the caller having to guess a sleep
+        duration and hope the background task finished in time."""
+        _mock_seal_session_prereqs()
+        client = self._make_async_client()
+
+        respx.post(_RECALL_URL).mock(return_value=_mock_recall([]))
+
+        gate = asyncio.Event()
+        completed = {"value": False}
+
+        async def gated_analyze(self, *args, **kwargs):
+            await gate.wait()
+            completed["value"] = True
+
+        with patch.object(MemWal, "analyze", gated_analyze):
+            smart = with_memwal_openai(
+                client, key=_KEY_HEX, account_id=_ACCOUNT_ID, server_url=_SERVER, auto_save=True
+            )
+            await smart.chat.completions.create(
+                model="gpt-4o",
+                messages=[{"role": "user", "content": "remember this"}],
+            )
+
+            # Fire-and-forget: the response came back, but the gated
+            # analyze() call has not completed yet.
+            assert completed["value"] is False
+
+            gate.set()
+            await smart.memwal_flush()
+
+            assert completed["value"] is True
+
     @respx.mock
     async def test_min_relevance_filter(self) -> None:
         """Memories with relevance below min_relevance are not injected."""
@@ -671,4 +727,70 @@ class TestWithMemWalOpenAI:
         assert "TV support appointment is 9 AM to noon" not in system_msgs[0]["content"]
         user_msgs = [m for m in captured if isinstance(m, dict) and m.get("role") == "user"]
         assert len(user_msgs) == 2
-        assert "TV support appointment is 9 AM to noon" in user_msgs[0]["content"]
+
+
+# ============================================================
+# _PendingSaves
+# ============================================================
+
+
+class TestPendingSaves:
+    """Tests for _PendingSaves, isolated from the middleware wrappers."""
+
+    def test_flush_sync_does_not_crash_on_a_task_from_a_different_loop(self) -> None:
+        """Reproduces the real bug: a caller that mixes an earlier async
+        call (which populates self._tasks on that call's event loop) with
+        flush_sync() from sync code afterward. asyncio.Tasks are bound to
+        the loop that created them and can't be awaited from a new one --
+        flush_sync() must not raise out of this, since it's typically
+        called during cleanup."""
+        pending = _PendingSaves()
+        task_holder: dict = {}
+        release = threading.Event()
+
+        def _run_other_loop() -> None:
+            async def _pending_forever() -> None:
+                await asyncio.get_event_loop().run_in_executor(None, release.wait)
+
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            task = loop.create_task(_pending_forever())
+            task_holder["task"] = task
+            pending.track_task(task)
+            # Keep this loop alive long enough for the main thread's
+            # flush_sync() to observe the still-pending, cross-loop task.
+            loop.run_until_complete(asyncio.sleep(0.3))
+
+        other_loop_thread = threading.Thread(target=_run_other_loop)
+        other_loop_thread.start()
+
+        # Give the other thread time to create its loop and schedule the task.
+        import time
+
+        time.sleep(0.05)
+        assert task_holder.get("task") is not None
+        assert not task_holder["task"].done()
+
+        try:
+            pending.flush_sync()  # must not raise
+        finally:
+            release.set()
+            other_loop_thread.join(timeout=2)
+
+    def test_completed_threads_are_untracked_without_flushing(self) -> None:
+        """A long-lived client that keeps using fire-and-forget saves but
+        never calls flush()/flush_sync() must not accumulate one Thread
+        object per save forever."""
+        pending = _PendingSaves()
+        done = threading.Event()
+
+        pending.spawn_thread(done.set)
+
+        assert done.wait(timeout=2), "background thread never ran"
+        # Give the thread's own cleanup a moment to run after done.set()
+        # returns (the Event is set from inside the tracked callable,
+        # microseconds before the thread function itself returns).
+        import time
+
+        time.sleep(0.05)
+        assert len(pending._threads) == 0