From 39a8fab36d88ba17a99f461d32c7b1105429630d Mon Sep 17 00:00:00 2001 From: Abhinav Maurya Date: Mon, 2 Nov 2026 23:52:37 +0530 Subject: [PATCH] fix(memory): retrieve all records when deleting namespace memories --- src/bedrock_agentcore/memory/session.py | 12 ++-- .../bedrock_agentcore/memory/test_session.py | 59 +++++++++++++++++++ 2 files changed, 66 insertions(+), 5 deletions(-) diff --git a/src/bedrock_agentcore/memory/session.py b/src/bedrock_agentcore/memory/session.py index 5f65198a..ef5b484c 100644 --- a/src/bedrock_agentcore/memory/session.py +++ b/src/bedrock_agentcore/memory/session.py @@ -965,7 +965,7 @@ def list_long_term_memory_records( self, namespace_prefix: Optional[str] = None, strategy_id: Optional[str] = None, - max_results: int = 10, + max_results: Optional[int] = 10, namespace: Optional[str] = None, namespace_path: Optional[str] = None, ) -> List[MemoryRecord]: @@ -980,7 +980,7 @@ def list_long_term_memory_records( Args: namespace_prefix: DEPRECATED. Use ``namespace`` or ``namespace_path`` instead. strategy_id: Optional strategy filter - max_results: Maximum records to return + max_results: Maximum records to return. Pass ``None`` to retrieve all records. namespace: Exact-match namespace (preserves pre-redesign behavior during the service grace period) namespace_path: Hierarchical path-prefix namespace @@ -1014,11 +1014,13 @@ def list_long_term_memory_records( all_records.extend([MemoryRecord(record) for record in memory_records]) # Stop if we've reached max_results - if len(all_records) >= max_results: + if max_results is not None and len(all_records) >= max_results: break logger.info(" ✅ Found a total of %d long-term records.", len(all_records)) - return all_records[:max_results] + if max_results is not None: + return all_records[:max_results] + return all_records except ClientError as e: logger.error(" ❌ Error listing long-term records: %s", e) @@ -1106,7 +1108,7 @@ def delete_all_long_term_memories_in_namespace(self, namespace: str) -> Dict[str # Retrieve all memory records in the specified namespace # TODO: revisit after full deprecation of namespace fields - memory_records = self.list_long_term_memory_records(namespace=namespace) + memory_records = self.list_long_term_memory_records(namespace=namespace, max_results=None) logger.info(" -> Found %d memory records to delete", len(memory_records)) if not memory_records: diff --git a/tests/bedrock_agentcore/memory/test_session.py b/tests/bedrock_agentcore/memory/test_session.py index 3108587a..b1d60573 100644 --- a/tests/bedrock_agentcore/memory/test_session.py +++ b/tests/bedrock_agentcore/memory/test_session.py @@ -1621,6 +1621,65 @@ def test_delete_all_long_term_memories_in_namespace_partial_failure(self): assert result["failedRecords"][0]["memoryRecordId"] == "rec-2" assert result["failedRecords"][0]["errorMessage"] == "Access denied" + def test_delete_all_long_term_memories_in_namespace_paginated(self): + """Test delete_all_long_term_memories_in_namespace deletes all records across pages.""" + with patch("boto3.Session") as mock_session_class: + mock_session = MagicMock() + mock_session.region_name = "us-west-2" + mock_client_instance = MagicMock() + mock_session.client.return_value = mock_client_instance + mock_session_class.return_value = mock_session + + manager = MemorySessionManager(memory_id="testMemory-1234567890", region_name="us-west-2") + + # Mock paginator returning 25 records over 3 pages (exceeding default max_results=10) + mock_paginator = MagicMock() + mock_client_instance.get_paginator.return_value = mock_paginator + records = [{"memoryRecordId": f"rec-{i:03d}"} for i in range(25)] + mock_paginator.paginate.return_value = [ + {"memoryRecordSummaries": records[0:10]}, + {"memoryRecordSummaries": records[10:20]}, + {"memoryRecordSummaries": records[20:25]}, + ] + mock_client_instance.batch_delete_memory_records.side_effect = lambda memoryId, records: { + "successfulRecords": [{"memoryRecordId": r["memoryRecordId"]} for r in records], + "failedRecords": [], + } + + result = manager.delete_all_long_term_memories_in_namespace("test/namespace/") + + assert len(result["successfulRecords"]) == 25 + assert len(result["failedRecords"]) == 0 + mock_client_instance.batch_delete_memory_records.assert_called_once_with( + memoryId="testMemory-1234567890", + records=[{"memoryRecordId": f"rec-{i:03d}"} for i in range(25)], + ) + + def test_list_long_term_memory_records_unlimited(self): + """Test list_long_term_memory_records with max_results=None returns all records.""" + with patch("boto3.Session") as mock_session_class: + mock_session = MagicMock() + mock_session.region_name = "us-west-2" + mock_client_instance = MagicMock() + mock_session.client.return_value = mock_client_instance + mock_session_class.return_value = mock_session + + manager = MemorySessionManager(memory_id="testMemory-1234567890", region_name="us-west-2") + + mock_paginator = MagicMock() + mock_client_instance.get_paginator.return_value = mock_paginator + records = [{"memoryRecordId": f"rec-{i:03d}"} for i in range(25)] + mock_paginator.paginate.return_value = [ + {"memoryRecordSummaries": records[0:10]}, + {"memoryRecordSummaries": records[10:20]}, + {"memoryRecordSummaries": records[20:25]}, + ] + + result = manager.list_long_term_memory_records(namespace="test/namespace/", max_results=None) + + assert len(result) == 25 + assert all(isinstance(record, MemoryRecord) for record in result) + def test_list_actor_sessions_success(self): """Test list_actor_sessions successful execution.""" with patch("boto3.Session") as mock_session_class: