Add drain and resume mutation for build-id rollout state

Operators could inspect per-build-id rollout state via
Client.list_task_queue_build_ids, but the async and sync clients had no
way to cut traffic off a build before deleting its workers. Without that
mutation, rollouts had to stop each worker process by hand, which is
easy to mis-sequence when rolling back a bad build or cutting over from
unversioned to versioned workers.

Add two coroutine methods on the async Client plus their sync mirrors
that wrap the new server endpoints:

- drain_task_queue_build_id(task_queue, build_id)
  POSTs to /task-queues/<tq>/build-ids/drain with {"build_id": ...}.
- resume_task_queue_build_id(task_queue, build_id)
  POSTs to /task-queues/<tq>/build-ids/resume with {"build_id": ...}.

Passing None as build_id targets the unversioned cohort (pre-rollout
default). Both calls are idempotent: repeated drains do not shift the
recorded drained_at timestamp, and repeated resumes leave the cohort
active.

Both methods return a TaskQueueBuildIdRolloutState dataclass carrying
the server's mutation receipt (namespace, task_queue, build_id,
drain_intent, drained_at) so operators can verify the outcome
programmatically. The raw server payload is preserved on the dataclass
for callers that need a byte-identical mirror.

Two new polyglot parity fixtures cover the cross-repo contract with the
CLI under tests/fixtures/control-plane/:

- task-queue-build-id-drain-parity.json
- task-queue-build-id-resume-parity.json

The matching CLI-side fixtures and parity test will land separately so
the cross-repo drift gate stays green.

Author: durable-workflow-ops

src/durable_workflow/__init__.py

@@ -34,6 +34,7 @@
     TaskQueueAdmission,
     TaskQueueBuildIdCohort,
     TaskQueueBuildIdRollout,
+    TaskQueueBuildIdRolloutState,
     TaskQueueDescription,
     TaskQueueList,
     TaskQueueQueryAdmission,
@@ -181,6 +182,7 @@
     "TaskQueueAdmission",
     "TaskQueueBuildIdCohort",
     "TaskQueueBuildIdRollout",
+    "TaskQueueBuildIdRolloutState",
     "TaskQueueDescription",
     "TaskQueueList",
     "TaskQueueQueryAdmission",

src/durable_workflow/client.py

@@ -534,6 +534,36 @@ def from_dict(cls, data: dict[str, Any]) -> TaskQueueBuildIdRollout:
         )
 
 
+@dataclass
+class TaskQueueBuildIdRolloutState:
+    """Operator-recorded drain intent for one ``(task_queue, build_id)`` cohort.
+
+    Returned by ``drain_task_queue_build_id`` and ``resume_task_queue_build_id``.
+    ``build_id`` is ``None`` for the unversioned cohort (workers registered
+    without a build identifier). ``drain_intent`` is ``"active"`` or
+    ``"draining"``. ``drained_at`` is set only when ``drain_intent`` is
+    ``"draining"``; repeated drains do not shift the timestamp.
+    """
+
+    namespace: str | None
+    task_queue: str
+    build_id: str | None
+    drain_intent: str
+    drained_at: str | None
+    raw: dict[str, Any] | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> TaskQueueBuildIdRolloutState:
+        return cls(
+            namespace=data.get("namespace"),
+            task_queue=str(data.get("task_queue") or ""),
+            build_id=data.get("build_id") if isinstance(data.get("build_id"), str) else None,
+            drain_intent=str(data.get("drain_intent") or ""),
+            drained_at=data.get("drained_at") if isinstance(data.get("drained_at"), str) else None,
+            raw=data,
+        )
+
+
 @dataclass
 class WorkerDescription:
     """Current server view of one registered worker."""
@@ -1392,6 +1422,67 @@ async def list_task_queue_build_ids(self, task_queue: str) -> TaskQueueBuildIdRo
             )
         return TaskQueueBuildIdRollout.from_dict(data)
 
+    async def drain_task_queue_build_id(
+        self,
+        task_queue: str,
+        build_id: str | None,
+    ) -> TaskQueueBuildIdRolloutState:
+        """Mark a build-id cohort as draining so it stops claiming new tasks.
+
+        Workers registered under ``build_id`` keep running their in-flight
+        work but are blocked from claiming fresh tasks, and future workers
+        that heartbeat under the same ``build_id`` land as draining too.
+        Pass ``None`` to drain the unversioned cohort (workers registered
+        without a build identifier). Idempotent: repeated drains do not
+        shift the recorded ``drained_at`` timestamp.
+        """
+        return await self._mutate_task_queue_build_id_rollout(
+            task_queue,
+            build_id,
+            action="drain",
+        )
+
+    async def resume_task_queue_build_id(
+        self,
+        task_queue: str,
+        build_id: str | None,
+    ) -> TaskQueueBuildIdRolloutState:
+        """Revert a previous drain so a build-id cohort can claim work again.
+
+        Resuming clears both ``drain_intent`` and ``drained_at`` for the
+        cohort and flips any still-running workers back to ``active``.
+        Pass ``None`` to resume the unversioned cohort. Idempotent:
+        resuming an already-active cohort is a no-op.
+        """
+        return await self._mutate_task_queue_build_id_rollout(
+            task_queue,
+            build_id,
+            action="resume",
+        )
+
+    async def _mutate_task_queue_build_id_rollout(
+        self,
+        task_queue: str,
+        build_id: str | None,
+        *,
+        action: str,
+    ) -> TaskQueueBuildIdRolloutState:
+        data = await self._request(
+            "POST",
+            f"/task-queues/{quote(task_queue, safe='')}/build-ids/{action}",
+            json={"build_id": build_id},
+            context=task_queue,
+        )
+        if not isinstance(data, dict):
+            raise ServerError(
+                200,
+                {
+                    "reason": f"invalid_task_queue_build_id_{action}_response",
+                    "message": f"expected JSON object, got {type(data).__name__}",
+                },
+            )
+        return TaskQueueBuildIdRolloutState.from_dict(data)
+
     # ── Search attributes ─────────────────────────────────────────────
     async def list_search_attributes(self) -> SearchAttributeList:
         """List system and custom search attribute definitions for this namespace."""

src/durable_workflow/sync.py

@@ -18,6 +18,7 @@
     ScheduleTriggerResult,
     StorageTestResult,
     TaskQueueBuildIdRollout,
+    TaskQueueBuildIdRolloutState,
     TaskQueueDescription,
     TaskQueueList,
     WorkflowCommandResult,
@@ -291,6 +292,26 @@ def list_task_queue_build_ids(self, task_queue: str) -> TaskQueueBuildIdRollout:
         )
         return result
 
+    def drain_task_queue_build_id(
+        self,
+        task_queue: str,
+        build_id: str | None,
+    ) -> TaskQueueBuildIdRolloutState:
+        result: TaskQueueBuildIdRolloutState = _run(
+            self._async.drain_task_queue_build_id(task_queue, build_id)
+        )
+        return result
+
+    def resume_task_queue_build_id(
+        self,
+        task_queue: str,
+        build_id: str | None,
+    ) -> TaskQueueBuildIdRolloutState:
+        result: TaskQueueBuildIdRolloutState = _run(
+            self._async.resume_task_queue_build_id(task_queue, build_id)
+        )
+        return result
+
     def list_namespaces(self) -> NamespaceList:
         result: NamespaceList = _run(self._async.list_namespaces())
         return result

tests/fixtures/control-plane/task-queue-build-id-drain-parity.json

@@ -0,0 +1,40 @@
+{
+  "schema": "durable-workflow.polyglot.control-plane-request-fixture",
+  "version": 1,
+  "operation": "task_queue.build_id.drain",
+  "request": {
+    "method": "POST",
+    "path": "/task-queues/orders-critical/build-ids/drain",
+    "body": {
+      "build_id": "build-2026.04.21-z9"
+    }
+  },
+  "semantic_body": {
+    "namespace": "orders-prod",
+    "task_queue": "orders-critical",
+    "build_id": "build-2026.04.21-z9",
+    "drain_intent": "draining",
+    "drained_at": "2026-04-22T09:45:00Z"
+  },
+  "response_body": {
+    "namespace": "orders-prod",
+    "task_queue": "orders-critical",
+    "build_id": "build-2026.04.21-z9",
+    "drain_intent": "draining",
+    "drained_at": "2026-04-22T09:45:00Z"
+  },
+  "cli": {
+    "argv": {
+      "task-queue": "orders-critical",
+      "--build-id": "build-2026.04.21-z9",
+      "--json": true
+    }
+  },
+  "sdk_python": {
+    "method": "drain_task_queue_build_id",
+    "args": {
+      "task_queue": "orders-critical",
+      "build_id": "build-2026.04.21-z9"
+    }
+  }
+}

tests/fixtures/control-plane/task-queue-build-id-resume-parity.json

@@ -0,0 +1,40 @@
+{
+  "schema": "durable-workflow.polyglot.control-plane-request-fixture",
+  "version": 1,
+  "operation": "task_queue.build_id.resume",
+  "request": {
+    "method": "POST",
+    "path": "/task-queues/orders-critical/build-ids/resume",
+    "body": {
+      "build_id": "build-2026.04.21-z9"
+    }
+  },
+  "semantic_body": {
+    "namespace": "orders-prod",
+    "task_queue": "orders-critical",
+    "build_id": "build-2026.04.21-z9",
+    "drain_intent": "active",
+    "drained_at": null
+  },
+  "response_body": {
+    "namespace": "orders-prod",
+    "task_queue": "orders-critical",
+    "build_id": "build-2026.04.21-z9",
+    "drain_intent": "active",
+    "drained_at": null
+  },
+  "cli": {
+    "argv": {
+      "task-queue": "orders-critical",
+      "--build-id": "build-2026.04.21-z9",
+      "--json": true
+    }
+  },
+  "sdk_python": {
+    "method": "resume_task_queue_build_id",
+    "args": {
+      "task_queue": "orders-critical",
+      "build_id": "build-2026.04.21-z9"
+    }
+  }
+}

