resolving comments

Author: Aayush Kataria

Docs/concepts.md

@@ -179,7 +179,9 @@ Set any value to `0` to disable that processing type. For example, setting `THRE
 
 | Container | Partition Key | Purpose |
 |-----------|---------------|---------|
-| `memories` | `/user_id`, `/thread_id` (hierarchical) | Existing memory store |
+| `memories` | `/user_id`, `/thread_id` (hierarchical) | Durable derived memories (`fact`, `episodic`, `procedural`) |
+| `memories_turns` | `/user_id`, `/thread_id` (hierarchical) | Raw conversation turns (`turn`) — append-only, TTL-pruned |
+| `memories_summaries` | `/user_id`, `/thread_id` (hierarchical) | Thread + user summaries (`thread_summary`, `user_summary`) |
 | `counter` | `/user_id`, `/thread_id` (hierarchical) | Message count tracking for automatic processing |
 | `leases` | `/id` | Change feed checkpointing container created by `create_memory_store()` |
 

agent_memory_toolkit/aio/cosmos_memory_client.py

@@ -831,8 +831,22 @@ async def process_now_and_wait(self, *, user_id: str, thread_id: str, timeout: f
         return False
 
     async def _summary_exists(self, *, user_id: str, thread_id: str) -> bool:
+        from azure.cosmos.exceptions import (
+            CosmosHttpResponseError,
+            CosmosResourceNotFoundError,
+        )
+
         try:
             results = await self.get_thread_summary(user_id=user_id, thread_id=thread_id, recent_k=1)
-        except Exception:
+        except CosmosResourceNotFoundError:
+            return False
+        except CosmosHttpResponseError as exc:
+            logger.warning(
+                "_summary_exists: Cosmos error user_id=%s thread_id=%s status=%s: %s",
+                user_id,
+                thread_id,
+                getattr(exc, "status_code", None),
+                exc,
+            )
             return False
         return bool(results)

agent_memory_toolkit/aio/store/memory_store.py

@@ -361,34 +361,56 @@ async def update(
         metadata: Optional[dict[str, Any]] = None,
     ) -> None:
         """Update a memory document via point read in the container for ``memory_type``."""
-        from azure.cosmos.exceptions import CosmosResourceNotFoundError
-
-        container = self._container_for_type(memory_type)
-        try:
-            doc = await container.read_item(item=memory_id, partition_key=[user_id, thread_id])
-        except CosmosResourceNotFoundError as exc:
-            raise MemoryNotFoundError(memory_id=memory_id, user_id=user_id, thread_id=thread_id) from exc
-        except Exception as exc:
-            raise _wrap_cosmos_exception(exc, message=f"async update read failed for {memory_id}: {exc}") from exc
-
-        actual_type = doc.get("type")
-        if actual_type != memory_type:
-            raise MemoryTypeMismatchError(memory_id=memory_id, expected=memory_type, actual=actual_type)
+        import random
 
-        if content is not None:
-            doc["content"] = content
-        if role is not None:
-            doc["role"] = role
-        if metadata is not None:
-            doc["metadata"] = metadata
-        doc["updated_at"] = datetime.now(timezone.utc).isoformat()
+        from azure.core import MatchConditions
+        from azure.cosmos.exceptions import (
+            CosmosAccessConditionFailedError,
+            CosmosResourceNotFoundError,
+        )
 
-        try:
-            await container.replace_item(item=doc["id"], body=doc)
-        except Exception as exc:
-            raise _wrap_cosmos_exception(exc, message=f"async update replace failed for {memory_id}: {exc}") from exc
+        container = self._container_for_type(memory_type)
+        max_attempts = 5
+        attempts = 0
+        while True:
+            try:
+                doc = await container.read_item(item=memory_id, partition_key=[user_id, thread_id])
+            except CosmosResourceNotFoundError as exc:
+                raise MemoryNotFoundError(memory_id=memory_id, user_id=user_id, thread_id=thread_id) from exc
+            except Exception as exc:
+                raise _wrap_cosmos_exception(exc, message=f"async update read failed for {memory_id}: {exc}") from exc
+
+            actual_type = doc.get("type")
+            if actual_type != memory_type:
+                raise MemoryTypeMismatchError(memory_id=memory_id, expected=memory_type, actual=actual_type)
+
+            if content is not None:
+                doc["content"] = content
+            if role is not None:
+                doc["role"] = role
+            if metadata is not None:
+                doc["metadata"] = metadata
+            doc["updated_at"] = datetime.now(timezone.utc).isoformat()
 
-        logger.info("Async updated record %s", memory_id)
+            kwargs: dict[str, Any] = {"item": doc["id"], "body": doc}
+            if etag := doc.get("_etag"):
+                kwargs.update(match_condition=MatchConditions.IfNotModified, etag=etag)
+            try:
+                await container.replace_item(**kwargs)
+                logger.info("Async updated record %s", memory_id)
+                return
+            except CosmosAccessConditionFailedError as exc:
+                attempts += 1
+                if attempts >= max_attempts:
+                    raise MemoryConflictError(
+                        f"update conflicted after {max_attempts} attempts for memory_id={memory_id!r}"
+                    ) from exc
+                base = 0.02 * (2 ** (attempts - 1))
+                await asyncio.sleep(base + random.uniform(0, base))
+            except Exception as exc:
+                raise _wrap_cosmos_exception(
+                    exc, message=f"async update replace failed for {memory_id}: {exc}"
+                ) from exc
 
     async def delete(
         self,
@@ -400,11 +422,16 @@ async def delete(
     ) -> None:
         """Delete a memory document from the container for ``memory_type``.
 
-        Reads the doc first to verify its ``type`` matches ``memory_type`` so a
-        wrong routing key cannot silently delete the wrong document (within
-        MEMORIES, fact/episodic/procedural share the same partition key).
+        Reads the doc first to verify its ``type`` matches ``memory_type`` and
+        then issues the delete with ``If-Match`` on the read ETag, so a
+        concurrent type mutation between read and delete is rejected (412)
+        rather than silently dropping the wrong document.
         """
-        from azure.cosmos.exceptions import CosmosResourceNotFoundError
+        from azure.core import MatchConditions
+        from azure.cosmos.exceptions import (
+            CosmosAccessConditionFailedError,
+            CosmosResourceNotFoundError,
+        )
 
         container = self._container_for_type(memory_type)
         try:
@@ -418,10 +445,17 @@ async def delete(
         if actual_type != memory_type:
             raise MemoryTypeMismatchError(memory_id=memory_id, expected=memory_type, actual=actual_type)
 
+        kwargs: dict[str, Any] = {"item": memory_id, "partition_key": [user_id, thread_id]}
+        if etag := doc.get("_etag"):
+            kwargs.update(match_condition=MatchConditions.IfNotModified, etag=etag)
         try:
-            await container.delete_item(item=memory_id, partition_key=[user_id, thread_id])
+            await container.delete_item(**kwargs)
         except CosmosResourceNotFoundError as exc:
             raise MemoryNotFoundError(memory_id=memory_id, user_id=user_id, thread_id=thread_id) from exc
+        except CosmosAccessConditionFailedError as exc:
+            raise MemoryConflictError(
+                f"delete conflicted for memory_id={memory_id!r} — document was modified after the type check"
+            ) from exc
         except Exception as exc:
             raise _wrap_cosmos_exception(exc, message=f"async delete failed for {memory_id}: {exc}") from exc
 
@@ -664,11 +698,13 @@ async def mark_superseded(
             )
             del exc
             return False
-        except CosmosHttpResponseError:
-            logger.exception(
-                "supersede transient failure id=%s superseder=%s",
+        except CosmosHttpResponseError as exc:
+            logger.warning(
+                "supersede failed id=%s superseder=%s status=%s: %s",
                 old_doc.get("id"),
                 superseder_id,
+                getattr(exc, "status_code", None),
+                exc,
             )
             return False
 

agent_memory_toolkit/cosmos_memory_client.py

@@ -809,8 +809,22 @@ def process_now_and_wait(self, *, user_id: str, thread_id: str, timeout: float =
         return False
 
     def _summary_exists(self, *, user_id: str, thread_id: str) -> bool:
+        from azure.cosmos.exceptions import (
+            CosmosHttpResponseError,
+            CosmosResourceNotFoundError,
+        )
+
         try:
             results = self.get_thread_summary(user_id=user_id, thread_id=thread_id, recent_k=1)
-        except Exception:
+        except CosmosResourceNotFoundError:
+            return False
+        except CosmosHttpResponseError as exc:
+            logger.warning(
+                "_summary_exists: Cosmos error user_id=%s thread_id=%s status=%s: %s",
+                user_id,
+                thread_id,
+                getattr(exc, "status_code", None),
+                exc,
+            )
             return False
         return bool(results)

agent_memory_toolkit/store/memory_store.py

@@ -272,15 +272,16 @@ def push(self, local_memory: list[dict[str, Any]], batch_size: int = 25) -> None
                     )
 
             bodies = [self._prepare_doc(body) for body in bodies]
-            for record, body in zip(batch, bodies):
+            for body in bodies:
                 memory_type = body.get("type")
                 if memory_type not in _CONTAINER_FOR_TYPE:
                     raise ValueError(
                         f"push: record id={body.get('id')!r} has invalid type={memory_type!r}. "
                         f"Set 'type' to one of {sorted(_CONTAINER_FOR_TYPE)} on every local "
                         f"memory before calling push_to_cosmos."
                     )
-                container = self._container_for_type(memory_type)
+            for record, body in zip(batch, bodies):
+                container = self._container_for_type(body.get("type"))
                 try:
                     container.upsert_item(body=body)
                 except Exception as exc:
@@ -375,34 +376,55 @@ def update(
         metadata: Optional[dict[str, Any]] = None,
     ) -> None:
         """Update a memory document via point read in the container for ``memory_type``."""
-        from azure.cosmos.exceptions import CosmosResourceNotFoundError
-
-        container = self._container_for_type(memory_type)
-        try:
-            doc = container.read_item(item=memory_id, partition_key=[user_id, thread_id])
-        except CosmosResourceNotFoundError as exc:
-            raise MemoryNotFoundError(memory_id=memory_id, user_id=user_id, thread_id=thread_id) from exc
-        except Exception as exc:
-            raise _wrap_cosmos_exception(exc, message=f"update read failed for {memory_id}: {exc}") from exc
-
-        actual_type = doc.get("type")
-        if actual_type != memory_type:
-            raise MemoryTypeMismatchError(memory_id=memory_id, expected=memory_type, actual=actual_type)
+        import random
+        import time
 
-        if content is not None:
-            doc["content"] = content
-        if role is not None:
-            doc["role"] = role
-        if metadata is not None:
-            doc["metadata"] = metadata
-        doc["updated_at"] = datetime.now(timezone.utc).isoformat()
+        from azure.core import MatchConditions
+        from azure.cosmos.exceptions import (
+            CosmosAccessConditionFailedError,
+            CosmosResourceNotFoundError,
+        )
 
-        try:
-            container.replace_item(item=doc["id"], body=doc)
-        except Exception as exc:
-            raise _wrap_cosmos_exception(exc, message=f"update replace failed for {memory_id}: {exc}") from exc
+        container = self._container_for_type(memory_type)
+        max_attempts = 5
+        attempts = 0
+        while True:
+            try:
+                doc = container.read_item(item=memory_id, partition_key=[user_id, thread_id])
+            except CosmosResourceNotFoundError as exc:
+                raise MemoryNotFoundError(memory_id=memory_id, user_id=user_id, thread_id=thread_id) from exc
+            except Exception as exc:
+                raise _wrap_cosmos_exception(exc, message=f"update read failed for {memory_id}: {exc}") from exc
+
+            actual_type = doc.get("type")
+            if actual_type != memory_type:
+                raise MemoryTypeMismatchError(memory_id=memory_id, expected=memory_type, actual=actual_type)
+
+            if content is not None:
+                doc["content"] = content
+            if role is not None:
+                doc["role"] = role
+            if metadata is not None:
+                doc["metadata"] = metadata
+            doc["updated_at"] = datetime.now(timezone.utc).isoformat()
 
-        logger.info("Updated record %s", memory_id)
+            kwargs: dict[str, Any] = {"item": doc["id"], "body": doc}
+            if etag := doc.get("_etag"):
+                kwargs.update(match_condition=MatchConditions.IfNotModified, etag=etag)
+            try:
+                container.replace_item(**kwargs)
+                logger.info("Updated record %s", memory_id)
+                return
+            except CosmosAccessConditionFailedError as exc:
+                attempts += 1
+                if attempts >= max_attempts:
+                    raise MemoryConflictError(
+                        f"update conflicted after {max_attempts} attempts for memory_id={memory_id!r}"
+                    ) from exc
+                base = 0.02 * (2 ** (attempts - 1))
+                time.sleep(base + random.uniform(0, base))
+            except Exception as exc:
+                raise _wrap_cosmos_exception(exc, message=f"update replace failed for {memory_id}: {exc}") from exc
 
     def delete(
         self,
@@ -414,11 +436,16 @@ def delete(
     ) -> None:
         """Delete a memory document from the container for ``memory_type``.
 
-        Reads the doc first to verify its ``type`` matches ``memory_type`` so a
-        wrong routing key cannot silently delete the wrong document (within
-        MEMORIES, fact/episodic/procedural share the same partition key).
+        Reads the doc first to verify its ``type`` matches ``memory_type`` and
+        then issues the delete with ``If-Match`` on the read ETag, so a
+        concurrent type mutation between read and delete is rejected (412)
+        rather than silently dropping the wrong document.
         """
-        from azure.cosmos.exceptions import CosmosResourceNotFoundError
+        from azure.core import MatchConditions
+        from azure.cosmos.exceptions import (
+            CosmosAccessConditionFailedError,
+            CosmosResourceNotFoundError,
+        )
 
         container = self._container_for_type(memory_type)
         try:
@@ -432,10 +459,17 @@ def delete(
         if actual_type != memory_type:
             raise MemoryTypeMismatchError(memory_id=memory_id, expected=memory_type, actual=actual_type)
 
+        kwargs: dict[str, Any] = {"item": memory_id, "partition_key": [user_id, thread_id]}
+        if etag := doc.get("_etag"):
+            kwargs.update(match_condition=MatchConditions.IfNotModified, etag=etag)
         try:
-            container.delete_item(item=memory_id, partition_key=[user_id, thread_id])
+            container.delete_item(**kwargs)
         except CosmosResourceNotFoundError as exc:
             raise MemoryNotFoundError(memory_id=memory_id, user_id=user_id, thread_id=thread_id) from exc
+        except CosmosAccessConditionFailedError as exc:
+            raise MemoryConflictError(
+                f"delete conflicted for memory_id={memory_id!r} — document was modified after the type check"
+            ) from exc
         except Exception as exc:
             raise _wrap_cosmos_exception(exc, message=f"delete failed for {memory_id}: {exc}") from exc
 
@@ -680,11 +714,13 @@ def mark_superseded(
             )
             del exc
             return False
-        except CosmosHttpResponseError:
-            logger.exception(
-                "supersede transient failure id=%s superseder=%s",
+        except CosmosHttpResponseError as exc:
+            logger.warning(
+                "supersede failed id=%s superseder=%s status=%s: %s",
                 old_doc.get("id"),
                 superseder_id,
+                getattr(exc, "status_code", None),
+                exc,
             )
             return False
 

tests/integration/test_processor_integration.py

@@ -163,9 +163,13 @@ def _fake_monotonic():
         assert client.get_thread_summary.call_count >= 1
 
     def test_timeout_swallows_search_errors(self, monkeypatch):
+        from azure.cosmos.exceptions import CosmosHttpResponseError
+
         client = _build_client(processor=DurableFunctionProcessor())
         client.get_thread = MagicMock(return_value=[])
-        client.get_thread_summary = MagicMock(side_effect=RuntimeError("transient"))
+        client.get_thread_summary = MagicMock(
+            side_effect=CosmosHttpResponseError(message="429 throttled", status_code=429)
+        )
 
         monkeypatch.setattr("time.sleep", lambda *_a, **_k: None)
 

tests/integration/test_processor_integration_async.py

@@ -165,9 +165,13 @@ async def _no_sleep(*_a, **_k):
 
     @pytest.mark.asyncio
     async def test_timeout_swallows_search_errors(self, monkeypatch):
+        from azure.cosmos.exceptions import CosmosHttpResponseError
+
         client = _build_client(processor=AsyncDurableFunctionProcessor())
         client.get_thread = AsyncMock(return_value=[])
-        client.get_thread_summary = AsyncMock(side_effect=RuntimeError("transient"))
+        client.get_thread_summary = AsyncMock(
+            side_effect=CosmosHttpResponseError(message="429 throttled", status_code=429)
+        )
 
         async def _no_sleep(*_a, **_k):
             return None