diff --git a/DependencyInjection/StompPhpStompExtension.php b/DependencyInjection/StompPhpStompExtension.php
index fb852f6..5183668 100644
--- a/DependencyInjection/StompPhpStompExtension.php
+++ b/DependencyInjection/StompPhpStompExtension.php
@@ -4,7 +4,7 @@
 
 namespace StompPhp\StompBundle\DependencyInjection;
 
-use Stomp\Client;
+use Stomp\StatefulStomp;
 use StompPhp\StompBundle\Stomp\ClientFactory;
 use StompPhp\StompBundle\Stomp\Subscription;
 use Symfony\Component\DependencyInjection\ContainerBuilder;
@@ -46,7 +46,7 @@ private function getClientDefinition(string $name, array $options): Definition
         unset($options['public']);
 
         return
-            (new Definition(Client::class))
+            (new Definition(StatefulStomp::class))
                 ->setFactory(
                     [
                         ClientFactory::class,
diff --git a/Stomp/ClientFactory.php b/Stomp/ClientFactory.php
index 965fb71..aed18a1 100644
--- a/Stomp/ClientFactory.php
+++ b/Stomp/ClientFactory.php
@@ -4,7 +4,6 @@
 
 namespace StompPhp\StompBundle\Stomp;
 
-use Stomp\Client;
 use Stomp\StatefulStomp;
 
 final class ClientFactory
@@ -12,7 +11,7 @@ final class ClientFactory
     /**
      * @param array $options
      *
-     * @return Client
+     * @return StatefulStomp
      *
      * @throws \Stomp\Exception\ConnectionException
      */