feat: inlay hints for auto-implicits (#6768)

This PR adds preliminary support for inlay hints, as well as support for
inlay hints that denote the auto-implicits of a function. Hovering over
an auto-implicit displays its type and double-clicking the auto-implicit
inserts it into the text document.

![Inlay hints for
auto-implicits](https://github.com/user-attachments/assets/fb204c42-5997-4f10-9617-c65f1042d732)

This PR is an extension of #3910.

### Known issues

- In VS Code, when inserting an inlay hint, the inlay hint may linger
for a couple of seconds before it disappears. This is a defect of the VS
Code implementation of inlay hints and cannot adequately be resolved by
us.
- When making a change to the document, it may take a couple of seconds
until the inlay hints respond to the change. This is deliberate and
intended to reduce the amount of inlay hint flickering while typing. VS
Code has a mechanism of its own for this, but in my experience it is
still far too sensitive without additional latency.
- Inserting an auto-implicit inlay hint that depends on an auto-implicit
meta-variable causes a "failed to infer binder type" error. We can't
display these meta-variables in the inlay hint because they don't have a
user-displayable name, so it is not clear how to resolve this problem.
- Inlay hints are currently always resolved eagerly, i.e. we do not
support the `textDocument/inlayHint/resolve` request yet. Implementing
support for this request is future work.

### Other changes
- Axioms did not support auto-implicits due to an oversight in the
implementation. This PR ensures they do.
- In order to reduce the amount of inlay hint flickering when making a
change to the document, the language server serves old inlay hints for
parts of the file that have not been processed yet. This requires LSP
request handler state (that sometimes must be invalidated on
`textDocument/didChange`), so this PR introduces the notion of a
stateful LSP request handler.
- The partial response mechanism that we use for semantic tokens, where
we simulate incremental LSP responses by periodically emitting refresh
requests to the client, is generalized to accommodate both inlay hints
and semantic tokens. Additionally, it is made more robust to ensure that
we never emit refresh requests while a corresponding request is in
flight, which causes VS Code to discard the respond of the request, as
well as to ensure that we keep prompting VS Code to send another request
if it spuriously decides not to respond to one of our refresh requests.
- The synthetic identifier of an `example` had the full declaration as
its (non-canonical synthetic) range. Since we need a reasonable position
for the identifier to insert an inlay hint for the auto-implicits of an
`example`, we change the (canonical synthetic) range of the synthetic
identifier to that of the `example` keyword.
- The semantic highlighting request handling is moved to a separate
file.

### Breaking changes
- The semantic highlighting request handler is not a pure request
handler anymore, but a stateful one. Notably, this means that clients
that extend the semantic highlighting of the Lean language server with
the `chainLspRequestHandler` function must now use the
`chainStatefulLspRequestHandler` function instead.

src/Init/Prelude.lean

@@ -2503,6 +2503,9 @@ instance (p₁ p₂ : String.Pos) : Decidable (LE.le p₁ p₂) :=
 instance (p₁ p₂ : String.Pos) : Decidable (LT.lt p₁ p₂) :=
   inferInstanceAs (Decidable (LT.lt p₁.byteIdx p₂.byteIdx))
 
+instance : Min String.Pos := minOfLe
+instance : Max String.Pos := maxOfLe
+
 /-- A `String.Pos` pointing at the end of this string. -/
 @[inline] def String.endPos (s : String) : String.Pos where
   byteIdx := utf8ByteSize s

src/Lean/Data/Lsp/Basic.lean

@@ -426,5 +426,9 @@ structure WorkDoneProgressOptions where
   workDoneProgress := false
   deriving ToJson, FromJson
 
+structure ResolveSupport where
+  properties : Array String
+  deriving FromJson, ToJson
+
 end Lsp
 end Lean

src/Lean/Data/Lsp/Capabilities.lean

@@ -30,6 +30,7 @@ structure CompletionClientCapabilities where
 structure TextDocumentClientCapabilities where
   completion? : Option CompletionClientCapabilities := none
   codeAction? : Option CodeActionClientCapabilities := none
+  inlayHint?  : Option InlayHintClientCapabilities  := none
   deriving ToJson, FromJson
 
 structure ShowDocumentClientCapabilities where
@@ -81,6 +82,7 @@ structure ServerCapabilities where
   foldingRangeProvider      : Bool                           := false
   semanticTokensProvider?   : Option SemanticTokensOptions   := none
   codeActionProvider?       : Option CodeActionOptions       := none
+  inlayHintProvider?        : Option InlayHintOptions        := none
   deriving ToJson, FromJson
 
 end Lsp

src/Lean/Data/Lsp/CodeActions.lean

@@ -129,10 +129,6 @@ structure CodeAction extends WorkDoneProgressParams, PartialResultParams where
   data?        : Option Json := none
   deriving ToJson, FromJson
 
-structure ResolveSupport where
-  properties : Array String
-  deriving FromJson, ToJson
-
 structure CodeActionLiteralSupportValueSet where
   /-- The code action kind values the client supports. When this
     property exists the client also guarantees that it will

src/Lean/Data/Lsp/LanguageFeatures.lean

@@ -445,5 +445,82 @@ structure RenameParams extends TextDocumentPositionParams where
 structure PrepareRenameParams extends TextDocumentPositionParams
   deriving FromJson, ToJson
 
+structure InlayHintParams extends WorkDoneProgressParams where
+  textDocument : TextDocumentIdentifier
+  range        : Range
+  deriving FromJson, ToJson
+
+inductive InlayHintTooltip
+  | plaintext (text : String)
+  | markdown (markup : MarkupContent)
+
+instance : FromJson InlayHintTooltip where
+  fromJson?
+    | .str s => .ok <| .plaintext s
+    | j@(.obj _) => do return .markdown (← fromJson? j)
+    | j => .error s!"invalid inlay hint tooltip {j}"
+
+instance : ToJson InlayHintTooltip where
+  toJson
+    | .plaintext text => toJson text
+    | .markdown markup => toJson markup
+
+structure InlayHintLabelPart where
+  value     : String
+  tooltip?  : Option InlayHintTooltip := none
+  location? : Option Location := none
+  command?  : Option Command := none
+  deriving FromJson, ToJson
+
+inductive InlayHintLabel
+  | name (n : String)
+  | parts (p : Array InlayHintLabelPart)
+
+instance : FromJson InlayHintLabel where
+  fromJson?
+    | .str s => .ok <| .name s
+    | j@(.arr _) => do return .parts (← fromJson? j)
+    | j => .error s!"invalid inlay hint label {j}"
+
+instance : ToJson InlayHintLabel where
+  toJson
+    | .name n => toJson n
+    | .parts p => toJson p
+
+inductive InlayHintKind where
+  | type
+  | parameter
+
+instance : FromJson InlayHintKind where
+  fromJson?
+    | 1 => .ok .type
+    | 2 => .ok .parameter
+    | j => .error s!"unknown inlay hint kind {j}"
+
+instance : ToJson InlayHintKind where
+  toJson
+    | .type => 1
+    | .parameter => 2
+
+structure InlayHint where
+  position      : Position
+  label         : InlayHintLabel
+  kind?         : Option InlayHintKind := none
+  textEdits?    : Option (Array TextEdit) := none
+  tooltip?      : Option (InlayHintTooltip) := none
+  paddingLeft?  : Option Bool := none
+  paddingRight? : Option Bool := none
+  data?         : Option Json := none
+  deriving FromJson, ToJson
+
+structure InlayHintClientCapabilities where
+  dynamicRegistration? : Option Bool := none
+  resolveSupport?      : Option ResolveSupport := none
+  deriving FromJson, ToJson
+
+structure InlayHintOptions extends WorkDoneProgressOptions where
+  resolveProvider? : Option Bool := none
+  deriving FromJson, ToJson
+
 end Lsp
 end Lean

src/Lean/Data/Lsp/Utf16.lean

@@ -89,15 +89,31 @@ def utf8PosToLspPos (text : FileMap) (pos : String.Pos) : Lsp.Position :=
 def utf8RangeToLspRange (text : FileMap) (range : String.Range) : Lsp.Range :=
   { start := text.utf8PosToLspPos range.start, «end» := text.utf8PosToLspPos range.stop }
 
+def lspRangeToUtf8Range (text : FileMap) (range : Lsp.Range) : String.Range :=
+  { start := text.lspPosToUtf8Pos range.start, stop := text.lspPosToUtf8Pos range.end }
+
 end FileMap
-end Lean
+
+def DeclarationRange.ofFilePositions (text : FileMap) (pos : Position) (endPos : Position)
+    : DeclarationRange := {
+  pos,
+  charUtf16 := text.leanPosToLspPos pos |>.character
+  endPos,
+  endCharUtf16 := text.leanPosToLspPos endPos |>.character
+}
+
+def DeclarationRange.ofStringPositions (text : FileMap) (pos : String.Pos) (endPos : String.Pos)
+    : DeclarationRange :=
+  .ofFilePositions text (text.toPosition pos) (text.toPosition endPos)
 
 /--
 Convert the Lean `DeclarationRange` to an LSP `Range` by turning the 1-indexed line numbering into a
 0-indexed line numbering and converting the character offset within the line to a UTF-16 indexed
 offset.
 -/
-def Lean.DeclarationRange.toLspRange (r : Lean.DeclarationRange) : Lsp.Range := {
+def DeclarationRange.toLspRange (r : DeclarationRange) : Lsp.Range := {
   start := ⟨r.pos.line - 1, r.charUtf16⟩
   «end» := ⟨r.endPos.line - 1, r.endCharUtf16⟩
 }
+
+end Lean

src/Lean/Elab/BuiltinCommand.lean

@@ -227,11 +227,16 @@ private def replaceBinderAnnotation (binder : TSyntax ``Parser.Term.bracketedBin
     return #[binder]
 
 @[builtin_command_elab «variable»] def elabVariable : CommandElab
-  | `(variable $binders*) => do
+  | `(variable%$tk $binders*) => do
     let binders ← binders.flatMapM replaceBinderAnnotation
     -- Try to elaborate `binders` for sanity checking
     runTermElabM fun _ => Term.withSynthesize <| Term.withAutoBoundImplicit <|
-      Term.elabBinders binders fun _ => pure ()
+      Term.elabBinders binders fun xs => do
+        -- Determine the set of auto-implicits for this variable command and add an inlay hint
+        -- for them. We will only actually add the auto-implicits to a type when the variables
+        -- declared here are used in some other declaration, but this is nonetheless the right
+        -- place to display the inlay hint.
+        let _ ← Term.addAutoBoundImplicits xs (tk.getTailPos? (canonicalOnly := true))
     -- Remark: if we want to produce error messages when variables shadow existing ones, here is the place to do it.
     for binder in binders do
       let varUIds ← (← getBracketedBinderIds binder) |>.mapM (withFreshMacroScope ∘ MonadQuotation.addMacroScope)

src/Lean/Elab/Declaration.lean

@@ -105,12 +105,13 @@ def elabAxiom (modifiers : Modifiers) (stx : Syntax) : CommandElabM Unit := do
     let scopeLevelNames ← Term.getLevelNames
     let ⟨shortName, declName, allUserLevelNames⟩ ← Term.expandDeclId (← getCurrNamespace) scopeLevelNames declId modifiers
     addDeclarationRangesForBuiltin declName modifiers.stx stx
+    Term.withAutoBoundImplicit do
     Term.withAutoBoundImplicitForbiddenPred (fun n => shortName == n) do
     Term.withDeclName declName <| Term.withLevelNames allUserLevelNames <| Term.elabBinders binders.getArgs fun xs => do
       Term.applyAttributesAt declName modifiers.attrs AttributeApplicationTime.beforeElaboration
       let type ← Term.elabType typeStx
       Term.synthesizeSyntheticMVarsNoPostponing
-      let xs ← Term.addAutoBoundImplicits xs
+      let xs ← Term.addAutoBoundImplicits xs (declId.getTailPos? (canonicalOnly := true))
       let type ← instantiateMVars type
       let type ← mkForallFVars xs type
       let type ← mkForallFVars vars type (usedOnly := true)

src/Lean/Elab/DeclarationRange.lean

@@ -15,15 +15,7 @@ def getDeclarationRange? [Monad m] [MonadFileMap m] (stx : Syntax) : m (Option D
   let some range := stx.getRange?
     | return none
   let fileMap ← getFileMap
-  --let range := fileMap.utf8RangeToLspRange
-  let pos    := fileMap.toPosition range.start
-  let endPos := fileMap.toPosition range.stop
-  return some {
-    pos          := pos
-    charUtf16    := fileMap.leanPosToLspPos pos |>.character
-    endPos       := endPos
-    endCharUtf16 := fileMap.leanPosToLspPos endPos |>.character
-  }
+  return some <| .ofStringPositions fileMap range.start range.stop
 
 /--
   For most builtin declarations, the selection range is just its name, which is stored in the second position.

src/Lean/Elab/DefView.lean

@@ -185,7 +185,7 @@ def mkDefViewOfOpaque (modifiers : Modifiers) (stx : Syntax) : CommandElabM DefV
 def mkDefViewOfExample (modifiers : Modifiers) (stx : Syntax) : DefView :=
   -- leading_parser "example " >> declSig >> declVal
   let (binders, type) := expandOptDeclSig stx[1]
-  let id              := mkIdentFrom stx `_example
+  let id              := mkIdentFrom stx[0] `_example (canonical := true)
   let declId          := mkNode ``Parser.Command.declId #[id, mkNullNode]
   { ref := stx, headerRef := mkNullNode stx.getArgs[:2], kind := DefKind.example, modifiers := modifiers,
     declId := declId, binders := binders, type? := type, value := stx[2] }

src/Lean/Elab/Inductive.lean

@@ -194,7 +194,7 @@ private def elabCtors (indFVars : Array Expr) (params : Array Expr) (r : ElabHea
             return type
         let type ← elabCtorType
         Term.synthesizeSyntheticMVarsNoPostponing
-        let ctorParams ← Term.addAutoBoundImplicits ctorParams
+        let ctorParams ← Term.addAutoBoundImplicits ctorParams (ctorView.declId.getTailPos? (canonicalOnly := true))
         let except (mvarId : MVarId) := ctorParams.any fun ctorParam => ctorParam.isMVar && ctorParam.mvarId! == mvarId
         /-
           We convert metavariables in the resulting type into extra parameters. Otherwise, we would not be able to elaborate

src/Lean/Elab/InfoTree/InlayHints.lean

@@ -0,0 +1,63 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Marc Huisinga
+-/
+prelude
+import Lean.Meta.Basic
+
+namespace Lean.Elab
+
+open Lean
+
+structure InlayHintLinkLocation where
+  module : Name
+  range  : String.Range
+
+structure InlayHintLabelPart where
+  value     : String
+  tooltip?  : Option String := none
+  location? : Option InlayHintLinkLocation := none
+
+inductive InlayHintLabel
+  | name (n : String)
+  | parts (p : Array InlayHintLabelPart)
+
+inductive InlayHintKind where
+  | type
+  | parameter
+
+structure InlayHintTextEdit where
+  range   : String.Range
+  newText : String
+
+structure InlayHintInfo where
+  position     : String.Pos
+  label        : InlayHintLabel
+  kind?        : Option InlayHintKind := none
+  textEdits    : Array InlayHintTextEdit := #[]
+  tooltip?     : Option String := none
+  paddingLeft  : Bool := false
+  paddingRight : Bool := false
+
+structure InlayHint extends InlayHintInfo where
+  lctx               : LocalContext
+  deferredResolution : InlayHintInfo → MetaM InlayHintInfo := fun i => .pure i
+  deriving TypeName
+
+namespace InlayHint
+
+def toCustomInfo (i : InlayHint) : CustomInfo := {
+  stx := .missing
+  value := .mk i
+}
+
+def ofCustomInfo? (c : CustomInfo) : Option InlayHint :=
+  c.value.get? InlayHint
+
+def resolveDeferred (i : InlayHint) : MetaM InlayHint := do
+  let info := i.toInlayHintInfo
+  let info ← i.deferredResolution info
+  return { i with toInlayHintInfo := info }
+
+end InlayHint

src/Lean/Elab/MutualDef.lean

@@ -197,7 +197,7 @@ private def elabHeaders (views : Array DefView)
               type ← cleanupOfNat type
             let (binderIds, xs) := xs.unzip
             -- TODO: add forbidden predicate using `shortDeclName` from `views`
-            let xs ← addAutoBoundImplicits xs
+            let xs ← addAutoBoundImplicits xs (view.declId.getTailPos? (canonicalOnly := true))
             type ← mkForallFVars' xs type
             type ← instantiateMVarsProfiling type
             let levelNames ← getLevelNames

src/Lean/Elab/MutualInductive.lean

@@ -288,7 +288,9 @@ private def elabHeadersAux (views : Array InductiveView) (i : Nat) (acc : Array
           let typeStx ← view.type?.getDM `(Sort _)
           let type ← Term.elabType typeStx
           Term.synthesizeSyntheticMVarsNoPostponing
-          let indices ← Term.addAutoBoundImplicits #[]
+          let inlayHintPos? := view.binders.getTailPos? (canonicalOnly := true)
+            <|> view.declId.getTailPos? (canonicalOnly := true)
+          let indices ← Term.addAutoBoundImplicits #[] inlayHintPos?
           let type ← mkForallFVars indices type
           if view.allowIndices then
             unless (← isTypeFormerType type) do
@@ -297,7 +299,7 @@ private def elabHeadersAux (views : Array InductiveView) (i : Nat) (acc : Array
             unless (← whnfD type).isSort do
               throwErrorAt typeStx "invalid resulting type, expecting 'Type _' or 'Prop'"
           return (type, indices.size)
-        let params ← Term.addAutoBoundImplicits params
+        let params ← Term.addAutoBoundImplicits params (view.declId.getTailPos? (canonicalOnly := true))
         trace[Elab.inductive] "header params: {params}, type: {type}"
         let levelNames ← Term.getLevelNames
         return acc.push { lctx := (← getLCtx), localInsts := (← getLocalInstances), levelNames, params, type, view }

src/Lean/Elab/Structure.lean

@@ -562,7 +562,7 @@ private def elabFieldTypeValue (view : StructFieldView) : TermElabM (Option Expr
       | some valStx =>
         Term.synthesizeSyntheticMVarsNoPostponing
         -- TODO: add forbidden predicate using `shortDeclName` from `view`
-        let params ← Term.addAutoBoundImplicits params
+        let params ← Term.addAutoBoundImplicits params (view.nameId.getTailPos? (canonicalOnly := true))
         let value ← Term.withoutAutoBoundImplicit <| Term.elabTerm valStx none
         registerFailedToInferFieldType view.name (← inferType value) view.nameId
         registerFailedToInferDefaultValue view.name value valStx
@@ -572,7 +572,7 @@ private def elabFieldTypeValue (view : StructFieldView) : TermElabM (Option Expr
       let type ← Term.elabType typeStx
       registerFailedToInferFieldType view.name type typeStx
       Term.synthesizeSyntheticMVarsNoPostponing
-      let params ← Term.addAutoBoundImplicits params
+      let params ← Term.addAutoBoundImplicits params (view.nameId.getTailPos? (canonicalOnly := true))
       match view.value? with
       | none        =>
         let type  ← mkForallFVars params type

src/Lean/Elab/Term.lean

@@ -14,7 +14,9 @@ import Lean.Elab.Config
 import Lean.Elab.Level
 import Lean.Elab.DeclModifiers
 import Lean.Elab.PreDefinition.TerminationHint
+import Lean.Elab.DeclarationRange
 import Lean.Language.Basic
+import Lean.Elab.InfoTree.InlayHints
 
 namespace Lean.Elab
 
@@ -1975,22 +1977,76 @@ where
         else
           go (mvarIdsNew.toList ++ mvarId :: mvarIds) result visited
 
+/--
+Adds an `InlayHintInfo` for the fvar auto implicits in `autos` at `inlayHintPos`.
+The inserted inlay hint has a hover that denotes the type of the auto-implicit (with meta-variables)
+and can be inserted at `inlayHintPos`.
+-/
+def addAutoBoundImplicitsInlayHint (autos : Array Expr) (inlayHintPos : String.Pos) : TermElabM Unit := do
+  -- If the list of auto-implicits contains a non-type fvar, then the list of auto-implicits will
+  -- also contain an mvar that denotes the type of the non-type fvar.
+  -- For example, the auto-implicit `x` in a type `Foo x` for `Foo.{u} {α : Sort u} (x : α) : Type`
+  -- also comes with an auto-implicit mvar denoting the type of `x`.
+  -- We have no way of displaying this mvar to the user in an inlay hint, as it doesn't have a name,
+  -- so we filter it.
+  -- This also means that inserting the inlay hint with the syntax displayed in the inlay hint will
+  -- cause a "failed to infer binder type" error, since we don't have a name to insert in the code.
+  let autos := autos.filter (· matches .fvar ..)
+  if autos.isEmpty then
+    return
+  let autoNames ← autos.mapM (·.fvarId!.getUserName)
+  let formattedHint := s!" \{{" ".intercalate <| Array.toList <| autoNames.map toString}}"
+  let autoLabelParts : List (InlayHintLabelPart × Option Expr) := Array.toList <| ← autos.mapM fun auto => do
+    let name := toString <| ← auto.fvarId!.getUserName
+    return ({ value := name }, some auto)
+  let p value : InlayHintLabelPart × Option Expr := ({ value }, none)
+  let labelParts := [p " ", p "{"] ++ [p " "].intercalate (autoLabelParts.map ([·])) ++ [p "}"]
+  let labelParts := labelParts.toArray
+  let deferredResolution ih := do
+    let .parts ps := ih.label
+      | return ih
+    let ps ← ps.mapIdxM fun i p => do
+      let some (part, some auto) := labelParts[i]?
+        | return p
+      let type := toString <| ← Meta.ppExpr <| ← instantiateMVars (← inferType auto)
+      return { p with
+        tooltip? := s!"{part.value} : {type}"
+      }
+    return { ih with label := .parts ps }
+  pushInfoLeaf <| .ofCustomInfo {
+      position := inlayHintPos
+      label := .parts <| labelParts.map (·.1)
+      textEdits := #[{
+        range := ⟨inlayHintPos, inlayHintPos⟩,
+        newText := formattedHint
+      }]
+      kind? := some .parameter
+      lctx := ← getLCtx
+      deferredResolution
+      : InlayHint
+    }.toCustomInfo
+
 /--
   Return `autoBoundImplicits ++ xs`
   This method throws an error if a variable in `autoBoundImplicits` depends on some `x` in `xs`.
   The `autoBoundImplicits` may contain free variables created by the auto-implicit feature, and unassigned free variables.
   It avoids the hack used at `autoBoundImplicitsOld`.
 
+  If `inlayHintPos?` is set, this function also inserts an inlay hint denoting `autoBoundImplicits`.
+  See `addAutoBoundImplicitsInlayHint` for more information.
+
   Remark: we cannot simply replace every occurrence of `addAutoBoundImplicitsOld` with this one because a particular
   use-case may not be able to handle the metavariables in the array being given to `k`.
 -/
-def addAutoBoundImplicits (xs : Array Expr) : TermElabM (Array Expr) := do
+def addAutoBoundImplicits (xs : Array Expr) (inlayHintPos? : Option String.Pos) : TermElabM (Array Expr) := do
   let autos := (← read).autoBoundImplicits
   go autos.toList #[]
 where
   go (todo : List Expr) (autos : Array Expr) : TermElabM (Array Expr) := do
     match todo with
     | [] =>
+      if let some inlayHintPos := inlayHintPos? then
+        addAutoBoundImplicitsInlayHint autos inlayHintPos
       for auto in autos do
         if auto.isFVar then
           let localDecl ← auto.fvarId!.getDecl
@@ -2009,7 +2065,7 @@ where
   We use this method to simplify the conversion of code using `autoBoundImplicitsOld` to `autoBoundImplicits`.
 -/
 def addAutoBoundImplicits' (xs : Array Expr) (type : Expr) (k : Array Expr → Expr → TermElabM α) : TermElabM α := do
-  let xs ← addAutoBoundImplicits xs
+  let xs ← addAutoBoundImplicits xs none
   if xs.all (·.isFVar) then
     k xs type
   else

src/Lean/Server/AsyncList.lean

@@ -97,18 +97,51 @@ partial def waitFind? (p : α → Bool) : AsyncList ε α → Task (Except ε (O
       | .ok tl   => tl.waitFind? p
       | .error e => .pure <| .error e
 
-/-- Retrieve the already-computed prefix of the list. If computation has finished with an error, return it as well. -/
-partial def getFinishedPrefix : AsyncList ε α → BaseIO (List α × Option ε)
+/--
+Retrieve the already-computed prefix of the list. If computation has finished with an error, return it as well.
+The returned boolean indicates whether the complete `AsyncList` was returned, or whether only a
+proper prefix was returned.
+-/
+partial def getFinishedPrefix : AsyncList ε α → BaseIO (List α × Option ε × Bool)
   | cons hd tl => do
-    let ⟨tl, e?⟩ ← tl.getFinishedPrefix
-    pure ⟨hd :: tl, e?⟩
-  | nil => pure ⟨[], none⟩
+    let ⟨tl, e?, isComplete⟩ ← tl.getFinishedPrefix
+    pure ⟨hd :: tl, e?, isComplete⟩
+  | nil => pure ⟨[], none, true⟩
   | delayed tl => do
     if (← hasFinished tl) then
       match tl.get with
       | Except.ok tl => tl.getFinishedPrefix
-      | Except.error e => pure ⟨[], some e⟩
-    else pure ⟨[], none⟩
+      | Except.error e => pure ⟨[], some e, true⟩
+    else pure ⟨[], none, false⟩
+
+partial def getFinishedPrefixWithTimeout (xs : AsyncList ε α) (timeoutMs : UInt32) : BaseIO (List α × Option ε × Bool) := do
+  let timeoutTask : Task (Unit ⊕ Except ε (AsyncList ε α)) ← BaseIO.asTask (prio := .dedicated) do
+    IO.sleep timeoutMs
+    return .inl ()
+  go timeoutTask xs
+where
+  go (timeoutTask : Task (Unit ⊕ Except ε (AsyncList ε α))) (xs : AsyncList ε α) : BaseIO (List α × Option ε × Bool) := do
+    match xs with
+    | cons hd tl =>
+      let ⟨tl, e?, isComplete⟩ ← go timeoutTask tl
+      return ⟨hd :: tl, e?, isComplete⟩
+    | nil => return ⟨[], none, true⟩
+    | delayed tl =>
+      let r ← IO.waitAny [
+        timeoutTask,
+        tl.map (sync := true) .inr
+      ]
+      match r with
+      | .inl _ => return ⟨[], none, false⟩ -- Timeout - stop waiting
+      | .inr (.ok tl) => go timeoutTask tl
+      | .inr (.error e) => return ⟨[], some e, true⟩
+
+partial def getFinishedPrefixWithConsistentLatency (xs : AsyncList ε α) (latencyMs : UInt32) : BaseIO (List α × Option ε × Bool) := do
+  let timestamp ← IO.monoMsNow
+  let r ← xs.getFinishedPrefixWithTimeout latencyMs
+  let passedTimeMs := (← IO.monoMsNow) - timestamp
+  IO.sleep <| (latencyMs.toNat - passedTimeMs).toUInt32
+  return r
 
 def waitHead? (as : AsyncList ε α) : Task (Except ε (Option α)) :=
   as.waitFind? fun _ => true

src/Lean/Server/FileSource.lean

@@ -100,4 +100,7 @@ instance : FileSource RpcKeepAliveParams where
 instance : FileSource CodeActionParams where
   fileSource p := fileSource p.textDocument
 
+instance : FileSource InlayHintParams where
+  fileSource p := fileSource p.textDocument
+
 end Lean.Lsp

src/Lean/Server/FileWorker.lean

@@ -59,27 +59,33 @@ open IO
 open Snapshots
 open JsonRpc
 
+structure RefreshInfo where
+  lastRefreshTimestamp      : Nat
+  successiveRefreshAttempts : Nat
+
+structure PartialHandlerInfo where
+  refreshMethod            : String
+  requestsInFlight         : Nat
+  pendingRefreshInfo?      : Option RefreshInfo
+  deriving Inhabited
+
 open Widget in
 
 structure WorkerContext where
   /-- Synchronized output channel for LSP messages. Notifications for outdated versions are
     discarded on read. -/
-  chanOut              : Std.Channel JsonRpc.Message
+  chanOut               : Std.Channel JsonRpc.Message
   /--
   Latest document version received by the client, used for filtering out notifications from
   previous versions.
   -/
-  maxDocVersionRef     : IO.Ref Int
-  freshRequestIdRef    : IO.Ref Int
-  /--
-  Channel that receives a message for every a `$/lean/fileProgress` notification, indicating whether
-  the notification suggests that the file is currently being processed.
-  -/
-  chanIsProcessing     : Std.Channel Bool
+  maxDocVersionRef      : IO.Ref Int
+  freshRequestIdRef     : IO.Ref Int
   /--
   Diagnostics that are included in every single `textDocument/publishDiagnostics` notification.
   -/
   stickyDiagnosticsRef : IO.Ref (Array InteractiveDiagnostic)
+  partialHandlersRef   : IO.Ref (RBMap String PartialHandlerInfo compare)
   hLog                 : FS.Stream
   initParams           : InitializeParams
   processor            : Parser.InputContext → BaseIO Lean.Language.Lean.InitialSnapshot
@@ -90,6 +96,23 @@ structure WorkerContext where
   -/
   cmdlineOpts          : Options
 
+def WorkerContext.modifyGetPartialHandler (ctx : WorkerContext) (method : String)
+    (f : PartialHandlerInfo → α × PartialHandlerInfo) : BaseIO α :=
+  ctx.partialHandlersRef.modifyGet fun partialHandlers => Id.run do
+    let h := partialHandlers.find! method
+    let (r, h) := f h
+    (r, partialHandlers.insert method h)
+
+def WorkerContext.modifyPartialHandler (ctx : WorkerContext) (method : String)
+    (f : PartialHandlerInfo → PartialHandlerInfo) : BaseIO Unit :=
+  ctx.partialHandlersRef.modify fun partialHandlers => Id.run do
+  let some h := partialHandlers.find? method
+    | return partialHandlers
+  partialHandlers.insert method <| f h
+
+def WorkerContext.updateRequestsInFlight (ctx : WorkerContext) (method : String) (f : Nat → Nat) : BaseIO Unit :=
+    ctx.modifyPartialHandler method fun h => { h with requestsInFlight := f h.requestsInFlight }
+
 /-! # Asynchronous snapshot elaboration -/
 
 section Elab
@@ -340,11 +363,16 @@ section Initialization
     let clientHasWidgets := initParams.initializationOptions?.bind (·.hasWidgets?) |>.getD false
     let maxDocVersionRef ← IO.mkRef 0
     let freshRequestIdRef ← IO.mkRef (0 : Int)
-    let chanIsProcessing ← Std.Channel.new
     let stickyDiagnosticsRef ← IO.mkRef ∅
-    let chanOut ← mkLspOutputChannel maxDocVersionRef chanIsProcessing
+    let chanOut ← mkLspOutputChannel maxDocVersionRef
     let srcSearchPathPromise ← IO.Promise.new
-
+    let partialHandlersRef ← IO.mkRef <| RBMap.fromArray (cmp := compare) <|
+      (← partialLspRequestHandlerMethods).map fun (method, refreshMethod) =>
+        (method, {
+          refreshMethod
+          requestsInFlight := 0
+          pendingRefreshInfo? := none
+        })
     let processor := Language.Lean.process (setupImports meta opts chanOut srcSearchPathPromise)
     let processor ← Language.mkIncrementalProcessor processor
     let initSnap ← processor meta.mkInputContext
@@ -358,9 +386,9 @@ section Initialization
       initParams
       processor
       clientHasWidgets
+      partialHandlersRef
       maxDocVersionRef
       freshRequestIdRef
-      chanIsProcessing
       cmdlineOpts := opts
       stickyDiagnosticsRef
     }
@@ -383,7 +411,7 @@ section Initialization
         the output FS stream after discarding outdated notifications. This is the only component of
         the worker with access to the output stream, so we can synchronize messages from parallel
         elaboration tasks here. -/
-    mkLspOutputChannel maxDocVersion chanIsProcessing : IO (Std.Channel JsonRpc.Message) := do
+    mkLspOutputChannel maxDocVersion : IO (Std.Channel JsonRpc.Message) := do
       let chanOut ← Std.Channel.new
       let _ ← chanOut.forAsync (prio := .dedicated) fun msg => do
         -- discard outdated notifications; note that in contrast to responses, notifications can
@@ -399,12 +427,10 @@ section Initialization
         if let some version := version? then
           if version < (← maxDocVersion.get) then
             return
+
           -- note that because of `server.reportDelayMs`, we cannot simply set `maxDocVersion` here
           -- as that would allow outdated messages to be reported until the delay is over
         o.writeLspMessage msg |>.catchExceptions (fun _ => pure ())
-        if let .notification "$/lean/fileProgress" (some params) := msg then
-          if let some (params : LeanFileProgressParams) := fromJson? (toJson params) |>.toOption then
-            chanIsProcessing.send (! params.processing.isEmpty)
       return chanOut
 
     getImportClosure? (snap : Language.Lean.InitialSnapshot) : Array Name := Id.run do
@@ -422,7 +448,7 @@ section ServerRequests
       (ctx    : WorkerContext)
       (method : String)
       (param  : α)
-      : IO Unit := do
+      : BaseIO Unit := do
     let freshRequestId ← ctx.freshRequestIdRef.modifyGet fun freshRequestId =>
       (freshRequestId, freshRequestId + 1)
     let r : JsonRpc.Request α := ⟨freshRequestId, method, param⟩
@@ -455,8 +481,18 @@ section NotificationHandling
   def handleDidChange (p : DidChangeTextDocumentParams) : WorkerM Unit := do
     let docId := p.textDocument
     let changes := p.contentChanges
+    let ctx ← read
+    let st ← get
     let oldDoc := (←get).doc
     let newVersion := docId.version?.getD 0
+    let _ ← IO.mapTask (t := st.srcSearchPathTask) fun srcSearchPath =>
+      let rc : RequestContext :=
+        { rpcSessions := st.rpcSessions
+          srcSearchPath
+          doc := oldDoc
+          hLog := ctx.hLog
+          initParams := ctx.initParams }
+      RequestM.runInIO (handleOnDidChange p) rc
     if ¬ changes.isEmpty then
       let newDocText := foldDocumentChanges changes oldDoc.meta.text
       updateDocument ⟨docId.uri, newVersion, newDocText, oldDoc.meta.dependencyBuildMode⟩
@@ -594,6 +630,11 @@ section MessageHandling
     let ctx ← read
     let st ← get
 
+    ctx.modifyPartialHandler method fun h => { h with
+      pendingRefreshInfo? := none
+      requestsInFlight := h.requestsInFlight + 1
+    }
+
     -- special cases
     try
       match method with
@@ -644,8 +685,17 @@ section MessageHandling
           IO.asTask do
             ctx.chanOut.send <| e.toLspResponseError id
         | Except.ok t => (IO.mapTask · t) fun
-          | Except.ok resp =>
-            ctx.chanOut.send <| .response id (toJson resp)
+          | Except.ok r => do
+            ctx.chanOut.send <| .response id (toJson r.response)
+            let timestamp ← IO.monoMsNow
+            ctx.modifyPartialHandler method fun h => { h with
+              requestsInFlight := h.requestsInFlight - 1
+              pendingRefreshInfo? :=
+                if r.isComplete then
+                  none
+                else
+                  some { lastRefreshTimestamp := timestamp, successiveRefreshAttempts := 0 }
+            }
           | Except.error e =>
             ctx.chanOut.send <| e.toLspResponseError id
     queueRequest id t
@@ -698,19 +748,78 @@ section MainLoop
     | _ => throwServerError "Got invalid JSON-RPC message"
 end MainLoop
 
-def runRefreshTask : WorkerM (Task (Except IO.Error Unit)) := do
+def runRefreshTasks : WorkerM (Array (Task Unit)) := do
+  let timeUntilRefreshMs := 2000
+  -- We limit the amount of successive refresh attempts in case the user has switched files,
+  -- in which case VS Code won't respond to any refresh request for the given file.
+  -- Since we don't want to spam the client with refresh requests for every single file that they
+  -- switched away from, we limit the amount of attempts.
+  let maxSuccessiveRefreshAttempts := 10
   let ctx ← read
-  IO.asTask (prio := Task.Priority.dedicated) do
-    while ! (←IO.checkCanceled) do
-      let pastProcessingStates ← ctx.chanIsProcessing.recvAllCurrent
-      if pastProcessingStates.isEmpty then
-        -- Processing progress has not changed since we last sent out a refresh request
-        -- => do not send out another one for now so that we do not make the client spam
-        --    semantic token requests while idle and already having received an up-to-date state
-        IO.sleep 1000
-        continue
-      sendServerRequest ctx "workspace/semanticTokens/refresh" (none : Option Nat)
-      IO.sleep 2000
+  let mut tasks := #[]
+  for (method, refreshMethod) in ← partialLspRequestHandlerMethods do
+    tasks := tasks.push <| ← BaseIO.asTask (prio := .dedicated) do
+      while true do
+        let lastRefreshTimestamp? ← ctx.modifyGetPartialHandler method fun h => Id.run do
+          let some info := h.pendingRefreshInfo?
+            | return (none, h)
+          if info.successiveRefreshAttempts >= maxSuccessiveRefreshAttempts then
+            return (none, { h with pendingRefreshInfo? := none })
+          return (some info.lastRefreshTimestamp, h)
+        let some lastRefreshTimestamp := lastRefreshTimestamp?
+          | let cancelled ← sleepWithCancellation timeUntilRefreshMs.toUInt32
+            if cancelled then
+              return
+            continue
+
+        let currentTimestamp ← IO.monoMsNow
+        let passedTimeMs := currentTimestamp - lastRefreshTimestamp
+        let remainingTimeMs := timeUntilRefreshMs - passedTimeMs
+        if remainingTimeMs > 0 then
+          let cancelled ← sleepWithCancellation remainingTimeMs.toUInt32
+          if cancelled then
+            return
+
+        let currentTimestamp ← IO.monoMsNow
+        let canRefresh := ← ctx.modifyGetPartialHandler method fun h => Id.run do
+          let some pendingRefreshInfo := h.pendingRefreshInfo?
+            | return (false, h)
+          -- If there is a request in flight and we emit a refresh request, VS Code will discard
+          -- the response for the request in flight.
+          -- To avoid this (especially for long-running requests), we only emit refresh requests
+          -- once there are no pending requests anymore.
+          if h.requestsInFlight > 0 then
+            return (false, h)
+          let h := { h with
+            pendingRefreshInfo? := some {
+              lastRefreshTimestamp := currentTimestamp
+              successiveRefreshAttempts := pendingRefreshInfo.successiveRefreshAttempts + 1
+            }
+          }
+          (true, h)
+        if ! canRefresh then
+          let cancelled ← sleepWithCancellation timeUntilRefreshMs.toUInt32
+          if cancelled then
+            return
+          continue
+
+        sendServerRequest ctx refreshMethod (none : Option Nat)
+  return tasks
+
+where
+
+  sleepWithCancellation (ms : UInt32) : BaseIO Bool := do
+    if (← IO.checkCanceled) then
+      return true
+    let napMs := 200
+    let mut remainingMs := ms
+    while remainingMs > 0 do
+      let remainingNapMs := if remainingMs < napMs then remainingMs else napMs
+      IO.sleep remainingNapMs
+      remainingMs := remainingMs - remainingNapMs
+      if (← IO.checkCanceled) then
+        return true
+    return false
 
 def initAndRunWorker (i o e : FS.Stream) (opts : Options) : IO Unit := do
   let i ← maybeTee "fwIn.txt" false i
@@ -730,9 +839,10 @@ def initAndRunWorker (i o e : FS.Stream) (opts : Options) : IO Unit := do
     throw err
   StateRefT'.run' (s := st) <| ReaderT.run (r := ctx) do
     try
-      let refreshTask ← runRefreshTask
+      let refreshTasks ← runRefreshTasks
       mainLoop i
-      IO.cancel refreshTask
+      for refreshTasks in refreshTasks do
+        IO.cancel refreshTasks
     catch err =>
       let st ← get
       writeErrorDiag st.doc.meta err

src/Lean/Server/FileWorker/InlayHints.lean

@@ -0,0 +1,180 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Marc Huisinga
+-/
+prelude
+import Lean.Server.GoTo
+import Lean.Server.Requests
+
+namespace Lean.Elab
+
+def InlayHintLinkLocation.toLspLocation (srcSearchPath : SearchPath) (text : FileMap)
+    (l : InlayHintLinkLocation) : IO (Option Lsp.Location) := do
+  let some uri ← Server.documentUriFromModule srcSearchPath l.module
+    | return none
+  return some {
+    uri
+    range := text.utf8RangeToLspRange l.range
+  }
+
+def InlayHintLabelPart.toLspInlayHintLabelPart (srcSearchPath : SearchPath) (text : FileMap)
+    (p : InlayHintLabelPart) : IO Lsp.InlayHintLabelPart := do
+    let location? ← p.location?.bindM fun loc => loc.toLspLocation srcSearchPath text
+    let tooltip? := do return .markdown { kind := .markdown, value := ← p.tooltip? }
+    return {
+      value := p.value
+      location?,
+      tooltip?
+    }
+
+def InlayHintLabel.toLspInlayHintLabel (srcSearchPath : SearchPath) (text : FileMap) : InlayHintLabel → IO Lsp.InlayHintLabel
+  | .name n => do return .name n
+  | .parts p => do return .parts <| ← p.mapM (·.toLspInlayHintLabelPart srcSearchPath text)
+
+def InlayHintKind.toLspInlayHintKind : InlayHintKind → Lsp.InlayHintKind
+  | .type => .type
+  | .parameter => .parameter
+
+def InlayHintTextEdit.toLspTextEdit (text : FileMap) (e : InlayHintTextEdit) : Lsp.TextEdit := {
+  range := text.utf8RangeToLspRange e.range
+  newText := e.newText
+}
+
+def InlayHintInfo.toLspInlayHint (srcSearchPath : SearchPath) (text : FileMap) (i : InlayHintInfo) : IO Lsp.InlayHint := do
+  return {
+    position := text.utf8PosToLspPos i.position
+    label := ← i.label.toLspInlayHintLabel srcSearchPath text
+    kind? := i.kind?.map (·.toLspInlayHintKind)
+    textEdits? := some <| i.textEdits.map (·.toLspTextEdit text)
+    tooltip? := do return .markdown { kind := .markdown, value := ← i.tooltip? }
+    paddingLeft? := i.paddingLeft
+    paddingRight? := i.paddingRight
+  }
+
+end Lean.Elab
+
+namespace Lean.Server.FileWorker
+open Lsp
+
+def applyEditToHint? (hintMod : Name) (ihi : Elab.InlayHintInfo) (range : String.Range) (newText : String) : Option Elab.InlayHintInfo := do
+  let isLabelLocAffectedByEdit :=
+    match ihi.label with
+    | .name _ => false
+    | .parts ps => ps.any (·.location?.any (fun loc => loc.module == hintMod && range.overlaps loc.range (includeFirstStop := true)))
+  -- We always include the stop of the edit range because insertion edit ranges may be empty,
+  -- but we must still invalidate the inlay hint in this case.
+  let isInlayHintInvalidatedByEdit :=
+    range.contains ihi.position (includeStop := true) ||
+      ihi.textEdits.any (range.overlaps ·.range (includeFirstStop := true)) ||
+      isLabelLocAffectedByEdit
+  if isInlayHintInvalidatedByEdit then
+    none
+  let byteOffset : Int := (newText.toSubstring.bsize : Int) - (range.bsize : Int)
+  let shift (p : String.Pos) : String.Pos :=
+    if range.stop < p then
+      ⟨p.byteIdx + byteOffset |>.toNat⟩
+    else if p < range.start then
+      p
+    else -- `range.start <= p && p <= range.stop`
+      panic! s!"Got position {p} that should have been invalidated by edit at range {range.start}-{range.stop}"
+  let shiftRange (r : String.Range) : String.Range := ⟨shift r.start, shift r.stop⟩
+  return { ihi with
+    position := shift ihi.position
+    textEdits := ihi.textEdits.map fun edit => { edit with range := shiftRange edit.range }
+    label :=
+      match ihi.label with
+      | .name s => .name s
+      | .parts ps => .parts <| ps.map fun p => { p with
+        location? := p.location?.map fun loc =>
+          if loc.module == hintMod then
+            { loc with range := shiftRange loc.range }
+          else
+            loc
+      }
+  }
+
+structure InlayHintState where
+  oldInlayHints : Array Elab.InlayHintInfo
+  deriving TypeName, Inhabited
+
+def InlayHintState.init : InlayHintState := {
+  oldInlayHints := #[]
+}
+
+def handleInlayHints (_ : InlayHintParams) (s : InlayHintState) :
+    RequestM (LspResponse (Array InlayHint) × InlayHintState) := do
+  let ctx ← read
+  let srcSearchPath := ctx.srcSearchPath
+  -- We delay sending inlay hints by 1000ms to avoid inlay hint flickering on the client.
+  -- VS Code already has a mechanism for this, but it is not sufficient.
+  -- Note that 1000ms of latency for this request are actually 2000ms of latency in VS Code after a
+  -- `textDocument/didChange` notification because VS Code (for some reason) emits two inlay hint
+  -- requests in succession after a change, immediately invalidating the result of the first.
+  -- Finally, for some stupid reason, VS Code doesn't remove the inlay hint when applying it,
+  -- so this additional latency causes the applied inlay hint to linger around for a bit.
+  let (snaps, _, isComplete) ← ctx.doc.cmdSnaps.getFinishedPrefixWithConsistentLatency 1000
+  let finishedRange? : Option String.Range := do
+    return ⟨⟨0⟩, ← List.max? <| snaps.map (fun s => s.endPos)⟩
+  let oldInlayHints :=
+    if let some finishedRange := finishedRange? then
+      s.oldInlayHints.filter fun (ihi : Elab.InlayHintInfo) =>
+        ! finishedRange.contains ihi.position
+    else
+      s.oldInlayHints
+  let newInlayHints : Array Elab.InlayHintInfo ← (·.2) <$> StateT.run (s := #[]) do
+    for s in snaps do
+      s.infoTree.visitM' (postNode := fun ci i _ => do
+        let .ofCustomInfo i := i
+          | return
+        let some ih := Elab.InlayHint.ofCustomInfo? i
+          | return
+        let ih ← ci.runMetaM ih.lctx ih.resolveDeferred
+        modify (·.push ih.toInlayHintInfo))
+  let inlayHints := newInlayHints ++ oldInlayHints
+  let lspInlayHints ← inlayHints.mapM (·.toLspInlayHint srcSearchPath ctx.doc.meta.text)
+  return ({ response := lspInlayHints, isComplete }, { s with oldInlayHints := inlayHints })
+
+def handleInlayHintsDidChange (p : DidChangeTextDocumentParams)
+    : StateT InlayHintState RequestM Unit := do
+  let meta := (← read).doc.meta
+  let text := meta.text
+  let srcSearchPath := (← read).srcSearchPath
+  let .ok (some modName) ← EIO.toBaseIO <| do
+      let some path := System.Uri.fileUriToPath? meta.uri
+        | return none
+      let some mod ← searchModuleNameOfFileName path srcSearchPath
+        | return some <| ← moduleNameOfFileName path none
+      return some mod
+    | return
+  let s ← get
+  let mut updatedOldInlayHints := #[]
+  for ihi in s.oldInlayHints do
+    let mut ihi := ihi
+    let mut inlayHintInvalidated := false
+    for c in p.contentChanges do
+      let .rangeChange changeRange newText := c
+        | set <| { s with oldInlayHints := #[] } -- `fullChange` => all old inlay hints invalidated
+          return
+      let changeRange := text.lspRangeToUtf8Range changeRange
+      let some ihi' := applyEditToHint? modName ihi changeRange newText
+        | -- Change in some position of inlay hint => inlay hint invalidated
+          inlayHintInvalidated := true
+          break
+      ihi := ihi'
+    if ! inlayHintInvalidated then
+      updatedOldInlayHints := updatedOldInlayHints.push ihi
+  set <| { s with oldInlayHints := updatedOldInlayHints }
+
+builtin_initialize
+  registerPartialStatefulLspRequestHandler
+    "textDocument/inlayHint"
+    "workspace/inlayHint/refresh"
+    InlayHintParams
+    (Array InlayHint)
+    InlayHintState
+    .init
+    handleInlayHints
+    handleInlayHintsDidChange
+
+end Lean.Server.FileWorker

src/Lean/Server/FileWorker/RequestHandling.lean

@@ -5,20 +5,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Wojciech Nawrocki, Marc Huisinga
 -/
 prelude
-import Lean.DeclarationRange
-
-import Lean.Data.Json
-import Lean.Data.Lsp
-
-import Lean.Parser.Tactic.Doc
-
-import Lean.Server.FileWorker.Utils
-import Lean.Server.Requests
+import Lean.Server.FileWorker.InlayHints
+import Lean.Server.FileWorker.SemanticHighlighting
 import Lean.Server.Completion
 import Lean.Server.References
-import Lean.Server.GoTo
 
-import Lean.Widget.InteractiveGoal
 import Lean.Widget.Diff
 
 namespace Lean.Server.FileWorker
@@ -443,175 +434,6 @@ where
             return toDocumentSymbols text stxs (syms.push sym) stack
         toDocumentSymbols text stxs syms stack
 
-/--
-`SyntaxNodeKind`s for which the syntax node and its children receive no semantic highlighting.
--/
-def noHighlightKinds : Array SyntaxNodeKind := #[
-  -- usually have special highlighting by the client
-  ``Lean.Parser.Term.sorry,
-  ``Lean.Parser.Term.type,
-  ``Lean.Parser.Term.prop,
-  -- not really keywords
-  `antiquotName,
-  ``Lean.Parser.Command.docComment,
-  ``Lean.Parser.Command.moduleDoc]
-
--- TODO: make extensible, or don't
-/-- Keywords for which a specific semantic token is provided. -/
-def keywordSemanticTokenMap : RBMap String SemanticTokenType compare :=
-  RBMap.empty
-    |>.insert "sorry" .leanSorryLike
-    |>.insert "admit" .leanSorryLike
-    |>.insert "stop" .leanSorryLike
-    |>.insert "#exit" .leanSorryLike
-
-/-- Semantic token information for a given `Syntax`. -/
-structure LeanSemanticToken where
-  /-- Syntax of the semantic token. -/
-  stx  : Syntax
-  /-- Type of the semantic token. -/
-  type : SemanticTokenType
-
-/-- Semantic token information with absolute LSP positions. -/
-structure AbsoluteLspSemanticToken where
-  /-- Start position of the semantic token. -/
-  pos     : Lsp.Position
-  /-- End position of the semantic token. -/
-  tailPos : Lsp.Position
-  /-- Start position of the semantic token. -/
-  type    : SemanticTokenType
-  deriving BEq, Hashable, FromJson, ToJson
-
-/--
-Given a set of `LeanSemanticToken`, computes the `AbsoluteLspSemanticToken` with absolute
-LSP position information for each token.
--/
-def computeAbsoluteLspSemanticTokens
-    (text     : FileMap)
-    (beginPos : String.Pos)
-    (endPos?  : Option String.Pos)
-    (tokens   : Array LeanSemanticToken)
-    : Array AbsoluteLspSemanticToken :=
-  tokens.filterMap fun ⟨stx, tokenType⟩ => do
-    let (pos, tailPos) := (← stx.getPos?, ← stx.getTailPos?)
-    guard <| beginPos <= pos && endPos?.all (pos < ·)
-    let (lspPos, lspTailPos) := (text.utf8PosToLspPos pos, text.utf8PosToLspPos tailPos)
-    return ⟨lspPos, lspTailPos, tokenType⟩
-
-/-- Filters all duplicate semantic tokens with the same `pos`, `tailPos` and `type`. -/
-def filterDuplicateSemanticTokens (tokens : Array AbsoluteLspSemanticToken) : Array AbsoluteLspSemanticToken :=
-  tokens.groupByKey id |>.toArray.map (·.1)
-
-/--
-Given a set of `AbsoluteLspSemanticToken`, computes the LSP `SemanticTokens` data with
-token-relative positioning.
-See https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_semanticTokens.
--/
-def computeDeltaLspSemanticTokens (tokens : Array AbsoluteLspSemanticToken) : SemanticTokens := Id.run do
-  let tokens := tokens.qsort fun ⟨pos1, tailPos1, _⟩ ⟨pos2, tailPos2, _⟩ =>
-    pos1 < pos2 || pos1 == pos2 && tailPos1 <= tailPos2
-  let mut data : Array Nat := Array.mkEmpty (5*tokens.size)
-  let mut lastPos : Lsp.Position := ⟨0, 0⟩
-  for ⟨pos, tailPos, tokenType⟩ in tokens do
-    let deltaLine := pos.line - lastPos.line
-    let deltaStart := pos.character - (if pos.line == lastPos.line then lastPos.character else 0)
-    let length := tailPos.character - pos.character
-    let tokenType := tokenType.toNat
-    let tokenModifiers := 0
-    data := data ++ #[deltaLine, deltaStart, length, tokenType, tokenModifiers]
-    lastPos := pos
-  return { data }
-
-/--
-Collects all semantic tokens that can be deduced purely from `Syntax`
-without elaboration information.
--/
-partial def collectSyntaxBasedSemanticTokens : (stx : Syntax) → Array LeanSemanticToken
-  | `($e.$id:ident)    =>
-    let tokens := collectSyntaxBasedSemanticTokens e
-    tokens.push ⟨id, SemanticTokenType.property⟩
-  | `($e |>.$field:ident) =>
-    let tokens := collectSyntaxBasedSemanticTokens e
-    tokens.push ⟨field, SemanticTokenType.property⟩
-  | stx => Id.run do
-    if noHighlightKinds.contains stx.getKind then
-      return #[]
-    let mut tokens :=
-      if stx.isOfKind choiceKind then
-        collectSyntaxBasedSemanticTokens stx[0]
-      else
-        stx.getArgs.map collectSyntaxBasedSemanticTokens |>.flatten
-    let Syntax.atom _ val := stx
-      | return tokens
-    let isRegularKeyword := val.length > 0 && val.front.isAlpha
-    let isHashKeyword := val.length > 1 && val.front == '#' && (val.get ⟨1⟩).isAlpha
-    if ! isRegularKeyword && ! isHashKeyword then
-      return tokens
-    return tokens.push ⟨stx, keywordSemanticTokenMap.findD val .keyword⟩
-
-/-- Collects all semantic tokens from the given `Elab.InfoTree`. -/
-def collectInfoBasedSemanticTokens (i : Elab.InfoTree) : Array LeanSemanticToken :=
-  List.toArray <| i.deepestNodes fun _ i _ => do
-    let .ofTermInfo ti := i
-      | none
-    let .original .. := ti.stx.getHeadInfo
-      | none
-    if let `($_:ident) := ti.stx then
-      if let Expr.fvar fvarId .. := ti.expr then
-        if let some localDecl := ti.lctx.find? fvarId then
-          -- Recall that `isAuxDecl` is an auxiliary declaration used to elaborate a recursive definition.
-          if localDecl.isAuxDecl then
-            if ti.isBinder then
-              return ⟨ti.stx, SemanticTokenType.function⟩
-          else if ! localDecl.isImplementationDetail then
-            return ⟨ti.stx, SemanticTokenType.variable⟩
-    if ti.stx.getKind == Parser.Term.identProjKind then
-      return ⟨ti.stx, SemanticTokenType.property⟩
-    none
-
-/-- Computes the semantic tokens in the range [beginPos, endPos?). -/
-def handleSemanticTokens (beginPos : String.Pos) (endPos? : Option String.Pos)
-    : RequestM (RequestTask SemanticTokens) := do
-  let doc ← readDoc
-  match endPos? with
-  | none =>
-    -- Only grabs the finished prefix so that we do not need to wait for elaboration to complete
-    -- for the full file before sending a response. This means that the response will be incomplete,
-    -- which we mitigate by regularly sending `workspace/semanticTokens/refresh` requests in the
-    -- `FileWorker` to tell the client to re-compute the semantic tokens.
-    let (snaps, _) ← doc.cmdSnaps.getFinishedPrefix
-    asTask <| run doc snaps
-  | some endPos =>
-    let t := doc.cmdSnaps.waitUntil (·.endPos >= endPos)
-    mapTask t fun (snaps, _) => run doc snaps
-where
-  run doc snaps : RequestM SemanticTokens := do
-    let mut leanSemanticTokens := #[]
-    for s in snaps do
-      if s.endPos <= beginPos then
-        continue
-      let syntaxBasedSemanticTokens := collectSyntaxBasedSemanticTokens s.stx
-      let infoBasedSemanticTokens := collectInfoBasedSemanticTokens s.infoTree
-      leanSemanticTokens := leanSemanticTokens ++ syntaxBasedSemanticTokens ++ infoBasedSemanticTokens
-    let absoluteLspSemanticTokens := computeAbsoluteLspSemanticTokens doc.meta.text beginPos endPos? leanSemanticTokens
-    let absoluteLspSemanticTokens := filterDuplicateSemanticTokens absoluteLspSemanticTokens
-    let semanticTokens := computeDeltaLspSemanticTokens absoluteLspSemanticTokens
-    return semanticTokens
-
-/-- Computes all semantic tokens for the document. -/
-def handleSemanticTokensFull (_ : SemanticTokensParams)
-    : RequestM (RequestTask SemanticTokens) := do
-  handleSemanticTokens 0 none
-
-/-- Computes the semantic tokens in the range provided by `p`. -/
-def handleSemanticTokensRange (p : SemanticTokensRangeParams)
-    : RequestM (RequestTask SemanticTokens) := do
-  let doc ← readDoc
-  let text := doc.meta.text
-  let beginPos := text.lspPosToUtf8Pos p.range.start
-  let endPos := text.lspPosToUtf8Pos p.range.end
-  handleSemanticTokens beginPos endPos
-
 partial def handleFoldingRange (_ : FoldingRangeParams)
   : RequestM (RequestTask (Array FoldingRange)) := do
   let doc ← readDoc
@@ -753,16 +575,6 @@ builtin_initialize
     DocumentSymbolParams
     DocumentSymbolResult
     handleDocumentSymbol
-  registerLspRequestHandler
-    "textDocument/semanticTokens/full"
-    SemanticTokensParams
-    SemanticTokens
-    handleSemanticTokensFull
-  registerLspRequestHandler
-    "textDocument/semanticTokens/range"
-    SemanticTokensRangeParams
-    SemanticTokens
-    handleSemanticTokensRange
   registerLspRequestHandler
     "textDocument/foldingRange"
     FoldingRangeParams

src/Lean/Server/FileWorker/SemanticHighlighting.lean

@@ -0,0 +1,211 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sebastian Ullrich, Marc Huisinga
+-/
+prelude
+import Lean.Server.Requests
+
+namespace Lean.Server.FileWorker
+open Lsp
+open RequestM
+
+/--
+`SyntaxNodeKind`s for which the syntax node and its children receive no semantic highlighting.
+-/
+def noHighlightKinds : Array SyntaxNodeKind := #[
+  -- usually have special highlighting by the client
+  ``Lean.Parser.Term.sorry,
+  ``Lean.Parser.Term.type,
+  ``Lean.Parser.Term.prop,
+  -- not really keywords
+  `antiquotName,
+  ``Lean.Parser.Command.docComment,
+  ``Lean.Parser.Command.moduleDoc]
+
+-- TODO: make extensible, or don't
+/-- Keywords for which a specific semantic token is provided. -/
+def keywordSemanticTokenMap : RBMap String SemanticTokenType compare :=
+  RBMap.empty
+    |>.insert "sorry" .leanSorryLike
+    |>.insert "admit" .leanSorryLike
+    |>.insert "stop" .leanSorryLike
+    |>.insert "#exit" .leanSorryLike
+
+/-- Semantic token information for a given `Syntax`. -/
+structure LeanSemanticToken where
+  /-- Syntax of the semantic token. -/
+  stx  : Syntax
+  /-- Type of the semantic token. -/
+  type : SemanticTokenType
+
+/-- Semantic token information with absolute LSP positions. -/
+structure AbsoluteLspSemanticToken where
+  /-- Start position of the semantic token. -/
+  pos     : Lsp.Position
+  /-- End position of the semantic token. -/
+  tailPos : Lsp.Position
+  /-- Start position of the semantic token. -/
+  type    : SemanticTokenType
+  deriving BEq, Hashable, FromJson, ToJson
+
+/--
+Given a set of `LeanSemanticToken`, computes the `AbsoluteLspSemanticToken` with absolute
+LSP position information for each token.
+-/
+def computeAbsoluteLspSemanticTokens
+    (text     : FileMap)
+    (beginPos : String.Pos)
+    (endPos?  : Option String.Pos)
+    (tokens   : Array LeanSemanticToken)
+    : Array AbsoluteLspSemanticToken :=
+  tokens.filterMap fun ⟨stx, tokenType⟩ => do
+    let (pos, tailPos) := (← stx.getPos?, ← stx.getTailPos?)
+    guard <| beginPos <= pos && endPos?.all (pos < ·)
+    let (lspPos, lspTailPos) := (text.utf8PosToLspPos pos, text.utf8PosToLspPos tailPos)
+    return ⟨lspPos, lspTailPos, tokenType⟩
+
+/-- Filters all duplicate semantic tokens with the same `pos`, `tailPos` and `type`. -/
+def filterDuplicateSemanticTokens (tokens : Array AbsoluteLspSemanticToken) : Array AbsoluteLspSemanticToken :=
+  tokens.groupByKey id |>.toArray.map (·.1)
+
+/--
+Given a set of `AbsoluteLspSemanticToken`, computes the LSP `SemanticTokens` data with
+token-relative positioning.
+See https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_semanticTokens.
+-/
+def computeDeltaLspSemanticTokens (tokens : Array AbsoluteLspSemanticToken) : SemanticTokens := Id.run do
+  let tokens := tokens.qsort fun ⟨pos1, tailPos1, _⟩ ⟨pos2, tailPos2, _⟩ =>
+    pos1 < pos2 || pos1 == pos2 && tailPos1 <= tailPos2
+  let mut data : Array Nat := Array.mkEmpty (5*tokens.size)
+  let mut lastPos : Lsp.Position := ⟨0, 0⟩
+  for ⟨pos, tailPos, tokenType⟩ in tokens do
+    let deltaLine := pos.line - lastPos.line
+    let deltaStart := pos.character - (if pos.line == lastPos.line then lastPos.character else 0)
+    let length := tailPos.character - pos.character
+    let tokenType := tokenType.toNat
+    let tokenModifiers := 0
+    data := data ++ #[deltaLine, deltaStart, length, tokenType, tokenModifiers]
+    lastPos := pos
+  return { data }
+
+/--
+Collects all semantic tokens that can be deduced purely from `Syntax`
+without elaboration information.
+-/
+partial def collectSyntaxBasedSemanticTokens : (stx : Syntax) → Array LeanSemanticToken
+  | `($e.$id:ident)    =>
+    let tokens := collectSyntaxBasedSemanticTokens e
+    tokens.push ⟨id, SemanticTokenType.property⟩
+  | `($e |>.$field:ident) =>
+    let tokens := collectSyntaxBasedSemanticTokens e
+    tokens.push ⟨field, SemanticTokenType.property⟩
+  | stx => Id.run do
+    if noHighlightKinds.contains stx.getKind then
+      return #[]
+    let mut tokens :=
+      if stx.isOfKind choiceKind then
+        collectSyntaxBasedSemanticTokens stx[0]
+      else
+        stx.getArgs.map collectSyntaxBasedSemanticTokens |>.flatten
+    let Syntax.atom _ val := stx
+      | return tokens
+    let isRegularKeyword := val.length > 0 && val.front.isAlpha
+    let isHashKeyword := val.length > 1 && val.front == '#' && (val.get ⟨1⟩).isAlpha
+    if ! isRegularKeyword && ! isHashKeyword then
+      return tokens
+    return tokens.push ⟨stx, keywordSemanticTokenMap.findD val .keyword⟩
+
+/-- Collects all semantic tokens from the given `Elab.InfoTree`. -/
+def collectInfoBasedSemanticTokens (i : Elab.InfoTree) : Array LeanSemanticToken :=
+  List.toArray <| i.deepestNodes fun _ i _ => do
+    let .ofTermInfo ti := i
+      | none
+    let .original .. := ti.stx.getHeadInfo
+      | none
+    if let `($_:ident) := ti.stx then
+      if let Expr.fvar fvarId .. := ti.expr then
+        if let some localDecl := ti.lctx.find? fvarId then
+          -- Recall that `isAuxDecl` is an auxiliary declaration used to elaborate a recursive definition.
+          if localDecl.isAuxDecl then
+            if ti.isBinder then
+              return ⟨ti.stx, SemanticTokenType.function⟩
+          else if ! localDecl.isImplementationDetail then
+            return ⟨ti.stx, SemanticTokenType.variable⟩
+    if ti.stx.getKind == Parser.Term.identProjKind then
+      return ⟨ti.stx, SemanticTokenType.property⟩
+    none
+
+/-- Computes the semantic tokens in the range [beginPos, endPos?). -/
+def handleSemanticTokens (beginPos : String.Pos) (endPos? : Option String.Pos)
+    : RequestM (RequestTask (LspResponse SemanticTokens)) := do
+  let doc ← readDoc
+  match endPos? with
+  | none =>
+    -- Only grabs the finished prefix so that we do not need to wait for elaboration to complete
+    -- for the full file before sending a response. This means that the response will be incomplete,
+    -- which we mitigate by regularly sending `workspace/semanticTokens/refresh` requests in the
+    -- `FileWorker` to tell the client to re-compute the semantic tokens.
+    let (snaps, _, isComplete) ← doc.cmdSnaps.getFinishedPrefixWithTimeout 2000
+    asTask <| do
+      return { response := ← run doc snaps, isComplete }
+  | some endPos =>
+    let t := doc.cmdSnaps.waitUntil (·.endPos >= endPos)
+    mapTask t fun (snaps, _) => do
+      return { response := ← run doc snaps, isComplete := true }
+where
+  run doc snaps : RequestM SemanticTokens := do
+    let mut leanSemanticTokens := #[]
+    for s in snaps do
+      if s.endPos <= beginPos then
+        continue
+      let syntaxBasedSemanticTokens := collectSyntaxBasedSemanticTokens s.stx
+      let infoBasedSemanticTokens := collectInfoBasedSemanticTokens s.infoTree
+      leanSemanticTokens := leanSemanticTokens ++ syntaxBasedSemanticTokens ++ infoBasedSemanticTokens
+    let absoluteLspSemanticTokens := computeAbsoluteLspSemanticTokens doc.meta.text beginPos endPos? leanSemanticTokens
+    let absoluteLspSemanticTokens := filterDuplicateSemanticTokens absoluteLspSemanticTokens
+    let semanticTokens := computeDeltaLspSemanticTokens absoluteLspSemanticTokens
+    return semanticTokens
+
+structure SemanticTokensState where
+  deriving TypeName, Inhabited
+
+/-- Computes all semantic tokens for the document. -/
+def handleSemanticTokensFull (_ : SemanticTokensParams) (_ : SemanticTokensState)
+    : RequestM (LspResponse SemanticTokens × SemanticTokensState) := do
+  let t ← handleSemanticTokens 0 none
+  match t.get with
+  | .error e => throw e
+  | .ok r => return (r, ⟨⟩)
+
+def handleSemanticTokensDidChange (_ : DidChangeTextDocumentParams)
+    : StateT SemanticTokensState RequestM Unit := do
+  return
+
+/-- Computes the semantic tokens in the range provided by `p`. -/
+def handleSemanticTokensRange (p : SemanticTokensRangeParams)
+    : RequestM (RequestTask SemanticTokens) := do
+  let doc ← readDoc
+  let text := doc.meta.text
+  let beginPos := text.lspPosToUtf8Pos p.range.start
+  let endPos := text.lspPosToUtf8Pos p.range.end
+  let t ← handleSemanticTokens beginPos endPos
+  return t.map fun r => r.map (·.response)
+
+builtin_initialize
+  registerLspRequestHandler
+    "textDocument/semanticTokens/range"
+    SemanticTokensRangeParams
+    SemanticTokens
+    handleSemanticTokensRange
+  registerPartialStatefulLspRequestHandler
+    "textDocument/semanticTokens/full"
+    "workspace/semanticTokens/refresh"
+    SemanticTokensParams
+    SemanticTokens
+    SemanticTokensState
+    ⟨⟩
+    handleSemanticTokensFull
+    handleSemanticTokensDidChange
+
+end Lean.Server.FileWorker

src/Lean/Server/Requests.lean

@@ -16,6 +16,8 @@ import Lean.Server.FileWorker.Utils
 
 import Lean.Server.Rpc.Basic
 
+import Std.Sync.Mutex
+
 namespace Lean.Language
 
 /--
@@ -115,6 +117,9 @@ namespace RequestM
 open FileWorker
 open Snapshots
 
+def runInIO (x : RequestM α) (ctx : RequestContext) : IO α := do
+  x.run ctx |>.adaptExcept (IO.userError ·.message)
+
 def readDoc [Monad m] [MonadReaderOf RequestContext m] : m EditableDocument := do
   let rc ← readThe RequestContext
   return rc.doc
@@ -343,17 +348,224 @@ def chainLspRequestHandler (method : String)
   else
     throw <| IO.userError s!"Failed to chain LSP request handler for '{method}': no initial handler registered"
 
-def routeLspRequest (method : String) (params : Json) : IO (Except RequestError DocumentUri) := do
-  match (← lookupLspRequestHandler method) with
-  | none => return Except.error <| RequestError.methodNotFound method
-  | some rh => return rh.fileSource params
+inductive RequestHandlerCompleteness where
+  | complete
+  /--
+  A request handler is partial if the LSP spec states that the request method implemented by
+  the handler should be responded to with the full state of the document, but our implementation
+  of the handler only returns a partial result for the document
+  (e.g. only for the processed regions of the document, to reduce latency after a `didChange`).
+  A request handler can only be partial if LSP also specifies a corresponding `refresh`
+  server-to-client request, e.g. `workspace/inlayHint/refresh` for `textDocument/inlayHint`.
+  This is necessary because we use the `refresh` request to prompt the client to re-request the
+  data for a partial request if we returned a partial response for that request in the past,
+  so that the client eventually converges to a complete set of information for the full document.
+  -/
+  | «partial» (refreshMethod : String)
 
-def handleLspRequest (method : String) (params : Json) : RequestM (RequestTask Json) := do
-  match (← lookupLspRequestHandler method) with
-  | none =>
-    throw <| .internalError
-      s!"request '{method}' routed through watchdog but unknown in worker; are both using the same plugins?"
-  | some rh => rh.handle params
+structure LspResponse (α : Type) where
+  response   : α
+  isComplete : Bool
+
+structure StatefulRequestHandler where
+  fileSource      : Json → Except RequestError Lsp.DocumentUri
+  /--
+  `handle` with explicit state management for chaining request handlers.
+  This function is pure w.r.t. `lastTaskMutex` and `stateRef`, but not `RequestM`.
+  -/
+  pureHandle      : Json → Dynamic → RequestM (LspResponse Json × Dynamic)
+  handle          : Json → RequestM (RequestTask (LspResponse Json))
+  /--
+  `onDidChange` with explicit state management for chaining request handlers.
+  This function is pure w.r.t. `lastTaskMutex` and `stateRef`, but not `RequestM`.
+  -/
+  pureOnDidChange : DidChangeTextDocumentParams → (StateT Dynamic RequestM) Unit
+  onDidChange     : DidChangeTextDocumentParams → RequestM Unit
+  lastTaskMutex   : Std.Mutex (Task Unit)
+  initState       : Dynamic
+  /--
+  `stateRef` is synchronized in `registerStatefulLspRequestHandler` by using `lastTaskMutex` to
+  ensure that stateful request tasks for the same handler are executed sequentially (in order of arrival).
+  -/
+  stateRef        : IO.Ref Dynamic
+  completeness    : RequestHandlerCompleteness
+
+builtin_initialize statefulRequestHandlers : IO.Ref (PersistentHashMap String StatefulRequestHandler) ←
+  IO.mkRef {}
+
+private def getState! (method : String) (state : Dynamic) stateType [TypeName stateType] : RequestM stateType := do
+  let some state := state.get? stateType
+    | throw <| .internalError s!"Got invalid state type in stateful LSP request handler for {method}"
+  return state
+
+private def getIOState! (method : String) (state : Dynamic) stateType [TypeName stateType] : IO stateType := do
+  let some state := state.get? stateType
+    | throw <| .userError s!"Got invalid state type in stateful LSP request handler for {method}"
+  return state
+
+private def overrideStatefulLspRequestHandler
+    (method : String) (completeness : RequestHandlerCompleteness)
+    paramType [FromJson paramType] [FileSource paramType]
+    respType [ToJson respType]
+    stateType [TypeName stateType]
+    (initState : stateType)
+    (handler : paramType → stateType → RequestM (LspResponse respType × stateType))
+    (onDidChange : DidChangeTextDocumentParams → StateT stateType RequestM Unit)
+    : IO Unit := do
+  if !(← Lean.initializing) then
+    throw <| IO.userError s!"Failed to register stateful LSP request handler for '{method}': only possible during initialization"
+  let fileSource := fun j =>
+    parseRequestParams paramType j |>.map Lsp.fileSource
+  let lastTaskMutex ← Std.Mutex.new <| Task.pure ()
+  let initState := Dynamic.mk initState
+  let stateRef ← IO.mkRef initState
+
+  let pureHandle : Json → Dynamic → RequestM (LspResponse Json × Dynamic) := fun param state => do
+    let param ← liftExcept <| parseRequestParams paramType param
+    let state ← getState! method state stateType
+    let (r, state') ← handler param state
+    return ({ r with response := toJson r.response }, Dynamic.mk state')
+
+  let handle : Json → RequestM (RequestTask (LspResponse Json)) := fun param => lastTaskMutex.atomically do
+    let lastTask ← get
+    let requestTask ← RequestM.mapTask lastTask fun () => do
+      let state ← stateRef.get
+      let (r, state') ← pureHandle param state
+      stateRef.set state'
+      return r
+    set <| requestTask.map <| fun _ => ()
+    return requestTask
+
+  let pureOnDidChange : DidChangeTextDocumentParams → (StateT Dynamic RequestM) Unit := fun param => do
+    let state ← getState! method (← get) stateType
+    let ((), state') ← onDidChange param |>.run state
+    set <| Dynamic.mk state'
+
+  let onDidChange : DidChangeTextDocumentParams → RequestM Unit := fun param => lastTaskMutex.atomically do
+      let lastTask ← get
+      let didChangeTask ← RequestM.mapTask (t := lastTask) fun () => do
+        let state ← stateRef.get
+        let ((), state') ← pureOnDidChange param |>.run state
+        stateRef.set state'
+      set <| didChangeTask.map <| fun _ => ()
+
+  statefulRequestHandlers.modify fun rhs => rhs.insert method {
+    fileSource,
+    pureHandle,
+    handle,
+    pureOnDidChange,
+    onDidChange,
+    lastTaskMutex,
+    initState,
+    stateRef,
+    completeness
+  }
+
+private def registerStatefulLspRequestHandler
+    (method : String) (completeness : RequestHandlerCompleteness)
+    paramType [FromJson paramType] [FileSource paramType]
+    respType [ToJson respType]
+    stateType [TypeName stateType]
+    (initState : stateType)
+    (handler : paramType → stateType → RequestM (LspResponse respType × stateType))
+    (onDidChange : DidChangeTextDocumentParams → StateT stateType RequestM Unit)
+    : IO Unit := do
+  if (← requestHandlers.get).contains method then
+    throw <| IO.userError s!"Failed to register stateful LSP request handler for '{method}': already registered"
+  overrideStatefulLspRequestHandler method completeness paramType respType stateType initState handler onDidChange
+
+def registerCompleteStatefulLspRequestHandler (method : String)
+    paramType [FromJson paramType] [FileSource paramType]
+    respType [ToJson respType]
+    stateType [TypeName stateType]
+    (initState : stateType)
+    (handler : paramType → stateType → RequestM (respType × stateType))
+    (onDidChange : DidChangeTextDocumentParams → StateT stateType RequestM Unit)
+    : IO Unit :=
+  let handler : paramType → stateType → RequestM (LspResponse respType × stateType) := fun p s => do
+    let (response, s) ← handler p s
+    return ({ response, isComplete := true }, s)
+  registerStatefulLspRequestHandler method .complete paramType respType stateType initState handler onDidChange
+
+def registerPartialStatefulLspRequestHandler (method refreshMethod : String)
+    paramType [FromJson paramType] [FileSource paramType]
+    respType [ToJson respType]
+    stateType [TypeName stateType]
+    (initState : stateType)
+    (handler : paramType → stateType → RequestM (LspResponse respType × stateType))
+    (onDidChange : DidChangeTextDocumentParams → StateT stateType RequestM Unit) :=
+  registerStatefulLspRequestHandler method (.partial refreshMethod) paramType respType stateType initState handler onDidChange
+
+def isStatefulLspRequestMethod (method : String) : BaseIO Bool := do
+  return (← statefulRequestHandlers.get).contains method
+
+def lookupStatefulLspRequestHandler (method : String) : BaseIO (Option StatefulRequestHandler) := do
+  return (← statefulRequestHandlers.get).find? method
+
+def partialLspRequestHandlerMethods : IO (Array (String × String)) := do
+  return (← statefulRequestHandlers.get).toArray.filterMap fun (method, h) => do
+    let .partial refreshMethod := h.completeness
+      | none
+    return (method, refreshMethod)
+
+def chainStatefulLspRequestHandler (method : String)
+    paramType [FromJson paramType] [ToJson paramType] [FileSource paramType]
+    respType [FromJson respType] [ToJson respType]
+    stateType [TypeName stateType]
+    (handler : paramType → LspResponse respType → stateType → RequestM (LspResponse respType × stateType))
+    (onDidChange : DidChangeTextDocumentParams → StateT stateType RequestM Unit) : IO Unit := do
+  if ! (← Lean.initializing) then
+    throw <| IO.userError s!"Failed to chain stateful LSP request handler for '{method}': only possible during initialization"
+  let some oldHandler ← lookupStatefulLspRequestHandler method
+    | throw <| IO.userError s!"Failed to chain stateful LSP request handler for '{method}': no initial handler registered"
+  let oldHandle := oldHandler.pureHandle
+  let oldOnDidChange := oldHandler.pureOnDidChange
+  let initState ← getIOState! method oldHandler.initState stateType
+  let handle (p : paramType) (s : stateType) : RequestM (LspResponse respType × stateType) := do
+    let (r, s) ← oldHandle (toJson p) (Dynamic.mk s)
+    let .ok response := fromJson? r.response
+      | throw <| RequestError.internalError "Failed to convert response of previous request handler when chaining stateful LSP request handlers"
+    let r := { r with response := response }
+    let s ← getState! method s stateType
+    handler p r s
+  let onDidChange (p : DidChangeTextDocumentParams) : StateT stateType RequestM Unit := do
+    let s ← get
+    let ((), s) ← oldOnDidChange p |>.run (Dynamic.mk s)
+    let s ← getState! method s stateType
+    let ((), s) ← onDidChange p |>.run s
+    set <| s
+  overrideStatefulLspRequestHandler method oldHandler.completeness paramType respType stateType initState
+    handle onDidChange
+
+def handleOnDidChange (p : DidChangeTextDocumentParams) : RequestM Unit := do
+  (← statefulRequestHandlers.get).forM fun _ handler => do
+    handler.onDidChange p
+
+def handleLspRequest (method : String) (params : Json) : RequestM (RequestTask (LspResponse Json)) := do
+  if ← isStatefulLspRequestMethod method then
+    match ← lookupStatefulLspRequestHandler method with
+    | none =>
+      throw <| .internalError
+        s!"request '{method}' routed through watchdog but unknown in worker; are both using the same plugins?"
+    | some rh => rh.handle params
+  else
+    match ← lookupLspRequestHandler method with
+    | none =>
+      throw <| .internalError
+        s!"request '{method}' routed through watchdog but unknown in worker; are both using the same plugins?"
+    | some rh =>
+      let t ← rh.handle params
+      return t.map (sync := true) fun r => r.map ({response := ·, isComplete := true })
+
+def routeLspRequest (method : String) (params : Json) : IO (Except RequestError DocumentUri) := do
+  if ← isStatefulLspRequestMethod method then
+    match ← lookupStatefulLspRequestHandler method with
+    | none => return Except.error <| RequestError.methodNotFound method
+    | some rh => return rh.fileSource params
+  else
+    match ← lookupLspRequestHandler method with
+    | none => return Except.error <| RequestError.methodNotFound method
+    | some rh => return rh.fileSource params
 
 end HandlerTable
 end Lean.Server

src/Lean/Server/Watchdog.lean

@@ -1096,6 +1096,9 @@ def mkLeanServerCapabilities : ServerCapabilities := {
     resolveProvider? := true,
     codeActionKinds? := some #["quickfix", "refactor"]
   }
+  inlayHintProvider? := some {
+    resolveProvider? := false
+  }
 }
 
 def initAndRunWatchdogAux : ServerM Unit := do

src/Lean/Syntax.lean

@@ -24,6 +24,14 @@ def String.Range.contains (r : String.Range) (pos : String.Pos) (includeStop :=
 def String.Range.includes (super sub : String.Range) : Bool :=
   super.start <= sub.start && super.stop >= sub.stop
 
+def String.Range.overlaps (first second : String.Range)
+    (includeFirstStop := false) (includeSecondStop := false) : Bool :=
+  (if includeFirstStop then second.start <= first.stop else second.start < first.stop) &&
+    (if includeSecondStop then first.start <= second.stop else first.start < second.stop)
+
+def String.Range.bsize (r : String.Range) : Nat :=
+  r.stop.byteIdx - r.start.byteIdx
+
 namespace Lean
 
 def SourceInfo.updateTrailing (trailing : Substring) : SourceInfo → SourceInfo

tests/lean/1018unknowMVarIssue.lean.expected.out

@@ -26,7 +26,8 @@ a : α
       • β : Type @ ⟨7, 33⟩-⟨7, 34⟩ @ Lean.Elab.Term.elabIdent
         • [.] β : some Sort.{?_uniq} @ ⟨7, 33⟩-⟨7, 34⟩
         • β : Type @ ⟨7, 33⟩-⟨7, 34⟩
-      • _example (isBinder := true) : {α β : Type} → α → Fam2 α β → β @ ⟨7, 0⟩†-⟨10, 19⟩†
+      • CustomInfo(Lean.Elab.InlayHint)
+      • _example (isBinder := true) : {α β : Type} → α → Fam2 α β → β @ ⟨7, 0⟩†!-⟨7, 7⟩†!
       • a (isBinder := true) : α @ ⟨7, 9⟩-⟨7, 10⟩
       • x (isBinder := true) : Fam2 α β @ ⟨7, 17⟩-⟨7, 18⟩
       • CustomInfo(Lean.Elab.Term.BodyInfo)