memwal 0.2.1__tar.gz → 0.2.2.dev0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (38) hide show
  1. {memwal-0.2.1 → memwal-0.2.2.dev0}/CHANGELOG.md +12 -0
  2. {memwal-0.2.1 → memwal-0.2.2.dev0}/PKG-INFO +11 -1
  3. {memwal-0.2.1 → memwal-0.2.2.dev0}/README.md +10 -0
  4. {memwal-0.2.1 → memwal-0.2.2.dev0}/memwal/__init__.py +1 -1
  5. {memwal-0.2.1 → memwal-0.2.2.dev0}/memwal/client.py +170 -23
  6. {memwal-0.2.1 → memwal-0.2.2.dev0}/memwal/mock.py +10 -3
  7. {memwal-0.2.1 → memwal-0.2.2.dev0}/memwal/types.py +7 -4
  8. {memwal-0.2.1 → memwal-0.2.2.dev0}/memwal/utils.py +23 -10
  9. {memwal-0.2.1 → memwal-0.2.2.dev0}/pyproject.toml +1 -1
  10. memwal-0.2.2.dev0/tests/test_destroy_and_cancel_shield.py +661 -0
  11. {memwal-0.2.1 → memwal-0.2.2.dev0}/.gitignore +0 -0
  12. {memwal-0.2.1 → memwal-0.2.2.dev0}/examples/.env.example +0 -0
  13. {memwal-0.2.1 → memwal-0.2.2.dev0}/examples/.gitignore +0 -0
  14. {memwal-0.2.1 → memwal-0.2.2.dev0}/examples/async_remember_demo.py +0 -0
  15. {memwal-0.2.1 → memwal-0.2.2.dev0}/examples/interactive_demo.py +0 -0
  16. {memwal-0.2.1 → memwal-0.2.2.dev0}/examples/verify_credentials.py +0 -0
  17. {memwal-0.2.1 → memwal-0.2.2.dev0}/memwal/compatibility.py +0 -0
  18. {memwal-0.2.1 → memwal-0.2.2.dev0}/memwal/middleware.py +0 -0
  19. {memwal-0.2.1 → memwal-0.2.2.dev0}/notebooks/walrus_memory_python_sdk.ipynb +0 -0
  20. {memwal-0.2.1 → memwal-0.2.2.dev0}/run_tests.py +0 -0
  21. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/__init__.py +0 -0
  22. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_auth_rejected_message.py +0 -0
  23. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_blank_text.py +0 -0
  24. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_client.py +0 -0
  25. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_empty_namespace.py +0 -0
  26. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_env_presets.py +0 -0
  27. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_integration.py +0 -0
  28. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_list_namespaces.py +0 -0
  29. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_middleware.py +0 -0
  30. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_mock.py +0 -0
  31. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_nonfinite_validation.py +0 -0
  32. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_normalize_server_url.py +0 -0
  33. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_polling_delay.py +0 -0
  34. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_recall_weights.py +0 -0
  35. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_remember_derived_key.py +0 -0
  36. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_retry_after_http_date.py +0 -0
  37. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_signing.py +0 -0
  38. {memwal-0.2.1 → memwal-0.2.2.dev0}/tests/test_transport_retry.py +0 -0
@@ -1,5 +1,17 @@
1
1
  # memwal
2
2
 
