diff --git a/rules-tests/Bundle230/Rector/Class_/AddAnnotationToRepositoryRector/Fixture/some_service_document_repository.php.inc b/rules-tests/Bundle230/Rector/Class_/AddAnnotationToRepositoryRector/Fixture/some_service_document_repository.php.inc
index d6afa5ea..7b3d7b14 100644
--- a/rules-tests/Bundle230/Rector/Class_/AddAnnotationToRepositoryRector/Fixture/some_service_document_repository.php.inc
+++ b/rules-tests/Bundle230/Rector/Class_/AddAnnotationToRepositoryRector/Fixture/some_service_document_repository.php.inc
@@ -25,7 +25,7 @@ use Doctrine\Persistence\ManagerRegistry;
 use App\Entity\SomeEntity;
 
 /**
- * @extends \Doctrine\Bundle\DoctrineBundle\Repository\ServiceEntityRepository<\App\Entity\SomeEntity>
+ * @extends \Doctrine\Bundle\MongoDBBundle\Repository\ServiceDocumentRepository<\App\Entity\SomeEntity>
  */
 final class SomeServiceDocumentRepository extends ServiceDocumentRepository
 {
diff --git a/rules/Bundle230/Rector/Class_/AddAnnotationToRepositoryRector.php b/rules/Bundle230/Rector/Class_/AddAnnotationToRepositoryRector.php
index 51677627..a033b5bb 100644
--- a/rules/Bundle230/Rector/Class_/AddAnnotationToRepositoryRector.php
+++ b/rules/Bundle230/Rector/Class_/AddAnnotationToRepositoryRector.php
@@ -86,7 +86,8 @@ public function getNodeTypes(): array
      */
     public function refactor(Node $node): ?Node
     {
-        if (! $this->isRepositoryClass($node)) {
+        $repositoryClass = $this->matchServiceRepositoryClass($node);
+        if ($repositoryClass === null) {
             return null;
         }
 
@@ -95,22 +96,25 @@ public function refactor(Node $node): ?Node
             return null;
         }
 
-        $this->addAnnotationToNode($node, $entityClass);
+        $this->addAnnotationToNode($node, $entityClass, $repositoryClass);
 
         return $node;
     }
 
-    private function isRepositoryClass(Class_ $class): bool
+    private function matchServiceRepositoryClass(Class_ $class): ?string
     {
         if (! $class->extends instanceof Name) {
-            return false;
+            return null;
         }
 
-        if ($this->isName($class->extends, DoctrineClass::SERVICE_ENTITY_REPOSITORY)) {
-            return true;
+        if (! $this->isNames(
+            $class->extends,
+            [DoctrineClass::SERVICE_ENTITY_REPOSITORY, DoctrineClass::SERVICE_DOCUMENT_REPOSITORY]
+        )) {
+            return null;
         }
 
-        return $this->isName($class->extends, DoctrineClass::SERVICE_DOCUMENT_REPOSITORY);
+        return $this->getName($class->extends);
     }
 
     private function getEntityClassFromConstructor(Class_ $class): ?string
@@ -149,11 +153,11 @@ private function getEntityClassFromConstructor(Class_ $class): ?string
         return null;
     }
 
-    private function addAnnotationToNode(Class_ $class, string $entityClass): void
+    private function addAnnotationToNode(Class_ $class, string $entityClass, string $repositoryClass): void
     {
         $phpDocInfo = $this->phpDocInfoFactory->createFromNodeOrEmpty($class);
 
-        $genericsAnnotation = sprintf('\%s<\%s>', DoctrineClass::SERVICE_ENTITY_REPOSITORY, $entityClass);
+        $genericsAnnotation = sprintf('\%s<\%s>', $repositoryClass, $entityClass);
         $phpDocInfo->addPhpDocTagNode(new PhpDocTagNode('@extends', new GenericTagValueNode($genericsAnnotation)));
 
         $this->docBlockUpdater->updateRefactoredNodeWithPhpDocInfo($class);