diff --git a/src/viewport-three/viewport-host.ts b/src/viewport-three/viewport-host.ts
index e2758c28..f1a97225 100644
--- a/src/viewport-three/viewport-host.ts
+++ b/src/viewport-three/viewport-host.ts
@@ -6541,15 +6541,20 @@ export class ViewportHost {
   }
 
   private getTerrainFoliageMaskPreviewLayerId(terrainId: string): string | null {
+    const brushState = this.currentTerrainBrushState;
+
+    if (brushState === null || brushState.terrainId !== terrainId) {
+      return null;
+    }
+
     if (
-      this.currentTerrainBrushState === null ||
-      this.currentTerrainBrushState.terrainId !== terrainId ||
-      !isTerrainFoliageLayerMaskTool(this.currentTerrainBrushState.tool)
+      brushState.tool !== "foliagePaint" &&
+      brushState.tool !== "foliageErase"
     ) {
       return null;
     }
 
-    return this.currentTerrainBrushState.foliageLayerId;
+    return brushState.foliageLayerId;
   }
 
   private isTerrainFoliageBlockerMaskPreviewActive(
@@ -9699,7 +9704,8 @@ export class ViewportHost {
       tool: toolState.tool,
       referenceHeight,
       layerIndex: toolState.tool === "paint" ? toolState.layerIndex : null,
-      foliageLayerId: isTerrainFoliageLayerMaskTool(toolState.tool)
+      foliageLayerId:
+        toolState.tool === "foliagePaint" || toolState.tool === "foliageErase"
         ? toolState.foliageLayerId
         : null
     });
diff --git a/tests/domain/terrain.command.test.ts b/tests/domain/terrain.command.test.ts
index 62da1db1..ac155ac5 100644
--- a/tests/domain/terrain.command.test.ts
+++ b/tests/domain/terrain.command.test.ts
@@ -177,7 +177,8 @@ describe("terrain commands", () => {
           terrainId: terrain.id,
           heightSamples: [{ index: 4, before: 0, after: 2 }],
           paintWeights: [{ index: 1, before: 0, after: 0.5 }],
-          foliageMaskValues: []
+          foliageMaskValues: [],
+          foliageBlockerMaskValues: []
         }
       })
     );
@@ -263,7 +264,8 @@ describe("terrain commands", () => {
               before: 0,
               after: 0.75
             }
-          ]
+          ],
+          foliageBlockerMaskValues: []
         }
       })
     );