diff --git a/packages/mermaid/src/diagrams/state/stateDb.js b/packages/mermaid/src/diagrams/state/stateDb.js
index cc44659eb..37cea295f 100644
--- a/packages/mermaid/src/diagrams/state/stateDb.js
+++ b/packages/mermaid/src/diagrams/state/stateDb.js
@@ -643,10 +643,17 @@ export class StateDB {
   }
 
   getDirection() {
+    const doc = this.rootDoc.find((doc) => doc.stmt === 'dir');
+    this.direction = doc ? doc.value : 'TB';
     return this.direction;
   }
   setDirection(dir) {
-    this.direction = dir;
+    let doc = this.rootDoc.find((doc) => doc.stmt === 'dir');
+    if (doc) {
+      doc.value = dir;
+    } else {
+      this.rootDoc.unshift({ stmt: 'dir', value: dir });
+    }
   }
 
   trimColon(str) {