feat: specify version in waitForDiagnostics

src/Lean/Data/Lsp/Extra.lean

@@ -8,20 +8,22 @@ import Lean.Data.Json
 import Lean.Data.JsonRpc
 import Lean.Data.Lsp.Basic
 
-/-
+/-!
 This file contains Lean-specific extensions to LSP.
 The following additional packets are supported:
-- "textDocument/waitForDiagnostics": Yields a response when all the diagnostics that were pending at the time of
-  arrival of the packet have been emitted. Exists for synchronization purposes, e.g. during testing or when external tools might want to use
-  our LSP server.
+- "textDocument/waitForDiagnostics": Yields a response when all the diagnostics for a version of the document
+  greater or equal to the specified one have been emitted. If the request specifies a version above the most
+  recently processed one, the server will delay the response until it does receive the specified version.
+  Exists for synchronization purposes, e.g. during testing or when external tools might want to use our LSP server.
 -/
 
 namespace Lean.Lsp
 
 open Json
 
-structure WaitForDiagnosticsParam where
-  uri : DocumentUri
+structure WaitForDiagnosticsParams where
+  uri     : DocumentUri
+  version : Nat
   deriving ToJson, FromJson
 
 structure WaitForDiagnostics

src/Lean/Data/Lsp/Ipc.lean

@@ -50,9 +50,9 @@ def waitForExit : IpcM UInt32 := do
 
 /-- Waits for the worker to emit all diagnostics for the current document version
 and returns them as a list. -/
-partial def collectDiagnostics (waitForDiagnosticsId : RequestID := 0) (target : DocumentUri)
+partial def collectDiagnostics (waitForDiagnosticsId : RequestID := 0) (target : DocumentUri) (version : Nat)
 : IpcM (List (Notification PublishDiagnosticsParams)) := do
-  writeRequest ⟨waitForDiagnosticsId, "textDocument/waitForDiagnostics", WaitForDiagnosticsParam.mk target⟩
+  writeRequest ⟨waitForDiagnosticsId, "textDocument/waitForDiagnostics", WaitForDiagnosticsParams.mk target version⟩
   let rec loop : IpcM (List (Notification PublishDiagnosticsParams)) := do
     match ←readMessage with
     | Message.response id _ =>

src/Lean/Server/FileSource.lean

@@ -52,7 +52,7 @@ instance DefinitionParams.hasFileSource : FileSource DefinitionParams :=
 instance TypeDefinitionParams.hasFileSource : FileSource TypeDefinitionParams :=
   ⟨fun h => fileSource h.toTextDocumentPositionParams⟩
 
-instance WaitForDiagnosticsParam.hasFileSource : FileSource WaitForDiagnosticsParam :=
+instance WaitForDiagnosticsParams.hasFileSource : FileSource WaitForDiagnosticsParams :=
   ⟨fun p => p.uri⟩
 
 instance DocumentSymbolParams.hasFileSource : FileSource DocumentSymbolParams :=

src/Lean/Server/FileWorker.lean

@@ -488,12 +488,23 @@ section RequestHandling
             children? := syms.toArray
           } :: syms', stxs'')
 
-  def handleWaitForDiagnostics (id : RequestID) (p : WaitForDiagnosticsParam)
+  partial def handleWaitForDiagnostics (id : RequestID) (p : WaitForDiagnosticsParams)
     : ServerM (Task (Except IO.Error (Except RequestError WaitForDiagnostics))) := do
     let st ← read
-    let doc ← st.docRef.get
-    let t ← doc.cmdSnaps.waitAll
-    t.map fun _ => Except.ok $ Except.ok WaitForDiagnostics.mk
+    let rec waitLoop : IO EditableDocument := do
+      let doc ← st.docRef.get
+      if p.version ≤ doc.meta.version then
+        return doc
+      else
+        IO.sleep 50
+        waitLoop
+    let t ← IO.asTask waitLoop
+    let t ← IO.bindTask t fun
+      | Except.error e => unreachable!
+      | Except.ok doc => do
+        let t₁ ← doc.cmdSnaps.waitAll
+        return t₁.map fun _ => Except.ok WaitForDiagnostics.mk
+    return t.map fun _ => Except.ok <| Except.ok WaitForDiagnostics.mk
 
 end RequestHandling
 
@@ -531,7 +542,7 @@ section MessageHandling
           st.hOut.writeLspResponseError { id := id, code := ErrorCode.internalError, message := toString e }
       queueRequest id t₁
     match method with
-    | "textDocument/waitForDiagnostics" => handle WaitForDiagnosticsParam WaitForDiagnostics handleWaitForDiagnostics
+    | "textDocument/waitForDiagnostics" => handle WaitForDiagnosticsParams WaitForDiagnostics handleWaitForDiagnostics
     | "textDocument/hover"              => handle HoverParams (Option Hover) handleHover
     | "textDocument/declaration"        => handle TextDocumentPositionParams (Array LocationLink) <| handleDefinition (goToType? := false)
     | "textDocument/definition"         => handle TextDocumentPositionParams (Array LocationLink) <| handleDefinition (goToType? := false)

src/Lean/Server/Watchdog.lean

@@ -379,7 +379,7 @@ section MessageHandling
       let uri := fileSource parsedParams
       tryWriteMessage uri ⟨id, method, parsedParams⟩ FileWorker.writeRequest
     match method with
-    | "textDocument/waitForDiagnostics" => handle WaitForDiagnosticsParam
+    | "textDocument/waitForDiagnostics" => handle WaitForDiagnosticsParams
     | "textDocument/hover"              => handle HoverParams
     | "textDocument/declaration"        => handle DeclarationParams
     | "textDocument/definition"         => handle DefinitionParams

tests/lean/server/diags.lean

@@ -11,7 +11,7 @@ open IO Lean Lsp
 
     hIn.write (←FS.readBinFile "open_content.log")
     hIn.flush
-    let diags ← Ipc.collectDiagnostics 1 "file:///test.lean"
+    let diags ← Ipc.collectDiagnostics 1 "file:///test.lean" 1
     if diags.isEmpty then
       throw $ userError "Test failed, no diagnostics received."
     else

tests/lean/server/edits.lean

@@ -12,12 +12,12 @@ open IO Lean Lsp
 
     hIn.write (←FS.readBinFile "open_content.log")
     hIn.flush
-    discard <| Ipc.collectDiagnostics 1 "file:///test.lean"
+    discard <| Ipc.collectDiagnostics 1 "file:///test.lean" 1
 
     hIn.write (←FS.readBinFile "content_changes.log")
     hIn.flush
 
-    let diags ← Ipc.collectDiagnostics 2 "file:///test.lean"
+    let diags ← Ipc.collectDiagnostics 2 "file:///test.lean" 7
     if diags.isEmpty then
       throw $ userError "Test failed, no diagnostics received."
     else

tests/lean/server/init_exit_worker.lean

@@ -9,7 +9,7 @@ open IO Lean Lsp
 
     hIn.write (←FS.readBinFile "open_empty.log")
     hIn.flush
-    let diags ← Ipc.collectDiagnostics 1 "file:///empty"
+    let diags ← Ipc.collectDiagnostics 1 "file:///empty" 1
     assert! diags.length == 1
 
     Ipc.writeNotification ⟨"exit", Json.null⟩