feat: bind WAC approval urls to a named wait_for_approval step (#10317)

* feat: bind WAC approval urls to a named wait_for_approval step Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix: reject duplicate WAC approval step keys instead of renaming them Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix: reject WAC approval links minted for a step that is not awaiting approval Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix: bind WAC approval links to the awaiting step and stop step key aliasing Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix: reject empty approval keys and scope minted-key writes to the workspace Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix: enforce WAC approval binding at consumption and reject colliding keys Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix: make WAC approval binding and collision checks atomic, harden TS step keys Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix: decrement WAC suspend atomically instead of from a pre-lock snapshot Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * chore: add sqlx cache entry for the atomic WAC suspend decrement Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix: omit empty approver param from python get_approval_urls Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * test: pin the suspend-snapshot decrement and the colliding-mint race Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * test: drop the suspend-snapshot interleave test, it cannot both be stable and discriminate Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix: reject step keys that cannot be minted as a URL path segment Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-08-18 08:01:26 +00:00 · 2026-07-25 11:41:48 +02:00
parent 15a6382c89
commit 9cef724ff2
29 changed files with 1676 additions and 228 deletions
@@ -3,7 +3,7 @@
 import asyncio
 import pytest

-from wmill.client import WorkflowCtx, _StepSuspend, TaskError, workflow, task, step, sleep, parallel, _run_workflow
+from wmill.client import WorkflowCtx, _StepSuspend, TaskError, workflow, task, step, sleep, parallel, wait_for_approval, _run_workflow


@task
@@ -1112,3 +1112,58 @@ class TestParallel:
        r = _run_workflow(wf, {}, {})
        assert r["type"] == "complete"
        assert r["result"] == []
+
+
+class TestApprovalKeys:
+    """`key` names the step that get_approval_urls() mints URLs against, so a
+    duplicate must fail rather than silently become `<key>_2` and leave the
+    caller holding a URL for the earlier step."""
+
+    def test_explicit_key_is_used_verbatim(self):
+        @workflow
+        async def wf():
+            return await wait_for_approval(key="manager")
+
+        assert _run_workflow(wf, {}, {})["key"] == "manager"
+
+    def test_duplicate_explicit_key_raises(self):
+        @workflow
+        async def wf():
+            await wait_for_approval(key="manager")
+            await wait_for_approval(key="manager")
+
+        with pytest.raises(RuntimeError, match="already used"):
+            _run_workflow(wf, {"completed_steps": {"manager": {"approved": True}}}, {})
+
+    def test_explicit_key_colliding_with_a_suffixed_step_key_raises(self):
+        """`step("dup")` twice yields `dup`/`dup_2`, so an approval explicitly named
+        `dup_2` would alias the second step's key."""
+
+        @workflow
+        async def wf():
+            await step("dup", lambda: 1)
+            await step("dup", lambda: 2)
+            await wait_for_approval(key="dup_2")
+
+        with pytest.raises(RuntimeError, match="already used"):
+            _run_workflow(wf, {"completed_steps": {"dup": 1, "dup_2": 2}}, {})
+
+    @pytest.mark.parametrize("bad", ["", "  ", ".", "..", "a/b"])
+    def test_unusable_key_raises(self, bad):
+        """The key travels as one path segment when its URLs are minted, so anything
+        `get_approval_urls` could not address must be refused here too."""
+
+        @workflow
+        async def wf():
+            await wait_for_approval(key=bad)
+
+        with pytest.raises(RuntimeError, match="non-empty step name"):
+            _run_workflow(wf, {}, {})
+
+    def test_unnamed_approvals_still_auto_number(self):
+        @workflow
+        async def wf():
+            await wait_for_approval()
+            await wait_for_approval()
+
+        assert _run_workflow(wf, {"completed_steps": {"approval": {}}}, {})["key"] == "approval_2"
@@ -1274,6 +1274,29 @@ class Windmill:
            params=params,
        ).json()

