feat: server -> client requests (#3271)

This PR adds support for requests from the server to the client in the
language server. It is based on #3014 and was developed during an
experiment for #3247 that unfortunately did not go anywhere.

src/Lean/Data/Lsp/Window.lean

@@ -0,0 +1,48 @@
+/-
+Copyright (c) 2024 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Marc Huisinga
+-/
+prelude
+import Lean.Data.Json
+
+open Lean
+
+inductive MessageType where
+  | error
+  | warning
+  | info
+  | log
+
+instance : FromJson MessageType where
+  fromJson?
+    | (1 : Nat) => .ok .error
+    | (2 : Nat) => .ok .warning
+    | (3 : Nat) => .ok .info
+    | (4 : Nat) => .ok .log
+    | _         => .error "Unknown MessageType ID"
+
+instance : ToJson MessageType where
+  toJson
+    | .error   => 1
+    | .warning => 2
+    | .info    => 3
+    | .log     => 4
+
+structure ShowMessageParams where
+  type    : MessageType
+  message : String
+  deriving FromJson, ToJson
+
+structure MessageActionItem where
+  title : String
+  deriving FromJson, ToJson
+
+structure ShowMessageRequestParams where
+  type     : MessageType
+  message  : String
+  actions? : Option (Array MessageActionItem)
+  deriving FromJson, ToJson
+
+def ShowMessageResponse := Option MessageActionItem
+  deriving FromJson, ToJson

src/Lean/Server/FileWorker.lean

@@ -222,6 +222,7 @@ structure WorkerState where
   /-- A map of RPC session IDs. We allow asynchronous elab tasks and request handlers
   to modify sessions. A single `Ref` ensures atomic transactions. -/
   rpcSessions        : RBMap UInt64 (IO.Ref RpcSession) compare
+  freshRequestID     : Nat
 
 abbrev WorkerM := ReaderT WorkerContext <| StateRefT WorkerState IO
 
@@ -315,6 +316,7 @@ section Initialization
       pendingRequests    := RBMap.empty
       rpcSessions        := RBMap.empty
       importCachingTask? := none
+      freshRequestID     := 0
     })
   where
     /-- Creates an LSP message output channel along with a reader that sends out read messages on
@@ -538,11 +540,15 @@ section MessageHandling
           | Except.error e =>
             ctx.chanOut.send <| e.toLspResponseError id
     queueRequest id t
+
+  def handleResponse (_ : RequestID) (result : Json) : WorkerM Unit :=
+    throwServerError s!"Unknown response kind: {result}"
+
 end MessageHandling
 
 section MainLoop
   variable (hIn : FS.Stream) in
-  partial def mainLoop : WorkerM Unit := do
+  partial def mainLoop : WorkerM UInt32 := do
     let mut st ← get
     let msg ← hIn.readLspMessage
     let filterFinishedTasks (acc : PendingRequestMap) (id : RequestID) (task : Task (Except IO.Error Unit))
@@ -569,10 +575,17 @@ section MainLoop
       handleRequest id method (toJson params)
       mainLoop
     | Message.notification "exit" none =>
-      return ()
+      return 0
     | Message.notification method (some params) =>
       handleNotification method (toJson params)
       mainLoop
+    | Message.response id result =>
+      handleResponse id result
+      mainLoop
+    | Message.responseError .. =>
+      -- Ignore all errors as we currently only handle a single request with an optional response
+      -- where failure is not an issue.
+      mainLoop
     | _ => throwServerError "Got invalid JSON-RPC message"
 end MainLoop
 

src/Lean/Server/Watchdog.lean

@@ -142,18 +142,49 @@ end FileWorker
 section ServerM
   abbrev FileWorkerMap := RBMap DocumentUri FileWorker compare
 
+  structure RequestIDTranslation where
+    sourceUri : DocumentUri
+    localID   : RequestID
+
+  abbrev PendingServerRequestMap := RBMap RequestID RequestIDTranslation compare
+
+  structure ServerRequestData where
+    pendingServerRequests : PendingServerRequestMap
+    freshServerRequestID  : Nat
+
+  def ServerRequestData.trackOutboundRequest
+      (data      : ServerRequestData)
+      (sourceUri : DocumentUri)
+      (localID   : RequestID)
+      : RequestID × ServerRequestData :=
+    let globalID := data.freshServerRequestID
+    let data := {
+      pendingServerRequests := data.pendingServerRequests.insert globalID ⟨sourceUri, localID⟩
+      freshServerRequestID  := globalID + 1
+    }
+    (globalID, data)
+
+  def ServerRequestData.translateInboundResponse
+      (data     : ServerRequestData)
+      (globalID : RequestID)
+      : Option RequestIDTranslation × ServerRequestData :=
+    let translation? := data.pendingServerRequests.find? globalID
+    let data := { data with pendingServerRequests := data.pendingServerRequests.erase globalID }
+    (translation?, data)
+
   structure ServerContext where
