fix: use properly random RPC session id

src/Lean/Data/Json/FromToJson.lean

@@ -65,19 +65,36 @@ instance : FromJson Name where
 instance : ToJson Name where
   toJson n := toString (repr n)
 
-/- Note that `USize`s are stored as strings because JavaScript
+/- Note that `USize`s and `UInt64`s are stored as strings because JavaScript
 cannot represent 64-bit numbers. -/
+def bignumFromJson? (j : Json) : Except String Nat := do
+  let s ← j.getStr?
+  let some v ← Syntax.decodeNatLitVal? s -- TODO maybe this should be in Std
+    | throw s!"expected a string-encoded number, got '{j}'"
+  return v
+  
+def bignumToJson (n : Nat) : Json :=
+  toString n
+
 instance : FromJson USize where
   fromJson? j := do
-    let s ← j.getStr?
-    let some v ← Syntax.decodeNatLitVal? s -- TODO maybe this should be in Std
-      | throw s!"expected a string-encoded `USize`, got '{j}'"
-    if v ≥ USize.size then
+    let n ← bignumFromJson? j
+    if n ≥ USize.size then
       throw "value '{j}' is too large for `USize`"
-    return USize.ofNat v
+    return USize.ofNat n
 
 instance : ToJson USize where
-  toJson v := toString (repr v)
+  toJson v := bignumToJson (USize.toNat v)
+
+instance : FromJson UInt64 where
+  fromJson? j := do
+    let n ← bignumFromJson? j
+    if n ≥ UInt64.size then
+      throw "value '{j}' is too large for `UInt64`"
+    return UInt64.ofNat n
+
+instance : ToJson UInt64 where
+  toJson v := bignumToJson (UInt64.toNat v)
 
 namespace Json
 

src/Lean/Data/Lsp/Extra.lean

@@ -84,7 +84,7 @@ structure RpcInitializeParams where
 
 structure RpcInitialized where
   uri : DocumentUri
-  sessionId : USize
+  sessionId : UInt64
   deriving FromJson, ToJson
 
 /-- A client request to execute a procedure previously bound for RPC.
@@ -92,7 +92,7 @@ structure RpcInitialized where
 Extending TDPP is weird. But in Lean, symbols exist in the context of a position
 within a source file. So we need this to refer to code in the env at that position. -/
 structure RpcCallParams extends TextDocumentPositionParams where
-  sessionId : USize
+  sessionId : UInt64
   /-- Procedure to invoke. Must be fully qualified. -/
   method : Name
   params : Json
@@ -103,7 +103,7 @@ when it no longer needs an `RpcRef` it has previously received from the server.
 Not doing so is safe but will leak memory. -/
 structure RpcReleaseParams where
   uri : DocumentUri
-  sessionId : USize
+  sessionId : UInt64
   ref : RpcRef
   deriving FromJson, ToJson
 

src/Lean/Server/FileWorker.lean

@@ -275,10 +275,12 @@ section NotificationHandling
   def handleRpcInitialize (p : RpcInitializeParams) : WorkerM Unit := do
     let ctx ← read
     let doc := (←get).doc
-    let newId ← USize.ofNat <$> IO.monoMsNow -- TODO may crash if called twice within the same millisecond?
+    /- We generate a random ID to ensure that session IDs do not repeat across re-initializations
+    and worker restarts. Otherwise, the client may attempt to use outdated references. -/
+    let newId ← ByteArray.toUInt64LE! <$> IO.getRandomBytes 8
     let newAliveRefs ← IO.mkRef Std.PersistentHashMap.empty
-    -- Just setting this should `dec` the previous session. Any associated references should
-    -- become collectable once all in-use references to the old session go out of scope.
+    /- Just setting this should `dec` the previous session. Any associated references
+    will then no longer be kept alive for the client. -/
     modify fun st => { st with rpcSesh := some { sessionId := newId, aliveRefs := newAliveRefs } }
     ctx.hOut.writeLspNotification
       <| { method := "$/lean/rpc/initialized"

src/Lean/Server/FileWorker/Rpc.lean

@@ -18,7 +18,7 @@ be initialized as follows:
 namespace Lean.Server
 
 private structure RpcProcedure where
-  wrapper : (sessionId : USize) → Json → RequestM (RequestTask Json)
+  wrapper : (sessionId : UInt64) → Json → RequestM (RequestTask Json)
 
 builtin_initialize rpcProcedures : IO.Ref (Std.PersistentHashMap Name RpcProcedure) ←
   IO.mkRef {}

src/Lean/Server/FileWorker/Utils.lean

@@ -54,7 +54,7 @@ structure EditableDocument where
   deriving Inhabited
 
 structure RpcSession where
-  sessionId : USize
+  sessionId : UInt64
   /-- Objects that are being kept alive for the RPC client, together with their type names,
   mapped to by their RPC reference.