feat: add WatchdogM wrapper for testing, collectDiagnostics and remove waitForRequests again after reconsideration

src/Lean/Data/JsonRpc.lean

@@ -21,6 +21,13 @@ inductive RequestID where
   | str (s : String)
   | num (n : JsonNumber)
   | null
+  deriving DecidableEq
+
+instance : ToString RequestID where
+  toString
+  | RequestID.str s => s!"\"s\""
+  | RequestID.num n => toString n
+  | RequestID.null => "null"
 
 /-- Error codes defined by JSON-RPC and LSP. -/
 inductive ErrorCode where
@@ -223,6 +230,18 @@ section
       else
         throw $ userError ("expected method '" ++ expectedMethod ++ "', got method '" ++ method ++ "'")
     | _ => throw $ userError "expected notification, got other type of message"
+
+  def readResponseAs (h : FS.Stream) (nBytes : Nat) (expectedID : RequestID) (α) [FromJson α] : IO (Response α) := do
+  let m ← h.readMessage nBytes
+  match m with
+  | Message.response id result =>
+    if id = expectedID then
+      match fromJson? result with
+      | some v => pure ⟨expectedID, v⟩
+      | none   => throw $ userError s!"unexpected result '{result.compress}'"
+    else
+      throw $ userError s!"expected id {expectedID}, got id {id}"
+  | _ => throw $ userError "expected response, got other type of message"
 end
 
 section

src/Lean/Data/Lsp/Communication.lean

@@ -56,6 +56,10 @@ section
   def readLspNotificationAs (h : FS.Stream) (expectedMethod : String) (α) [FromJson α] : IO (Notification α) := do
     let nBytes ← readLspHeader h
     h.readNotificationAs nBytes expectedMethod α
+
+  def readLspResponseAs (h : FS.Stream) (expectedID : RequestID) (α) [FromJson α] : IO (Response α) := do
+    let nBytes ← readLspHeader h
+    h.readResponseAs nBytes expectedID α
 end
 
 section
@@ -69,7 +73,7 @@ section
     h.putStr (header ++ j)
     h.flush
 
-  def writeLspRequest (h : FS.Stream) (r : Request α): IO Unit :=
+  def writeLspRequest (h : FS.Stream) (r : Request α) : IO Unit :=
     h.writeLspMessage r
 
   def writeLspNotification (h : FS.Stream) (n : Notification α) : IO Unit :=

src/Lean/Data/Lsp/Extra.lean

@@ -12,11 +12,8 @@ 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.
-- "textDocument/waitForResponses": Yields a response when all the requests that were pending at the time of
-  arrival of the packet have been emitted.
-Both of these exist for synchronization purposes, e.g. during testing or when external tools might want to use
-our LSP server.
+  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.
 -/
 
 namespace Lean.Lsp
@@ -42,23 +39,4 @@ instance : FromJson WaitForDiagnostics :=
 instance : ToJson WaitForDiagnostics :=
   ⟨fun o => mkObj []⟩
 
-structure WaitForResponsesParam where
-  uri : DocumentUri
-
-instance : FromJson WaitForResponsesParam :=
-  ⟨fun j => do
-    let id ← j.getObjValAs? DocumentUri "uri"
-    pure ⟨id⟩⟩
-
-instance : ToJson WaitForResponsesParam :=
-  ⟨fun o => mkObj [⟨"uri", toJson o.uri⟩]⟩
-
-structure WaitForResponses
-
-instance : FromJson WaitForResponses :=
-  ⟨fun j => WaitForResponses.mk⟩
-
-instance : ToJson WaitForResponses :=
-  ⟨fun o => mkObj []⟩
-
 end Lean.Lsp

src/Lean/Server/FileSource.lean

@@ -46,8 +46,5 @@ instance HoverParams.hasFileSource : FileSource HoverParams :=
 instance WaitForDiagnosticsParam.hasFileSource : FileSource WaitForDiagnosticsParam :=
   ⟨fun p => p.uri⟩
 
-instance WaitForResponsesParam.hasFileSource : FileSource WaitForResponsesParam :=
-  ⟨fun p => p.uri⟩
-
 end Lsp
 end Lean

src/Lean/Server/FileWorker.lean

@@ -220,16 +220,6 @@ section RequestHandling
     let e ← st.docRef.get
     discard $ e.cmdSnaps.waitAll
     st.hOut.writeLspResponse ⟨id, WaitForDiagnostics.mk⟩
-
-  /- The signature of this request is somewhat special, since it is a request that performs
-     an action related to requests. Hence we need to ensure that this handler does not
-     accidentally reference itself, by passing in the `PendingRequestMap` from before
-     the request was added. -/
-  def handleWaitForResponses (id : RequestID) (pr : PendingRequestMap) : ServerM Unit := do
-    for ⟨_, task⟩ in pr do
-      discard $ IO.wait task
-    (←read).hOut.writeLspResponse ⟨id, WaitForResponses.mk⟩
-
 end RequestHandling
 
 section MessageHandling
@@ -257,7 +247,6 @@ section MessageHandling
                       (handler : RequestID → paramType → ServerM Unit) =>
       parseParams paramType params >>= queueRequest id handler
     match method with
-    | "textDocument/waitForResponses"   => queueRequest id handleWaitForResponses (←(←read).pendingRequestsRef.get)
     | "textDocument/waitForDiagnostics" => handle WaitForDiagnosticsParam handleWaitForDiagnostics
     | "textDocument/hover"              => handle HoverParams handleHover
     | _ => throwServerError $ "Got unsupported request: " ++ method ++

src/Lean/Server/Watchdog.lean

@@ -342,7 +342,6 @@ section MessageHandling
       let uri := fileSource parsedParams
       tryWriteMessage uri ⟨id, method, parsedParams⟩ FileWorker.writeRequest
     match method with
-    | "textDocument/waitForResponses"   => handle WaitForResponsesParam
     | "textDocument/waitForDiagnostics" => handle WaitForDiagnosticsParam
     | "textDocument/hover"              => handle HoverParams
     | _ => throwServerError $ "Got unsupported request: " ++ method ++ " params: " ++ toString params
@@ -433,7 +432,7 @@ def initAndRunWatchdogAux : ServerM Unit := do
      throw err
 
 def initAndRunWatchdog (i o e : FS.Stream) : IO Unit := do
-  let workerPath ← match (← IO.getEnv "LEAN_WORKER_PATH") with
+  let workerPath ← match (←IO.getEnv "LEAN_WORKER_PATH") with
     | none   => IO.appPath
     | some p => p
   let fileWorkersRef ← IO.mkRef (RBMap.empty : FileWorkerMap)
@@ -464,14 +463,55 @@ def initAndRunWatchdog (i o e : FS.Stream) : IO Unit := do
 
 namespace Test
 
-def runWatchdogWithInputFile (fn : String) (searchPath : Option String) : IO Unit := do
-  let o ← IO.getStdout
-  let e ← IO.getStderr
-  FS.withFile fn FS.Mode.read $ fun hFile => do
-    Lean.initSearchPath searchPath
-    try initAndRunWatchdog (FS.Stream.ofHandle hFile) o e
-    catch err => e.putStrLn (toString err)
+def watchdogCfg : Process.StdioConfig where
+  stdin  := Process.Stdio.piped
+  stdout := Process.Stdio.piped
+  stderr := Process.Stdio.piped
 
+abbrev WatchdogM := ReaderT (Process.Child watchdogCfg) IO
+
+variable [ToJson α]
+
+def stdin : WatchdogM FS.Stream := do
+  FS.Stream.ofHandle (←read).stdin
+
+def stdout : WatchdogM FS.Stream := do
+  FS.Stream.ofHandle (←read).stdout
+
+def writeRequest (r : Request α) : WatchdogM Unit := do
+  (←stdin).writeLspRequest r
+
+def writeNotification (n : Notification α) : WatchdogM Unit := do
+  (←stdin).writeLspNotification n
+
+def readResponseAs (expectedID : RequestID) (α) [FromJson α] : WatchdogM (Response α) := do
+  (←stdout).readLspResponseAs expectedID α
+
+partial def collectDiagnostics (waitForDiagnosticsId : RequestID := 0) (target : DocumentUri)
+: WatchdogM (Array (Notification PublishDiagnosticsParams)) := do
+  writeRequest ⟨waitForDiagnosticsId, "textDocument/waitForDiagnostics", WaitForDiagnosticsParam.mk target⟩
+  let rec loop : WatchdogM (Array (Notification PublishDiagnosticsParams)) := do
+    let error : IO (Array (Notification PublishDiagnosticsParams)) := throw $ userError "got unexpected packet while collecting diagnostics"
+    match ←(←stdout).readLspMessage with
+    | Message.response id _ =>
+      if id = waitForDiagnosticsId then
+        #[]
+      else
+        error
+    | Message.notification "textDocument/publishDiagnostics" (some param) =>
+      match fromJson? (toJson param) with
+      | some diagnosticParam => (←loop).push ⟨"textDocument/publishDiagnostics", diagnosticParam⟩
+      | none                 => error
+    | _ => error
+  loop
+
+def runWatchdogTest (test : WatchdogM Unit) : IO Unit := do
+  let proc ← Process.spawn {
+    toStdioConfig := watchdogCfg
+    cmd           := ←IO.appPath
+    args          := #["--server"]
+  }
+  ReaderT.run test proc
 end Test
 
 @[export lean_server_watchdog_main]