memwal 0.1.11.dev6__tar.gz → 0.1.11.dev8__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 (33) hide show
  1. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/PKG-INFO +1 -1
  2. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/memwal/__init__.py +1 -1
  3. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/memwal/client.py +87 -11
  4. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/memwal/middleware.py +50 -15
  5. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/pyproject.toml +1 -1
  6. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/test_middleware.py +129 -7
  7. memwal-0.1.11.dev8/tests/test_retry_after_http_date.py +160 -0
  8. memwal-0.1.11.dev8/tests/test_transport_retry.py +195 -0
  9. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/.gitignore +0 -0
  10. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/CHANGELOG.md +0 -0
  11. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/README.md +0 -0
  12. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/examples/.env.example +0 -0
  13. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/examples/.gitignore +0 -0
  14. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/examples/async_remember_demo.py +0 -0
  15. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/examples/interactive_demo.py +0 -0
  16. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/examples/verify_credentials.py +0 -0
  17. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/memwal/compatibility.py +0 -0
  18. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/memwal/mock.py +0 -0
  19. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/memwal/types.py +0 -0
  20. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/memwal/utils.py +0 -0
  21. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/notebooks/walrus_memory_python_sdk.ipynb +0 -0
  22. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/run_tests.py +0 -0
  23. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/__init__.py +0 -0
  24. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/test_auth_rejected_message.py +0 -0
  25. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/test_blank_text.py +0 -0
  26. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/test_client.py +0 -0
  27. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/test_env_presets.py +0 -0
  28. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/test_integration.py +0 -0
  29. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/test_list_namespaces.py +0 -0
  30. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/test_mock.py +0 -0
  31. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/test_normalize_server_url.py +0 -0
  32. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/test_polling_delay.py +0 -0
  33. {memwal-0.1.11.dev6 → memwal-0.1.11.dev8}/tests/test_signing.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: memwal
3
- Version: 0.1.11.dev6
3
+ Version: 0.1.11.dev8
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
@@ -128,4 +128,4 @@ __all__ = [
128
128
  "RecallManualResult",
129
129
  ]
130
130
 
131
- __version__ = "0.1.11.dev6"
131
+ __version__ = "0.1.11.dev8"
@@ -29,11 +29,13 @@ import asyncio
29
29
  import base64
30
30
  import json
31
31
  import logging
32
+ import math
32
33
  import random
33
34
  import re
34
35
  import time
35
36
  import uuid
36
37
  from datetime import datetime, timezone
38
+ from email.utils import parsedate_to_datetime
37
39
  from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, TypeVar, Union
38
40
  from urllib.parse import ParseResult, urlencode, urlparse
39
41
 
@@ -199,18 +201,51 @@ def _positive_retry_ms(raw: object) -> Optional[int]:
199
201
  seconds = float(raw) # type: ignore[arg-type]
200
202
  except (TypeError, ValueError):
201
203
  return None
202
- if seconds <= 0 or seconds != seconds or seconds == float("inf"):
204
+ if seconds <= 0 or math.isnan(seconds) or math.isinf(seconds):
205
+ return None
206
+ try:
207
+ ms = seconds * 1000.0
208
+ if math.isinf(ms) or math.isnan(ms):
209
+ return None
210
+ return int(ms)
211
+ except (OverflowError, ValueError):
212
+ return None
213
+
214
+
215
+ def _http_date_retry_ms(raw: object) -> Optional[int]:
216
+ if not isinstance(raw, str):
217
+ return None
218
+ try:
219
+ dt = parsedate_to_datetime(raw)
220
+ if dt is None:
221
+ return None
222
+ if dt.tzinfo is None:
223
+ dt = dt.replace(tzinfo=timezone.utc)
224
+ target_ts = dt.timestamp()
225
+ # HTTP-date is absolute wall-clock time, so we deliberately compare against
226
+ # time.time() rather than monotonic _now_ms().
227
+ now_ts = time.time()
228
+ diff = target_ts - now_ts
229
+ if diff <= 0 or math.isnan(diff) or math.isinf(diff):
230
+ return None
231
+ ms = diff * 1000.0
232
+ if math.isinf(ms) or math.isnan(ms):
233
+ return None
234
+ return int(ms)
235
+ except Exception:
203
236
  return None
204
- return int(seconds * 1000)
205
237
 
206
238
 
207
239
  def _retry_after_ms(err: "_HttpStatusError") -> int:
208
- """Retry-After seconds, else JSON ``retry_after_seconds``, in ms."""
240
+ """Retry-After seconds or HTTP-date, else JSON ``retry_after_seconds``, in ms."""
209
241
 
210
242
  if err.retry_after:
211
243
  header_ms = _positive_retry_ms(err.retry_after)
212
244
  if header_ms is not None:
213
245
  return header_ms
246
+ date_ms = _http_date_retry_ms(err.retry_after)
247
+ if date_ms is not None:
248
+ return date_ms
214
249
  if not err.body:
215
250
  return 0
216
251
  try:
@@ -463,6 +498,10 @@ class MemWal:
463
498
  ) -> RememberAcceptedResult:
464
499
  return await self.remember(text, namespace, idempotency_key)
465
500
 
501
+ def _poll_deadline_s(self, deadline_ms: int) -> float:
502
+ remaining_s = (deadline_ms - _now_ms()) / 1000.0
503
+ return max(0.001, min(60.0, remaining_s))
504
+
466
505
  async def wait_for_remember_job(
467
506
  self,
468
507
  job_id: str,
@@ -498,12 +537,25 @@ class MemWal:
498
537
  retry_after_ms = 0
499
538
 
500
539
  try:
501
- data = await self._signed_request(
502
- "GET",
503
- f"/api/remember/{job_id}",
504
- {},
505
- accepted_statuses=(200, 404),
506
- )
540
+ poll_timeout = self._poll_deadline_s(deadline_ms)
541
+ try:
542
+ data = await self._signed_request(
543
+ "GET",
544
+ f"/api/remember/{job_id}",
545
+ {},
546
+ accepted_statuses=(200, 404),
547
+ timeout=poll_timeout,
548
+ )
549
+ except TypeError as err:
550
+ if "unexpected keyword argument" in str(err) and "timeout" in str(err):
551
+ data = await self._signed_request(
552
+ "GET",
553
+ f"/api/remember/{job_id}",
554
+ {},
555
+ accepted_statuses=(200, 404),
556
+ )
557
+ else:
558
+ raise
507
559
  except _HttpStatusError as err:
508
560
  if _is_transient_polling_status(err.status):
509
561
  last_refusal = err.status
@@ -512,6 +564,11 @@ class MemWal:
512
564
  retry_after_ms = _clamped_retry_after_ms(err, deadline_ms)
513
565
  continue
514
566
  raise
567
+ except httpx.TransportError:
568
+ # Network drops retry as status 0; continuing next loop iteration
569
+ # increments `attempt` so exponential backoff delay continues to grow.
570
+ last_refusal = 0
571
+ continue
515
572
 
516
573
  status_str = data.get("status")
517
574
  if status_str is None or status_str == "not_found":
@@ -632,7 +689,9 @@ class MemWal:
632
689
  return await self.remember_bulk_async(items)
633
690
 
634
691
  async def get_remember_bulk_status(
635
- self, job_ids: Sequence[str]
692
+ self,
693
+ job_ids: Sequence[str],
694
+ timeout: Optional[float] = None,
636
695
  ) -> RememberBulkStatusResult:
637
696
  """Poll the bulk-status endpoint for a batch of job_ids.
638
697
 
@@ -644,6 +703,7 @@ class MemWal:
644
703
  "POST",
645
704
  "/api/remember/bulk/status",
646
705
  {"job_ids": list(job_ids)},
706
+ timeout=timeout,
647
707
  )
648
708
  return RememberBulkStatusResult(
649
709
  results=[
@@ -726,7 +786,14 @@ class MemWal:
726
786
  retry_after_ms = 0
727
787
 
728
788
  try:
729
- batch = await self.get_remember_bulk_status(pending)
789
+ poll_timeout = self._poll_deadline_s(deadline_ms)
790
+ try:
791
+ batch = await self.get_remember_bulk_status(pending, timeout=poll_timeout)
792
+ except TypeError as err:
793
+ if "unexpected keyword argument" in str(err) and "timeout" in str(err):
794
+ batch = await self.get_remember_bulk_status(pending)
795
+ else:
796
+ raise
730
797
  except _HttpStatusError as err:
731
798
  if _is_transient_polling_status(err.status):
732
799
  last_refusal = err.status
@@ -735,6 +802,11 @@ class MemWal:
735
802
  retry_after_ms = _clamped_retry_after_ms(err, deadline_ms)
736
803
  continue
737
804
  raise
805
+ except httpx.TransportError:
806
+ # Network drops retry as status 0; continuing next loop iteration
807
+ # increments `attempt` so exponential backoff delay continues to grow.
808
+ last_refusal = 0
809
+ continue
738
810
  answered = True
739
811
 
740
812
  terminal_ids: Set[str] = {
@@ -757,6 +829,8 @@ class MemWal:
757
829
  retry_after = f"{wait_ms / 1000:g}" if wait_ms else None
758
830
  raise MemWalRateLimited(job_ids, retry_after) from err
759
831
  # Any other probe failure is not evidence about the jobs.
832
+ except httpx.TransportError:
833
+ pass
760
834
  else:
761
835
  answered = True
762
836
  settled = {
@@ -1515,6 +1589,7 @@ class MemWal:
1515
1589
  body: Dict[str, Any],
1516
1590
  accepted_statuses: tuple = (200,),
1517
1591
  include_seal_session: bool = True,
1592
+ timeout: Optional[float] = None,
1518
1593
  ) -> Dict[str, Any]:
1519
1594
  """Make a signed request to the server.
1520
1595
 
@@ -1575,6 +1650,7 @@ class MemWal:
1575
1650
  url=url,
1576
1651
  headers=headers,
1577
1652
  content=None if method_upper == "GET" else body_str,
1653
+ timeout=timeout,
1578
1654
  )
