diff --git a/src/Context/Environment/InitializedSymfonyExtensionEnvironment.php b/src/Context/Environment/InitializedSymfonyExtensionEnvironment.php
index 1a38591..8a40dd2 100644
--- a/src/Context/Environment/InitializedSymfonyExtensionEnvironment.php
+++ b/src/Context/Environment/InitializedSymfonyExtensionEnvironment.php
@@ -27,7 +27,10 @@ final class InitializedSymfonyExtensionEnvironment implements SymfonyExtensionEn
     /** @var Suite */
     private $suite;
 
-    /** @var Context[] */
+    /**
+     * @var array<class-string<Context>, Context>
+     * @psalm-var class-string-map<T as Context, T>
+     */
     private $contexts = [];
 
     public function __construct(Suite $suite)
@@ -74,6 +77,12 @@ public function hasContextClass($class): bool
     /**
      * @see http://behat.org/en/latest/cookbooks/accessing_contexts_from_each_other.html
      *
+     * @template T of Context
+     *
+     * @param class-string<T> $class
+     *
+     * @return T
+     *
      * @throws ContextNotFoundException
      */
     public function getContext(string $class): Context