From d2ca37e63eb9ab75131071af81854059ff2f66c1 Mon Sep 17 00:00:00 2001 From: Richard McDaniel <1130888+rmcdaniel@users.noreply.github.com> Date: Thu, 5 Nov 2026 19:47:47 +0000 Subject: [PATCH] Link handler and HTTP observations to scoped workflow task identity --- .../worker_startup_reconstruction/README.md | 10 +++ tests/test_worker_observation.py | 87 +++++++++++++++++++ tests/worker_observation.py | 44 ++++++++-- 3 files changed, 134 insertions(+), 7 deletions(-) diff --git a/tests/fixtures/worker_startup_reconstruction/README.md b/tests/fixtures/worker_startup_reconstruction/README.md index 321c167..379ce70 100644 --- a/tests/fixtures/worker_startup_reconstruction/README.md +++ b/tests/fixtures/worker_startup_reconstruction/README.md @@ -66,6 +66,16 @@ incompatible harness is a setup finding, not permission to patch the SDK. in the historical Worker. `workflow_handler_enter` is later. The intervening interval includes history loading, decode and replay. It is not a unique hydration-cause measurement. +- Handler milestones and existing HTTP calls inside that interceptor now inherit + only its allowlisted task/run/attempt identity through an async task context. + Concurrent tasks retain distinct contexts, and scope exit/error/cancellation + restores the previous context. The request path's bounded task ID is recorded + separately; it cannot replace or silently reconcile a conflicting interceptor + identity. Neither source reads request bodies or private claim tokens. Requests + made after interceptor exit have no inherited attempt, even if they address the + same task; the recorder does not substitute the most recently observed task. + These links do not establish a durable claim, successful decode/replay or a + publication receipt. Exact private/SQL/API linkage is still required. - The activity interceptor runs after argument decode. An absent activity handler milestone does not by itself identify a decode failure. - Submission metadata such as `recorded=false` must remain false. A returned diff --git a/tests/test_worker_observation.py b/tests/test_worker_observation.py index affe8f8..2d58333 100644 --- a/tests/test_worker_observation.py +++ b/tests/test_worker_observation.py @@ -17,6 +17,7 @@ ObservationInterceptor, WorkerObservation, observe_and_drain, + request_task_identity, result_or_worker_error, stop_and_drain, ) @@ -552,6 +553,92 @@ async def fail(context): assert "SECRET" not in json.dumps(report) +@pytest.mark.asyncio +async def test_handler_and_http_keep_scoped_task_and_request_identities_separate(monkeypatch: Any) -> None: + observation = WorkerObservation() + observation.install_http(monkeypatch) + interceptor = ObservationInterceptor(observation) + task = {"task_id": "original", "run_id": "run-original", "workflow_task_attempt": 0, + "task_token": "SECRET", "input": "SECRET"} + context = WorkflowTaskInterceptorContext(worker_id="worker", task_queue="queue", task=task) + expected = [{"type": "complete_workflow"}] + async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, json={}))) as client: + async def execute(context): + observation.note("workflow_handler_enter") + await client.post("https://example.test/api/worker/workflow-tasks/different/history?token=SECRET") + return expected + + assert await interceptor.execute_workflow_task(context, execute) is expected + observation.note("outside_interceptor") + await client.post("https://example.test/api/worker/workflow-tasks/original/complete") + report = observation.report() + linked = report["milestones"]["workflow_handler_enter"]["first"] + assert {key: linked[key] for key in ("task_id", "run_id", "workflow_task_attempt")} == { + "task_id": "original", "run_id": "run-original", "workflow_task_attempt": 0, + } + history, completion = report["recent"] + assert history["request_task"] == {"family": "workflow", "task_id": "different"} + assert history["interceptor_task"] == {"task_id": "original", "run_id": "run-original", "workflow_task_attempt": 0} + assert completion["request_task"] == {"family": "workflow", "task_id": "original"} + assert "interceptor_task" not in completion, "a completed scope must not supply the last attempt to a later request" + assert "task_id" not in report["milestones"]["outside_interceptor"]["first"] + assert "SECRET" not in json.dumps(report) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_kind", [None, RuntimeError, asyncio.CancelledError]) +async def test_task_context_isolated_across_concurrent_tasks_and_reset_on_exit( + failure_kind: type[BaseException] | None, +) -> None: + observation = WorkerObservation() + interceptor = ObservationInterceptor(observation) + failure = None if failure_kind is None else failure_kind("SECRET") + caught = [] + returned = [] + + async def run(index: int): + task = {"task_id": f"original-{index}", "run_id": f"run-{index}", "workflow_task_attempt": index} + context = WorkflowTaskInterceptorContext(worker_id="worker", task_queue="queue", task=task) + + async def execute(context): + await asyncio.sleep(0) + observation.note("workflow_handler_enter") + if index == 0 and failure is not None: + raise failure + return [{"type": "complete_workflow"}] + + try: + returned.append(await interceptor.execute_workflow_task(context, execute)) + except BaseException as error: + caught.append(error) + observation.note("outside_interceptor") + + await asyncio.gather(run(0), run(1)) + assert caught == ([] if failure is None else [failure]) + assert len(returned) == (2 if failure is None else 1) + report = observation.report() + handler = report["milestones"]["workflow_handler_enter"] + assert {handler[where]["task_id"] for where in ("first", "last")} == {"original-0", "original-1"} + for where in ("first", "last"): + identity = handler[where] + assert identity["run_id"] == f"run-{identity['workflow_task_attempt']}" + assert "task_id" not in report["milestones"]["outside_interceptor"][where] + assert "SECRET" not in json.dumps(report) + + +@pytest.mark.parametrize(("path", "expected"), [ + ("/api/worker/activity-tasks/original/complete", {"family": "activity", "task_id": "original"}), + ("/api/worker/query-tasks/original/complete", {"family": "query", "task_id": "original"}), + ("/api/worker/workflow-tasks/poll", {}), + ("/api/worker/workflow-tasks//history", {}), + ("/api/worker/workflow-tasks/original/other/history", {}), + ("/api/worker/workflow-tasks/original%2Fother/history", {}), + ("/api/worker/workflow-tasks/" + "x" * 161 + "/history", {}), +]) +def test_request_task_identity_is_bounded_and_ignores_non_task_routes(path: str, expected: dict[str, str]) -> None: + assert request_task_identity(httpx.Request("POST", "https://example.test" + path + "?token=SECRET")) == expected + + @pytest.mark.parametrize(("method", "path", "expected"), [ ("POST", "/api/worker/register", "registration"), ("POST", "/api/worker/registrations/worker/deregister", "deregistration"), diff --git a/tests/worker_observation.py b/tests/worker_observation.py index 90983f4..2edff9f 100644 --- a/tests/worker_observation.py +++ b/tests/worker_observation.py @@ -4,11 +4,13 @@ import asyncio import json +import re import threading import time import weakref from collections import Counter, deque from collections.abc import Callable, Coroutine +from contextvars import ContextVar from copy import deepcopy from typing import Any @@ -62,6 +64,19 @@ def task_identity(task: Any) -> dict[str, str | int]: } +def request_task_identity(request: httpx.Request) -> dict[str, str]: + path = request.url.path.rstrip("/") + for family in ("workflow", "activity", "query"): + marker = f"/worker/{family}-tasks/" + if marker not in path: + continue + task_id, separator, operation = path.rsplit(marker, 1)[1].partition("/") + if (separator and re.fullmatch(r"[A-Za-z0-9_-]{1,160}", task_id) + and re.fullmatch(r"[A-Za-z0-9_-]{1,80}", operation)): + return {"family": family, "task_id": task_id} + return {} + + class WorkerObservation: def __init__(self, *, clock: Callable[[], float] = time.monotonic, limit: int = 32) -> None: if limit < 1: @@ -86,8 +101,13 @@ def __init__(self, *, clock: Callable[[], float] = time.monotonic, limit: int = self._client_instances: dict[int, tuple[weakref.ReferenceType[httpx.AsyncClient], str]] = {} self.client_instances: dict[str, dict[str, Any]] = {} self.unattributed_client_requests = 0 + self._task_context: ContextVar[dict[str, str | int] | None] = ContextVar( + "worker_observation_task", default=None, + ) def note(self, phase: str, task: Any = None) -> None: + if task is None: + task = self._task_context.get() with self._lock: if phase not in self.milestones and len(self.milestones) >= self.limit: self.omitted_milestone_events += 1 @@ -124,6 +144,8 @@ async def observed_send( route = route_name(request) method = request.method if request.method in {"GET", "POST", "PUT", "PATCH", "DELETE"} else "OTHER" key = f"{method} {route}" + request_task = request_task_identity(request) + interceptor_task = task_identity(self._task_context.get()) with self._lock: self.total += 1 request_id = self.total @@ -136,6 +158,10 @@ async def observed_send( "method": method, "route": route, "started": started - self.started, "response_metadata_state": "missing", } + if request_task: + entry["request_task"] = request_task + if interceptor_task: + entry["interceptor_task"] = interceptor_task if len(self.active) == self.limit: self.active.pop(next(iter(self.active))) self.evicted_active += 1 @@ -328,14 +354,18 @@ def __init__(self, observation: WorkerObservation) -> None: async def execute_workflow_task( self, context: WorkflowTaskInterceptorContext, next: WorkflowTaskHandler, ) -> list[dict[str, Any]] | None: - self.observation.note("workflow_task_enter", context.task) + token = self.observation._task_context.set(task_identity(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 + 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 + finally: + self.observation._task_context.reset(token) async def execute_activity(self, context: ActivityInterceptorContext, next: ActivityHandler) -> Any: self.observation.note("activity_handler_enter", context.task)