From 205c69c0aacf6ad5b3b1da8457f02b62c82d5754 Mon Sep 17 00:00:00 2001 From: Richard McDaniel <1130888+rmcdaniel@users.noreply.github.com> Date: Thu, 5 Nov 2026 19:08:10 +0000 Subject: [PATCH] Observe worker state before and after shutdown without masking errors --- .../worker_startup_reconstruction/README.md | 15 ++-- .../integration/test_worker_task_families.py | 15 +--- tests/test_worker_observation.py | 87 +++++++++++++++++++ tests/worker_observation.py | 55 ++++++++++++ 4 files changed, 153 insertions(+), 19 deletions(-) diff --git a/tests/fixtures/worker_startup_reconstruction/README.md b/tests/fixtures/worker_startup_reconstruction/README.md index bce669be..ddd88ef2 100644 --- a/tests/fixtures/worker_startup_reconstruction/README.md +++ b/tests/fixtures/worker_startup_reconstruction/README.md @@ -76,12 +76,15 @@ incompatible harness is a setup finding, not permission to patch the SDK. while retaining that limit on per-client attribution. A missing cloned-client trace cannot count as full coverage. -The current driver emits a full before-shutdown observation on failure only. -The next observer-source revision must add the same bounded snapshot on success -before calling `stop_and_drain`, bind API/database snapshots to that boundary, -and preserve the existing primary exception if observation or shutdown fails. -Do not run the current sources and claim those additions were present. Return -the revised source hashes for review rather than carrying these source pins. +The driver now captures the same bounded observation before and after shutdown +on success and failure. Each capture has its own monotonic bracket; capture +failure is explicitly invalid. Reporting or shutdown failures cannot replace an +existing primary exception, and reporting failure still runs the original +stop/drain. Capture availability does not establish phase/SQL/API completeness. +No new probes, waits, timeout extensions or outer cleanup deadline are supplied +by this slice. Binding API/database snapshots to these boundaries and enforcing +the future ten-second overall cleanup bound remain pending. Return revised +source hashes for review rather than carrying the historical source pins. ## Server comparator review diff --git a/tests/integration/test_worker_task_families.py b/tests/integration/test_worker_task_families.py index 9e6a82c1..bdee744d 100644 --- a/tests/integration/test_worker_task_families.py +++ b/tests/integration/test_worker_task_families.py @@ -18,8 +18,8 @@ from tests.worker_observation import ( ObservationInterceptor, WorkerObservation, + observe_and_drain, result_or_worker_error, - stop_and_drain, ) current_observation: ContextVar[WorkerObservation] = ContextVar("task_family_observation") @@ -128,21 +128,10 @@ async def test_workers_execute_without_polling_empty_task_families( 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: 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 + await observe_and_drain(observation, identity, workers, runners, primary_error=primary_error) finally: current_observation.reset(context_token) diff --git a/tests/test_worker_observation.py b/tests/test_worker_observation.py index 22291571..9cbe5b1a 100644 --- a/tests/test_worker_observation.py +++ b/tests/test_worker_observation.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import itertools import json from typing import Any @@ -12,6 +13,7 @@ from tests.worker_observation import ( ObservationInterceptor, WorkerObservation, + observe_and_drain, result_or_worker_error, stop_and_drain, ) @@ -303,6 +305,91 @@ async def run(): assert drained == [True] +@pytest.mark.asyncio +@pytest.mark.parametrize("primary_kind", [None, RuntimeError, asyncio.CancelledError]) +@pytest.mark.parametrize("failure_phase", [None, "report", "emit", "shutdown"]) +async def test_boundary_observation_always_drains_and_preserves_primary_error( + monkeypatch: pytest.MonkeyPatch, primary_kind: type[BaseException] | None, failure_phase: str | None, +) -> None: + observation = WorkerObservation() + primary = None if primary_kind is None else primary_kind("SECRET primary") + secondary = ValueError("SECRET secondary") + stopped = asyncio.Event() + drained = [] + records: list[dict[str, Any]] = [] + original_report = observation.report + report_calls = 0 + + def report(): + nonlocal report_calls + report_calls += 1 + if failure_phase == "report" and report_calls == 1: + raise secondary + return original_report() + + def emit(record): + records.append(record) + if failure_phase == "emit" and len(records) == 1: + raise secondary + + class Worker: + async def stop(self): + observation.note("stop_requested") + stopped.set() + if failure_phase == "shutdown": + raise secondary + + async def run(): + await stopped.wait() + drained.append(True) + + monkeypatch.setattr(observation, "report", report) + runner = asyncio.create_task(run()) + + async def case(): + try: + if primary is not None: + raise primary + finally: + await observe_and_drain(observation, {"run_id": "original-run"}, [Worker()], [runner], + primary_error=primary, emit=emit) + + expected = primary if primary is not None else secondary if failure_phase is not None else None + if expected is None: + await case() + else: + with pytest.raises(type(expected)) as raised: + await case() + assert raised.value is expected + assert stopped.is_set() and runner.done() and drained == [True] + assert records[0]["phase"] == "before_shutdown" and records[-1]["phase"] == "after_shutdown" + assert "stop_requested" not in records[0].get("observation", {}).get("milestones", {}) + assert records[-1]["observation"]["milestones"]["stop_requested"]["count"] == 1 + assert all(record["run_id"] == "original-run" for record in records) + assert records[0]["observation_capture_state"] == ("invalid" if failure_phase == "report" else "complete") + assert "SECRET" not in json.dumps(records) + + +@pytest.mark.asyncio +async def test_shutdown_snapshots_bracket_cleanup_and_do_not_share_mutable_state() -> None: + ticks = itertools.count() + observation = WorkerObservation(clock=lambda: float(next(ticks))) + records: list[dict[str, Any]] = [] + + class Worker: + async def stop(self): + observation.note("stop_requested") + + await observe_and_drain(observation, {}, [Worker()], [], emit=records.append) + before, after = records + stop_at = after["observation"]["milestones"]["stop_requested"]["first"]["at"] + observation.started + before_report_at = before["observation"]["elapsed"] + observation.started + assert before["monotonic_before"] <= before_report_at <= before["monotonic_after"] + assert before["monotonic_after"] <= stop_at <= after["monotonic_before"] <= after["monotonic_after"] + assert "stop_requested" not in before["observation"]["milestones"] + assert before["observation_capture_state"] == after["observation_capture_state"] == "complete" + + @pytest.mark.asyncio async def test_failed_runner_does_not_skip_drain_of_its_peer() -> None: release = asyncio.Event() diff --git a/tests/worker_observation.py b/tests/worker_observation.py index 18992d28..a8907c42 100644 --- a/tests/worker_observation.py +++ b/tests/worker_observation.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import json import time from collections import Counter, deque from collections.abc import Callable, Coroutine @@ -229,6 +230,60 @@ async def stop_and_drain(workers: list[Any], runners: list[asyncio.Task[Any]]) - raise outcome +def print_observation(record: dict[str, Any]) -> None: + print(json.dumps(record), flush=True) + + +async def observe_and_drain( + observation: WorkerObservation, identity: dict[str, Any], + workers: list[Any], runners: list[asyncio.Task[Any]], *, + primary_error: BaseException | None = None, + emit: Callable[[dict[str, Any]], None] = print_observation, +) -> None: + """Bracket the existing shutdown with bounded reports, retaining the primary error. + + Capture availability is separate from report completeness and phase linkage. + Capture is synchronous; API/SQL probes and timeout extensions are absent. + """ + secondary_error: BaseException | None = None + + def capture(phase: str, error: BaseException | None) -> None: + nonlocal secondary_error + record: dict[str, Any] = {**identity, "phase": phase} + if error is not None: + record["exception"] = type(error).__name__ + try: + record["monotonic_before"] = observation.clock() + report = observation.report() + record["monotonic_after"] = observation.clock() + record["observation"] = report + record["observation_capture_state"] = "complete" + except BaseException as capture_error: + record["observation_capture_state"] = "invalid" + record["observation_exception"] = type(capture_error).__name__ + if secondary_error is None: + secondary_error = capture_error + try: + emit(record) + except BaseException as emit_error: + if secondary_error is None: + secondary_error = emit_error + + capture("before_shutdown", primary_error) + shutdown_error: BaseException | None = None + try: + await stop_and_drain(workers, runners) + except BaseException as error: + shutdown_error = error + if secondary_error is None: + secondary_error = error + capture("shutdown", error) + finally: + capture("after_shutdown", shutdown_error if shutdown_error is not None else primary_error) + if primary_error is None and secondary_error is not None: + raise secondary_error + + class ObservationInterceptor(PassthroughWorkerInterceptor): def __init__(self, observation: WorkerObservation) -> None: self.observation = observation