diff --git a/packages/mermaid/src/diagrams/state/stateDb.js b/packages/mermaid/src/diagrams/state/stateDb.js
index 2da0719ca..50785587e 100644
--- a/packages/mermaid/src/diagrams/state/stateDb.js
+++ b/packages/mermaid/src/diagrams/state/stateDb.js
@@ -130,6 +130,7 @@ export class StateDB {
     log.info('Setting root doc', o);
     // rootDoc = { id: 'root', doc: o };
     this.rootDoc = o;
+    this.extract(o);
   }
 
   getRootDoc() {
diff --git a/packages/mermaid/src/diagrams/state/stateRenderer-v3-unified.ts b/packages/mermaid/src/diagrams/state/stateRenderer-v3-unified.ts
index 109417c03..2998c8173 100644
--- a/packages/mermaid/src/diagrams/state/stateRenderer-v3-unified.ts
+++ b/packages/mermaid/src/diagrams/state/stateRenderer-v3-unified.ts
@@ -36,7 +36,6 @@ export const getClasses = function (
   text: string,
   diagramObj: any
 ): Map<string, DiagramStyleClassDef> {
-  diagramObj.db.extract(diagramObj.db.getRootDocV2());
   return diagramObj.db.getClasses();
 };