1579
1655
 
1580
1656
  if response.status_code not in accepted_statuses:
@@ -188,7 +188,7 @@ class _PendingSaves:
188
188
 
189
189
  try:
190
190
  asyncio.run(_drain(list(self._tasks)))
191
- except RuntimeError:
191
+ except (RuntimeError, ValueError):
192
192
  logger.warning(
193
193
  "Walrus Memory flush_sync() could not await %d pending "
194
194
  "task(s) bound to a different event loop (mixing sync "
@@ -366,12 +366,32 @@ def with_memwal_langchain(
366
366
  last_human_idx = i
367
367
  break
368
368
 
369
- guard_msg = SystemMessage(content=_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION)
370
369
  memory_msg = HumanMessage(content=memory_context)
371
370
  if last_human_idx > 0:
372
- result[last_human_idx:last_human_idx] = [guard_msg, memory_msg]
371
+ result.insert(last_human_idx, memory_msg)
373
372
  else:
374
- result[0:0] = [guard_msg, memory_msg]
373
+ result.insert(0, memory_msg)
374
+
375
+ if result and isinstance(result[0], SystemMessage):
376
+ existing_content = result[0].content
377
+ if isinstance(existing_content, list):
378
+ new_content = list(existing_content) + [
379
+ {"type": "text", "text": f"\n\n{_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION}"}
380
+ ]
381
+ else:
382
+ text_str = str(existing_content or "")
383
+ new_content = f"{text_str}\n\n{_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION}".strip()
384
+ if hasattr(result[0], "model_copy"):
385
+ result[0] = result[0].model_copy(update={"content": new_content})
386
+ else:
387
+ result[0] = SystemMessage(
388
+ content=new_content,
389
+ additional_kwargs=getattr(result[0], "additional_kwargs", {}),
390
+ id=getattr(result[0], "id", None),
391
+ name=getattr(result[0], "name", None),
392
+ )
393
+ else:
394
+ result.insert(0, SystemMessage(content=_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION))
375
395
 
376
396
  return result
377
397
  except Exception as e:
@@ -625,21 +645,36 @@ def _inject_openai_memory(
625
645
  messages: List[Dict[str, Any]],
626
646
  memory_context: str,
627
647
  ) -> List[Dict[str, Any]]:
628
- """Insert a fixed system guard plus memory content in a user message."""
648
+ """Insert a fixed system guard at index 0 and memory content in a user message."""
649
+ result = [dict(m) if isinstance(m, dict) else m for m in messages]
629
650
  last_user_idx = -1
630
- for i in range(len(messages) - 1, -1, -1):
631
- if isinstance(messages[i], dict) and messages[i].get("role") == "user":
651
+ for i in range(len(result) - 1, -1, -1):
652
+ if isinstance(result[i], dict) and result[i].get("role") == "user":
632
653
  last_user_idx = i
633
654
  break
634
655
 
635
- memory_messages: List[Dict[str, Any]] = [
636
- {"role": "system", "content": _UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION},
637
- {"role": "user", "content": memory_context},
638
- ]
639
-
656
+ memory_message = {"role": "user", "content": memory_context}
640
657
  if last_user_idx > 0:
641
- messages[last_user_idx:last_user_idx] = memory_messages
658
+ result.insert(last_user_idx, memory_message)
659
+ else:
660
+ result.insert(0, memory_message)
661
+
662
+ if result and isinstance(result[0], dict) and result[0].get("role") == "system":
663
+ existing_content = result[0].get("content")
664
+ if isinstance(existing_content, list):
665
+ result[0] = {
666
+ **result[0],
667
+ "content": existing_content + [
668
+ {"type": "text", "text": f"\n\n{_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION}"}
669
+ ],
670
+ }
671
+ else:
672
+ text_str = str(existing_content or "")
673
+ result[0] = {
674
+ **result[0],
675
+ "content": f"{text_str}\n\n{_UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION}".strip(),
676
+ }
642
677
  else:
643
- messages[0:0] = memory_messages
678
+ result.insert(0, {"role": "system", "content": _UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION})
644
679
 
645
- return messages
680
+ return result
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "memwal"
7
- version = "0.1.11.dev6"
7
+ version = "0.1.11.dev8"
8
8
  description = "Python SDK for Walrus Memory — Privacy-first AI memory with Ed25519 signing"
9
9
  readme = "README.md"
10
10
  license = "MIT"
@@ -240,27 +240,53 @@ class TestFormatMemories:
240
240
  class TestInjectOpenAIMemory:
241
241
  """Tests for _inject_openai_memory."""
242
242
 
243
- def test_inserts_before_last_user_message(self) -> None:
243
+ def test_folds_guard_into_system_message_and_inserts_memory_before_user_message(self) -> None:
244
244
  messages = [
245
245
  {"role": "system", "content": "You are helpful."},
246
246
  {"role": "user", "content": "question"},
247
247
  ]
248
248
  result = _inject_openai_memory(messages, "memory context")
249
- # Should be: existing-system | fixed-guard | memory-user | user
249
+ # Should be: existing-system-with-guard | memory-user | user
250
+ assert len(result) == 3
250
251
  assert result[0]["role"] == "system"
251
- assert result[1]["role"] == "system"
252
- assert "untrusted data" in result[1]["content"]
253
- assert result[2] == {"role": "user", "content": "memory context"}
254
- assert result[3]["role"] == "user"
252
+ assert "You are helpful." in result[0]["content"]
253
+ assert "untrusted data" in result[0]["content"]
254
+ assert result[1] == {"role": "user", "content": "memory context"}
255
+ assert result[2]["role"] == "user"
256
+ assert result[2]["content"] == "question"
257
+
258
+ def test_multi_turn_never_injects_system_message_after_index_0(self) -> None:
259
+ messages = [
260
+ {"role": "system", "content": "You are helpful."},
261
+ {"role": "user", "content": "first"},
262
+ {"role": "assistant", "content": "reply"},
263
+ {"role": "user", "content": "second"},
264
+ ]
265
+ result = _inject_openai_memory(messages, "recalled fact")
266
+ # System message only at index 0
267
+ assert result[0]["role"] == "system"
268
+ assert "You are helpful." in result[0]["content"]
269
+ assert "untrusted data" in result[0]["content"]
270
+
271
+ # No system message after index 0
272
+ assert not any(m.get("role") == "system" for m in result[1:])
273
+
274
+ # Memory is injected as user message right before "second"
275
+ assert result[1] == {"role": "user", "content": "first"}
276
+ assert result[2] == {"role": "assistant", "content": "reply"}
277
+ assert result[3] == {"role": "user", "content": "recalled fact"}
278
+ assert result[4] == {"role": "user", "content": "second"}
255
279
 
256
280
  def test_inserts_at_start_when_only_user_message(self) -> None:
257
281
  messages = [{"role": "user", "content": "question"}]
258
282
  result = _inject_openai_memory(messages, "context")
259
- # last_user_idx == 0, so prepend guard + memory-data messages
283
+ # When no system message exists, prepends guard system message at index 0
260
284
  assert result[0]["role"] == "system"
285
+ assert "untrusted data" in result[0]["content"]
261
286
  assert result[1]["role"] == "user"
262
287
  assert result[1]["content"] == "context"
263
288
  assert result[2]["role"] == "user"
289
+ assert result[2]["content"] == "question"
264
290
 
265
291
  def test_does_not_mutate_original(self) -> None:
266
292
  original = [{"role": "user", "content": "q"}]
@@ -335,6 +361,102 @@ class TestWithMemWalLangChain:
335
361
  assert "User loves coffee" in human_msgs[0].content
336
362
  assert "BEGIN_UNTRUSTED_WALRUS_MEMORY_" in human_msgs[0].content
337
363
 
364
+ @respx.mock
365
+ async def test_multi_turn_never_injects_system_message_after_index_0(self) -> None:
366
+ """In a multi-turn conversation, SystemMessage stays strictly at index 0."""
367
+ _mock_seal_session_prereqs()
368
+ from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
369
+ from langchain_core.outputs import ChatGeneration, ChatResult
370
+
371
+ llm = self._make_llm()
372
+ captured: list = []
373
+
374
+ async def _capture_agenerate(messages_batch, *a, **kw):
375
+ captured.extend(messages_batch)
376
+ return ChatResult(
377
+ generations=[ChatGeneration(message=AIMessage(content="ok"))]
378
+ )
379
+
380
+ llm._agenerate = _capture_agenerate
381
+
382
+ respx.post(_RECALL_URL).mock(
383
+ return_value=_mock_recall([
384
+ {"blob_id": "b1", "text": "User loves coffee", "distance": 0.05}
385
+ ])
386
+ )
387
+ respx.post(_ANALYZE_URL).mock(return_value=_mock_analyze())
388
+
389
+ smart_llm = with_memwal_langchain(
390
+ llm, key=_KEY_HEX, account_id=_ACCOUNT_ID, server_url=_SERVER, auto_save=False
391
+ )
392
+ messages = [
393
+ SystemMessage(content="You are helpful."),
394
+ HumanMessage(content="I like coffee"),
395
+ AIMessage(content="Noted!"),
396
+ HumanMessage(content="What do I like?"),
397
+ ]
398
+ await smart_llm._agenerate([messages])
399
+
400
+ assert len(captured) == 1
401
+ injected = captured[0]
402
+ # SystemMessage strictly at index 0, with guard folded in
403
+ assert isinstance(injected[0], SystemMessage)
404
+ assert "You are helpful." in injected[0].content
405
+ assert "untrusted data" in injected[0].content
406
+
407
+ # No SystemMessage after index 0
408
+ assert not any(isinstance(m, SystemMessage) for m in injected[1:])
409
+
410
+ # The recalled memory is in a HumanMessage before the final query
411
+ assert isinstance(injected[3], HumanMessage)
412
+ assert "User loves coffee" in injected[3].content
413
+ assert isinstance(injected[4], HumanMessage)
414
+ assert injected[4].content == "What do I like?"
415
+
416
+ @respx.mock
417
+ async def test_multimodal_system_message_appends_block_and_preserves_metadata(self) -> None:
418
+ """A list-content SystemMessage appends a text block and retains its metadata."""
419
+ _mock_seal_session_prereqs()
420
+ from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
421
+ from langchain_core.outputs import ChatGeneration, ChatResult
422
+
423
+ llm = self._make_llm()
424
+ captured: list = []
425
+
426
+ async def _capture_agenerate(messages_batch, *a, **kw):
427
+ captured.extend(messages_batch)
428
+ return ChatResult(
429
+ generations=[ChatGeneration(message=AIMessage(content="ok"))]
430
+ )
431
+
432
+ llm._agenerate = _capture_agenerate
433
+
434
+ respx.post(_RECALL_URL).mock(
435
+ return_value=_mock_recall([
436
+ {"blob_id": "b1", "text": "User loves coffee", "distance": 0.05}
437
+ ])
438
+ )
439
+ respx.post(_ANALYZE_URL).mock(return_value=_mock_analyze())
440
+
441
+ smart_llm = with_memwal_langchain(
442
+ llm, key=_KEY_HEX, account_id=_ACCOUNT_ID, server_url=_SERVER, auto_save=False
443
+ )
444
+ blocks = [
445
+ {"type": "text", "text": "Base instructions", "cache_control": {"type": "ephemeral"}}
446
+ ]
447
+ sys_msg = SystemMessage(content=blocks, id="sys-123", additional_kwargs={"source": "test"})
448
+ await smart_llm._agenerate([[sys_msg, HumanMessage("What do I like?")]])
449
+
450
+ assert len(captured) == 1
451
+ injected = captured[0]
452
+ assert isinstance(injected[0], SystemMessage)
453
+ assert injected[0].id == "sys-123"
454
+ assert injected[0].additional_kwargs == {"source": "test"}
455
+ assert isinstance(injected[0].content, list)
456
+ assert len(injected[0].content) == 2
457
+ assert injected[0].content[0] == blocks[0]
458
+ assert "untrusted data" in injected[0].content[1]["text"]
459
+
338
460
  @respx.mock
339
461
  async def test_no_memories_found_no_injection(self) -> None:
340
462
  """When no memories match, the message list is passed unchanged."""
@@ -0,0 +1,160 @@
1
+ """Tests for HTTP-date Retry-After header parsing and arithmetic overflow protection."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from datetime import datetime, timedelta, timezone
6
+ from email.utils import format_datetime
7
+ from typing import Optional
8
+
9
+ import pytest
10
+
11
+ from memwal.client import (
12
+ MemWal,
13
+ _clamped_retry_after_ms,
14
+ _HttpStatusError,
15
+ _positive_retry_ms,
16
+ _retry_after_ms,
17
+ )
18
+
19
+
20
+ def _err(
21
+ retry_after: Optional[str] = None,
22
+ body: Optional[str] = None,
23
+ ) -> _HttpStatusError:
24
+ return _HttpStatusError(429, body or "{}", retry_after=retry_after)
25
+
26
+
27
+ def test_positive_retry_ms_delta_seconds() -> None:
28
+ assert _positive_retry_ms("45") == 45000
29
+ assert _positive_retry_ms(45) == 45000
30
+ assert _positive_retry_ms("0.5") == 500
31
+ assert _positive_retry_ms(0.25) == 250
32
+
33
+
34
+ def test_positive_retry_ms_invalid_or_non_positive() -> None:
35
+ assert _positive_retry_ms(0) is None
36
+ assert _positive_retry_ms("0") is None
37
+ assert _positive_retry_ms("-5") is None
38
+ assert _positive_retry_ms("invalid") is None
39
+ assert _positive_retry_ms(None) is None
40
+ assert _positive_retry_ms(float("nan")) is None
41
+ assert _positive_retry_ms(float("inf")) is None
42
+
43
+
44
+ def test_positive_retry_ms_overflow_protection() -> None:
45
+ # 1e308 is finite float, but 1e308 * 1000 overflows to infinity
46
+ assert _positive_retry_ms("1e308") is None
47
+ assert _positive_retry_ms(1e308) is None
48
+
49
+
50
+ def test_retry_after_http_date_in_future() -> None:
51
+ future_time = datetime.now(timezone.utc) + timedelta(seconds=45)
52
+ future_header = format_datetime(future_time, usegmt=True)
53
+
54
+ err = _err(retry_after=future_header)
55
+ delay = _retry_after_ms(err)
56
+ # Delay should be approximately 45,000 ms (allowing slight test execution variance)
57
+ assert 40000 <= delay <= 46000
58
+
59
+
60
+ def test_retry_after_http_date_takes_priority_over_json() -> None:
61
+ future_time = datetime.now(timezone.utc) + timedelta(seconds=60)
62
+ future_header = format_datetime(future_time, usegmt=True)
63
+ body = '{"retry_after_seconds": 10}'
64
+
65
+ err = _err(retry_after=future_header, body=body)
66
+ delay = _retry_after_ms(err)
67
+ assert 55000 <= delay <= 61000
68
+
69
+
70
+ def test_retry_after_http_date_in_past_falls_back_to_json() -> None:
71
+ past_time = datetime.now(timezone.utc) - timedelta(seconds=10)
72
+ past_header = format_datetime(past_time, usegmt=True)
73
+ body = '{"retry_after_seconds": 15}'
74
+
75
+ err = _err(retry_after=past_header, body=body)
76
+ assert _retry_after_ms(err) == 15000
77
+
78
+
79
+ def test_retry_after_http_date_in_past_returns_zero_when_no_json() -> None:
80
+ past_time = datetime.now(timezone.utc) - timedelta(seconds=10)
81
+ past_header = format_datetime(past_time, usegmt=True)
82
+
83
+ err = _err(retry_after=past_header)
84
+ assert _retry_after_ms(err) == 0
85
+
86
+
87
+ def test_retry_after_malformed_header_falls_back_to_json() -> None:
88
+ err = _err(retry_after="Not A Valid Date Or Number", body='{"retry_after_seconds": 25}')
89
+ assert _retry_after_ms(err) == 25000
90
+
91
+
92
+ def test_retry_after_malformed_header_without_json_returns_zero() -> None:
93
+ err = _err(retry_after="invalid-header")
94
+ assert _retry_after_ms(err) == 0
95
+
96
+
97
+ def test_retry_after_overflowing_number_falls_back_to_json() -> None:
98
+ err = _err(retry_after="1e308", body='{"retry_after_seconds": 30}')
99
+ assert _retry_after_ms(err) == 30000
100
+
101
+
102
+ def test_clamped_retry_after_ms_with_http_date(monkeypatch: pytest.MonkeyPatch) -> None:
103
+ fixed_now = 1_000_000
104
+ monkeypatch.setattr("memwal.client._now_ms", lambda: fixed_now)
105
+
106
+ # Future date roughly 50s ahead
107
+ future_time = datetime.now(timezone.utc) + timedelta(seconds=50)
108
+ future_header = format_datetime(future_time, usegmt=True)
109
+ err = _err(retry_after=future_header)
110
+
111
+ # If remaining time is only 20s (deadline 1,020,000), clamped to 20,000
112
+ clamped = _clamped_retry_after_ms(err, deadline_ms=fixed_now + 20_000)
113
+ assert clamped == 20_000
114
+
115
+ # If remaining time is 100s, delay (~50,000 ms) is used
116
+ clamped_large = _clamped_retry_after_ms(err, deadline_ms=fixed_now + 100_000)
117
+ assert 45_000 <= clamped_large <= 51_000
118
+
119
+
120
+ @pytest.mark.asyncio
121
+ async def test_wait_for_remember_job_honors_http_date_retry_after(
122
+ monkeypatch: pytest.MonkeyPatch,
123
+ ) -> None:
124
+ fixed_ts = 1700000000.0
125
+ monkeypatch.setattr("time.time", lambda: fixed_ts)
126
+ clock = {"now": int(fixed_ts * 1000)}
127
+ slept: list[int] = []
128
+
129
+ def fake_now_ms() -> int:
130
+ return clock["now"]
131
+
132
+ async def fake_sleep_ms(ms: int) -> None:
133
+ slept.append(ms)
134
+ clock["now"] += ms
135
+
136
+ monkeypatch.setattr("memwal.client._now_ms", fake_now_ms)
137
+ monkeypatch.setattr("memwal.client._sleep_ms", fake_sleep_ms)
138
+
139
+ client = MemWal.create(key="11" * 32, account_id="0x1", server_url="http://localhost:8000")
140
+ calls = {"n": 0}
141
+
142
+ # 429 Retry-After is 2 seconds into future
143
+ future_dt = datetime.fromtimestamp(fixed_ts, tz=timezone.utc) + timedelta(seconds=2)
144
+ http_date = format_datetime(future_dt, usegmt=True)
145
+
146
+ async def status(*_args: object, **_kwargs: object) -> dict:
147
+ calls["n"] += 1
148
+ if calls["n"] == 1:
149
+ raise _HttpStatusError(429, "{}", retry_after=http_date)
150
+ return {"job_id": "job-1", "status": "done", "blob_id": "b-1"}
151
+
152
+ monkeypatch.setattr(client, "_signed_request", status)
153
+
154
+ res = await client.wait_for_remember_job("job-1", timeout_ms=10_000)
155
+ assert res.blob_id == "b-1"
156
+ assert calls["n"] == 2
157
+ # First sleep is attempt 0 (0 ms), second sleep should be 2000 ms from Retry-After
158
+ assert len(slept) == 2
159
+ assert slept[0] == 0
160
+ assert slept[1] == 2000
@@ -0,0 +1,195 @@
1
+ """Tests for network transport error retry and deadline bounding.
2
+
3
+ Covers WALM-696 and GH #1000.
4
+ """
5
+
6
+ from __future__ import annotations
7
+
8
+ import httpx
9
+ import pytest
10
+
11
+ from memwal.client import (
12
+ MemWal,
13
+ MemWalRememberJobTimeout,
14
+ RememberBulkOptions,
15
+ )
16
+
17
+
18
+ def _client() -> MemWal:
19
+ return MemWal.create(key="11" * 32, account_id="0x1", server_url="http://localhost:8000")
20
+
21
+
22
+ def _patch_clock(monkeypatch: pytest.MonkeyPatch) -> tuple[dict[str, int], list[int]]:
23
+ clock = {"now": 0}
24
+ slept: list[int] = []
25
+
26
+ def now_ms() -> int:
27
+ return clock["now"]
28
+
29
+ async def sleep_ms(ms: int) -> None:
30
+ slept.append(ms)
31
+ clock["now"] += ms
32
+
33
+ monkeypatch.setattr("memwal.client._now_ms", now_ms)
34
+ monkeypatch.setattr("memwal.client._sleep_ms", sleep_ms)
35
+ return clock, slept
36
+
37
+
38
+ @pytest.mark.asyncio
39
+ async def test_wait_for_remember_job_retries_connect_error(monkeypatch: pytest.MonkeyPatch) -> None:
40
+ _patch_clock(monkeypatch)
41
+ client = _client()
42
+ calls = {"n": 0}
43
+
44
+ async def status(*_args: object, **_kwargs: object) -> dict:
45
+ calls["n"] += 1
46
+ if calls["n"] == 1:
47
+ raise httpx.ConnectError("simulated connection drop")
48
+ return {
49
+ "job_id": "job-1",
50
+ "status": "done",
51
+ "blob_id": "blob-1",
52
+ "owner": "0x1",
53
+ "namespace": "default",
54
+ }
55
+
56
+ monkeypatch.setattr(client, "_signed_request", status)
57
+
58
+ done = await client.wait_for_remember_job("job-1", poll_interval_ms=100, timeout_ms=5000)
59
+ assert done.blob_id == "blob-1"
60
+ assert calls["n"] == 2
61
+
62
+
63
+ @pytest.mark.asyncio
64
+ async def test_wait_for_remember_job_retries_read_timeout(monkeypatch: pytest.MonkeyPatch) -> None:
65
+ _patch_clock(monkeypatch)
66
+ client = _client()
67
+ calls = {"n": 0}
68
+
69
+ async def status(*_args: object, **_kwargs: object) -> dict:
70
+ calls["n"] += 1
71
+ if calls["n"] == 1:
72
+ raise httpx.ReadTimeout("simulated socket read timeout")
73
+ return {
74
+ "job_id": "job-2",
75
+ "status": "done",
76
+ "blob_id": "blob-2",
77
+ "owner": "0x1",
78
+ "namespace": "default",
79
+ }
80
+
81
+ monkeypatch.setattr(client, "_signed_request", status)
82
+
83
+ done = await client.wait_for_remember_job("job-2", poll_interval_ms=100, timeout_ms=5000)
84
+ assert done.blob_id == "blob-2"
85
+ assert calls["n"] == 2
86
+
87
+
88
+ @pytest.mark.asyncio
89
+ async def test_wait_for_remember_job_times_out_with_http_0_message(
90
+ monkeypatch: pytest.MonkeyPatch,
91
+ ) -> None:
92
+ clock, slept = _patch_clock(monkeypatch)
93
+ client = _client()
94
+ calls = {"n": 0}
95
+
96
+ async def failing(*_args: object, **_kwargs: object) -> dict:
97
+ calls["n"] += 1
98
+ clock["now"] += 100
99
+ raise httpx.ConnectError("network unreachable")
100
+
101
+ monkeypatch.setattr(client, "_signed_request", failing)
102
+
103
+ with pytest.raises(MemWalRememberJobTimeout) as exc_info:
104
+ await client.wait_for_remember_job("job-fail", poll_interval_ms=100, timeout_ms=500)
105
+
106
+ err = exc_info.value
107
+ assert err.last_refusal == 0
108
+ assert err.status == 504
109
+ assert "last poll: HTTP 0" in str(err)
110
+ assert calls["n"] > 1
111
+
112
+
113
+ @pytest.mark.asyncio
114
+ async def test_wait_for_remember_job_bounds_poll_timeout(monkeypatch: pytest.MonkeyPatch) -> None:
115
+ clock, _ = _patch_clock(monkeypatch)
116
+ client = _client()
117
+ captured_timeouts: list[float] = []
118
+
119
+ async def capture(*_args: object, **kwargs: object) -> dict:
120
+ timeout = kwargs.get("timeout")
121
+ if timeout is not None:
122
+ captured_timeouts.append(timeout)
123
+ return {
124
+ "job_id": "job-bound",
125
+ "status": "done",
126
+ "blob_id": "blob-bound",
127
+ "owner": "0x1",
128
+ "namespace": "default",
129
+ }
130
+
131
+ monkeypatch.setattr(client, "_signed_request", capture)
132
+
133
+ clock["now"] = 1000
134
+ # timeout_ms=2000 -> deadline=3000 -> remaining=2.0s
135
+ await client.wait_for_remember_job("job-bound", timeout_ms=2000)
136
+ assert len(captured_timeouts) == 1
137
+ assert captured_timeouts[0] == pytest.approx(2.0, rel=1e-2)
138
+
139
+
140
+ @pytest.mark.asyncio
141
+ async def test_wait_for_remember_jobs_bulk_retries_transport_error(
142
+ monkeypatch: pytest.MonkeyPatch,
143
+ ) -> None:
144
+ _patch_clock(monkeypatch)
145
+ client = _client()
146
+ calls = {"n": 0}
147
+
148
+ async def bulk_status(job_ids: list[str], timeout: float | None = None):
149
+ from memwal.types import RememberBulkStatusItem, RememberBulkStatusResult
150
+
151
+ calls["n"] += 1
152
+ if calls["n"] == 1:
153
+ raise httpx.ConnectError("simulated bulk network drop")
154
+ return RememberBulkStatusResult(
155
+ results=[
156
+ RememberBulkStatusItem(
157
+ job_id="bulk-1",
158
+ status="done",
159
+ blob_id="blob-bulk-1",
160
+ )
161
+ ]
162
+ )
163
+
164
+ monkeypatch.setattr(client, "get_remember_bulk_status", bulk_status)
165
+
166
+ res = await client.wait_for_remember_jobs(
167
+ ["bulk-1"],
168
+ opts=RememberBulkOptions(poll_interval_ms=100, timeout_ms=5000),
169
+ )
170
+ assert len(res.results) == 1
171
+ assert res.results[0].status == "done"
172
+ assert res.results[0].blob_id == "blob-bulk-1"
173
+ assert calls["n"] == 2
174
+
175
+
176
+ @pytest.mark.asyncio
177
+ async def test_wait_for_remember_jobs_bulk_timeout_on_persistent_transport_error(
178
+ monkeypatch: pytest.MonkeyPatch,
179
+ ) -> None:
180
+ clock, _ = _patch_clock(monkeypatch)
181
+ client = _client()
182
+
183
+ async def bulk_failing(job_ids: list[str], timeout: float | None = None):
184
+ clock["now"] += 100
185
+ raise httpx.ConnectError("persistent bulk network drop")
186
+
187
+ monkeypatch.setattr(client, "get_remember_bulk_status", bulk_failing)
188
+
189
+ res = await client.wait_for_remember_jobs(
190
+ ["bulk-dead"],
191
+ opts=RememberBulkOptions(poll_interval_ms=100, timeout_ms=500),
192
+ )
193
+ assert len(res.results) == 1
194
+ assert res.results[0].status == "timeout"
195
+ assert "last poll: HTTP 0" in (res.results[0].error or "")
File without changes
File without changes
File without changes
File without changes