diff --git a/paper-server/patches/sources/net/minecraft/world/level/Level.java.patch b/paper-server/patches/sources/net/minecraft/world/level/Level.java.patch
index 2e46fe2b7de..142ee3dcfc2 100644
--- a/paper-server/patches/sources/net/minecraft/world/level/Level.java.patch
+++ b/paper-server/patches/sources/net/minecraft/world/level/Level.java.patch
@@ -84,7 +84,7 @@
 +    public boolean preventPoiUpdated = false; // CraftBukkit - SPIGOT-5710
 +    public boolean captureBlockStates = false;
 +    public boolean captureTreeGeneration = false;
-+    public Map<BlockPos, CapturedBlockState> capturedBlockStates = new java.util.LinkedHashMap<>();
++    public Map<BlockPos, org.bukkit.craftbukkit.block.CraftBlockState> capturedBlockStates = new java.util.LinkedHashMap<>(); // Paper
 +    public Map<BlockPos, BlockEntity> capturedTileEntities = new HashMap<>();
 +    public List<ItemEntity> captureDrops;
 +    public final it.unimi.dsi.fastutil.objects.Object2LongOpenHashMap<SpawnCategory> ticksPerSpawnCategory = new it.unimi.dsi.fastutil.objects.Object2LongOpenHashMap<>();
@@ -237,7 +237,7 @@
  
 +    // Paper start - if loaded
      @Nullable
-     @Override
++    @Override
 +    public final ChunkAccess getChunkIfLoadedImmediately(int x, int z) {
 +        return ((ServerLevel)this).chunkSource.getChunkAtIfLoadedImmediately(x, z);
 +    }
@@ -290,7 +290,7 @@
 +        return getWorldBorder().isWithinBounds(blockposition) ? getBlockStateIfLoaded(blockposition) : null;
 +    }
 +
-+    @Override
+     @Override
      public ChunkAccess getChunk(int chunkX, int chunkZ, ChunkStatus leastStatus, boolean create) {
 +        // Paper end
          ChunkAccess ichunkaccess = this.getChunkSource().getChunk(chunkX, chunkZ, leastStatus, create);
@@ -302,7 +302,7 @@
      public boolean setBlock(BlockPos pos, BlockState state, int flags, int maxUpdateDepth) {
 +        // CraftBukkit start - tree generation
 +        if (this.captureTreeGeneration) {
-+            CapturedBlockState blockstate = this.capturedBlockStates.get(pos);
++            CraftBlockState blockstate = this.capturedBlockStates.get(pos);
 +            if (blockstate == null) {
 +                blockstate = CapturedBlockState.getTreeBlockState(this, pos, flags);
 +                this.capturedBlockStates.put(pos.immutable(), blockstate);
@@ -315,7 +315,7 @@
          if (this.isOutsideBuildHeight(pos)) {
              return false;
          } else if (!this.isClientSide && this.isDebug()) {
-@@ -214,44 +400,123 @@
+@@ -214,44 +400,124 @@
          } else {
              LevelChunk chunk = this.getChunkAt(pos);
              Block block = state.getBlock();
@@ -324,7 +324,8 @@
 +            // CraftBukkit start - capture blockstates
 +            boolean captured = false;
 +            if (this.captureBlockStates && !this.capturedBlockStates.containsKey(pos)) {
-+                CapturedBlockState blockstate = CapturedBlockState.getBlockState(this, pos, flags);
++                CraftBlockState blockstate = (CraftBlockState) world.getBlockAt(pos.getX(), pos.getY(), pos.getZ()).getState(); // Paper - use CB getState to get a suitable snapshot
++                blockstate.setFlag(flags); // Paper - set flag
 +                this.capturedBlockStates.put(pos.immutable(), blockstate);
 +                captured = true;
 +            }
@@ -454,13 +455,13 @@
  
      public void onBlockStateChange(BlockPos pos, BlockState oldBlock, BlockState newBlock) {}
  
-@@ -340,10 +605,18 @@
+@@ -340,10 +606,18 @@
  
      @Override
      public BlockState getBlockState(BlockPos pos) {
 +        // CraftBukkit start - tree generation
 +        if (this.captureTreeGeneration) {
-+            CapturedBlockState previous = this.capturedBlockStates.get(pos);
++            CraftBlockState previous = this.capturedBlockStates.get(pos); // Paper
 +            if (previous != null) {
 +                return previous.getHandle();
 +            }
@@ -474,7 +475,7 @@
  
              return chunk.getBlockState(pos);
          }
-@@ -446,32 +719,44 @@
+@@ -446,32 +720,44 @@
              this.pendingBlockEntityTickers.clear();
          }
  
@@ -528,7 +529,7 @@
          }
      }
  
-@@ -510,13 +795,29 @@
+@@ -510,13 +796,29 @@
      @Nullable
      @Override
      public BlockEntity getBlockEntity(BlockPos pos) {
@@ -559,7 +560,7 @@
              this.getChunkAt(blockposition).addAndRegisterBlockEntity(blockEntity);
          }
      }
-@@ -643,7 +944,7 @@
+@@ -643,7 +945,7 @@
  
                  for (int k = 0; k < j; ++k) {
                      EnderDragonPart entitycomplexpart = aentitycomplexpart[k];
@@ -568,7 +569,7 @@
  
                      if (t0 != null && predicate.test(t0)) {
                          result.add(t0);
-@@ -912,7 +1213,7 @@
+@@ -912,7 +1214,7 @@
  
      public static enum ExplosionInteraction implements StringRepresentable {