fix: propagate parent cancel token to asTask subtasks

`CancelToken` now supports an optional parent pointer. When `isSet` is
called, it checks both the token itself and its parent (recursively).
This ensures that server-level cancellation propagates to subtasks
spawned via `TacticM.asTask'` (as used by `TacticM.par` in `try?`).

Previously, `Core.wrapAsync` replaced the cancel token with a fresh one,
so subtasks were completely disconnected from the server's cancellation.
This meant that when a file was re-elaborated, `try?`'s parallel subtasks
would continue running until they finished or hit heartbeat limits.

Now `CoreM.asTask` creates a child token with the parent's token as its
parent, so `Core.checkInterrupted` in the subtask also detects when the
parent has been cancelled.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

src/Init/System/IO.lean

@@ -1681,21 +1681,32 @@ tasks.
 -/
 structure CancelToken where
   private ref : IO.Ref Bool
+  /-- Optional parent token. When set, `isSet` also checks the parent (recursively). -/
+  parent? : Option CancelToken := none
 deriving Nonempty
 
 namespace CancelToken
 
 /-- Creates a new cancellation token. -/
-def new : BaseIO CancelToken :=
-  CancelToken.mk <$> IO.mkRef false
+def new : BaseIO CancelToken := do
+  return { ref := (← IO.mkRef false) }
 
-/-- Activates a cancellation token. Idempotent. -/
+/-- Creates a new cancellation token with a parent. The child token is considered set
+    whenever the parent is set (in addition to being set directly). -/
+def newWithParent (parent : CancelToken) : BaseIO CancelToken := do
+  return { ref := (← IO.mkRef false), parent? := some parent }
+
+/-- Activates a cancellation token. Idempotent. Does not activate the parent. -/
 def set (tk : CancelToken) : BaseIO Unit :=
   tk.ref.set true
 
-/-- Checks whether the cancellation token has been activated. -/
-def isSet (tk : CancelToken) : BaseIO Bool :=
-  tk.ref.get
+/-- Checks whether the cancellation token has been activated, either directly or
+    via a parent token. -/
+partial def isSet (tk : CancelToken) : BaseIO Bool := do
+  if (← tk.ref.get) then return true
+  if let some parent := tk.parent? then
+    return (← parent.isSet)
+  return false
 
 -- separate definition as otherwise no unboxed version is generated
 @[export lean_io_cancel_token_is_set]

src/Init/Try.lean

@@ -52,6 +52,20 @@ namespace Lean.Parser.Tactic
 
 syntax (name := tryTrace) "try?" optConfig : tactic
 
+/--
+`try? => tac` runs `tac` through `try?`'s suggestion engine (`evalSuggest`) without the
+automatic tactic generation phase. This is useful for testing `evalSuggest` directly with
+explicit tactic scripts using `try?`'s internal combinators (`attempt_all`, `attempt_all_par`,
+`first_par`, etc.).
+
+Example:
+```
+example : 1 = 1 := by
+  try? => first | rfl | simp
+```
+-/
+syntax (name := tryTraceWith) "try?" optConfig " => " tacticSeq : tactic
+
 /-- Helper internal tactic for implementing the tactic `try?`. -/
 syntax (name := attemptAll) "attempt_all " withPosition((ppDedent(ppLine) colGe "| " tacticSeq)+) : tactic
 

src/Lean/Elab/Tactic/Try.lean

