From 3d75c2ce2b18185342f018a8b36cb4635fb6d566 Mon Sep 17 00:00:00 2001
From: Marc Huisinga <mhuisi@protonmail.com>
Date: Fri, 3 Oct 2025 12:54:36 +0200
Subject: [PATCH] fix: eliminate potential source of inlay hint flakiness
 (#10664)

This PR fixes one potential source of inlay hint flakiness.

In the old `IO.waitAny` implementation, we could rely on the fact that
if all tasks in the list were finished, `IO.waitAny` would pick the
first finished one. In the new implementation (#9732), this isn't the
case anymore for fairness reasons, but this also means that in
`IO.AsyncList.getFinishedPrefixWithTimeout`, it can happen that we don't
scan the full finished command snapshot prefix because we pick the
timeout task before the finished snapshot task. This is likely the cause
of a flaky test failure
[here](https://github.com/leanprover/lean4/actions/runs/18215430028/job/51863870111),
where the inlay hint test yielded no result (the timeout task has an
edit delay of 0ms in the first inlay hint request that is emitted,
finishes immediately and can thus immediately cause the finished prefix
to be skipped with the new `waitAny` implementation).

This PR fixes this issue by adding a `hasFinished` check before the
`waitAny` to ensure that we always scan the finished prefix and don't
need to rely on a brittle invariant that doesn't hold anymore. It also
converts some `Task.get`s to `IO.wait` for safety so that the compiler
can't re-order them.
---
 src/Lean/Server/AsyncList.lean  | 22 +++++++++++++---------
 src/Lean/Server/ServerTask.lean |  2 ++
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/Lean/Server/AsyncList.lean b/src/Lean/Server/AsyncList.lean
index 3fade41b13..bdb4f7de7d 100644
--- a/src/Lean/Server/AsyncList.lean
+++ b/src/Lean/Server/AsyncList.lean
@@ -78,7 +78,7 @@ partial def getFinishedPrefix : AsyncList ε α → BaseIO (List α × Option ε
   | nil => pure ⟨[], none, true⟩
   | delayed tl => do
     if ← tl.hasFinished then
-      match tl.get with
+      match ← tl.wait with
       | Except.ok tl => tl.getFinishedPrefix
       | Except.error e => pure ⟨[], some e, true⟩
     else pure ⟨[], none, false⟩
@@ -102,14 +102,18 @@ where
       return ⟨hd :: tl, e?, isComplete⟩
     | nil => return ⟨[], none, true⟩
     | delayed tl =>
-      let tl : ServerTask (Except ε (AsyncList ε α)) := tl
-      let tl := tl.mapCheap .inr
-      let cancelTks := cancelTks.map (·.mapCheap .inl)
-      let r ← ServerTask.waitAny (tl :: cancelTks ++ [timeoutTask])
-      match r with
-      | .inl _ => return ⟨[], none, false⟩ -- Timeout or cancellation - stop waiting
-      | .inr (.ok tl) => go timeoutTask tl
-      | .inr (.error e) => return ⟨[], some e, true⟩
+      if ← tl.hasFinished then
+        match ← tl.wait with
+        | .ok tl => go timeoutTask tl
+        | .error e => return ⟨[], some e, true⟩
+      else
+        let tl := tl.mapCheap .inr
+        let cancelTks := cancelTks.map (·.mapCheap .inl)
+        let r ← ServerTask.waitAny (tl :: cancelTks ++ [timeoutTask])
+        match r with
+        | .inl _ => return ⟨[], none, false⟩ -- Timeout or cancellation - stop waiting
+        | .inr (.ok tl) => go timeoutTask tl
+        | .inr (.error e) => return ⟨[], some e, true⟩
 
 partial def getFinishedPrefixWithConsistentLatency (xs : AsyncList ε α) (latencyMs : UInt32)
     (cancelTks : List (ServerTask Unit) := []) : BaseIO (List α × Option ε × Bool) := do
diff --git a/src/Lean/Server/ServerTask.lean b/src/Lean/Server/ServerTask.lean
index 11c5f1538f..d755c14b39 100644
--- a/src/Lean/Server/ServerTask.lean
+++ b/src/Lean/Server/ServerTask.lean
@@ -54,6 +54,8 @@ def pure (x : α) : ServerTask α := Task.pure x
 
 def get (t : ServerTask α) : α := t.task.get
 
+def wait (t : ServerTask α) : BaseIO α := IO.wait t.task
+
 def mapCheap (f : α → β) (t : ServerTask α) : ServerTask β :=
   t.task.map f (sync := true)