From 46d912d0ce70b8238be31936778a5b12ac031b61 Mon Sep 17 00:00:00 2001
From: Greg Oledzki <greg.oledzki@moderne.io>
Date: Tue, 24 Mar 2026 15:50:41 +0100
Subject: [PATCH] Fix comment handling in DeleteKey

---
 .../java/org/openrewrite/toml/DeleteKey.java  | 46 +++++++++++++++++++
 .../org/openrewrite/toml/DeleteKeyTest.java   |  2 -
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/rewrite-toml/src/main/java/org/openrewrite/toml/DeleteKey.java b/rewrite-toml/src/main/java/org/openrewrite/toml/DeleteKey.java
index fa162fb606f..d18b6f1e428 100644
--- a/rewrite-toml/src/main/java/org/openrewrite/toml/DeleteKey.java
+++ b/rewrite-toml/src/main/java/org/openrewrite/toml/DeleteKey.java
@@ -22,7 +22,13 @@
 import org.openrewrite.Option;
 import org.openrewrite.Recipe;
 import org.openrewrite.TreeVisitor;
+import org.openrewrite.toml.tree.Comment;
+import org.openrewrite.toml.tree.Space;
 import org.openrewrite.toml.tree.Toml;
+import org.openrewrite.toml.tree.TomlRightPadded;
+
+import java.util.ArrayList;
+import java.util.List;
 
 @Value
 @EqualsAndHashCode(callSuper = false)
@@ -45,6 +51,46 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
                 Toml.KeyValue kv = super.visitKeyValue(keyValue, ctx);
                 return matcher.matches(getCursor()) ? null : kv;
             }
+
+            @Override
+            public Toml.Table visitTable(Toml.Table table, ExecutionContext ctx) {
+                Toml.Table t = table;
+                t = t.withPrefix(visitSpace(t.getPrefix(), ctx));
+                t = t.withMarkers(visitMarkers(t.getMarkers(), ctx));
+
+                List<TomlRightPadded<Toml>> original = t.getPadding().getValues();
+                List<TomlRightPadded<Toml>> result = new ArrayList<>(original.size());
+                boolean changed = false;
+                boolean prevDeleted = false;
+
+                for (TomlRightPadded<Toml> rp : original) {
+                    Toml visited = visit(rp.getElement(), ctx);
+                    if (visited == null) {
+                        prevDeleted = changed = true;
+                        continue;
+                    }
+                    if (prevDeleted) {
+                        visited = visited.withPrefix(
+                                stripInlineCommentFromDeletedLine(visited.getPrefix()));
+                        prevDeleted = false;
+                    }
+                    TomlRightPadded<Toml> newRp = visited == rp.getElement() ? rp : rp.withElement(visited);
+                    if (newRp != rp) changed = true;
+                    result.add(newRp);
+                }
+                return changed ? t.getPadding().withValues(result) : t;
+            }
+
+            private Space stripInlineCommentFromDeletedLine(Space prefix) {
+                List<Comment> comments = prefix.getComments();
+                if (prefix.getWhitespace().contains("\n") || comments.isEmpty()) {
+                    return prefix;
+                }
+                return Space.build(
+                        comments.get(0).getSuffix(),
+                        new ArrayList<>(comments.subList(1, comments.size()))
+                );
+            }
         };
     }
 }
diff --git a/rewrite-toml/src/test/java/org/openrewrite/toml/DeleteKeyTest.java b/rewrite-toml/src/test/java/org/openrewrite/toml/DeleteKeyTest.java
index e630b960a39..ad573c2cef4 100644
--- a/rewrite-toml/src/test/java/org/openrewrite/toml/DeleteKeyTest.java
+++ b/rewrite-toml/src/test/java/org/openrewrite/toml/DeleteKeyTest.java
@@ -15,7 +15,6 @@
  */
 package org.openrewrite.toml;
 
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.openrewrite.DocumentExample;
 import org.openrewrite.test.RewriteTest;
@@ -225,7 +224,6 @@ void deleteKeyWithArrayValue() {
         );
     }
 
-    @Disabled
     @Test
     void preserveComments() {
         rewriteRun(