3
+ ## 0.2.2
4
+
5
+ Unreleased.
6
+
7
+ ### Added
8
+
9
+ - `MemWal.destroy()` and `MemWalSync.destroy()` (plus `close(clear_keys=True)`) zero the client's internal `bytearray` seed copy in place, clear `_private_key_hex`, drop `_signing_key` and cached SEAL session credentials, cancel in-flight initialization tasks, and reject any subsequent or in-flight operation with `MemWalError`. `MemWal.create()` and `MemWalSync.create()` also accept 32-byte `bytes` / `bytearray` keys and copy caller `bytearray` buffers so `destroy()` leaves the caller's buffer intact. (#975)
10
+
11
+ ### Fixed
12
+
13
+ - Shield shared `_build_seal_session` and `_resolve_owner` initialization tasks with `asyncio.shield`, clear the shared task slot via a completion callback that consumes exceptions on orphaned tasks, and cancel/drain in-flight initialization tasks in `close()` so cancelling one concurrent caller does not abort sibling callers or leave unretrieved task exceptions. (#1029)
14
+
3
15
  ## 0.2.1
4
16
 
5
17
  Unreleased.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: memwal
3
- Version: 0.2.1
3
+ Version: 0.2.2.dev0
4
4
  Summary: Python SDK for Walrus Memory — Privacy-first AI memory with Ed25519 signing
5
5
  Project-URL: Homepage, https://memory.walrus.xyz
6
6
  Project-URL: Documentation, https://memory.walrus.xyz
@@ -259,6 +259,16 @@ Create a new async client.
259
259
  | `await remember_manual(opts)` | Store encrypted payload + pre-computed vector |
260
260
  | `await recall_manual(opts)` | Search with pre-computed vector |
261
261
  | `await get_public_key_hex()` | Get Ed25519 public key |
262
+ | `destroy()` | Zero internal key seed buffer, drop signing credentials/session cache, and permanently invalidate the client |
263
+ | `await close(clear_keys=False)` | Close the underlying HTTP client (pass `clear_keys=True` to also run `destroy()`) |
264
+
265
+ ### Key Lifecycle & `destroy()`
266
+
267
+ For parity with the TypeScript SDK's `destroy()`, `MemWal` and `MemWalSync` expose `destroy()` and `close(clear_keys=True)`:
268
+
269
+ - `key` may be passed as a hex string, `suiprivkey1...` string, or a 32-byte `bytes` / `bytearray`. When passed as `bytearray`, the SDK copies the bytes into its own internal buffer so `destroy()` zeroes only the client's copy and leaves the caller's `bytearray` untouched.
270
+ - Calling `destroy()` (or `await client.close(clear_keys=True)` / `sync_client.close(clear_keys=True)`) zero-fills the client's internal `bytearray` seed copy in place, clears `_private_key_hex`, drops the `_signing_key` and cached SEAL session references, cancels any in-flight initialization tasks, and causes subsequent or in-flight calls to raise `MemWalError`.
271
+ - **CPython memory semantics**: Python `str` and `bytes` objects (and PyNaCl's internal C-level `SigningKey` buffers) are immutable in CPython and cannot be overwritten in place from pure Python. `destroy()` zeroes the mutable `bytearray` copy and drops all references to immutable key objects so they become eligible for garbage collection once the caller also releases its references. To minimize plaintext `str` copies in process memory, pass a 32-byte `bytearray` to `MemWal.create()` and zero your own buffer after creation.
262
272
 
263
273
  ## Authentication
264
274
 
@@ -218,6 +218,16 @@ Create a new async client.
218
218
  | `await remember_manual(opts)` | Store encrypted payload + pre-computed vector |
219
219
  | `await recall_manual(opts)` | Search with pre-computed vector |
220
220
  | `await get_public_key_hex()` | Get Ed25519 public key |
221
+ | `destroy()` | Zero internal key seed buffer, drop signing credentials/session cache, and permanently invalidate the client |
222
+ | `await close(clear_keys=False)` | Close the underlying HTTP client (pass `clear_keys=True` to also run `destroy()`) |
223
+
224
+ ### Key Lifecycle & `destroy()`
225
+
226
+ For parity with the TypeScript SDK's `destroy()`, `MemWal` and `MemWalSync` expose `destroy()` and `close(clear_keys=True)`:
227
+
228
+ - `key` may be passed as a hex string, `suiprivkey1...` string, or a 32-byte `bytes` / `bytearray`. When passed as `bytearray`, the SDK copies the bytes into its own internal buffer so `destroy()` zeroes only the client's copy and leaves the caller's `bytearray` untouched.
229
+ - Calling `destroy()` (or `await client.close(clear_keys=True)` / `sync_client.close(clear_keys=True)`) zero-fills the client's internal `bytearray` seed copy in place, clears `_private_key_hex`, drops the `_signing_key` and cached SEAL session references, cancels any in-flight initialization tasks, and causes subsequent or in-flight calls to raise `MemWalError`.
230
+ - **CPython memory semantics**: Python `str` and `bytes` objects (and PyNaCl's internal C-level `SigningKey` buffers) are immutable in CPython and cannot be overwritten in place from pure Python. `destroy()` zeroes the mutable `bytearray` copy and drops all references to immutable key objects so they become eligible for garbage collection once the caller also releases its references. To minimize plaintext `str` copies in process memory, pass a 32-byte `bytearray` to `MemWal.create()` and zero your own buffer after creation.
221
231
 
222
232
  ## Authentication
223
233
 
@@ -128,4 +128,4 @@ __all__ = [
128
128
  "RecallManualResult",
129
129
  ]
130
130
 
131
- __version__ = "0.2.1"
131
+ __version__ = "0.2.2.dev0"
@@ -437,6 +437,12 @@ def _check_vector(vector: Any) -> None:
437
437
  ):
438
438
  raise ValueError("vector must be a non-empty list of finite numbers")
439
439
 
440
+ def _destroyed_client_error() -> "MemWalError":
441
+ return MemWalError(
442
+ "Walrus Memory client was destroyed; its keys are zeroed. Create a new one."
443
+ )
444
+
445
+
440
446
  class MemWal:
441
447
  """Async-native Walrus Memory client.
442
448
 
@@ -445,8 +451,13 @@ class MemWal:
445
451
  """
446
452
 
447
453
  def __init__(self, config: MemWalConfig) -> None:
448
- self._private_key_hex = normalize_private_key(config.key)
449
- self._signing_key = build_signing_key(self._private_key_hex)
454
+ self._destroyed: bool = False
455
+ self._private_key_hex: str = normalize_private_key(config.key)
456
+ signing_key = build_signing_key(self._private_key_hex)
457
+ self._signing_key: Optional[nacl.signing.SigningKey] = signing_key
458
+ # Mutable copy of the 32-byte seed so destroy() can zero-fill it in
459
+ # place without mutating any caller-supplied bytearray (#975).
460
+ self._private_key_bytes: bytearray = bytearray(bytes(signing_key))
450
461
  self._account_id = config.account_id
451
462
  self._server_url = normalize_server_url(config.server_url)
452
463
  self._namespace = config.namespace
@@ -473,7 +484,7 @@ class MemWal:
473
484
  @classmethod
474
485
  def create(
475
486
  cls,
476
- key: str,
487
+ key: Union[str, bytes, bytearray],
477
488
  account_id: str,
478
489
  server_url: Optional[str] = None,
479
490
  namespace: str = "default",
@@ -482,7 +493,10 @@ class MemWal:
482
493
  """Create a new Walrus Memory client instance.
483
494
 
484
495
  Args:
485
- key: Ed25519 private key hex string (the delegate key).
496
+ key: Ed25519 private key (hex string, ``suiprivkey1...``, or 32-byte
497
+ ``bytes`` / ``bytearray``). When passed as ``bytearray``, the
498
+ bytes are copied into an internal buffer so :meth:`destroy`
499
+ zeroes only the client's copy and leaves the caller's buffer intact.
486
500
  account_id: Walrus Memory account object ID on Sui.
487
501
  server_url: Server URL (default: ``https://relayer.memory.walrus.xyz``).
488
502
  namespace: Default namespace for memory isolation (default: ``"default"``).
@@ -510,8 +524,65 @@ class MemWal:
510
524
  self._client = httpx.AsyncClient(timeout=60.0)
511
525
  return self._client
512
526
 
513
- async def close(self) -> None:
514
- """Close the underlying HTTP client."""
527
+ def destroy(self) -> None:
528
+ """Wipe internal key buffers, drop signing credentials, and invalidate the client.
529
+
530
+ Provides parity with the TypeScript SDK's ``destroy()`` method (#975):
531
+
532
+ - Zeroes ``_private_key_bytes`` (the client's internal ``bytearray`` seed
533
+ copy) in place.
534
+ - Clears ``_private_key_hex``, drops the ``_signing_key`` reference, and
535
+ clears cached SEAL session tokens, server config, owner address, and
536
+ pending idempotency keys.
537
+ - Marks the client as destroyed so any subsequent or in-flight operation
538
+ raises :class:`MemWalError`.
539
+
540
+ Note on CPython memory semantics:
541
+ Python ``str`` and ``bytes`` objects (and PyNaCl's internal C-level
542
+ ``SigningKey`` buffers) are immutable in CPython and cannot be zeroed
543
+ in place from pure Python. ``destroy()`` zeroes the client's mutable
544
+ ``bytearray`` seed copy in place and drops all references to immutable
545
+ key objects so they become eligible for garbage collection once the
546
+ caller also releases its references. Callers that want to minimize
547
+ long-lived plaintext ``str`` keys in memory should pass ``key`` as a
548
+ 32-byte ``bytearray`` and zero their own buffer after ``create()``.
549
+ """
550
+ self._destroyed = True
551
+ if self._private_key_bytes:
552
+ self._private_key_bytes[:] = b"\x00" * len(self._private_key_bytes)
553
+ self._private_key_hex = ""
554
+ self._signing_key = None
555
+ self._session_cache = None
556
+ self._server_config = None
557
+ self._relayer_version_metadata = None
558
+ self._owner_address = None
559
+ self._pending_remember_keys.clear()
560
+ for attr in ("_session_build_task", "_owner_task"):
561
+ task = getattr(self, attr, None)
562
+ setattr(self, attr, None)
563
+ if task is not None and not task.done():
564
+ task.cancel()
565
+
566
+ async def close(self, clear_keys: bool = False) -> None:
567
+ """Close the underlying HTTP client and cancel any in-flight initialization tasks.
568
+
569
+ Args:
570
+ clear_keys: If ``True``, also call :meth:`destroy` to wipe internal
571
+ key buffers, drop signing credentials, and invalidate the client.
572
+ """
573
+ for attr in ("_session_build_task", "_owner_task"):
574
+ task = getattr(self, attr, None)
575
+ setattr(self, attr, None)
576
+ if task is not None and not task.done():
577
+ task.cancel()
578
+ try:
579
+ await task
580
+ except (asyncio.CancelledError, Exception):
581
+ # Ignore cancellation or transport errors from in-flight
582
+ # session tasks during teardown.
583
+ pass
584
+ if clear_keys:
585
+ self.destroy()
515
586
  if self._client is not None and not self._client.is_closed:
516
587
  await self._client.aclose()
517
588
  self._client = None
@@ -1443,7 +1514,11 @@ class MemWal:
1443
1514
  Returns:
1444
1515
  :class:`HealthResult` with status and version.
1445
1516
  """
1517
+ if self._destroyed:
1518
+ raise _destroyed_client_error()
1446
1519
  response = await self._http.get(f"{self._server_url}/health")
1520
+ if self._destroyed:
1521
+ raise _destroyed_client_error()
1447
1522
  if response.status_code != 200:
1448
1523
  raise MemWalError(f"Health check failed: {response.status_code}")
1449
1524
  data = response.json()
@@ -1462,7 +1537,8 @@ class MemWal:
1462
1537
 
1463
1538
  async def compatibility(self) -> Dict[str, Any]:
1464
1539
  """Fetch and validate the relayer compatibility contract."""
1465
-
1540
+ if self._destroyed:
1541
+ raise _destroyed_client_error()
1466
1542
  return await self._ensure_compatible_relayer()
1467
1543
 
1468
1544
  # ============================================================
@@ -1545,6 +1621,8 @@ class MemWal:
1545
1621
  Returns:
1546
1622
  Public key hex string.
1547
1623
  """
1624
+ if self._destroyed or self._signing_key is None:
1625
+ raise _destroyed_client_error()
1548
1626
  return bytes_to_hex(bytes(self._signing_key.verify_key))
1549
1627
 
1550
1628
  # ============================================================
@@ -1552,6 +1630,8 @@ class MemWal:
1552
1630
  # ============================================================
1553
1631
 
1554
1632
  async def _ensure_compatible_relayer(self) -> Dict[str, Any]:
1633
+ if self._destroyed:
1634
+ raise _destroyed_client_error()
1555
1635
  if self._relayer_version_metadata is not None:
1556
1636
  return self._relayer_version_metadata
1557
1637
 
@@ -1559,14 +1639,20 @@ class MemWal:
1559
1639
  self._compatibility_lock = asyncio.Lock()
1560
1640
 
1561
1641
  async with self._compatibility_lock:
1642
+ if self._destroyed:
1643
+ raise _destroyed_client_error()
1562
1644
  if self._relayer_version_metadata is not None:
1563
1645
  return self._relayer_version_metadata
1564
1646
 
1565
1647
  version_response = await self._http.get(f"{self._server_url}/version")
1648
+ if self._destroyed:
1649
+ raise _destroyed_client_error()
1566
1650
  if version_response.status_code == 200:
1567
1651
  metadata = version_response.json()
1568
1652
  elif version_response.status_code in (404, 405):
1569
1653
  health_response = await self._http.get(f"{self._server_url}/health")
1654
+ if self._destroyed:
1655
+ raise _destroyed_client_error()
1570
1656
  if health_response.status_code != 200:
1571
1657
  raise MemWalError(
1572
1658
  "Walrus Memory compatibility check failed: "
@@ -1684,10 +1770,16 @@ class MemWal:
1684
1770
  )
1685
1771
 
1686
1772
  async def _build_seal_session_inner(self) -> str:
1773
+ if self._destroyed or self._signing_key is None:
1774
+ raise _destroyed_client_error()
1687
1775
  cfg = await self._fetch_server_config()
1776
+ if self._destroyed or self._signing_key is None:
1777
+ raise _destroyed_client_error()
1688
1778
  await self._assert_first_package_version(
1689
1779
  cfg["network"], cfg["suiRpcUrl"], cfg["packageId"]
1690
1780
  )
1781
+ if self._destroyed or self._signing_key is None:
1782
+ raise _destroyed_client_error()
1691
1783
 
1692
1784
  session_signing_key = nacl.signing.SigningKey.generate()
1693
1785
  session_public_key = bytes(session_signing_key.verify_key)
@@ -1716,6 +1808,8 @@ class MemWal:
1716
1808
  separators=(",", ":"),
1717
1809
  )
1718
1810
  session_bytes = base64.b64encode(json_str.encode("utf-8")).decode("utf-8")
1811
+ if self._destroyed:
1812
+ raise _destroyed_client_error()
1719
1813
  self._session_cache = (
1720
1814
  session_bytes,
1721
1815
  int(time.time() * 1000) + SEAL_SESSION_TTL_MIN * 60_000 - SEAL_SESSION_SAFETY_MARGIN_MS,
@@ -1723,22 +1817,40 @@ class MemWal:
1723
1817
  return session_bytes
1724
1818
 
1725
1819
  async def _build_seal_session(self) -> str:
1820
+ if self._destroyed:
1821
+ raise _destroyed_client_error()
1726
1822
  now_ms = int(time.time() * 1000)
1727
1823
  if self._session_cache is not None:
1728
1824
  cached_bytes, expires_at_ms = self._session_cache
1729
1825
  if now_ms < expires_at_ms:
1730
1826
  return cached_bytes
1731
1827
 
1732
- if self._session_build_task is not None:
1733
- return await self._session_build_task
1828
+ task = self._session_build_task
1829
+ if task is None or task.done():
1830
+ task = asyncio.create_task(self._build_seal_session_inner())
1831
+ self._session_build_task = task
1832
+
1833
+ def _on_session_build_done(done_task: "asyncio.Task[str]") -> None:
1834
+ if self._session_build_task is done_task:
1835
+ self._session_build_task = None
1836
+ if not done_task.cancelled():
1837
+ done_task.exception()
1838
+
1839
+ task.add_done_callback(_on_session_build_done)
1734
1840
 
1735
- self._session_build_task = asyncio.create_task(self._build_seal_session_inner())
1736
1841
  try:
1737
- return await self._session_build_task
1738
- finally:
1739
- self._session_build_task = None
1842
+ result = await asyncio.shield(task)
1843
+ except asyncio.CancelledError:
1844
+ if self._destroyed:
1845
+ raise _destroyed_client_error() from None
1846
+ raise
1847
+ if self._destroyed:
1848
+ raise _destroyed_client_error()
1849
+ return result
1740
1850
 
1741
1851
  async def _resolve_owner_inner(self) -> str:
1852
+ if self._destroyed:
1853
+ raise _destroyed_client_error()
1742
1854
  # POST /api/stats authenticates with the same delegate scheme and
1743
1855
  # returns the owner the relayer resolved from our key. Same approach
1744
1856
  # as the TypeScript SDK's resolveOwner().
@@ -1748,6 +1860,8 @@ class MemWal:
1748
1860
  {"namespace": self._namespace},
1749
1861
  include_seal_session=False,
1750
1862
  )
1863
+ if self._destroyed:
1864
+ raise _destroyed_client_error()
1751
1865
  owner = data.get("owner")
1752
1866
  if not owner:
1753
1867
  raise MemWalError(
@@ -1763,17 +1877,33 @@ class MemWal:
1763
1877
  The owner-scoped read routes take the address in the path, but the
1764
1878
  client is configured with only a delegate key and account id.
1765
1879
  """
1880
+ if self._destroyed:
1881
+ raise _destroyed_client_error()
1766
1882
  if self._owner_address is not None:
1767
1883
  return self._owner_address
1768
1884
 
1769
- if self._owner_task is not None:
1770
- return await self._owner_task
1885
+ task = self._owner_task
1886
+ if task is None or task.done():
1887
+ task = asyncio.create_task(self._resolve_owner_inner())
1888
+ self._owner_task = task
1889
+
1890
+ def _on_owner_done(done_task: "asyncio.Task[str]") -> None:
1891
+ if self._owner_task is done_task:
1892
+ self._owner_task = None
1893
+ if not done_task.cancelled():
1894
+ done_task.exception()
1895
+
1896
+ task.add_done_callback(_on_owner_done)
1771
1897
 
1772
- self._owner_task = asyncio.create_task(self._resolve_owner_inner())
1773
1898
  try:
1774
- return await self._owner_task
1775
- finally:
1776
- self._owner_task = None
1899
+ result = await asyncio.shield(task)
1900
+ except asyncio.CancelledError:
1901
+ if self._destroyed:
1902
+ raise _destroyed_client_error() from None
1903
+ raise
1904
+ if self._destroyed:
1905
+ raise _destroyed_client_error()
1906
+ return result
1777
1907
 
1778
1908
  async def _signed_request(
1779
1909
  self,
@@ -1803,8 +1933,12 @@ class MemWal:
1803
1933
  - ``x-account-id``: Walrus Memory account object ID
1804
1934
  - ``Content-Type``: application/json
1805
1935
  """
1936
+ if self._destroyed or self._signing_key is None:
1937
+ raise _destroyed_client_error()
1806
1938
 
1807
1939
  await self._ensure_compatible_relayer()
1940
+ if self._destroyed or self._signing_key is None:
1941
+ raise _destroyed_client_error()
1808
1942
 
1809
1943
  method_upper = method.upper()
1810
1944
  timestamp = str(int(time.time()))
@@ -1838,6 +1972,8 @@ class MemWal:
1838
1972
  }
1839
1973
  if include_seal_session:
1840
1974
  headers["x-seal-session"] = await self._build_seal_session()
1975
+ if self._destroyed or self._signing_key is None:
1976
+ raise _destroyed_client_error()
1841
1977
 
1842
1978
  response = await self._http.request(
1843
1979
  method=method_upper,
@@ -1846,6 +1982,8 @@ class MemWal:
1846
1982
  content=None if method_upper == "GET" else body_str,
1847
1983
  timeout=timeout,
1848
1984
  )
1985
+ if self._destroyed or self._signing_key is None:
1986
+ raise _destroyed_client_error()
1849
1987
 
1850
1988
  if response.status_code not in accepted_statuses:
1851
1989
  err_text = response.text
@@ -2073,7 +2211,7 @@ class MemWalSync:
2073
2211
  @classmethod
2074
2212
  def create(
2075
2213
  cls,
2076
- key: str,
2214
+ key: Union[str, bytes, bytearray],
2077
2215
  account_id: str,
2078
2216
  server_url: Optional[str] = None,
2079
2217
  namespace: str = "default",
@@ -2278,9 +2416,18 @@ class MemWalSync:
2278
2416
  """Synchronous version of :meth:`MemWal.get_public_key_hex`."""
2279
2417
  return self._run(self._inner.get_public_key_hex())
2280
2418
 
2281
- def close(self) -> None:
2282
- """Close the underlying HTTP client."""
2283
- self._run(self._inner.close())
2419
+ def destroy(self) -> None:
2420
+ """Synchronous version of :meth:`MemWal.destroy`."""
2421
+ self._inner.destroy()
2422
+
2423
+ def close(self, clear_keys: bool = False) -> None:
2424
+ """Close the underlying HTTP client.
2425
+
2426
+ Args:
2427
+ clear_keys: If ``True``, also call :meth:`destroy` to wipe internal
2428
+ key buffers, drop signing credentials, and invalidate the client.
2429
+ """
2430
+ self._run(self._inner.close(clear_keys=clear_keys))
2284
2431
 
2285
2432
  def __enter__(self) -> "MemWalSync":
2286
2433
  return self
@@ -129,9 +129,13 @@ class MemWalMock:
129
129
  ) -> "MemWalMock":
130
130
  return cls(namespace, owner, initial_memories)
131
131
 
132
- async def close(self) -> None:
132
+ def destroy(self) -> None:
133
133
  """No-op for API parity with :class:`MemWal`."""
134
134
 
135
+ async def close(self, clear_keys: bool = False) -> None:
136
+ """No-op for API parity with :class:`MemWal`."""
137
+ del clear_keys
138
+
135
139
  async def __aenter__(self) -> "MemWalMock":
136
140
  return self
137
141
 
@@ -714,8 +718,11 @@ class MemWalMockSync:
714
718
  def clear(self, namespace: Optional[str] = None) -> int:
715
719
  return self._inner.clear(namespace)
716
720
 
717
- def close(self) -> None:
718
- self._run(self._inner.close())
721
+ def destroy(self) -> None:
722
+ self._inner.destroy()
723
+
724
+ def close(self, clear_keys: bool = False) -> None:
725
+ self._run(self._inner.close(clear_keys=clear_keys))
719
726
 
720
727
  def __enter__(self) -> "MemWalMockSync":
721
728
  return self
@@ -9,7 +9,7 @@ the Walrus Memory Rust server (TEE).
9
9
  from __future__ import annotations
10
10
 
11
11
  from dataclasses import dataclass
12
- from typing import Any, Dict, List, Optional
12
+ from typing import Any, Dict, List, Optional, Union
13
13
 
14
14
  # ============================================================
15
15
  # Config
@@ -34,8 +34,11 @@ class MemWalConfig:
34
34
  """Configuration for creating a Walrus Memory client.
35
35
 
36
36
  Attributes:
37
- key: Ed25519 private key (hex string). This is the delegate key from the
38
- Walrus Memory dashboard.
37
+ key: Ed25519 private key (hex string, ``suiprivkey1...``, or 32-byte
38
+ ``bytes`` / ``bytearray``). This is the delegate key from the
39
+ Walrus Memory dashboard. When passed as ``bytearray``, the client
40
+ copies the bytes into its own internal buffer so :meth:`MemWal.destroy`
41
+ zeroes only the client's copy.
39
42
  account_id: Walrus Memory account object ID on Sui.
40
43
  server_url: Server URL. ``None`` means omitted. Any explicit URL,
41
44
  including the hosted default, wins over ``env``.
@@ -46,7 +49,7 @@ class MemWalConfig:
46
49
  > ``DEFAULT_SERVER_URL``.
47
50
  """
48
51
 
49
- key: str
52
+ key: Union[str, bytes, bytearray]
50
53
  account_id: str
51
54
  server_url: Optional[str] = None
52
55
  namespace: str = "default"
@@ -10,7 +10,7 @@ from __future__ import annotations
10
10
  import base64
11
11
  import hashlib
12
12
  from datetime import datetime, timezone
13
- from typing import Tuple
13
+ from typing import Tuple, Union
14
14
 
15
15
  import nacl.signing
16
16
 
@@ -51,7 +51,7 @@ def hex_to_bytes(hex_str: str) -> bytes:
51
51
  return bytes.fromhex(clean)
52
52
 
53
53
 
54
- def bytes_to_hex(b: bytes) -> str:
54
+ def bytes_to_hex(b: Union[bytes, bytearray]) -> str:
55
55
  """Convert bytes to a lowercase hex string (no ``0x`` prefix).
56
56
 
57
57
  Args:
@@ -75,11 +75,14 @@ def sha256_hex(data: str) -> str:
75
75
  return hashlib.sha256(data.encode("utf-8")).hexdigest()
76
76
 
77
77
 
78
- def build_signing_key(private_key_hex: str) -> nacl.signing.SigningKey:
79
- """Build a PyNaCl ``SigningKey`` from an Ed25519 private key hex string.
78
+ def build_signing_key(
79
+ private_key: Union[str, bytes, bytearray]
80
+ ) -> nacl.signing.SigningKey:
81
+ """Build a PyNaCl ``SigningKey`` from an Ed25519 private key hex string or 32-byte buffer.
80
82
 
81
83
  Args:
82
- private_key_hex: 32-byte Ed25519 seed as hex (64 hex chars), optionally ``0x``-prefixed.
84
+ private_key: 32-byte Ed25519 seed as hex (64 hex chars, optionally ``0x``-prefixed)
85
+ or raw 32-byte ``bytes`` / ``bytearray``.
83
86
 
84
87
  Returns:
85
88
  A ``nacl.signing.SigningKey`` instance.
@@ -87,7 +90,10 @@ def build_signing_key(private_key_hex: str) -> nacl.signing.SigningKey:
87
90
  Raises:
88
91
  ValueError: If the decoded seed is not exactly 32 bytes.
89
92
  """
90
- seed_bytes = hex_to_bytes(private_key_hex)
93
+ if isinstance(private_key, (bytes, bytearray)):
94
+ seed_bytes = bytes(private_key)
95
+ else:
96
+ seed_bytes = hex_to_bytes(private_key)
91
97
  if len(seed_bytes) != 32:
92
98
  raise ValueError(
93
99
  f"Ed25519 seed must be exactly 32 bytes, got {len(seed_bytes)}"
@@ -290,12 +296,19 @@ def decode_sui_private_key(encoded: str) -> bytes:
290
296
  return seed
291
297
 
292
298
 
293
- def normalize_private_key(key: str) -> str:
294
- """Accept either a hex seed or a Sui ``suiprivkey1...`` string, return hex.
299
+ def normalize_private_key(key: Union[str, bytes, bytearray]) -> str:
300
+ """Accept a hex seed, ``suiprivkey1...`` string, or 32-byte ``bytes``/``bytearray``; return hex.
295
301
 
296
- The TypeScript SDK takes both, and the dashboard hands out the bech32 form,
297
- so hex-only input would reject a key a user reasonably expects to work.
302
+ The TypeScript SDK accepts hex, ``suiprivkey1...``, and 32-byte ``Uint8Array``.
298
303
  """
304
+ if isinstance(key, (bytes, bytearray)):
305
+ if len(key) != 32:
306
+ raise ValueError(f"Ed25519 seed must be exactly 32 bytes, got {len(key)}")
307
+ return bytes_to_hex(key)
308
+ if not isinstance(key, str):
309
+ raise TypeError(
310
+ f"key must be str, bytes, or bytearray, got {type(key).__name__}"
311
+ )
299
312
  candidate = key.strip()
300
313
  if candidate.lower().startswith("suiprivkey1"):
301
314
  return bytes_to_hex(decode_sui_private_key(candidate))
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "memwal"
7
- version = "0.2.1"
7
+ version = "0.2.2.dev0"
8
8
  description = "Python SDK for Walrus Memory — Privacy-first AI memory with Ed25519 signing"
9
9
  readme = "README.md"
10
10
  license = "MIT"
@@ -0,0 +1,661 @@
1
+ """Tests for MemWal.destroy() key-wipe (GH #975) and concurrent
2
+ initialization cancellation shielding (GH #1029)."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import asyncio
7
+ import json
8
+ from typing import Dict, List
9
+
10
+ import httpx
11
+ import pytest
12
+
13
+ from memwal import (
14
+ MemWal,
15
+ MemWalError,
16
+ MemWalMock,
17
+ MemWalMockSync,
18
+ MemWalSync,
19
+ RecallParams,
20
+ )
21
+
22
+ PRIVATE_KEY_HEX = "a" * 64
23
+ ACCOUNT_ID = "0x" + "1" * 64
24
+ PACKAGE_ID = "0x" + "2" * 64
25
+ SERVER_URL = "https://relayer.example.com"
26
+
27
+
28
+ # ============================================================
29
+ # GH #975: destroy(), close(clear_keys=True), bytes/bytearray keys
30
+ # ============================================================
31
+
32
+
33
+ class TestDestroyAndKeyWipe:
34
+ @pytest.mark.asyncio
35
+ async def test_accepts_bytes_and_bytearray_keys_and_clones_caller_buffer(self):
36
+ seed_bytes = bytes.fromhex(PRIVATE_KEY_HEX)
37
+ caller_buf = bytearray(seed_bytes)
38
+
39
+ from_hex = MemWal.create(
40
+ key=PRIVATE_KEY_HEX,
41
+ account_id=ACCOUNT_ID,
42
+ server_url=SERVER_URL,
43
+ )
44
+ from_bytes = MemWal.create(
45
+ key=seed_bytes,
46
+ account_id=ACCOUNT_ID,
47
+ server_url=SERVER_URL,
48
+ )
49
+ from_bytearray = MemWal.create(
50
+ key=caller_buf,
51
+ account_id=ACCOUNT_ID,
52
+ server_url=SERVER_URL,
53
+ )
54
+
55
+ expected_pub = await from_hex.get_public_key_hex()
56
+ assert await from_bytes.get_public_key_hex() == expected_pub
57
+ assert await from_bytearray.get_public_key_hex() == expected_pub
58
+
59
+ # Mutating caller's bytearray after create() does not affect client
60
+ caller_buf[:] = b"\xff" * 32
61
+ assert from_bytearray._private_key_bytes == bytearray(seed_bytes)
62
+ assert await from_bytearray.get_public_key_hex() == expected_pub
63
+
64
+ await from_hex.close()
65
+ await from_bytes.close()
66
+ await from_bytearray.close()
67
+
68
+ def test_rejects_wrong_length_bytes_key(self):
69
+ with pytest.raises(ValueError, match="32 bytes"):
70
+ MemWal.create(
71
+ key=b"\xaa" * 16,
72
+ account_id=ACCOUNT_ID,
73
+ server_url=SERVER_URL,
74
+ )
75
+
76
+ @pytest.mark.asyncio
77
+ async def test_destroy_zeroes_internal_buffer_without_mutating_caller_bytearray(self):
78
+ caller_buf = bytearray(bytes.fromhex(PRIVATE_KEY_HEX))
79
+ client = MemWal.create(
80
+ key=caller_buf,
81
+ account_id=ACCOUNT_ID,
82
+ server_url=SERVER_URL,
83
+ )
84
+ internal_ref = client._private_key_bytes
85
+ assert internal_ref == bytearray(bytes.fromhex(PRIVATE_KEY_HEX))
86
+
87
+ client.destroy()
88
+ # Idempotent second call
89
+ client.destroy()
90
+
91
+ # Internal bytearray was zeroed in place
92
+ assert internal_ref == bytearray(32)
93
+ assert all(b == 0 for b in client._private_key_bytes)
94
+ # Caller's bytearray was NOT mutated
95
+ assert caller_buf == bytearray(bytes.fromhex(PRIVATE_KEY_HEX))
96
+ # Credentials and caches dropped
97
+ assert client._private_key_hex == ""
98
+ assert client._signing_key is None
99
+ assert client._session_cache is None
100
+
101
+ with pytest.raises(MemWalError, match="Walrus Memory client was destroyed"):
102
+ await client.get_public_key_hex()
103
+ with pytest.raises(MemWalError, match="Walrus Memory client was destroyed"):
104
+ await client.remember("secret")
105
+ with pytest.raises(MemWalError, match="Walrus Memory client was destroyed"):
106
+ await client.recall("secret")
107
+ with pytest.raises(MemWalError, match="Walrus Memory client was destroyed"):
108
+ await client.health()
109
+ with pytest.raises(MemWalError, match="Walrus Memory client was destroyed"):
110
+ await client.compatibility()
111
+ with pytest.raises(MemWalError, match="Walrus Memory client was destroyed"):
112
+ await client.restore("default")
113
+
114
+ await client.close()
115
+
116
+ @pytest.mark.asyncio
117
+ async def test_destroy_mid_flight_aborts_pending_request(self):
118
+ config_entered = asyncio.Event()
119
+ release_config = asyncio.Event()
120
+
121
+ async def handler(request: httpx.Request) -> httpx.Response:
122
+ path = request.url.path
123
+ if path == "/version":
124
+ return httpx.Response(
125
+ 200,
126
+ json={
127
+ "apiVersion": "1.0.0",
128
+ "relayerVersion": "0.1.0",
129
+ "minSupportedSdk": {
130
+ "typescript": "0.0.0",
131
+ "python": "0.0.0",
132
+ "mcp": "0.0.0",
133
+ },
134
+ },
135
+ )
136
+ if path == "/config":
137
+ config_entered.set()
138
+ await release_config.wait()
139
+ return httpx.Response(
140
+ 200,
141
+ json={
142
+ "packageId": PACKAGE_ID,
143
+ "network": "testnet",
144
+ "suiRpcUrl": "https://rpc.testnet.sui.io",
145
+ },
146
+ )
147
+ return httpx.Response(500, text=f"unexpected {path}")
148
+
149
+ client = MemWal.create(
150
+ key=PRIVATE_KEY_HEX,
151
+ account_id=ACCOUNT_ID,
152
+ server_url=SERVER_URL,
153
+ )
154
+ client._client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
155
+
156
+ recall_task = asyncio.create_task(client.recall(RecallParams(query="test")))
157
+ await config_entered.wait()
158
+
159
+ client.destroy()
160
+ release_config.set()
161
+
162
+ with pytest.raises(MemWalError, match="Walrus Memory client was destroyed"):
163
+ _ = await recall_task
164
+ assert client._session_cache is None
165
+ await client.close()
166
+
167
+ @pytest.mark.asyncio
168
+ async def test_destroy_while_signed_request_and_health_in_flight_rejects_response(self):
169
+ recall_entered = asyncio.Event()
170
+ release_recall = asyncio.Event()
171
+ health_entered = asyncio.Event()
172
+ release_health = asyncio.Event()
173
+
174
+ async def handler(request: httpx.Request) -> httpx.Response:
175
+ path = request.url.path
176
+ if path == "/version":
177
+ return httpx.Response(
178
+ 200,
179
+ json={
180
+ "apiVersion": "1.0.0",
181
+ "relayerVersion": "0.1.0",
182
+ "minSupportedSdk": {
183
+ "typescript": "0.0.0",
184
+ "python": "0.0.0",
185
+ "mcp": "0.0.0",
186
+ },
187
+ },
188
+ )
189
+ if path == "/config":
190
+ return httpx.Response(
191
+ 200,
192
+ json={
193
+ "packageId": PACKAGE_ID,
194
+ "network": "testnet",
195
+ "suiRpcUrl": "https://rpc.testnet.sui.io",
196
+ },
197
+ )
198
+ if request.url.host == "graphql.testnet.sui.io":
199
+ return httpx.Response(200, json={"data": {"object": {"version": 1}}})
200
+ if path == "/api/recall":
201
+ recall_entered.set()
202
+ await release_recall.wait()
203
+ return httpx.Response(
204
+ 200,
205
+ json={
206
+ "results": [
207
+ {
208
+ "id": "mem-1",
209
+ "blob_id": "blob-1",
210
+ "text": "match",
211
+ "distance": 0.1,
212
+ "namespace": "default",
213
+ }
214
+ ],
215
+ "total": 1,
216
+ },
217
+ )
218
+ if path == "/health":
219
+ health_entered.set()
220
+ await release_health.wait()
221
+ return httpx.Response(
222
+ 200,
223
+ json={"status": "ok", "version": "0.1.0"},
224
+ )
225
+ return httpx.Response(500, text=f"unexpected {path}")
226
+
227
+ client = MemWal.create(
228
+ key=PRIVATE_KEY_HEX,
229
+ account_id=ACCOUNT_ID,
230
+ server_url=SERVER_URL,
231
+ )
232
+ client._client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
233
+
234
+ recall_task = asyncio.create_task(client.recall(RecallParams(query="test")))
235
+ health_task = asyncio.create_task(client.health())
236
+ await recall_entered.wait()
237
+ await health_entered.wait()
238
+
239
+ client.destroy()
240
+ release_recall.set()
241
+ release_health.set()
242
+
243
+ with pytest.raises(MemWalError, match="Walrus Memory client was destroyed"):
244
+ _ = await recall_task
245
+ with pytest.raises(MemWalError, match="Walrus Memory client was destroyed"):
246
+ _ = await health_task
247
+ await client.close()
248
+
249
+ @pytest.mark.asyncio
250
+ async def test_close_with_clear_keys_destroys_client(self):
251
+ client = MemWal.create(
252
+ key=PRIVATE_KEY_HEX,
253
+ account_id=ACCOUNT_ID,
254
+ server_url=SERVER_URL,
255
+ )
256
+ # Default close(clear_keys=False) preserves keys
257
+ await client.close()
258
+ assert client._destroyed is False
259
+ assert await client.get_public_key_hex() != ""
260
+
261
+ # close(clear_keys=True) wipes keys and marks destroyed
262
+ await client.close(clear_keys=True)
263
+ assert client._destroyed is True
264
+ assert all(b == 0 for b in client._private_key_bytes)
265
+ with pytest.raises(MemWalError, match="Walrus Memory client was destroyed"):
266
+ await client.get_public_key_hex()
267
+
268
+ def test_sync_client_destroy_and_close_clear_keys(self):
269
+ caller_buf = bytearray(bytes.fromhex(PRIVATE_KEY_HEX))
270
+ sync_client = MemWalSync.create(
271
+ key=caller_buf,
272
+ account_id=ACCOUNT_ID,
273
+ server_url=SERVER_URL,
274
+ )
275
+ assert sync_client.get_public_key_hex() != ""
276
+
277
+ sync_client.close(clear_keys=True)
278
+ assert caller_buf == bytearray(bytes.fromhex(PRIVATE_KEY_HEX))
279
+ with pytest.raises(MemWalError, match="Walrus Memory client was destroyed"):
280
+ sync_client.get_public_key_hex()
281
+
282
+ @pytest.mark.asyncio
283
+ async def test_mock_parity_destroy_and_close_clear_keys(self):
284
+ mock = MemWalMock.create()
285
+ mock.destroy()
286
+ await mock.close(clear_keys=True)
287
+
288
+ sync_mock = MemWalMockSync.create()
289
+ sync_mock.destroy()
290
+ sync_mock.close(clear_keys=True)
291
+
292
+
293
+ # ============================================================
294
+ # GH #1029: Concurrent initialization cancellation shielding
295
+ # ============================================================
296
+
297
+
298
+ class Gate:
299
+ """Deterministic barrier when N coroutines enter, plus explicit release."""
300
+
301
+ def __init__(self, parties: int = 1) -> None:
302
+ self._parties = parties
303
+ self._count = 0
304
+ self.entered = asyncio.Event()
305
+ self.release_event = asyncio.Event()
306
+
307
+ async def wait(self) -> None:
308
+ self._count += 1
309
+ if self._count >= self._parties:
310
+ self.entered.set()
311
+ await self.release_event.wait()
312
+
313
+ def release(self) -> None:
314
+ self.release_event.set()
315
+
316
+
317
+ class FakeTransport:
318
+ """Counted transport for /version, /config, GraphQL package version, and /api/recall."""
319
+
320
+ def __init__(
321
+ self,
322
+ *,
323
+ config_gate: Gate | None = None,
324
+ fail_config_times: int = 0,
325
+ ) -> None:
326
+ self.config_gate = config_gate
327
+ self.fail_config_times = fail_config_times
328
+ self.counts: Dict[str, int] = {
329
+ "version": 0,
330
+ "config": 0,
331
+ "graphql": 0,
332
+ "recall": 0,
333
+ "stats": 0,
334
+ "restore": 0,
335
+ }
336
+ self.recall_headers: List[Dict[str, str]] = []
337
+
338
+ async def handle(self, request: httpx.Request) -> httpx.Response:
339
+ path = request.url.path
340
+ if path == "/version":
341
+ self.counts["version"] += 1
342
+ return httpx.Response(
343
+ 200,
344
+ json={
345
+ "apiVersion": "1.0.0",
346
+ "relayerVersion": "0.1.0",
347
+ "minSupportedSdk": {
348
+ "typescript": "0.0.0",
349
+ "python": "0.0.0",
350
+ "mcp": "0.0.0",
351
+ },
352
+ },
353
+ )
354
+
355
+ if path == "/config":
356
+ self.counts["config"] += 1
357
+ if self.config_gate is not None:
358
+ await self.config_gate.wait()
359
+ if self.counts["config"] <= self.fail_config_times:
360
+ return httpx.Response(503, text="config temporary failure")
361
+ return httpx.Response(
362
+ 200,
363
+ json={
364
+ "packageId": PACKAGE_ID,
365
+ "network": "testnet",
366
+ "suiRpcUrl": "https://rpc.testnet.sui.io:443",
367
+ },
368
+ )
369
+
370
+ if request.url.host == "graphql.testnet.sui.io":
371
+ self.counts["graphql"] += 1
372
+ return httpx.Response(200, json={"data": {"object": {"version": 1}}})
373
+
374
+ if path == "/api/recall":
375
+ self.counts["recall"] += 1
376
+ self.recall_headers.append(dict(request.headers))
377
+ payload = json.loads(request.content.decode("utf-8"))
378
+ return httpx.Response(
379
+ 200,
380
+ json={
381
+ "results": [
382
+ {
383
+ "id": f"mem-{self.counts['recall']}",
384
+ "blob_id": f"blob-{self.counts['recall']}",
385
+ "text": f"match for {payload['query']}",
386
+ "distance": 0.1,
387
+ "namespace": payload["namespace"],
388
+ }
389
+ ],
390
+ "total": 1,
391
+ },
392
+ )
393
+
394
+ if path == "/api/stats":
395
+ self.counts["stats"] += 1
396
+ if self.config_gate is not None:
397
+ await self.config_gate.wait()
398
+ return httpx.Response(200, json={"owner": "0x" + "9" * 64})
399
+
400
+ if path.startswith("/v1/owners/") and path.endswith("/namespaces"):
401
+ self.counts["restore"] += 1
402
+ return httpx.Response(
403
+ 200,
404
+ json={"namespaces": [], "has_more": False, "snapshot_version": 2},
405
+ )
406
+
407
+ return httpx.Response(500, text=f"unexpected request {request.method} {request.url}")
408
+
409
+
410
+ def make_client(transport: FakeTransport) -> MemWal:
411
+ client = MemWal.create(
412
+ key=PRIVATE_KEY_HEX,
413
+ account_id=ACCOUNT_ID,
414
+ server_url=SERVER_URL,
415
+ )
416
+ client._client = httpx.AsyncClient(transport=httpx.MockTransport(transport.handle))
417
+ return client
418
+
419
+
420
+ async def wait_until(predicate, *, turns: int = 20) -> None:
421
+ for _ in range(turns):
422
+ if predicate():
423
+ return
424
+ await asyncio.sleep(0)
425
+ assert predicate()
426
+
427
+
428
+ @pytest.mark.asyncio
429
+ async def test_normal_concurrent_and_cached_retry() -> None:
430
+ gate = Gate()
431
+ transport = FakeTransport(config_gate=gate)
432
+ client = make_client(transport)
433
+
434
+ first = asyncio.create_task(client.recall(RecallParams(query="first")))
435
+ await gate.entered.wait()
436
+
437
+ second = asyncio.create_task(client.recall(RecallParams(query="second")))
438
+ await wait_until(lambda: client._session_build_task is not None)
439
+ gate.release()
440
+
441
+ r1, r2 = await asyncio.gather(first, second)
442
+ r3 = await client.recall(RecallParams(query="third"))
443
+
444
+ assert [r1.total, r2.total, r3.total] == [1, 1, 1]
445
+ assert transport.counts == {
446
+ "version": 1,
447
+ "config": 1,
448
+ "graphql": 1,
449
+ "recall": 3,
450
+ "stats": 0,
451
+ "restore": 0,
452
+ }
453
+ sessions = {h["x-seal-session"] for h in transport.recall_headers}
454
+ assert len(sessions) == 1
455
+
456
+ await client.close()
457
+
458
+
459
+ @pytest.mark.asyncio
460
+ async def test_cancel_creator_preserves_waiter_and_late_join() -> None:
461
+ gate = Gate()
462
+ transport = FakeTransport(config_gate=gate)
463
+ client = make_client(transport)
464
+
465
+ creator = asyncio.create_task(client.recall(RecallParams(query="creator")))
466
+ await gate.entered.wait()
467
+
468
+ waiter = asyncio.create_task(client.recall(RecallParams(query="waiter")))
469
+ await asyncio.sleep(0)
470
+
471
+ creator.cancel()
472
+ with pytest.raises(asyncio.CancelledError):
473
+ _ = await creator
474
+
475
+ # Join after creator cancelled while the shared session build is still in flight.
476
+ late = asyncio.create_task(client.recall(RecallParams(query="late")))
477
+ await asyncio.sleep(0)
478
+
479
+ gate.release()
480
+ waiter_result, late_result = await asyncio.gather(waiter, late)
481
+ retry_result = await client.recall(RecallParams(query="retry"))
482
+
483
+ assert waiter_result.results[0].text == "match for waiter"
484
+ assert late_result.results[0].text == "match for late"
485
+ assert retry_result.results[0].text == "match for retry"
486
+ assert transport.counts["config"] == 1
487
+ assert transport.counts["graphql"] == 1
488
+ assert transport.counts["recall"] == 3
489
+
490
+ await client.close()
491
+
492
+
493
+ @pytest.mark.asyncio
494
+ async def test_cancel_waiter_preserves_creator_and_late_join() -> None:
495
+ gate = Gate()
496
+ transport = FakeTransport(config_gate=gate)
497
+ client = make_client(transport)
498
+
499
+ creator = asyncio.create_task(client.recall(RecallParams(query="creator")))
500
+ await gate.entered.wait()
501
+
502
+ waiter = asyncio.create_task(client.recall(RecallParams(query="waiter")))
503
+ await asyncio.sleep(0)
504
+
505
+ waiter.cancel()
506
+ with pytest.raises(asyncio.CancelledError):
507
+ _ = await waiter
508
+
509
+ late = asyncio.create_task(client.recall(RecallParams(query="late")))
510
+ await asyncio.sleep(0)
511
+
512
+ gate.release()
513
+ creator_result, late_result = await asyncio.gather(creator, late)
514
+ assert creator_result.results[0].text == "match for creator"
515
+ assert late_result.results[0].text == "match for late"
516
+ assert transport.counts["config"] == 1
517
+ assert transport.counts["graphql"] == 1
518
+ assert transport.counts["recall"] == 2
519
+
520
+ await client.close()
521
+
522
+
523
+ @pytest.mark.asyncio
524
+ async def test_real_initialization_failure_reaches_both_then_retry_succeeds() -> None:
525
+ gate = Gate()
526
+ transport = FakeTransport(config_gate=gate, fail_config_times=1)
527
+ client = make_client(transport)
528
+
529
+ first = asyncio.create_task(client.recall(RecallParams(query="first")))
530
+ await gate.entered.wait()
531
+
532
+ second = asyncio.create_task(client.recall(RecallParams(query="second")))
533
+ await asyncio.sleep(0)
534
+ gate.release()
535
+
536
+ outcomes = await asyncio.gather(first, second, return_exceptions=True)
537
+ assert all(isinstance(outcome, MemWalError) for outcome in outcomes)
538
+ assert transport.counts["config"] == 1
539
+ assert transport.counts["recall"] == 0
540
+ assert client._session_build_task is None
541
+
542
+ recovered = await client.recall(RecallParams(query="recovered"))
543
+ assert recovered.results[0].text == "match for recovered"
544
+ assert transport.counts["config"] == 2
545
+ assert transport.counts["graphql"] == 1
546
+ assert transport.counts["recall"] == 1
547
+
548
+ await client.close()
549
+
550
+
551
+ @pytest.mark.asyncio
552
+ async def test_all_cancelled_success_populates_cache() -> None:
553
+ gate = Gate()
554
+ transport = FakeTransport(config_gate=gate)
555
+ client = make_client(transport)
556
+
557
+ first = asyncio.create_task(client.recall(RecallParams(query="first")))
558
+ await gate.entered.wait()
559
+ second = asyncio.create_task(client.recall(RecallParams(query="second")))
560
+ await asyncio.sleep(0)
561
+
562
+ first.cancel()
563
+ second.cancel()
564
+ for task in (first, second):
565
+ with pytest.raises(asyncio.CancelledError):
566
+ _ = await task
567
+
568
+ gate.release()
569
+ await wait_until(lambda: client._session_build_task is None)
570
+
571
+ later = await client.recall(RecallParams(query="later"))
572
+ assert later.results[0].text == "match for later"
573
+ assert transport.counts["config"] == 1
574
+ assert transport.counts["graphql"] == 1
575
+ assert transport.counts["recall"] == 1
576
+
577
+ await client.close()
578
+
579
+
580
+ @pytest.mark.asyncio
581
+ async def test_all_cancelled_failure_has_no_unretrieved_exception() -> None:
582
+ loop = asyncio.get_running_loop()
583
+ unhandled: List[Dict[str, object]] = []
584
+ previous_handler = loop.get_exception_handler()
585
+
586
+ def handler(_loop: asyncio.AbstractEventLoop, context: Dict[str, object]) -> None:
587
+ unhandled.append(context)
588
+
589
+ loop.set_exception_handler(handler)
590
+ try:
591
+ gate = Gate()
592
+ transport = FakeTransport(config_gate=gate, fail_config_times=1)
593
+ client = make_client(transport)
594
+
595
+ first = asyncio.create_task(client.recall(RecallParams(query="first")))
596
+ await gate.entered.wait()
597
+ first.cancel()
598
+ with pytest.raises(asyncio.CancelledError):
599
+ _ = await first
600
+
601
+ gate.release()
602
+ await wait_until(lambda: client._session_build_task is None)
603
+
604
+ later = await client.recall(RecallParams(query="later"))
605
+ assert later.results[0].text == "match for later"
606
+ assert transport.counts["config"] == 2
607
+ await client.close()
608
+ finally:
609
+ loop.set_exception_handler(previous_handler)
610
+
611
+ assert unhandled == []
612
+
613
+
614
+ @pytest.mark.asyncio
615
+ async def test_close_cancels_and_drains_orphaned_initialization() -> None:
616
+ gate = Gate()
617
+ transport = FakeTransport(config_gate=gate)
618
+ client = make_client(transport)
619
+
620
+ first = asyncio.create_task(client.recall(RecallParams(query="first")))
621
+ await gate.entered.wait()
622
+ build_task = client._session_build_task
623
+ assert build_task is not None
624
+
625
+ first.cancel()
626
+ with pytest.raises(asyncio.CancelledError):
627
+ _ = await first
628
+
629
+ await client.close()
630
+ assert build_task.cancelled()
631
+ assert client._session_build_task is None
632
+ assert client._session_cache is None
633
+
634
+
635
+ @pytest.mark.asyncio
636
+ async def test_cancel_owner_resolution_creator_preserves_waiter() -> None:
637
+ gate = Gate()
638
+ transport = FakeTransport(config_gate=gate)
639
+ client = make_client(transport)
640
+
641
+ creator = asyncio.create_task(client.list_namespaces())
642
+ await gate.entered.wait()
643
+
644
+ waiter = asyncio.create_task(client.list_namespaces())
645
+ await asyncio.sleep(0)
646
+
647
+ creator.cancel()
648
+ with pytest.raises(asyncio.CancelledError):
649
+ _ = await creator
650
+
651
+ late = asyncio.create_task(client.list_namespaces())
652
+ await asyncio.sleep(0)
653
+
654
+ gate.release()
655
+ waiter_res, late_res = await asyncio.gather(waiter, late)
656
+ assert waiter_res.namespaces == []
657
+ assert late_res.namespaces == []
658
+ assert transport.counts["stats"] == 1
659
+ assert transport.counts["restore"] == 2
660
+
661
+ await client.close()
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes