fix: Clear Diagnostics when file is closed

src/Lean/Data/Lsp/Ipc.lean

@@ -39,6 +39,23 @@ def writeRequest (r : Request α) : IpcM Unit := do
 def writeNotification (n : Notification α) : IpcM Unit := do
   (←stdin).writeLspNotification n
 
+def shutdown (requestNo : Nat) : IpcM Unit := do
+  let hIn ← stdout
+  let hOut ← stdin
+  hOut.writeLspRequest ⟨requestNo, "shutdown", Json.null⟩
+  repeat
+    let shutMsg ← hIn.readLspMessage
+    match shutMsg with
+    | Message.response id result =>
+      assert! result.isNull
+      if id != requestNo then
+        throw <| IO.userError s!"Expected id {requestNo}, got id {id}"
+
+      hOut.writeLspNotification ⟨"exit", Json.null⟩
+      break
+    | _ =>  -- ignore other messages in between.
+      pure ()
+
 def readMessage : IpcM JsonRpc.Message := do
   (←stdout).readLspMessage
 

src/Lean/Server/FileWorker.lean

@@ -421,6 +421,16 @@ section MessageHandling
 end MessageHandling
 
 section MainLoop
+
+  /-- called when document is closed so we can clear any diagnostics and cancel pending work -/
+  def exitWorker : WorkerM Unit := do
+    let ctx ← read
+    let state ← get
+    let doc := state.doc
+    doc.cancelTk.set
+    -- this is acting like an acknowledgement back to the Watchdog that exit was in fact received.
+    publishDiagnostics doc.meta #[] ctx.hOut
+
   partial def mainLoop : WorkerM Unit := do
     let ctx ← read
     let mut st ← get
@@ -449,8 +459,7 @@ section MainLoop
       handleRequest id method (toJson params)
       mainLoop
     | Message.notification "exit" none =>
-      let doc := (←get).doc
-      doc.cancelTk.set
+      exitWorker
       return ()
     | Message.notification method (some params) =>
       handleNotification method (toJson params)

src/Lean/Server/Watchdog.lean

@@ -99,6 +99,7 @@ section Utils
     worker arrives. -/
     | crashed (queuedMsgs : Array JsonRpc.Message)
     | running
+    | exiting
 
   abbrev PendingRequestMap := RBMap RequestID JsonRpc.Message compare
 
@@ -123,7 +124,7 @@ section FileWorker
     doc                : OpenDocument
     proc               : Process.Child workerCfg
     commTask           : Task WorkerEvent
-    state              : WorkerState
+    state              : IO.Ref WorkerState
     -- This should not be mutated outside of namespace FileWorker, as it is used as shared mutable state
     /-- The pending requests map contains all requests
     that have been received from the LSP client, but were not answered yet.
@@ -224,7 +225,8 @@ section ServerM
   /-- Creates a Task which forwards a worker's messages into the output stream until an event
   which must be handled in the main watchdog thread (e.g. an I/O error) happens. -/
   private partial def forwardMessages (fw : FileWorker) : ServerM (Task WorkerEvent) := do
-    let o := (←read).hOut
+    let ctx ← read
+    let o := ctx.hOut
     let rec loop : ServerM WorkerEvent := do
       try
         let msg ← fw.stdout.readLspMessage
@@ -246,14 +248,26 @@ section ServerM
               if let Except.ok params := FromJson.fromJson? <| ToJson.toJson params then
                 handleIleanInfoFinal fw params
           | _ => o.writeLspMessage msg
+
+        match (← fw.state.get) with
+        | .exiting =>
+          -- the above may have just sent the response to the "exit" command,
+          -- and now we ensure that clear diagnostics is the last thing we should do
+          -- so we never leave stale diagnostics around for this file.
+          let meta := fw.doc.meta
+          dbg_trace "sending clear diagnostics from Watchdog!"
+          publishDiagnostics meta #[] o
+          return WorkerEvent.terminated
+        | _ => pure
+
       catch err =>
         -- If writeLspMessage from above errors we will block here, but the main task will
         -- quit eventually anyways if that happens
         let exitCode ← fw.proc.wait
         if exitCode = 0 then
           -- Worker was terminated
           fw.errorPendingRequests o ErrorCode.contentModified
-            ("The file worker has been terminated. Either the header has changed,"
+            (s!"The file worker for {fw.doc.meta.uri} has been terminated. Either the header has changed,"
             ++ " or the file was closed, or the server is shutting down.")
           return WorkerEvent.terminated
         else
@@ -263,6 +277,7 @@ section ServerM
           publishProgressAtPos fw.doc.meta 0 o (kind := LeanFileProgressKind.fatalError)
           return WorkerEvent.crashed err
       loop
+
     let task ← IO.asTask (loop $ ←read) Task.Priority.dedicated
     return task.map fun
       | Except.ok ev   => ev
@@ -283,7 +298,7 @@ section ServerM
       doc                := ⟨m, headerAst⟩
       proc               := workerProc
       commTask           := Task.pure WorkerEvent.terminated
-      state              := WorkerState.running
+      state              := ← IO.mkRef WorkerState.running
       pendingRequestsRef := pendingRequestsRef
       groupedEditsRef    := ← IO.mkRef none
     }
@@ -306,6 +321,7 @@ section ServerM
   def terminateFileWorker (uri : DocumentUri) : ServerM Unit := do
     let fw ← findFileWorker! uri
     try
+      fw.state.modify (λ _ => WorkerState.exiting)
       fw.stdin.writeLspMessage (Message.notification "exit" none)
     catch _ =>
       /- The file worker must have crashed just when we were about to terminate it!
@@ -318,7 +334,7 @@ section ServerM
     eraseFileWorker uri
 
   def handleCrash (uri : DocumentUri) (queuedMsgs : Array JsonRpc.Message) : ServerM Unit := do
-    updateFileWorkers { ←findFileWorker! uri with state := WorkerState.crashed queuedMsgs }
+    (←findFileWorker! uri).state.modify (λ _ => WorkerState.crashed queuedMsgs)
 
   /-- Tries to write a message, sets the state of the FileWorker to `crashed` if it does not succeed
       and restarts the file worker if the `crashed` flag was already set. Just logs an error if there
@@ -336,7 +352,7 @@ section ServerM
       | none    => (false, none)
     if pendingEdit then
       return
-    match fw.state with
+    match (← fw.state.get) with
     | WorkerState.crashed queuedMsgs =>
       let mut queuedMsgs := queuedMsgs
       if queueFailedMessage then
@@ -356,7 +372,7 @@ section ServerM
           crashedMsgs := crashedMsgs.push msg
       if ¬ crashedMsgs.isEmpty then
         handleCrash uri crashedMsgs
-    | WorkerState.running =>
+    | _ =>
       let initialQueuedMsgs :=
         if queueFailedMessage then
           #[msg]
@@ -587,7 +603,7 @@ section MainLoop
     let workers ← st.fileWorkersRef.get
     let mut workerTasks := #[]
     for (_, fw) in workers do
-      if let WorkerState.running := fw.state then
+      if let WorkerState.running := (← fw.state.get) then
         workerTasks := workerTasks.push <| fw.commTask.map (ServerEvent.workerEvent fw)
         if let some ge ← fw.groupedEditsRef.get then
           workerTasks := workerTasks.push <| ge.signalTask.map (ServerEvent.workerEvent fw)

tests/lean/interactive/run.lean

@@ -178,8 +178,6 @@ partial def main (args : List String) : IO Unit := do
       | _ =>
         lastActualLineNo := lineNo
       lineNo := lineNo + 1
-    Ipc.writeRequest ⟨requestNo, "shutdown", Json.null⟩
-    let shutResp ← Ipc.readResponseAs requestNo Json
-    assert! shutResp.result.isNull
-    Ipc.writeNotification ⟨"exit", Json.null⟩
+
+    Ipc.shutdown requestNo
     discard <| Ipc.waitForExit

tests/lean/server/diags.lean

@@ -25,8 +25,5 @@ def main : IO Unit := do
       else
         throw $ userError "Failed parsing test file."
 
-      Ipc.writeRequest ⟨2, "shutdown", Json.null⟩
-      let shutResp ← Ipc.readResponseAs 2 Json
-      assert! shutResp.result.isNull
-      Ipc.writeNotification ⟨"exit", Json.null⟩
+      Ipc.shutdown 2
       discard $ Ipc.waitForExit

tests/lean/server/edits.lean

@@ -30,8 +30,5 @@ def main : IO Unit := do
       else
         throw $ userError "Failed parsing test file."
 
-      Ipc.writeRequest ⟨3, "shutdown", Json.null⟩
-      let shutResp ← Ipc.readResponseAs 3 Json
-      assert! shutResp.result.isNull
-      Ipc.writeNotification ⟨"exit", Json.null⟩
+      Ipc.shutdown 3
       discard $ Ipc.waitForExit

tests/lean/server/init_exit.lean

@@ -8,10 +8,7 @@ def main : IO Unit := do
     hIn.flush
     let initResp ← Ipc.readResponseAs 0 InitializeResult
     let regWatchReq ← Ipc.readRequestAs "client/registerCapability" Json
-    Ipc.writeNotification ⟨"initialized", InitializedParams.mk⟩ 
+    Ipc.writeNotification ⟨"initialized", InitializedParams.mk⟩
 
-    Ipc.writeRequest ⟨1, "shutdown", Json.null⟩
-    let shutdownResp ← Ipc.readResponseAs 1 Json
-    assert! shutdownResp.result.isNull
-    Ipc.writeNotification ⟨"exit", Json.null⟩
+    Ipc.shutdown 1
     discard Ipc.waitForExit