@@ -1042,4 +1042,17 @@ private def evalAndSuggestWithBy (tk : Syntax) (tac : TSyntax `tactic) (original
         evalAndSuggest tk stx originalMaxHeartbeats config
   | _ => throwUnsupportedSyntax
 
+@[builtin_tactic Lean.Parser.Tactic.tryTraceWith] def evalTryTraceWith : Tactic := fun stx => do
+  match stx with
+  | `(tactic| try?%$tk $config:optConfig => $tac:tacticSeq) => Tactic.focus do withMainContext do
+    let config ← elabTryConfig config
+    let originalMaxHeartbeats ← getMaxHeartbeats
+    let tac ← `(tactic| ($tac:tacticSeq))
+    withUnlimitedHeartbeats do
+      if config.wrapWithBy then
+        evalAndSuggestWithBy tk tac originalMaxHeartbeats config
+      else
+        evalAndSuggest tk tac originalMaxHeartbeats config
+  | _ => throwUnsupportedSyntax
+
 end Lean.Elab.Tactic.Try

src/Lean/Elab/Task.lean

@@ -36,17 +36,20 @@ returning
 * a cancellation hook and
 * a monadic value with the cached result (and subsequent state as it was after running).
 
-The task is run with a fresh `CancelToken` in its context, so it can detect cancellation
-via `Core.checkInterrupted`. The cancellation hook sets this token.
+The task is run with a child `CancelToken` in its context, so it can detect cancellation
+via `Core.checkInterrupted`. The cancellation hook sets this token. If a parent cancel token
+exists in the current context, the child token is linked to it via `CancelToken.newWithParent`,
+so that server-level cancellation propagates to subtasks.
 
 Uses `Core.wrapAsync` internally to properly handle name generators and heartbeats.
-
-Note: We only set the cancel token and don't call `IO.cancel task`. We're uncertain whether
-`IO.cancel` is also necessary - it may be required for tasks that use `IO.checkCanceled`
-instead of `Core.checkInterrupted`.
 -/
 def asTask (t : CoreM α) : CoreM (BaseIO Unit × Task (CoreM α)) := do
-  let cancelToken ← IO.CancelToken.new
+  -- Create a child cancel token that also checks the parent's token.
+  -- This ensures that server-level cancellation propagates to subtasks.
+  let parentToken? := (← read).cancelTk?
+  let cancelToken ← match parentToken? with
+    | some parent => IO.CancelToken.newWithParent parent
+    | none => IO.CancelToken.new
   -- Use wrapAsync to properly handle name generators and heartbeats,
   -- but modify it to return both the result and state
   let wrappedAct ← Core.wrapAsync (fun () => do let a ← t; let s ← get; return (a, s)) (some cancelToken)

src/Lean/Server/Test/Cancel.lean

@@ -9,8 +9,10 @@ module
 prelude
 public import Lean.Elab.Command
 public import Lean.Elab.Tactic.Basic
+public import Lean.Elab.Task
 public meta import Lean.Elab.Command
 public meta import Lean.Elab.Tactic.Basic
+public meta import Lean.Elab.Task
 
 public section
 
@@ -184,6 +186,66 @@ elab_rules : tactic
   dbg_trace "blocked!"
   log "blocked"
 
+/-! ## Helpers for testing `asTask` subtask cancellation -/
+
+meta initialize parOnceRef : IO.Ref (Option (Task Unit)) ← IO.mkRef none
+
+/--
+Like `wait_for_cancel_once`, but spawns a parallel subtask via `TacticM.asTask'` that waits for
+cancellation. This tests whether subtasks spawned via `asTask'` (as used by `TacticM.par` in `try?`)
+inherit their parent's cancellation token.
+
+The subtask gets a fresh cancel token from `asTask'` (via `Core.wrapAsync`), so the parent's
+cancellation token is NOT propagated. The test expects to see "leaked!" if the bug exists,
+or "subtask-cancelled!" if the fix is in place.
+-/
+scoped syntax "wait_for_cancel_once_par_subtask" : tactic
+@[incremental]
+elab_rules : tactic
+| `(tactic| wait_for_cancel_once_par_subtask) => do
+  let prom ← IO.Promise.new
+  if let some t := (← parOnceRef.modifyGet (fun old => (old, old.getD prom.result!))) then
+    IO.wait t
+    return
+
+  -- Send "blocked" diagnostic (like wait_for_cancel_once)
+  dbg_trace "blocked!"
+  log "blocked"
+  let ctx ← readThe Elab.Term.Context
+  let some tacSnap := ctx.tacSnap? | unreachable!
+  tacSnap.new.resolve {
+    diagnostics := (← Language.Snapshot.Diagnostics.ofMessageLog (← Core.getMessageLog))
+    stx := default
+    finished := default
+  }
+
+  -- Now spawn a subtask via asTask' (same mechanism used by TacticM.par in try?)
+  -- This subtask gets a FRESH cancel token, disconnected from the server's token
+  let subtask ← Elab.Tactic.TacticM.asTask' do
+    -- Poll for cancellation (checks the FRESH token, not the server's)
+    let mut cancelled := false
+    for _ in List.range 100 do -- 100 * 50ms = 5s max
+      IO.sleep 50
+      let ctx ← readThe Core.Context
+      if let some cancelTk := ctx.cancelTk? then
+        if (← cancelTk.isSet) then
+          cancelled := true
+          break
+    if cancelled then
+      IO.eprintln "subtask-cancelled!"
+    else
+      IO.eprintln "leaked!"
+    prom.resolve ()
+
+  -- Block the main thread waiting for the subtask
+  -- The server's cancellation sets the PARENT token, but the main thread is blocked in IO.wait
+  -- and can only check checkInterrupted after the subtask completes
+  discard <| IO.wait subtask
+
+  -- If we reach here (subtask didn't throw), resolve and check interrupted
+  prom.resolve ()
+  Core.checkInterrupted
+
 meta initialize cmdOnceRef : IO.Ref (Option (Task Unit)) ← IO.mkRef none
 
 /--

tests/elab/delabConst.lean

@@ -222,15 +222,15 @@ end PrvTest
 The name `IO.CancelToken.ref✝` is a private imported name.
 -/
 /--
-info: def IO.CancelToken.isSet : IO.CancelToken → BaseIO Bool :=
-fun tk => ST.Ref.get (IO.CancelToken.ref✝ tk)
+info: def IO.CancelToken.set : IO.CancelToken → BaseIO Unit :=
+fun tk => ST.Ref.set (IO.CancelToken.ref✝ tk) true
 -/
-#guard_msgs in #print IO.CancelToken.isSet
+#guard_msgs in #print IO.CancelToken.set
 /-!
 Even if `IO` is opened, it won't print as `CancelToken.ref✝`, but the full name.
 -/
 /--
-info: def IO.CancelToken.isSet : CancelToken → BaseIO Bool :=
-fun tk => ST.Ref.get (IO.CancelToken.ref✝ tk)
+info: def IO.CancelToken.set : CancelToken → BaseIO Unit :=
+fun tk => ST.Ref.set (IO.CancelToken.ref✝ tk) true
 -/
-#guard_msgs in open IO in #print IO.CancelToken.isSet
+#guard_msgs in open IO in #print IO.CancelToken.set

tests/elab/printStructure.lean

@@ -29,15 +29,16 @@ constructor:
 
 /-! Structure with private field, imported -/
 /--
-info: structure IO.CancelToken : Type
-number of parameters: 0
+info: structure IO.Promise (α : Type) : Type
+number of parameters: 1
 fields:
-  private IO.CancelToken.ref✝ : IO.Ref Bool
+  private IO.Promise.prom✝ : IO.PromisePointed✝.type
+  private IO.Promise.h✝ : Nonempty α
 constructor:
-  private IO.CancelToken.mk✝ (ref : IO.Ref Bool) : IO.CancelToken
+  private IO.Promise.mk✝ {α : Type} (prom : IO.PromisePointed✝.type) (h : Nonempty α) : IO.Promise α
 -/
 #guard_msgs in
-#print IO.CancelToken
+#print IO.Promise
 
 /-! Structure with private field, current module -/
 structure PrivField where

tests/elab/try_eval_suggest.lean

@@ -0,0 +1,51 @@
+/-!
+# Tests for `try? => tac` syntax
+
+These tests verify that the `try? => tac` syntax correctly runs a user-supplied tactic
+through `evalSuggest` and reports suggestions.
+-/
+
+/-- info: Try this:
+  [apply] rfl -/
+#guard_msgs (info) in
+example : 1 = 1 := by
+  try? => rfl
+
+/-- info: Try this:
+  [apply] rfl -/
+#guard_msgs (info) in
+example : 1 = 1 := by
+  try? => first | assumption | rfl
+
+/-- info: Try these:
+  [apply] rfl
+  [apply] simp_all -/
+#guard_msgs (info) in
+example : 1 = 1 := by
+  try? => attempt_all | rfl | simp_all
+
+/-- info: Try these:
+  [apply] rfl
+  [apply] simp_all -/
+#guard_msgs (info) in
+example : 1 = 1 := by
+  try? => attempt_all_par | rfl | simp_all
+
+-- first_par returns whichever finishes first; just test it doesn't error
+#guard_msgs (drop info) in
+example : 1 = 1 := by
+  try? => first_par | rfl | simp_all
+
+/-- info: Try these:
+  [apply] assumption
+  [apply] rfl -/
+#guard_msgs (info) in
+example (h : 1 = 1) : 1 = 1 := by
+  try? => attempt_all | assumption | rfl
+
+-- `max` config option should limit suggestions
+/-- info: Try this:
+  [apply] rfl -/
+#guard_msgs (info) in
+example : 1 = 1 := by
+  try? (max := 1) => attempt_all | rfl | simp_all

tests/server_interactive/cancellation_par_subtask.lean

@@ -0,0 +1,27 @@
+import Lean.Server.Test.Cancel
+open Lean.Server.Test.Cancel
+
+/-!
+Test that subtasks spawned via `TacticM.asTask'` (as used by `TacticM.par` in `try?`)
+are cancelled when the parent command is cancelled.
+
+The `wait_for_cancel_once_par_subtask` tactic:
+1. Sends "blocked" diagnostic
+2. Spawns a subtask via `asTask'` that polls its cancel token for 5 seconds
+3. Blocks the main thread waiting for the subtask
+
+When the edit invalidates this declaration, the server sets the command's cancel token.
+But the subtask has a FRESH cancel token (from `asTask'`/`Core.wrapAsync`), so:
+- BUG: the subtask runs for the full 5 seconds and prints "leaked!"
+- FIX: the subtask detects cancellation and terminates early
+-/
+
+example : True := by
+  trivial
+       --^ waitFor: blocked
+       --^ insert: "; skip"
+       --^ collectDiagnostics
+
+theorem t : True := by
+  wait_for_cancel_once_par_subtask
+  trivial

tests/server_interactive/cancellation_par_subtask.lean.out.expected

@@ -0,0 +1,25 @@
+blocked!
+subtask-cancelled!
+{"version": 2,
+ "uri": "file:///cancellation_par_subtask.lean",
+ "diagnostics":
+ [{"source": "Lean 4",
+   "severity": 3,
+   "range":
+   {"start": {"line": 18, "character": 0}, "end": {"line": 19, "character": 0}},
+   "message": "Goals accomplished!",
+   "leanTags": [2],
+   "isSilent": true,
+   "fullRange":
+   {"start": {"line": 18, "character": 0},
+    "end": {"line": 19, "character": 15}}},
+  {"source": "Lean 4",
+   "severity": 3,
+   "range":
+   {"start": {"line": 24, "character": 0}, "end": {"line": 25, "character": 0}},
+   "message": "Goals accomplished!",
+   "leanTags": [2],
+   "isSilent": true,
+   "fullRange":
+   {"start": {"line": 24, "character": 0},
+    "end": {"line": 26, "character": 9}}}]}