-    hIn            : FS.Stream
-    hOut           : FS.Stream
-    hLog           : FS.Stream
+    hIn               : FS.Stream
+    hOut              : FS.Stream
+    hLog              : FS.Stream
     /-- Command line arguments. -/
-    args           : List String
-    fileWorkersRef : IO.Ref FileWorkerMap
+    args              : List String
+    fileWorkersRef    : IO.Ref FileWorkerMap
     /-- We store these to pass them to workers. -/
-    initParams     : InitializeParams
-    workerPath     : System.FilePath
-    srcSearchPath  : System.SearchPath
-    references     : IO.Ref References
+    initParams        : InitializeParams
+    workerPath        : System.FilePath
+    srcSearchPath     : System.SearchPath
+    references        : IO.Ref References
+    serverRequestData : IO.Ref ServerRequestData
 
   abbrev ServerM := ReaderT ServerContext IO
 
@@ -210,6 +241,10 @@ section ServerM
           | Message.responseError id _ _ _ => do
             fw.erasePendingRequest id
             o.writeLspMessage msg
+          | Message.request id method params? =>
+            let globalID ← (←read).serverRequestData.modifyGet
+              (·.trackOutboundRequest fw.doc.uri id)
+            o.writeLspMessage (Message.request globalID method params?)
           | Message.notification "$/lean/ileanInfoUpdate" params =>
             if let some params := params then
               if let Except.ok params := FromJson.fromJson? <| ToJson.toJson params then
@@ -767,6 +802,21 @@ section MessageHandling
       -- implementation-dependent notifications can be safely ignored
       if !"$/".isPrefixOf method then
         (←read).hLog.putStrLn s!"Got unsupported notification: {method}"
+
+  def handleResponse (id : RequestID) (result : Json) : ServerM Unit := do
+    let some translation ← (← read).serverRequestData.modifyGet (·.translateInboundResponse id)
+      | return
+    tryWriteMessage translation.sourceUri (Response.mk translation.localID result)
+
+  def handleResponseError
+      (id      : RequestID)
+      (code    : ErrorCode)
+      (message : String)
+      (data?   : Option Json)
+      : ServerM Unit := do
+    let some translation ← (← read).serverRequestData.modifyGet (·.translateInboundResponse id)
+      | return
+    tryWriteMessage translation.sourceUri (ResponseError.mk translation.localID code message data?)
 end MessageHandling
 
 section MainLoop
@@ -811,11 +861,12 @@ section MainLoop
       | Message.request id method (some params) =>
         handleRequest id method (toJson params)
         mainLoop (←runClientTask)
-      | Message.response .. =>
-        -- TODO: handle client responses
+      | Message.response id result =>
+        handleResponse id result
+        mainLoop (←runClientTask)
+      | Message.responseError id code message data? =>
+        handleResponseError id code message data?
         mainLoop (←runClientTask)
-      | Message.responseError _ _ e .. =>
-        throwServerError s!"Unhandled response error: {e}"
       | Message.notification method (some params) =>
         handleNotification method (toJson params)
         mainLoop (←runClientTask)
@@ -929,6 +980,10 @@ def initAndRunWatchdog (args : List String) (i o e : FS.Stream) : IO Unit := do
   let references ← IO.mkRef .empty
   startLoadingReferences references
   let fileWorkersRef ← IO.mkRef (RBMap.empty : FileWorkerMap)
+  let serverRequestData ← IO.mkRef {
+    pendingServerRequests := RBMap.empty
+    freshServerRequestID  := 0
+  }
   let i ← maybeTee "wdIn.txt" false i
   let o ← maybeTee "wdOut.txt" true o
   let e ← maybeTee "wdErr.txt" true e
@@ -967,6 +1022,7 @@ def initAndRunWatchdog (args : List String) (i o e : FS.Stream) : IO Unit := do
     workerPath
     srcSearchPath
     references
+    serverRequestData
     : ServerContext
   }