diff --git a/src/viewport-three/viewport-host.ts b/src/viewport-three/viewport-host.ts
index 4eb2c492..d40a6be7 100644
--- a/src/viewport-three/viewport-host.ts
+++ b/src/viewport-three/viewport-host.ts
@@ -9699,10 +9699,9 @@ export class ViewportHost {
       tool: toolState.tool,
       referenceHeight,
       layerIndex: toolState.tool === "paint" ? toolState.layerIndex : null,
-      foliageLayerId:
-        toolState.tool === "foliagePaint" || toolState.tool === "foliageErase"
-          ? toolState.foliageLayerId
-          : null
+      foliageLayerId: isTerrainFoliageLayerMaskTool(toolState.tool)
+        ? toolState.foliageLayerId
+        : null
     });
   }