diff --git a/format/format.go b/format/format.go
index c55d7998d1..2376c94db3 100644
--- a/format/format.go
+++ b/format/format.go
@@ -300,8 +300,10 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment)
 
 	if len(rule.Head.Args) > 0 {
 		closeLoc = closingLoc('(', ')', '{', '}', rule.Location)
-	} else {
+	} else if rule.Head.Key != nil {
 		closeLoc = closingLoc('[', ']', '{', '}', rule.Location)
+	} else {
+		closeLoc = closingLoc(0, 0, '{', '}', rule.Location)
 	}
 
 	comments = w.insertComments(comments, closeLoc)
diff --git a/format/testfiles/test_end_of_rule_comment.rego b/format/testfiles/test_end_of_rule_comment.rego
new file mode 100644
index 0000000000..0c3bd75bee
--- /dev/null
+++ b/format/testfiles/test_end_of_rule_comment.rego
@@ -0,0 +1,7 @@
+package foo
+
+bar {
+	# before
+	input.bar
+	# after
+}
diff --git a/format/testfiles/test_end_of_rule_comment.rego.formatted b/format/testfiles/test_end_of_rule_comment.rego.formatted
new file mode 100644
index 0000000000..0c3bd75bee
--- /dev/null
+++ b/format/testfiles/test_end_of_rule_comment.rego.formatted
@@ -0,0 +1,7 @@
+package foo
+
+bar {
+	# before
+	input.bar
+	# after
+}
diff --git a/format/testfiles/test_issue_2299.rego b/format/testfiles/test_issue_2299.rego
index 5f1c5b74f0..5af69145d9 100644
--- a/format/testfiles/test_issue_2299.rego
+++ b/format/testfiles/test_issue_2299.rego
@@ -25,7 +25,8 @@ else = z {
 # Mixed compact and newline separated
 p = x {
     foo == "bar"
-} # some special case
+}
+# some special case
 # with lots of comments
 # describing it
 else = y {