diff --git a/model/map/src/main/java/org/keycloak/models/map/authSession/MapAuthenticationSessionEntity.java b/model/map/src/main/java/org/keycloak/models/map/authSession/MapAuthenticationSessionEntity.java
index 9df9b2131d..1adcb9aa84 100644
--- a/model/map/src/main/java/org/keycloak/models/map/authSession/MapAuthenticationSessionEntity.java
+++ b/model/map/src/main/java/org/keycloak/models/map/authSession/MapAuthenticationSessionEntity.java
@@ -31,6 +31,9 @@ import java.util.Set;
 @DeepCloner.Root
 public interface MapAuthenticationSessionEntity extends UpdatableEntity {
 
+    String getTabId();
+    void setTabId(String tabId);
+
     Map<String, String> getUserSessionNotes();
     void setUserSessionNotes(Map<String, String> userSessionNotes);
     void setUserSessionNote(String name, String value);
diff --git a/model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionAdapter.java b/model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionAdapter.java
index a0b7523d89..96428376f2 100644
--- a/model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionAdapter.java
+++ b/model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionAdapter.java
@@ -61,9 +61,8 @@ public class MapRootAuthenticationSessionAdapter extends AbstractRootAuthenticat
 
     @Override
     public Map<String, AuthenticationSessionModel> getAuthenticationSessions() {
-        return Optional.ofNullable(entity.getAuthenticationSessions()).orElseGet(Collections::emptyMap).entrySet().stream()
-                .collect(Collectors.toMap(Map.Entry::getKey,
-                        entry -> new MapAuthenticationSessionAdapter(session, this, entry.getKey(), (MapAuthenticationSessionEntity) entry.getValue())));
+        return Optional.ofNullable(entity.getAuthenticationSessions()).orElseGet(Collections::emptySet).stream()
+                .collect(Collectors.toMap(MapAuthenticationSessionEntity::getTabId, this::toAdapter));
     }
 
     @Override
@@ -72,13 +71,7 @@ public class MapRootAuthenticationSessionAdapter extends AbstractRootAuthenticat
             return null;
         }
 
-        AuthenticationSessionModel authSession = getAuthenticationSessions().get(tabId);
-        if (authSession != null && client.equals(authSession.getClient())) {
-            session.getContext().setAuthenticationSession(authSession);
-            return authSession;
-        } else {
-            return null;
-        }
+        return entity.getAuthenticationSession(tabId).map(this::toAdapter).map(this::setAuthContext).orElse(null);
     }
 
     @Override
@@ -90,16 +83,15 @@ public class MapRootAuthenticationSessionAdapter extends AbstractRootAuthenticat
 
         int timestamp = Time.currentTime();
         authSessionEntity.setTimestamp(timestamp);
+        String tabId = generateTabId();
+        authSessionEntity.setTabId(tabId);
 
-        String tabId =  generateTabId();
-        entity.setAuthenticationSession(tabId, authSessionEntity);
+        entity.addAuthenticationSession(authSessionEntity);
 
         // Update our timestamp when adding new authenticationSession
         entity.setTimestamp(timestamp);
 
-        MapAuthenticationSessionAdapter authSession = new MapAuthenticationSessionAdapter(session, this, tabId, entity.getAuthenticationSessions().get(tabId));
-        session.getContext().setAuthenticationSession(authSession);
-        return authSession;
+        return entity.getAuthenticationSession(tabId).map(this::toAdapter).map(this::setAuthContext).orElse(null);
     }
 
     @Override
@@ -123,4 +115,13 @@ public class MapRootAuthenticationSessionAdapter extends AbstractRootAuthenticat
     private String generateTabId() {
         return Base64Url.encode(SecretGenerator.getInstance().randomBytes(8));
     }
+
+    private MapAuthenticationSessionAdapter toAdapter(MapAuthenticationSessionEntity entity) {
+        return new MapAuthenticationSessionAdapter(session, this, entity.getTabId(), entity);
+    }
+
+    private MapAuthenticationSessionAdapter setAuthContext(MapAuthenticationSessionAdapter adapter) {
+        session.getContext().setAuthenticationSession(adapter);
+        return adapter;
+    }
 }
diff --git a/model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionEntity.java b/model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionEntity.java
index f8192ce932..182256d554 100644
--- a/model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionEntity.java
+++ b/model/map/src/main/java/org/keycloak/models/map/authSession/MapRootAuthenticationSessionEntity.java
@@ -23,8 +23,9 @@ import org.keycloak.models.map.common.DeepCloner;
 import org.keycloak.models.map.common.UpdatableEntity;
 
 import java.util.Collections;
-import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 
 /**
  * @author <a href="mailto:mkanis@redhat.com">Martin Kanis</a>
@@ -51,16 +52,32 @@ public interface MapRootAuthenticationSessionEntity extends AbstractEntity, Upda
             this.updated |= id != null;
         }
 
+        @Override
+        public Optional<MapAuthenticationSessionEntity> getAuthenticationSession(String tabId) {
+            Set<MapAuthenticationSessionEntity> authenticationSessions = getAuthenticationSessions();
+            if (authenticationSessions == null || authenticationSessions.isEmpty()) return Optional.empty();
+
+            return authenticationSessions.stream().filter(as -> Objects.equals(as.getTabId(), tabId)).findFirst();
+        }
+
+        @Override
+        public Boolean removeAuthenticationSession(String tabId) {
+            Set<MapAuthenticationSessionEntity> authenticationSessions = getAuthenticationSessions();
+            boolean removed = authenticationSessions != null && authenticationSessions.removeIf(c -> Objects.equals(c.getTabId(), tabId));
+            this.updated |= removed;
+            return removed;
+        }
+
         @Override
         public boolean isUpdated() {
             return this.updated ||
-                    Optional.ofNullable(getAuthenticationSessions()).orElseGet(Collections::emptyMap).values().stream().anyMatch(MapAuthenticationSessionEntity::isUpdated);
+                    Optional.ofNullable(getAuthenticationSessions()).orElseGet(Collections::emptySet).stream().anyMatch(MapAuthenticationSessionEntity::isUpdated);
         }
 
         @Override
         public void clearUpdatedFlag() {
             this.updated = false;
-            Optional.ofNullable(getAuthenticationSessions()).orElseGet(Collections::emptyMap).values().forEach(UpdatableEntity::clearUpdatedFlag);
+            Optional.ofNullable(getAuthenticationSessions()).orElseGet(Collections::emptySet).forEach(UpdatableEntity::clearUpdatedFlag);
         }
     }
 
@@ -70,8 +87,9 @@ public interface MapRootAuthenticationSessionEntity extends AbstractEntity, Upda
     Integer getTimestamp();
     void setTimestamp(Integer timestamp);
 
-    Map<String, MapAuthenticationSessionEntity> getAuthenticationSessions();
-    void setAuthenticationSessions(Map<String, MapAuthenticationSessionEntity> authenticationSessions);
-    void setAuthenticationSession(String tabId, MapAuthenticationSessionEntity entity);
+    Set<MapAuthenticationSessionEntity> getAuthenticationSessions();
+    void setAuthenticationSessions(Set<MapAuthenticationSessionEntity> authenticationSessions);
+    Optional<MapAuthenticationSessionEntity> getAuthenticationSession(String tabId);
+    void addAuthenticationSession(MapAuthenticationSessionEntity authenticationSession);
     Boolean removeAuthenticationSession(String tabId);
 }