chore: make rpc/connect a request

src/Lean/Data/Lsp/Extra.lean

@@ -26,6 +26,7 @@ structure WaitForDiagnosticsParams where
   version : Nat
   deriving FromJson, ToJson
 
+/-- `textDocument/waitForDiagnostics` client<-server reply. -/
 structure WaitForDiagnostics
 
 instance : FromJson WaitForDiagnostics :=
@@ -53,6 +54,7 @@ Otherwise returns `null`. -/
 structure PlainGoalParams extends TextDocumentPositionParams
   deriving FromJson, ToJson
 
+/-- `$/lean/plainGoal` client<-server reply. -/
 structure PlainGoal where
   /-- The goals as pretty-printed Markdown, or something like "no goals" if accomplished. -/
   rendered : String
@@ -66,6 +68,7 @@ Returns the expected type at the specified position, pretty-printed as a string.
 structure PlainTermGoalParams extends TextDocumentPositionParams
   deriving FromJson, ToJson
 
+/-- `$/lean/plainTermGoal` client<-server reply. -/
 structure PlainTermGoal where
   goal : String
   range : Range
@@ -81,21 +84,22 @@ structure RpcRef where
 instance : ToString RpcRef where
   toString r := toString r.p
 
-/-- `$/lean/rpc/connect` client->server notification.
+/-- `$/lean/rpc/connect` client->server request.
 
-A notification to connect to the RPC session at the given file's worker. Should be sent:
-- exactly once whenever RPC is first needed (e.g. on client startup)
-- if an `RpcNeedsReconnect` error is received in an RPC request -/
+Starts an RPC session at the given file's worker, replying with the new session ID.
+Multiple sessions may be started and operating concurrently.
+
+A session may be destroyed by the server at any time (e.g. due to a crash), in which case further
+RPC requests for that session will reply with `RpcNeedsReconnect` errors. The client should discard
+references held from that session and `connect` again. -/
 structure RpcConnectParams where
   uri : DocumentUri
   deriving FromJson, ToJson
 
-/-- `$/lean/rpc/connected` client<-server notification.
+/-- `$/lean/rpc/connect` client<-server reply.
 
-Indicates that an RPC connection had been made. On receiving this, the client should discard any
-references it may still be holding from previous RPC sessions. -/
+Indicates that an RPC connection had been made and a session started for it. -/
 structure RpcConnected where
-  uri : DocumentUri
   sessionId : UInt64
   deriving FromJson, ToJson
 
@@ -123,7 +127,7 @@ structure RpcReleaseParams where
   refs : Array RpcRef
   deriving FromJson, ToJson
 
-/-- `$/lean/rpc/keepAlive` client -> server notification.
+/-- `$/lean/rpc/keepAlive` client->server notification.
 
 The client must send an RPC notification every 10s in order to keep the RPC session alive.
 This is the simplest one. On not seeing any notifications for three 10s periods, the server

src/Lean/Server/FileWorker.lean

@@ -275,24 +275,6 @@ section NotificationHandling
   def handleCancelRequest (p : CancelParams) : WorkerM Unit := do
     updatePendingRequests (fun pendingRequests => pendingRequests.erase p.id)
 
