sched: fence scheduled resumes at shutdown so none outlive ownership

src/wf_api/runs.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from dataclasses import asdict
 from typing import Any, Protocol
 
@@ -198,14 +199,20 @@ class WorkflowRunApi:
         # Shared execution slot for schedule-owned resumes: the resume
         # gate holds the scheduler's own capacity accounting (no second
         # semaphore) behind the durable executing mark. Acquisition runs
-        # BEFORE the ACTIVE attempt mark, so a busy rejection leaves no
-        # fake ACTIVE attempt for work that never dispatched. Manual runs
-        # (no schedule admission, or no live scheduler) skip the gate.
+        # BEFORE the ACTIVE attempt mark, so a busy or draining rejection
+        # leaves no fake ACTIVE attempt for work that never dispatched.
+        # The caller's task is bound atomically with acquisition, so the
+        # shutdown drain can cancel and join it after the grace deadline.
+        # Manual runs (no schedule admission, or no live scheduler) skip
+        # the gate.
         gate = self.resume_slot_gate
-        slot = gate.acquire(run_id) if gate is not None else None
+        if gate is not None:
+            slot = gate.acquire(run_id, owner_task=asyncio.current_task())
+        else:
+            slot = None
         slot_held = slot is not None
         try:
-            return await self._resume_scheduled_or_manual(
+            result = await self._resume_scheduled_or_manual(
                 run_id=run_id,
                 resume_payload=resume_payload,
                 resume_outcome=resume_outcome,
@@ -213,9 +220,29 @@ class WorkflowRunApi:
                 trace_values=trace_values,
                 store=store,
             )
-        finally:
+        except asyncio.CancelledError:
+            # Shutdown drain cancelled the execution mid-flight (or the
+            # caller went away): fence, don't release. The durable
+            # executing mark and the ACTIVE attempt stay exactly as a
+            # crash mid-resume would leave them, so restart recovery
+            # fails the run closed instead of presenting a half-resumed
+            # run as safe to retry. Cancellation is never swallowed.
+            if slot_held and gate is not None:
+                gate.fence(run_id)
+            raise
+        except BaseException:
+            # Pre-persist errors, torn persists, and validation failures:
+            # release is truthful here because either nothing executed
+            # (the marks were only ever ours) or the durable ACTIVE
+            # attempt / stopped result already captures the ambiguity for
+            # recovery to reconcile by attempt identity.
             if slot_held and gate is not None:
                 gate.release(run_id)
+            raise
+        else:
+            if slot_held and gate is not None:
+                gate.release(run_id)
+            return result
 
     async def _resume_scheduled_or_manual(
         self,