diff --git a/packages/mermaid/src/diagrams/state/stateRenderer-v2.js b/packages/mermaid/src/diagrams/state/stateRenderer-v2.js
index 990048ec7..1c9b2d1d3 100644
--- a/packages/mermaid/src/diagrams/state/stateRenderer-v2.js
+++ b/packages/mermaid/src/diagrams/state/stateRenderer-v2.js
@@ -84,6 +84,7 @@ export const setConf = function (cnf) {
  * @returns {object} ClassDef styles (a Map with keys = strings, values = )
  */
 export const getClasses = function (text, diagramObj) {
+  diagramObj.db.extract(diagramObj.db.getRootDocV2());
   return diagramObj.db.getClasses();
 };