:LspFold doesn't fold the right set of lines

autoload/lsp/lspserver.vim

@@ -1710,6 +1710,9 @@ def FoldRange(lspserver: dict<any>, fname: string)
     return
   endif
 
+  # Remove all the current folds
+  :normal! zE
+
   # interface FoldingRangeParams
   # interface TextDocumentIdentifier
   var params = {textDocument: {uri: util.LspFileToUri(fname)}}
@@ -1720,13 +1723,14 @@ def FoldRange(lspserver: dict<any>, fname: string)
 
   # result: FoldingRange[]
   var end_lnum: number
-  var last_lnum: number = line('$')
   for foldRange in reply.result
+    var start_lnum = foldRange.startLine + 1
     end_lnum = foldRange.endLine + 1
-    if end_lnum < foldRange.startLine + 2
-      end_lnum = foldRange.startLine + 2
+
+    if end_lnum < start_lnum
+      end_lnum = start_lnum
     endif
-    exe $':{foldRange.startLine + 2}, {end_lnum}fold'
+    exe $':{start_lnum}, {end_lnum}fold'
     # Open all the folds, otherwise the subsequently created folds are not
     # correct.
     :silent! foldopen!

test/gopls_tests.vim

@@ -38,7 +38,7 @@ def g:Test_LspGoto()
     func main() {
     }
   END
-  
+
   setline(1, lines)
   :redraw!
   g:WaitForServerFileLoad(0)
@@ -112,6 +112,42 @@ def g:Test_LspGoto()
   bw!
 enddef
 
+# Test for :LspFold command
+def g:Test_LspFold()
+  :silent! edit XLspFold1.go
+  sleep 200m
+  var lines =<< trim END
+    package main
+
+    // Some comment
+    // Some other comment
+
+    func plus(a int, b int) int {
+	return a + b
+    }
+
+    func main() {
+    }
+  END
+
+  setline(1, lines)
+  :redraw!
+  g:WaitForServerFileLoad(0)
+
+  :LspFold
+  sleep 50m
+  assert_equal(1, foldlevel(3))
+  var r = [0, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0]
+  assert_equal(r, range(1, 11)->map((_, v) => foldlevel(v)))
+  r = [-1, -1, 3, 3, -1, 6, 6, -1, -1, -1, -1]
+  assert_equal(r, range(1, 11)->map((_, v) => foldclosed(v)))
+  r = [-1, -1, 4, 4, -1, 7, 7, -1, -1, -1, -1]
+  assert_equal(r, range(1, 11)->map((_, v) => foldclosedend(v)))
+
+  v:errmsg = ''
+  bw!
+enddef
+
 # Start the gopls language server.  Returns true on success and false on
 # failure.
 def g:StartLangServer(): bool