-  def handleRpcConnect (p : RpcConnectParams) : WorkerM Unit := do
-    let st ← get
-    let doc := st.doc
-    let rpcSesh' ←
-      if !st.rpcSesh.clientConnected then
-        -- First client connection.
-        let s := { st.rpcSesh with clientConnected := true }
-        s.keptAlive
-      else
-        /- Client has restarted. Just setting the state should `dec` the previous session.
-        Any associated references will then no longer be kept alive for the client. -/
-        RpcSession.new true
-    set { st with rpcSesh := rpcSesh' }
-    (←read).hOut.writeLspNotification
-      <| { method := "$/lean/rpc/connected"
-           param := { uri := doc.meta.uri
-                      sessionId := rpcSesh'.sessionId : RpcConnected } }
-
 def handleRpcRelease (p : Lsp.RpcReleaseParams) : WorkerM Unit := do
   let st ← get
   if p.sessionId ≠ st.rpcSesh.sessionId then
@@ -310,6 +292,28 @@ def handleRpcKeepAlive (p : Lsp.RpcKeepAliveParams) : WorkerM Unit := do
 
 end NotificationHandling
 
+/-! Requests here are handled synchronously rather than in the asynchronous `RequestM`. -/
+section RequestHandling
+
+def handleRpcConnect (p : RpcConnectParams) : WorkerM RpcConnected := do
+  let st ← get
+  let rpcSesh' ←
+    if !st.rpcSesh.clientConnected then
+      -- First client connection.
+      let s := { st.rpcSesh with clientConnected := true }
+      s.keptAlive
+    else
+      /- Client has restarted. Just setting the state should `dec` the previous session.
+      Any associated references will then no longer be kept alive for the client. -/
+      RpcSession.new true
+  set { st with rpcSesh := rpcSesh' }
+
+  return {
+    sessionId := rpcSesh'.sessionId : RpcConnected
+  }
+
+end RequestHandling
+
 section MessageHandling
   def parseParams (paramType : Type) [FromJson paramType] (params : Json) : WorkerM paramType :=
     match fromJson? params with
@@ -322,7 +326,6 @@ section MessageHandling
     match method with
     | "textDocument/didChange" => handle DidChangeTextDocumentParams handleDidChange
     | "$/cancelRequest"        => handle CancelParams handleCancelRequest
-    | "$/lean/rpc/connect"     => handle RpcConnectParams handleRpcConnect
     | "$/lean/rpc/release"     => handle RpcReleaseParams handleRpcRelease
     | "$/lean/rpc/keepAlive"   => handle RpcKeepAliveParams handleRpcKeepAlive
     | _                        => throwServerError s!"Got unsupported notification method: {method}"
@@ -335,6 +338,19 @@ section MessageHandling
       : WorkerM Unit := do
     let ctx ← read
     let st ← get
+
+    if method == "$/lean/rpc/connect" then
+      try
+        let ps ← parseParams RpcConnectParams params
+        let resp ← handleRpcConnect ps
+        ctx.hOut.writeLspResponse ⟨id, resp⟩
+      catch e =>
+        ctx.hOut.writeLspResponseError
+          { id
+            code := ErrorCode.internalError
+            message := toString e }
+      return
+
     let rc : RequestContext :=
       { rpcSesh := st.rpcSesh
         srcSearchPath := ctx.srcSearchPath

src/Lean/Server/Requests.lean

@@ -132,7 +132,7 @@ private structure RequestHandler where
 builtin_initialize requestHandlers : IO.Ref (Std.PersistentHashMap String RequestHandler) ←
   IO.mkRef {}
 
-/-- NB: This method may only be called in `initialize`/`builtin_initialize` blocks.
+/-- NB: This method may only be called in `builtin_initialize` blocks.
 
 A registration consists of:
 - a type of JSON-parsable request data `paramType`

src/Lean/Server/Watchdog.lean

@@ -380,10 +380,16 @@ section MessageHandling
       | Except.error inner => throwServerError s!"Got param with wrong structure: {params.compress}\n{inner}"
 
   def handleRequest (id : RequestID) (method : String) (params : Json) : ServerM Unit := do
-    match (← routeLspRequest method params) with
+    let uri: DocumentUri ←
+      -- This request is handled specially.
+      if method == "$/lean/rpc/connect" then
+        let ps ← parseParams Lsp.RpcConnectParams params
+        fileSource ps
+      else match (← routeLspRequest method params) with
       | Except.error e =>
         (←read).hOut.writeLspResponseError <| e.toLspResponseError id
-      | Except.ok uri =>
+        return
+      | Except.ok uri => uri
         let fw ← try
           findFileWorker uri
         catch _ =>