+    def get_approval_urls(self, step_key: str = "approval", approver: str = None) -> dict:
+        """Get the resume URLs bound to one ``wait_for_approval`` step of this workflow.
+
+        Args:
+            step_key: Checkpoint key of the approval step, as passed to
+                ``wait_for_approval(key=...)``
+            approver: Optional approver name
+
+        Returns:
+            Dictionary with approvalPage, resume, and cancel URLs
+        """
+        from urllib.parse import quote
+
+        _assert_usable_step_key(step_key, "get_approval_urls step_key")
+        job_id = os.environ.get("WM_JOB_ID") or "NO_ID"
+        # Omit rather than send `approver=`: an empty value is echoed into the
+        # returned URLs and recorded as the approver instead of "anonymous".
+        params = {"approver": approver} if approver is not None else {}
+        return self.get(
+            f"/w/{self.workspace}/jobs/wac_approval_urls/{job_id}/{quote(step_key, safe='')}",
+            params=params,
+        ).json()
+
    def request_interactive_slack_approval(
        self,
        slack_resource_path: str,
@@ -2027,6 +2050,33 @@ def get_resume_urls(approver: str = None, flow_level: bool = None) -> dict:
    return _client.get_resume_urls(approver, flow_level)


+@init_global_client
+def get_approval_urls(step_key: str = "approval", approver: str = None) -> dict:
+    """Get the resume/cancel/approval-page URLs bound to one ``wait_for_approval`` step.
+
+    Unlike :func:`get_resume_urls`, which signs a random nonce, these address the
+    very ``resume_job`` record the step's built-in approval buttons use, so they
+    are stable across replays and safe to embed in a custom notification.
+
+    Args:
+        step_key: Checkpoint key of the approval step, as passed to
+            ``wait_for_approval(key=...)``. Keys must be unique within a workflow;
+            reusing one raises rather than silently renaming it. The URL only
+            resumes while that step is awaiting approval; used at any other moment
+            it is rejected rather than banking a row a different approval would
+            consume. Send it ahead of time — approvers just cannot act before the
+            workflow reaches the step.
+            ``resume`` and ``cancel`` are step-bound; ``approvalPage`` is not — it
+            opens the job's approval page, which acts on whichever approval is
+            pending when it is used.
+        approver: Optional approver name
+
+    Returns:
+        Dictionary with approvalPage, resume, and cancel URLs
+    """
+    return _client.get_approval_urls(step_key, approver)
+
+
@init_global_client
 def request_interactive_slack_approval(
    slack_resource_path: str,
@@ -2606,6 +2656,15 @@ import asyncio as _asyncio
 import contextvars as _contextvars


+def _assert_usable_step_key(key: str, what: str) -> None:
+    """A step key travels as one path segment when its URLs are minted, so it must be
+    non-empty and free of ``/`` and dot segments — otherwise ``wait_for_approval``
+    would accept a key ``get_approval_urls`` can never address."""
+    k = key.strip()
+    if not k or k in (".", "..") or "/" in key or "\\" in key:
+        raise RuntimeError(f"{what} must be a non-empty step name without `/` or dot segments")
+
+
 class _StepSuspend(BaseException):
    """Raised to suspend workflow execution. Inherits from BaseException
    so it is not caught by bare `except Exception:` blocks."""
@@ -2645,6 +2704,8 @@ class WorkflowCtx:
        checkpoint = checkpoint or {}
        self._completed: dict = checkpoint.get("completed_steps", {})
        self._counters: dict[str, int] = {}
+        # Every key handed out by _alloc_key, so distinct names can't alias one key.
+        self._used_keys: set[str] = set()
        self._pending: list = []
        self._executing_key: str | None = checkpoint.get("_executing_key")
        # Reuse a single httpx.AsyncClient across all fast-path step() calls
@@ -2666,10 +2727,20 @@ class WorkflowCtx:
        self._inline_lock: "_asyncio.Lock | None" = None

    def _alloc_key(self, name: str = "step") -> str:
-        """Name-based key: ``double`` for first call, ``double_2``, ``double_3`` for subsequent."""
+        """Name-based key: ``double`` for first call, ``double_2``, ``double_3`` for subsequent.
+
+        Suffixing alone can alias — a second ``step("x")`` and a first ``step("x_2")``
+        both want ``x_2`` — so keep bumping past keys already handed out. Allocation
+        order is fixed by the workflow body, so replays reproduce the same keys.
+        """
        n = self._counters.get(name, 0) + 1
+        key = name if n == 1 else f"{name}_{n}"
+        while key in self._used_keys:
+            n += 1
+            key = f"{name}_{n}"
        self._counters[name] = n
-        return name if n == 1 else f"{name}_{n}"
+        self._used_keys.add(key)
+        return key

    def _next_step(self, name: str, script: str, func=None, dispatch_type: str = "inline", _task_options: Optional[dict] = None, **kwargs):
        """Return an awaitable that either resolves from cache or suspends."""
@@ -2724,9 +2795,25 @@ class WorkflowCtx:
        )

    async def _wait_for_approval(
-        self, timeout: int = 1800, form: dict | None = None, self_approval: bool = True
+        self,
+        timeout: int = 1800,
+        form: dict | None = None,
+        self_approval: bool = True,
+        key: str | None = None,
    ):
-        key = self._alloc_key("approval")
+        if key is not None:
+            _assert_usable_step_key(key, "wait_for_approval key")
+        requested_key, key = key, self._alloc_key(key or "approval")
+
+        # An explicit key is an identifier callers mint URLs against, so silently
+        # renaming a duplicate to ``<key>_2`` would hand them a URL for the *first*
+        # step — which then fails with "resume request already sent" and parks the
+        # workflow until timeout. Unnamed approvals keep auto-numbering.
+        if requested_key and key != requested_key:
+            raise RuntimeError(
+                f'WAC step key "{requested_key}" is already used in this workflow. '
+                "Give each wait_for_approval() its own key so get_approval_urls() can address it."
+            )

        if key in self._completed:
            return self._completed[key]
@@ -3073,11 +3160,13 @@ async def wait_for_approval(
    timeout: int = 1800,
    form: dict | None = None,
    self_approval: bool = True,
+    key: str | None = None,
 ) -> dict:
    """Suspend the workflow and wait for an external approval.

-    Use ``get_resume_urls()`` (wrapped in ``step()``) to obtain
-    resume/cancel/approval URLs before calling this function.
+    Pass ``key`` to name the step, then ``get_approval_urls(key)`` yields the URLs
+    that resume exactly this approval — route them through your own channel.
+    Without a key the steps are named ``approval``, ``approval_2``, ...

    Returns a dict with ``value`` (form data), ``approver``, and ``approved``.

@@ -3085,16 +3174,19 @@ async def wait_for_approval(
        timeout: Approval timeout in seconds (default 1800).
        form: Optional form schema for the approval page.
        self_approval: Whether the user who triggered the flow can approve it (default True).
+        key: Optional checkpoint key naming this approval step.

    Example::

-        urls = await step("urls", lambda: get_resume_urls())
-        await step("notify", lambda: send_email(urls["approvalPage"]))
-        result = await wait_for_approval(timeout=3600)
+        urls = await step("urls", lambda: get_approval_urls("manager"))
+        await step("notify", lambda: send_email(urls["resume"], urls["cancel"]))
+        result = await wait_for_approval(key="manager", timeout=3600)
    """
    ctx: WorkflowCtx | None = _workflow_ctx.get(None)
    if ctx is not None:
-        return await ctx._wait_for_approval(timeout=timeout, form=form, self_approval=self_approval)
+        return await ctx._wait_for_approval(
+            timeout=timeout, form=form, self_approval=self_approval, key=key
+        )
    raise RuntimeError("wait_for_approval can only be called inside a @workflow")