From 4c2865f5f87bf128d86f5c6e717a7d7e47194d2b Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Sat, 17 Oct 2026 06:37:35 +0000 Subject: [PATCH 1/4] Capture bounded worker startup observations in integration tests --- .../integration/test_worker_task_families.py | 59 ++++- tests/test_worker_observation.py | 233 ++++++++++++++++++ tests/worker_observation.py | 203 +++++++++++++++ 3 files changed, 489 insertions(+), 6 deletions(-) create mode 100644 tests/test_worker_observation.py create mode 100644 tests/worker_observation.py diff --git a/tests/integration/test_worker_task_families.py b/tests/integration/test_worker_task_families.py index 93a826d..929621a 100644 --- a/tests/integration/test_worker_task_families.py +++ b/tests/integration/test_worker_task_families.py @@ -7,12 +7,21 @@ import json import uuid from collections import Counter +from contextvars import ContextVar from typing import Any import httpx import pytest from durable_workflow import Client, Worker, activity, workflow +from tests.worker_observation import ( + ObservationInterceptor, + WorkerObservation, + result_or_worker_error, + stop_and_drain, +) + +current_observation: ContextVar[WorkerObservation] = ContextVar("task_family_observation") @activity.defn(name="tests.task-family.echo") @@ -23,12 +32,16 @@ async def echo(value: str) -> str: @workflow.defn(name="tests.task-family.activity") class ActivityWorkflow: def run(self, ctx: Any, value: str) -> Any: + if observation := current_observation.get(None): + observation.note("workflow_handler_enter") return (yield ctx.schedule_activity("tests.task-family.echo", [value])) @workflow.defn(name="tests.task-family.timer") class TimerWorkflow: def run(self, ctx: Any, value: str) -> Any: + if observation := current_observation.get(None): + observation.note("workflow_handler_enter") yield ctx.sleep(2) return f"timer:{value}" @@ -49,8 +62,10 @@ async def record(request: httpx.Request) -> None: @pytest.mark.parametrize("mode", ["split", "mixed", "timer"]) @pytest.mark.parametrize("entry_point", ["run", "run_until"]) async def test_workers_execute_without_polling_empty_task_families( - server_url: str, server_token: str, mode: str, entry_point: str, + server_url: str, server_token: str, mode: str, entry_point: str, monkeypatch: pytest.MonkeyPatch, ) -> None: + observation = WorkerObservation() + observation.install_http(monkeypatch) suffix = uuid.uuid4().hex[:12] queue = f"task-family-{suffix}" async with ( @@ -58,6 +73,10 @@ async def test_workers_execute_without_polling_empty_task_families( Client(server_url, token=server_token, namespace="default") as workflow_client, Client(server_url, token=server_token, namespace="default") as activity_client, ): + observation.clients.update({ + id(control._http): "control", id(workflow_client._http): "workflow", + id(activity_client._http): "activity", + }) workflow_polls = record_polls(workflow_client) activity_polls = record_polls(activity_client) workflow_definition = TimerWorkflow if mode == "timer" else ActivityWorkflow @@ -66,32 +85,60 @@ async def test_workers_execute_without_polling_empty_task_families( workflows=[workflow_definition], activities=[echo] if mode == "mixed" else [], poll_timeout=0.5, shutdown_timeout=10, + interceptors=[ObservationInterceptor(observation)], ) workers = [worker] if mode == "split": workers.append(Worker( activity_client, task_queue=queue, workflows=[], activities=[echo], poll_timeout=0.5, shutdown_timeout=10, + interceptors=[ObservationInterceptor(observation)], )) handle = await control.start_workflow( workflow_type="tests.task-family.timer" if mode == "timer" else "tests.task-family.activity", task_queue=queue, workflow_id=f"task-family-{suffix}", input=[suffix], ) + context_token = current_observation.set(observation) runners = [asyncio.create_task( selected.run() if entry_point == "run" else selected.run_until( workflow_id=handle.workflow_id, timeout=30, poll_interval=0.05, ) ) for selected in workers] + identity = { + "sdk_version": importlib.metadata.version("durable-workflow"), + "mode": mode, "entry_point": entry_point, "workflow_id": handle.workflow_id, + "run_id": handle.run_id, "task_queue": queue, + "workers": [selected.worker_id for selected in workers], + } + primary_error: BaseException | None = None try: expected = f"timer:{suffix}" if mode == "timer" else f"echo:{suffix}" - assert await handle.result(timeout=30, poll_interval=0.1) == expected + assert await result_or_worker_error(handle.result(timeout=30, poll_interval=0.1), runners) == expected if entry_point == "run_until": terminal = await asyncio.wait_for(asyncio.gather(*runners), timeout=10) assert all((result.status or "").lower() == "completed" for result in terminal) + assert observation.milestones["workflow_handler_enter"]["count"] > 0 + assert ("activity_handler_enter" in observation.milestones) is (mode != "timer") + except BaseException as error: + primary_error = error + print(json.dumps({ + **identity, "phase": "before_shutdown", "exception": type(error).__name__, + "observation": observation.report(), + }), flush=True) + raise finally: - await asyncio.gather(*(selected.stop() for selected in workers)) - await asyncio.wait_for(asyncio.gather(*runners), timeout=10) + try: + await stop_and_drain(workers, runners) + except BaseException as error: + print(json.dumps({ + **identity, "phase": "shutdown", "exception": type(error).__name__, + "observation": observation.report(), + }), flush=True) + if primary_error is None: + raise + finally: + current_observation.reset(context_token) assert workflow_polls["workflow"] > 0 assert (workflow_polls["activity"] > 0) is (mode == "mixed") @@ -100,8 +147,8 @@ async def test_workers_execute_without_polling_empty_task_families( assert activity_polls["workflow"] == 0 assert activity_polls["query"] == 0 print(json.dumps({ - "sdk_version": importlib.metadata.version("durable-workflow"), - "mode": mode, "entry_point": entry_point, "workflow_id": handle.workflow_id, + **identity, "workflow_worker_polls": dict(workflow_polls), "activity_worker_polls": dict(activity_polls), "outcome": "pass", + "observation": observation.report(include_recent=False), }), flush=True) diff --git a/tests/test_worker_observation.py b/tests/test_worker_observation.py new file mode 100644 index 0000000..11a3248 --- /dev/null +++ b/tests/test_worker_observation.py @@ -0,0 +1,233 @@ +from __future__ import annotations + +import asyncio +import json +from typing import Any + +import httpx +import pytest + +from durable_workflow.interceptors import ActivityInterceptorContext, WorkflowTaskInterceptorContext +from tests.worker_observation import ( + ObservationInterceptor, + WorkerObservation, + result_or_worker_error, + stop_and_drain, +) + + +@pytest.mark.asyncio +async def test_http_observation_covers_additional_clients_without_exposing_payloads(monkeypatch: Any) -> None: + observation = WorkerObservation() + observation.install_http(monkeypatch) + requests = [] + + async def handle(request: httpx.Request) -> httpx.Response: + requests.append(await request.aread()) + return httpx.Response(503, json={ + "task": {"task_id": "task-original", "workflow_task_attempt": 0, "task_token": "SECRET"}, + "reason": "backend_lock_pressure", "error": {"code": "backend_lock_pressure", "message": "SECRET"}, + "recorded": False, "input": "SECRET", + }) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as first: + observation.clients[id(first)] = "workflow" + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as additional: + for client, family in ((first, "workflow"), (additional, "query")): + response = await client.post( + f"https://example.test/api/worker/{family}-tasks/poll?token=SECRET", + headers={"Authorization": "Bearer SECRET"}, content=b"SECRET", + ) + assert response.json()["input"] == "SECRET" + + report = observation.report() + assert report["counts"] == {"POST workflow_poll": 1, "POST query_poll": 1} + assert requests == [b"SECRET", b"SECRET"] + assert report["active_count"] == 0 + assert report["recent"][1]["client"] == "additional" + assert report["recent"][1]["workflow_task_attempt"] == 0 + assert report["recent"][1]["reason"] == "backend_lock_pressure" + assert report["recent"][1]["recorded"] is False + assert "SECRET" not in json.dumps(report) + + +@pytest.mark.asyncio +async def test_snapshot_keeps_inflight_registration_before_cancellation(monkeypatch: Any) -> None: + observation = WorkerObservation() + observation.install_http(monkeypatch) + started = asyncio.Event() + + async def held(request: httpx.Request) -> httpx.Response: + started.set() + await asyncio.Event().wait() + raise AssertionError("unreachable") + + async with httpx.AsyncClient(transport=httpx.MockTransport(held)) as client: + pending = asyncio.create_task(client.post("https://example.test/api/worker/registrations")) + await asyncio.wait_for(started.wait(), 5) + before = observation.report() + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + assert before["active_count"] == 1 + assert before["active"][0]["route"] == "registration" + assert "exception" not in before["active"][0] + after = observation.report() + assert after["active_count"] == 0 + assert after["recent"][0]["exception"] == "CancelledError" + + +@pytest.mark.asyncio +async def test_streamed_response_is_not_consumed_by_observer(monkeypatch: Any) -> None: + observation = WorkerObservation() + observation.install_http(monkeypatch) + reads = [] + + class Stream(httpx.AsyncByteStream): + async def __aiter__(self): + reads.append(True) + yield b'{"task":null}' + + async with httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, stream=Stream(), headers={"Content-Type": "application/json"}), + )) as client: + response = await client.send(client.build_request("GET", "https://example.test/stream"), stream=True) + assert reads == [] + assert "task_state" not in observation.report()["recent"][0] + assert await response.aread() == b'{"task":null}' + await response.aclose() + + +@pytest.mark.asyncio +async def test_observation_bounds_completed_records_and_milestones(monkeypatch: Any) -> None: + observation = WorkerObservation(limit=2) + observation.install_http(monkeypatch) + async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200))) as client: + for index in range(5): + await client.get(f"https://example.test/record/{index}") + observation.note(f"phase{index}", {"task_id": f"task{index}", "task_token": "SECRET"}) + report = observation.report() + assert report["requests"] == 5 + assert len(report["recent"]) == 2 + assert len(report["milestones"]) == 2 + assert "SECRET" not in json.dumps(report) + + +@pytest.mark.asyncio +async def test_active_record_limit_keeps_the_actual_inflight_count(monkeypatch: Any) -> None: + observation = WorkerObservation(limit=2) + observation.install_http(monkeypatch) + started = asyncio.Event() + entered = 0 + + async def held(request: httpx.Request) -> httpx.Response: + nonlocal entered + entered += 1 + if entered == 5: + started.set() + await asyncio.Event().wait() + raise AssertionError("unreachable") + + async with httpx.AsyncClient(transport=httpx.MockTransport(held)) as client: + runners = [asyncio.create_task(client.get("https://example.test/held")) for _ in range(5)] + try: + await asyncio.wait_for(started.wait(), 5) + report = observation.report() + assert report["active_count"] == 5 + assert len(report["active"]) == 2 + assert report["omitted_active"] == 3 + finally: + for runner in runners: + runner.cancel() + await asyncio.gather(*runners, return_exceptions=True) + assert observation.report()["active_count"] == 0 + + +@pytest.mark.asyncio +async def test_fatal_worker_is_reported_before_result_wait_and_cancels_wait() -> None: + cancelled = asyncio.Event() + started = asyncio.Event() + failure = RuntimeError("registration failed") + + async def result(): + try: + started.set() + await asyncio.Event().wait() + finally: + cancelled.set() + + async def worker(): + await started.wait() + raise failure + + runner = asyncio.create_task(worker()) + with pytest.raises(RuntimeError) as raised: + await asyncio.wait_for(result_or_worker_error(result(), [runner]), 5) + assert raised.value is failure + assert cancelled.is_set() + + +@pytest.mark.asyncio +async def test_normal_worker_return_does_not_replace_original_result() -> None: + worker_returned = asyncio.Event() + expected = object() + + async def result(): + await worker_returned.wait() + return expected + + async def worker(): + worker_returned.set() + + runner = asyncio.create_task(worker()) + assert await result_or_worker_error(result(), [runner]) is expected + + +@pytest.mark.asyncio +async def test_stop_failure_still_drains_original_runner() -> None: + stopped = asyncio.Event() + drained = [] + failure = RuntimeError("registration delete failed") + + class Worker: + async def stop(self): + stopped.set() + raise failure + + async def run(): + await stopped.wait() + drained.append(True) + + with pytest.raises(RuntimeError) as raised: + await stop_and_drain([Worker()], [asyncio.create_task(run())]) + assert raised.value is failure + assert drained == [True] + + +@pytest.mark.asyncio +async def test_interceptors_preserve_results_and_exceptions() -> None: + observation = WorkerObservation() + interceptor = ObservationInterceptor(observation) + task = {"task_id": "original", "workflow_task_attempt": 1, "task_token": "SECRET"} + context = WorkflowTaskInterceptorContext(worker_id="worker", task_queue="queue", task=task) + expected = [{"type": "complete_workflow"}] + + async def execute(context): + return expected + + assert await interceptor.execute_workflow_task(context, execute) is expected + activity = ActivityInterceptorContext( + worker_id="worker", task_queue="queue", task=task, activity_type="echo", args=("SECRET",), + ) + failure = RuntimeError("SECRET") + + async def fail(context): + raise failure + + with pytest.raises(RuntimeError) as raised: + await interceptor.execute_activity(activity, fail) + assert raised.value is failure + report = observation.report() + assert report["milestones"]["activity_handler_error"]["count"] == 1 + assert report["milestones"]["workflow_task_enter"]["first"]["task_id"] == "original" + assert "SECRET" not in json.dumps(report) diff --git a/tests/worker_observation.py b/tests/worker_observation.py new file mode 100644 index 0000000..1963d95 --- /dev/null +++ b/tests/worker_observation.py @@ -0,0 +1,203 @@ +"""Bounded, payload-free diagnostics for the ordinary task-family integration test.""" + +from __future__ import annotations + +import asyncio +import time +from collections import Counter, deque +from collections.abc import Callable, Coroutine +from copy import deepcopy +from typing import Any + +import httpx + +from durable_workflow.interceptors import ( + ActivityHandler, + ActivityInterceptorContext, + PassthroughWorkerInterceptor, + WorkflowTaskHandler, + WorkflowTaskInterceptorContext, +) + + +def route_name(request: httpx.Request) -> str: + path = request.url.path.rstrip("/") + for family in ("workflow", "activity", "query"): + if path.endswith(f"/worker/{family}-tasks/poll"): + return f"{family}_poll" + if f"/worker/{family}-tasks/" in path: + return f"{family}_submission" + if "/worker/registrations" in path: + return "registration" + if path.endswith("/cluster/info"): + return "discovery" + return "control" + + +def task_identity(task: Any) -> dict[str, str | int]: + if not isinstance(task, dict): + return {} + fields = ( + "task_id", "workflow_id", "run_id", "workflow_run_id", "activity_id", + "workflow_task_attempt", "activity_task_attempt", "attempt", "lease_owner", + "namespace", "task_queue", "workflow_type", + "lease_expires_at", "task_deadline_at", + ) + return { + key: value for key in fields + if (value := task.get(key)) is not None + and (type(value) is int or isinstance(value, str) and len(value) <= 160) + } + + +class WorkerObservation: + def __init__(self, *, clock: Callable[[], float] = time.monotonic, limit: int = 32) -> None: + if limit < 1: + raise ValueError("observation limit must be positive") + self.clock = clock + self.started = clock() + self.limit = limit + self.recent: deque[dict[str, Any]] = deque(maxlen=limit) + self.active: dict[int, dict[str, Any]] = {} + self.active_count = 0 + self.total = 0 + self.counts: Counter[str] = Counter() + self.milestones: dict[str, dict[str, Any]] = {} + self.clients: dict[int, str] = {} + + def note(self, phase: str, task: Any = None) -> None: + if phase not in self.milestones and len(self.milestones) >= self.limit: + return + entry = {"at": self.clock() - self.started, **task_identity(task)} + milestone = self.milestones.setdefault(phase, {"count": 0, "first": entry}) + milestone["count"] += 1 + milestone["last"] = entry + + def install_http(self, monkeypatch: Any) -> None: + original_send = httpx.AsyncClient.send + + async def observed_send( + client: httpx.AsyncClient, request: httpx.Request, *args: Any, **kwargs: Any, + ) -> httpx.Response: + self.total += 1 + request_id = self.total + started = self.clock() + route = route_name(request) + method = request.method if request.method in {"GET", "POST", "PUT", "PATCH", "DELETE"} else "OTHER" + self.counts[f"{method} {route}"] += 1 + entry: dict[str, Any] = { + "request": request_id, "client": self.clients.get(id(client), "additional"), + "method": method, "route": route, "started": started - self.started, + } + if len(self.active) == self.limit: + self.active.pop(next(iter(self.active))) + self.active[request_id] = entry + self.active_count += 1 + try: + response = await original_send(client, request, *args, **kwargs) + entry["status"] = response.status_code + if response.is_stream_consumed: + self._response_metadata(entry, response) + return response + except BaseException as error: + entry["exception"] = type(error).__name__ + raise + finally: + entry["elapsed"] = self.clock() - started + self.active_count -= 1 + self.active.pop(request_id, None) + self.recent.append(entry) + + monkeypatch.setattr(httpx.AsyncClient, "send", observed_send) + + @staticmethod + def _response_metadata(entry: dict[str, Any], response: httpx.Response) -> None: + if not response.headers.get("content-type", "").startswith("application/json"): + return + try: + if len(response.content) > 32_768: + return + body = response.json() + except (ValueError, httpx.ResponseNotRead): + return + if not isinstance(body, dict): + return + if "task" in body: + entry["task_state"] = ( + "present" if isinstance(body["task"], dict) else "null" if body["task"] is None else "invalid" + ) + entry.update(task_identity(body["task"])) + if type(body.get("recorded")) is bool: + entry["recorded"] = body["recorded"] + error = body.get("error") + if isinstance(error, dict): + code = error.get("code") + if isinstance(code, str) and len(code) <= 80 and code.replace("_", "").isalnum(): + entry["error_code"] = code + reason = body.get("reason") + if isinstance(reason, str) and len(reason) <= 80 and reason.replace("_", "").isalnum(): + entry["reason"] = reason + + def report(self, *, include_recent: bool = True) -> dict[str, Any]: + return deepcopy({ + "elapsed": self.clock() - self.started, "requests": self.total, + "counts": dict(self.counts), "milestones": self.milestones, + "active_count": self.active_count, "active": list(self.active.values()), + "omitted_active": self.active_count - len(self.active), + "recent": list(self.recent) if include_recent else [], + }) + + +async def result_or_worker_error(result: Coroutine[Any, Any, Any], runners: list[asyncio.Task[Any]]) -> Any: + """Surface a fatal worker error while the original bounded result wait is pending.""" + result_task = asyncio.create_task(result) + pending = {*runners, result_task} + try: + while True: + done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED) + for runner in done - {result_task}: + if runner.cancelled(): + raise asyncio.CancelledError + error = runner.exception() + if error is not None: + raise error + if result_task in done: + return await result_task + finally: + if not result_task.done(): + result_task.cancel() + await asyncio.gather(result_task, return_exceptions=True) + + +async def stop_and_drain(workers: list[Any], runners: list[asyncio.Task[Any]]) -> None: + try: + await asyncio.gather(*(worker.stop() for worker in workers)) + finally: + await asyncio.wait_for(asyncio.gather(*runners), timeout=10) + + +class ObservationInterceptor(PassthroughWorkerInterceptor): + def __init__(self, observation: WorkerObservation) -> None: + self.observation = observation + + async def execute_workflow_task( + self, context: WorkflowTaskInterceptorContext, next: WorkflowTaskHandler, + ) -> list[dict[str, Any]] | None: + self.observation.note("workflow_task_enter", context.task) + try: + result = await next(context) + self.observation.note("workflow_task_return", context.task) + return result + except BaseException: + self.observation.note("workflow_task_error", context.task) + raise + + async def execute_activity(self, context: ActivityInterceptorContext, next: ActivityHandler) -> Any: + self.observation.note("activity_handler_enter", context.task) + try: + result = await next(context) + self.observation.note("activity_handler_return", context.task) + return result + except BaseException: + self.observation.note("activity_handler_error", context.task) + raise From 8eaf4e87594017784e9926ab13a877c41f5b1d15 Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Sat, 17 Oct 2026 06:41:45 +0000 Subject: [PATCH 2/4] Drain every runner and retain bounded startup timing summaries --- tests/test_worker_observation.py | 47 ++++++++++++++++++++++++++++++++ tests/worker_observation.py | 21 ++++++++++---- 2 files changed, 62 insertions(+), 6 deletions(-) diff --git a/tests/test_worker_observation.py b/tests/test_worker_observation.py index 11a3248..c5e27da 100644 --- a/tests/test_worker_observation.py +++ b/tests/test_worker_observation.py @@ -77,6 +77,25 @@ async def held(request: httpx.Request) -> httpx.Response: assert after["recent"][0]["exception"] == "CancelledError" +@pytest.mark.asyncio +async def test_registration_timing_survives_recent_record_eviction(monkeypatch: Any) -> None: + now = [0.0] + observation = WorkerObservation(clock=lambda: now[0], limit=1) + observation.install_http(monkeypatch) + + async def reply(request: httpx.Request) -> httpx.Response: + now[0] += 3.5 if request.method == "POST" else 0.25 + return httpx.Response(200, json={"status": "completed", "outcome": "completed"}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(reply)) as client: + await client.post("https://example.test/api/worker/registrations") + await client.get("https://example.test/result") + report = observation.report(include_recent=False) + assert report["recent"] == [] + assert report["max_elapsed"]["POST registration"] == 3.5 + assert report["max_elapsed"]["GET control"] == 0.25 + + @pytest.mark.asyncio async def test_streamed_response_is_not_consumed_by_observer(monkeypatch: Any) -> None: observation = WorkerObservation() @@ -204,6 +223,34 @@ async def run(): assert drained == [True] +@pytest.mark.asyncio +async def test_failed_runner_does_not_skip_drain_of_its_peer() -> None: + release = asyncio.Event() + drained = [] + failure = RuntimeError("worker startup failed") + + async def failed(): + raise failure + + async def peer(): + await release.wait() + drained.append(True) + + runners = [asyncio.create_task(failed()), asyncio.create_task(peer())] + cleanup = asyncio.create_task(stop_and_drain([], runners)) + try: + for _ in range(10): + await asyncio.sleep(0) + release.set() + with pytest.raises(RuntimeError) as raised: + await cleanup + assert raised.value is failure + assert drained == [True] + finally: + release.set() + await asyncio.gather(*runners, cleanup, return_exceptions=True) + + @pytest.mark.asyncio async def test_interceptors_preserve_results_and_exceptions() -> None: observation = WorkerObservation() diff --git a/tests/worker_observation.py b/tests/worker_observation.py index 1963d95..24bb77e 100644 --- a/tests/worker_observation.py +++ b/tests/worker_observation.py @@ -62,6 +62,7 @@ def __init__(self, *, clock: Callable[[], float] = time.monotonic, limit: int = self.active_count = 0 self.total = 0 self.counts: Counter[str] = Counter() + self.max_elapsed: dict[str, float] = {} self.milestones: dict[str, dict[str, Any]] = {} self.clients: dict[int, str] = {} @@ -104,6 +105,8 @@ async def observed_send( raise finally: entry["elapsed"] = self.clock() - started + key = f"{method} {route}" + self.max_elapsed[key] = max(self.max_elapsed.get(key, 0), entry["elapsed"]) self.active_count -= 1 self.active.pop(request_id, None) self.recent.append(entry) @@ -134,14 +137,16 @@ def _response_metadata(entry: dict[str, Any], response: httpx.Response) -> None: code = error.get("code") if isinstance(code, str) and len(code) <= 80 and code.replace("_", "").isalnum(): entry["error_code"] = code - reason = body.get("reason") - if isinstance(reason, str) and len(reason) <= 80 and reason.replace("_", "").isalnum(): - entry["reason"] = reason + for field in ("reason", "status", "poll_status", "outcome"): + value = body.get(field) + if isinstance(value, str) and len(value) <= 80 and value.replace("_", "").isalnum(): + entry[field] = value def report(self, *, include_recent: bool = True) -> dict[str, Any]: return deepcopy({ "elapsed": self.clock() - self.started, "requests": self.total, "counts": dict(self.counts), "milestones": self.milestones, + "max_elapsed": self.max_elapsed, "active_count": self.active_count, "active": list(self.active.values()), "omitted_active": self.active_count - len(self.active), "recent": list(self.recent) if include_recent else [], @@ -170,10 +175,14 @@ async def result_or_worker_error(result: Coroutine[Any, Any, Any], runners: list async def stop_and_drain(workers: list[Any], runners: list[asyncio.Task[Any]]) -> None: + outcomes = await asyncio.gather(*(worker.stop() for worker in workers), return_exceptions=True) try: - await asyncio.gather(*(worker.stop() for worker in workers)) - finally: - await asyncio.wait_for(asyncio.gather(*runners), timeout=10) + outcomes += await asyncio.wait_for(asyncio.gather(*runners, return_exceptions=True), timeout=10) + except Exception as error: + outcomes.append(error) + for outcome in outcomes: + if isinstance(outcome, BaseException): + raise outcome class ObservationInterceptor(PassthroughWorkerInterceptor): From 1df8875c63a0215432de37ae4ae8eb8af7e847aa Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Sat, 17 Oct 2026 06:46:46 +0000 Subject: [PATCH 3/4] Observe the published worker registration route --- .../integration/test_worker_task_families.py | 13 +++++++----- tests/test_worker_observation.py | 20 +++++++++++++++++-- tests/worker_observation.py | 6 +++++- 3 files changed, 31 insertions(+), 8 deletions(-) diff --git a/tests/integration/test_worker_task_families.py b/tests/integration/test_worker_task_families.py index 929621a..1682d04 100644 --- a/tests/integration/test_worker_task_families.py +++ b/tests/integration/test_worker_task_families.py @@ -140,12 +140,15 @@ async def test_workers_execute_without_polling_empty_task_families( finally: current_observation.reset(context_token) - assert workflow_polls["workflow"] > 0 - assert (workflow_polls["activity"] > 0) is (mode == "mixed") + assert workflow_polls["workflow"] > 0, observation.report() + assert (workflow_polls["activity"] > 0) is (mode == "mixed"), observation.report() + assert observation.successful["POST registration"] >= len(workers), observation.report() + assert observation.successful["DELETE registration"] >= len(workers), observation.report() + assert observation.active_count == 0, observation.report() if mode == "split": - assert activity_polls["activity"] > 0 - assert activity_polls["workflow"] == 0 - assert activity_polls["query"] == 0 + assert activity_polls["activity"] > 0, observation.report() + assert activity_polls["workflow"] == 0, observation.report() + assert activity_polls["query"] == 0, observation.report() print(json.dumps({ **identity, "workflow_worker_polls": dict(workflow_polls), diff --git a/tests/test_worker_observation.py b/tests/test_worker_observation.py index c5e27da..b460acd 100644 --- a/tests/test_worker_observation.py +++ b/tests/test_worker_observation.py @@ -7,6 +7,7 @@ import httpx import pytest +from durable_workflow import Client from durable_workflow.interceptors import ActivityInterceptorContext, WorkflowTaskInterceptorContext from tests.worker_observation import ( ObservationInterceptor, @@ -51,6 +52,21 @@ async def handle(request: httpx.Request) -> httpx.Response: assert "SECRET" not in json.dumps(report) +@pytest.mark.asyncio +async def test_observer_follows_the_actual_client_registration_route(monkeypatch: Any) -> None: + observation = WorkerObservation() + observation.install_http(monkeypatch) + async with Client("https://example.test", token="SECRET", namespace="default") as client: + original = client._http + client._http = httpx.AsyncClient( + base_url=original.base_url, headers=original.headers, + transport=httpx.MockTransport(lambda request: httpx.Response(201, json={"worker_id": "worker"})), + ) + await original.aclose() + await client.register_worker(worker_id="worker", task_queue="queue", capability_manifest={}) + assert observation.report()["counts"].get("POST registration") == 1 + + @pytest.mark.asyncio async def test_snapshot_keeps_inflight_registration_before_cancellation(monkeypatch: Any) -> None: observation = WorkerObservation() @@ -63,7 +79,7 @@ async def held(request: httpx.Request) -> httpx.Response: raise AssertionError("unreachable") async with httpx.AsyncClient(transport=httpx.MockTransport(held)) as client: - pending = asyncio.create_task(client.post("https://example.test/api/worker/registrations")) + pending = asyncio.create_task(client.post("https://example.test/api/worker/register")) await asyncio.wait_for(started.wait(), 5) before = observation.report() pending.cancel() @@ -88,7 +104,7 @@ async def reply(request: httpx.Request) -> httpx.Response: return httpx.Response(200, json={"status": "completed", "outcome": "completed"}) async with httpx.AsyncClient(transport=httpx.MockTransport(reply)) as client: - await client.post("https://example.test/api/worker/registrations") + await client.post("https://example.test/api/worker/register") await client.get("https://example.test/result") report = observation.report(include_recent=False) assert report["recent"] == [] diff --git a/tests/worker_observation.py b/tests/worker_observation.py index 24bb77e..5e090f9 100644 --- a/tests/worker_observation.py +++ b/tests/worker_observation.py @@ -27,7 +27,7 @@ def route_name(request: httpx.Request) -> str: return f"{family}_poll" if f"/worker/{family}-tasks/" in path: return f"{family}_submission" - if "/worker/registrations" in path: + if path.endswith("/worker/register") or "/worker/registrations" in path: return "registration" if path.endswith("/cluster/info"): return "discovery" @@ -62,6 +62,7 @@ def __init__(self, *, clock: Callable[[], float] = time.monotonic, limit: int = self.active_count = 0 self.total = 0 self.counts: Counter[str] = Counter() + self.successful: Counter[str] = Counter() self.max_elapsed: dict[str, float] = {} self.milestones: dict[str, dict[str, Any]] = {} self.clients: dict[int, str] = {} @@ -97,6 +98,8 @@ async def observed_send( try: response = await original_send(client, request, *args, **kwargs) entry["status"] = response.status_code + if 200 <= response.status_code < 300: + self.successful[f"{method} {route}"] += 1 if response.is_stream_consumed: self._response_metadata(entry, response) return response @@ -146,6 +149,7 @@ def report(self, *, include_recent: bool = True) -> dict[str, Any]: return deepcopy({ "elapsed": self.clock() - self.started, "requests": self.total, "counts": dict(self.counts), "milestones": self.milestones, + "successful": dict(self.successful), "max_elapsed": self.max_elapsed, "active_count": self.active_count, "active": list(self.active.values()), "omitted_active": self.active_count - len(self.active), From e13e687e1b46c8eaef768483e3b7106ffc3557cb Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Sat, 17 Oct 2026 06:58:53 +0000 Subject: [PATCH 4/4] Separate exception-test setup from the failing operation --- tests/test_worker_observation.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_worker_observation.py b/tests/test_worker_observation.py index b460acd..34b2422 100644 --- a/tests/test_worker_observation.py +++ b/tests/test_worker_observation.py @@ -196,8 +196,9 @@ async def worker(): raise failure runner = asyncio.create_task(worker()) + pending_result = asyncio.wait_for(result_or_worker_error(result(), [runner]), 5) with pytest.raises(RuntimeError) as raised: - await asyncio.wait_for(result_or_worker_error(result(), [runner]), 5) + await pending_result assert raised.value is failure assert cancelled.is_set() @@ -233,8 +234,10 @@ async def run(): await stopped.wait() drained.append(True) + workers = [Worker()] + runners = [asyncio.create_task(run())] with pytest.raises(RuntimeError) as raised: - await stop_and_drain([Worker()], [asyncio.create_task(run())]) + await stop_and_drain(workers, runners) assert raised.value is failure assert drained == [True]