tests/test_client.py

@@ -1111,6 +1111,116 @@ async def test_list_task_queue_build_ids_surfaces_cohort_worker_counts(
         assert unversioned.last_heartbeat_at is None
         assert unversioned.first_seen_at is None
 
+    @pytest.mark.asyncio
+    async def test_drain_task_queue_build_id_matches_polyglot_fixture(self, client: Client) -> None:
+        fixture_path = (
+            Path(__file__).parent
+            / "fixtures"
+            / "control-plane"
+            / "task-queue-build-id-drain-parity.json"
+        )
+        fixture = json.loads(fixture_path.read_text())
+        assert fixture["operation"] == "task_queue.build_id.drain"
+        sdk = fixture["sdk_python"]
+        resp = _mock_response(200, fixture["response_body"])
+
+        with patch.object(
+            client._http, "request", new_callable=AsyncMock, return_value=resp
+        ) as mock:
+            result = await client.drain_task_queue_build_id(**sdk["args"])
+
+        assert mock.call_args.args[0] == fixture["request"]["method"]
+        assert mock.call_args.args[1] == f"/api{fixture['request']['path']}"
+        body = mock.call_args.kwargs.get("json")
+        assert body == fixture["request"]["body"]
+
+        semantic = fixture["semantic_body"]
+        assert result.namespace == semantic["namespace"]
+        assert result.task_queue == semantic["task_queue"]
+        assert result.build_id == semantic["build_id"]
+        assert result.drain_intent == semantic["drain_intent"]
+        assert result.drained_at == semantic["drained_at"]
+
+    @pytest.mark.asyncio
+    async def test_resume_task_queue_build_id_matches_polyglot_fixture(self, client: Client) -> None:
+        fixture_path = (
+            Path(__file__).parent
+            / "fixtures"
+            / "control-plane"
+            / "task-queue-build-id-resume-parity.json"
+        )
+        fixture = json.loads(fixture_path.read_text())
+        assert fixture["operation"] == "task_queue.build_id.resume"
+        sdk = fixture["sdk_python"]
+        resp = _mock_response(200, fixture["response_body"])
+
+        with patch.object(
+            client._http, "request", new_callable=AsyncMock, return_value=resp
+        ) as mock:
+            result = await client.resume_task_queue_build_id(**sdk["args"])
+
+        assert mock.call_args.args[0] == fixture["request"]["method"]
+        assert mock.call_args.args[1] == f"/api{fixture['request']['path']}"
+        body = mock.call_args.kwargs.get("json")
+        assert body == fixture["request"]["body"]
+
+        semantic = fixture["semantic_body"]
+        assert result.namespace == semantic["namespace"]
+        assert result.task_queue == semantic["task_queue"]
+        assert result.build_id == semantic["build_id"]
+        assert result.drain_intent == semantic["drain_intent"]
+        assert result.drained_at is None
+
+    @pytest.mark.asyncio
+    async def test_drain_task_queue_build_id_targets_unversioned_cohort_with_null_body(
+        self, client: Client
+    ) -> None:
+        resp = _mock_response(
+            200,
+            {
+                "namespace": "default",
+                "task_queue": "orders",
+                "build_id": None,
+                "drain_intent": "draining",
+                "drained_at": "2026-04-22T09:50:00Z",
+            },
+        )
+
+        with patch.object(
+            client._http, "request", new_callable=AsyncMock, return_value=resp
+        ) as mock:
+            result = await client.drain_task_queue_build_id("orders", None)
+
+        assert mock.call_args.args[0] == "POST"
+        assert mock.call_args.args[1] == "/api/task-queues/orders/build-ids/drain"
+        assert mock.call_args.kwargs.get("json") == {"build_id": None}
+        assert result.build_id is None
+        assert result.drain_intent == "draining"
+        assert result.drained_at == "2026-04-22T09:50:00Z"
+
+    @pytest.mark.asyncio
+    async def test_resume_task_queue_build_id_clears_drained_at(self, client: Client) -> None:
+        resp = _mock_response(
+            200,
+            {
+                "namespace": "default",
+                "task_queue": "orders",
+                "build_id": "build-alpha",
+                "drain_intent": "active",
+                "drained_at": None,
+            },
+        )
+
+        with patch.object(
+            client._http, "request", new_callable=AsyncMock, return_value=resp
+        ) as mock:
+            result = await client.resume_task_queue_build_id("orders", "build-alpha")
+
+        assert mock.call_args.args[1] == "/api/task-queues/orders/build-ids/resume"
+        assert mock.call_args.kwargs.get("json") == {"build_id": "build-alpha"}
+        assert result.drain_intent == "active"
+        assert result.drained_at is None
+
     @pytest.mark.asyncio
     async def test_list_task_queues_parses_admission(self, client: Client) -> None:
         resp = _mock_response(200, {