durable-workflow 2.4.3__tar.gz → 2.5.0__tar.gz

{durable_workflow-2.4.3/src/durable_workflow.egg-info → durable_workflow-2.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durable-workflow
-Version: 2.4.3
+Version: 2.5.0
 Summary: Python client and worker SDK for Durable Workflow Cloud and self-hosted Server
 Author: Durable Workflow Contributors
 License-Expression: MIT

{durable_workflow-2.4.3 → durable_workflow-2.5.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "durable-workflow"
-version = "2.4.3"
+version = "2.5.0"
 description = "Python client and worker SDK for Durable Workflow Cloud and self-hosted Server"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -71,8 +71,8 @@ durable-workflow-replay-conformance = "durable_workflow.replay_conformance:main"
 durable-workflow-workflow-updates-conformance = "durable_workflow.workflow_updates_conformance:main"
 
 [tool.durable-workflow]
-product-train = "2.4.3"
-registry-version = "2.4.3"
+product-train = "2.5.0"
+registry-version = "2.5.0"
 supported-server-versions = "2.5.0"
 worker-protocol-version = "1.19"
 control-plane-version = "2"

durable_workflow-2.5.0/src/durable_workflow/_sticky_workflow_cache.py

@@ -0,0 +1,103 @@
+"""Bounded snapshots of durable wire history, never live workflow instances."""
+
+from __future__ import annotations
+
+import json
+import time
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Any
+
+CacheKey = tuple[str, str, str]
+
+
+def complete_history(history: list[dict[str, Any]]) -> bool:
+    if not history or any(not isinstance(event, dict) for event in history):
+        return False
+    types = [event.get("event_type") for event in history[:2]]
+    if types[0] != "WorkflowStarted" and types != ["StartAccepted", "WorkflowStarted"]:
+        return False
+    return all(type(event.get("sequence")) is int and event["sequence"] == index
+               for index, event in enumerate(history, 1))
+
+
+@dataclass(frozen=True)
+class _Entry:
+    encoded: bytes
+    expires_at: float
+    resume_token: str | None
+    resume_offset: int
+
+
+class StickyWorkflowCache:
+    def __init__(self, capacity: int, max_bytes: int, ttl_seconds: int) -> None:
+        for name, value, minimum in (("capacity", capacity, 0), ("max_bytes", max_bytes, 1),
+                                     ("ttl_seconds", ttl_seconds, 1)):
+            if type(value) is not int or value < minimum:
+                raise ValueError(f"sticky_cache_{name} must be an integer of at least {minimum}")
+        if ttl_seconds > 3600:
+            raise ValueError("sticky_cache_ttl_seconds must be at most 3600")
+        self.capacity = capacity
+        self.max_bytes = max_bytes
+        self.ttl_seconds = ttl_seconds
+        self._entries: OrderedDict[CacheKey, _Entry] = OrderedDict()
+        self._bytes = 0
+        self._metrics = {"hit": 0, "miss": 0, "eviction": 0, "forced_cold_replay": 0}
+
+    @property
+    def enabled(self) -> bool:
+        return self.capacity > 0
+
+    def _expire(self) -> None:
+        now = time.monotonic()
+        for key, entry in list(self._entries.items()):
+            if entry.expires_at <= now:
+                self.discard(key)
+
+    def discard(self, key: CacheKey) -> None:
+        entry = self._entries.pop(key, None)
+        if entry is not None:
+            self._bytes -= len(entry.encoded)
+
+    def lookup(self, key: CacheKey) -> tuple[list[dict[str, Any]], str | None, int] | None:
+        self._expire()
+        entry = self._entries.get(key)
+        if entry is None:
+            return None
+        self._entries.move_to_end(key)
+        # Decoding isolates the retained snapshot from replay and handler mutation.
+        history: list[dict[str, Any]] = json.loads(entry.encoded)
+        return history, entry.resume_token, entry.resume_offset
+
+    def remember(self, key: CacheKey, history: list[dict[str, Any]], *,
+                 resume_token: str | None = None, resume_offset: int = 0) -> bool:
+        self._expire()
+        self.discard(key)
+        if not self.enabled or not complete_history(history):
+            return False
+        try:
+            encoded = json.dumps(history, ensure_ascii=False, separators=(",", ":"), allow_nan=False).encode("utf-8")
+        except (TypeError, ValueError):
+            return False
+        if len(encoded) > self.max_bytes:
+            return False
+        while self._entries and (len(self._entries) >= self.capacity or self._bytes + len(encoded) > self.max_bytes):
+            oldest = next(iter(self._entries))
+            self.discard(oldest)
+            self._metrics["eviction"] += 1
+        self._entries[key] = _Entry(encoded, time.monotonic() + self.ttl_seconds, resume_token, resume_offset)
+        self._bytes += len(encoded)
+        return True
+
+    def record_replay(self, *, hit: bool, forced: bool = False) -> None:
+        self._metrics["hit" if hit else "miss"] += 1
+        if forced:
+            self._metrics["forced_cold_replay"] += 1
+
+    def metrics(self) -> dict[str, int]:
+        self._expire()
+        return {**self._metrics, "entries": len(self._entries), "history_bytes": self._bytes}
+
+    def clear(self) -> None:
+        self._entries.clear()
+        self._bytes = 0

{durable_workflow-2.4.3 → durable_workflow-2.5.0}/src/durable_workflow/client.py

@@ -87,7 +87,7 @@ PORTABLE_WORKER_AFFINITY_CAPABILITY_MANIFEST: dict[str, dict[str, str | bool]] =
     "sticky_execution": {
         "supported": False,
         "minimum_protocol_version": "1.18",
-        "reason": "python_worker_uses_complete_durable_history_replay",
+        "reason": "sticky_cache_disabled",
     },
 }
 _MESSAGE_STREAMS_CAPABILITY = "message_streams"
@@ -5112,6 +5112,7 @@ class Client:
         commands: list[dict[str, Any]],
         message_stream_cursors: list[dict[str, Any]] | None = None,
         message_stream_waits: list[dict[str, Any]] | None = None,
+        sticky_cache: dict[str, Any] | None = None,
     ) -> Any:
         """Report successful execution of a workflow task with its emitted commands.
 
@@ -5143,6 +5144,8 @@ class Client:
             body["message_stream_cursors"] = message_stream_cursors
         if message_stream_waits:
             body["message_stream_waits"] = message_stream_waits
+        if sticky_cache is not None:
+            body["sticky_cache"] = sticky_cache
         return await self._request(
             "POST", f"/worker/workflow-tasks/{task_id}/complete", worker=True, json=body
         )

{durable_workflow-2.4.3 → durable_workflow-2.5.0}/src/durable_workflow/worker.py

@@ -41,6 +41,7 @@ from ._activity_process import CallbackFailure, CallbackInvocation, CallbackProc
 from ._cancellation_scope_history import CancellationScopeBudget, CancellationScopeDeliveryReceipt
 from ._cooperative_cancellation import CancellationRequest, read_cancellation_history
 from ._prepared_local_activity import PreparedAttempt, PreparedCancellationObserved, PreparedLocalRunner
+from ._sticky_workflow_cache import CacheKey, StickyWorkflowCache, complete_history
 from .activity import ActivityContext, ActivityInfo, _set_context
 from .auth_composition import (
     AUTH_COMPOSITION_CONTRACT_SCHEMA,
@@ -53,6 +54,7 @@ from .client import (
     CONTROL_PLANE_REQUEST_CONTRACT_SCHEMA,
     CONTROL_PLANE_REQUEST_CONTRACT_VERSION,
     CONTROL_PLANE_VERSION,
+    DEFAULT_SDK_VERSION,
     PORTABLE_WORKER_AFFINITY_CAPABILITY_MANIFEST,
     PROTOCOL_VERSION,
     Client,
@@ -1008,6 +1010,9 @@ class Worker:
         max_concurrent_workflow_tasks: int = 10,
         max_concurrent_activity_tasks: int = 10,
         max_concurrent_worker_sessions: int = 10,
+        sticky_cache_capacity: int = 0,
+        sticky_cache_max_bytes: int = 16 * 1024 * 1024,
+        sticky_cache_ttl_seconds: int = 300,
         shutdown_timeout: float = 30.0,
         heartbeat_interval: float = 60.0,
         metrics: MetricsRecorder | None = None,
@@ -1029,6 +1034,11 @@ class Worker:
         }
         self.activities = {_activity_name(a): a for a in activities}
         self.capabilities = tuple(dict.fromkeys(capability.strip() for capability in capabilities))
+        self._sticky_cache = StickyWorkflowCache(
+            sticky_cache_capacity, sticky_cache_max_bytes, sticky_cache_ttl_seconds,
+        )
+        if "sticky_execution" in self.capabilities and not self._sticky_cache.enabled:
+            raise ValueError("sticky_execution requires a positive sticky_cache_capacity")
         self._cooperative_cancellation_supported = False
         self._prepared_local_activities_supported = False
         self._prepared_local_activity_groups_supported = False
@@ -1211,6 +1221,13 @@ class Worker:
         _validate_server_compatibility(info)
         protocol = info.get("worker_protocol")
         server_capabilities = protocol.get("server_capabilities") if isinstance(protocol, Mapping) else None
+        sticky_support = (
+            server_capabilities.get("sticky_execution") if isinstance(server_capabilities, Mapping) else None
+        )
+        if self._sticky_cache.enabled and (
+            not isinstance(sticky_support, Mapping) or sticky_support.get("supported") is not True
+        ):
+            raise RuntimeError("sticky execution requires Server's advertised sticky_execution support")
         self._cooperative_cancellation_supported = (
             "cooperative_cancellation" in self.capabilities
             and isinstance(protocol, Mapping)
@@ -1292,6 +1309,15 @@ class Worker:
             capabilities.append("prepared_local_activity_cancellation_policies")
         if PORTABLE_WORKER_AFFINITY_CAPABILITY_MANIFEST["worker_sessions"]["supported"]:
             capabilities.append("worker_sessions")
+        affinity_manifest = {name: dict(entry) for name, entry in PORTABLE_WORKER_AFFINITY_CAPABILITY_MANIFEST.items()}
+        if self._sticky_cache.enabled:
+            if "sticky_execution" not in capabilities:
+                capabilities.append("sticky_execution")
+            affinity_manifest["sticky_execution"] = {
+                "supported": True,
+                "minimum_protocol_version": "1.18",
+                "implementation": "bounded_durable_history_cache",
+            }
 
         ack = await self.client.register_worker(
             worker_id=self.worker_id,
@@ -1306,7 +1332,7 @@ class Worker:
             build_id=self.build_id,
             capabilities=capabilities,
             capability_manifest={
-                **PORTABLE_WORKER_AFFINITY_CAPABILITY_MANIFEST,
+                **affinity_manifest,
                 **({"prepared_local_activities": {
                     "supported": True, "minimum_protocol_version": "1.20",
                     "implementation": "durable_sequential_admission",
@@ -1514,8 +1540,65 @@ class Worker:
     async def _load_workflow_claim_history(
         self, task: dict[str, Any], *, first_page_token: str | None = None,
     ) -> list[dict[str, Any]]:
-        history = [] if first_page_token is not None else list(task.get("history_events", []))
-        token = first_page_token if first_page_token is not None else task.get("next_history_page_token")
+        if first_page_token is not None:
+            return await self._fetch_workflow_claim_history_pages(task, [], first_page_token)
+        inline = list(task.get("history_events", []))
+        key = self._sticky_cache_key(task)
+        mode = task.get("sticky_replay_mode")
+        cached = self._sticky_cache.lookup(key) if key is not None and mode == "sticky_hit_expected" else None
+        if cached is not None and key is not None:
+            prefix, resume_token, resume_offset = cached
+            overlap = min(len(inline), len(prefix))
+            # Published Server responses also express the boundary as an event count.
+            # Retained histories are contiguous from sequence 1, so the two agree.
+            raw_last_sequence = task.get("last_history_sequence", task.get("total_history_events"))
+            last_sequence = raw_last_sequence if isinstance(raw_last_sequence, int) else -1
+            anchored = (
+                complete_history(inline) and inline[:overlap] == prefix[:overlap]
+                and type(raw_last_sequence) is int and last_sequence >= len(prefix)
+            )
+            if anchored:
+                if resume_token is not None and resume_offset >= len(inline) and len(inline) < len(prefix):
+                    try:
+                        tail = await self._fetch_workflow_claim_history_pages(task, [], resume_token)
+                    except (InvalidArgument, LocalActivityExecutionAborted):
+                        tail = []
+                    except ServerError as error:
+                        if error.status != 400 or error.reason() != "invalid_page_token":
+                            raise
+                        tail = []
+                    if tail[:len(prefix) - resume_offset] == prefix[resume_offset:]:
+                        history = prefix[:resume_offset] + tail
+                        if complete_history(history) and len(history) >= last_sequence:
+                            # Paging measured the offset within the tail; retain its absolute boundary.
+                            task["_sticky_resume_offset"] += resume_offset
+                            self._sticky_cache.record_replay(hit=True)
+                            return history
+                else:
+                    history = await self._fetch_workflow_claim_history_pages(
+                        task, inline, task.get("next_history_page_token"),
+                    )
+                    if complete_history(history) and history[:len(prefix)] == prefix and len(history) >= last_sequence:
+                        self._sticky_cache.record_replay(hit=True)
+                        return history
+            # A cached prefix must never replace an inconsistent authoritative page.
+            self._sticky_cache.discard(key)
+        if key is not None:
+            self._sticky_cache.record_replay(
+                hit=False, forced=mode in {"sticky_hit_expected", "forced_cold_replay"},
+            )
+        if mode in {"sticky_hit_expected", "forced_cold_replay"} and not complete_history(inline):
+            history = await self._fetch_workflow_claim_history_pages(task, [], "MA==")
+            if not complete_history(history):
+                raise LocalActivityExecutionAborted("cold replay lacks a complete canonical history prefix")
+            return history
+        return await self._fetch_workflow_claim_history_pages(task, inline, task.get("next_history_page_token"))
+
+    async def _fetch_workflow_claim_history_pages(
+        self, task: dict[str, Any], history: list[dict[str, Any]], token: Any,
+    ) -> list[dict[str, Any]]:
+        task["_sticky_resume_token"] = None
+        task["_sticky_resume_offset"] = 0
         seen: set[str] = set()
         while token is not None:
             if not isinstance(token, str) or not token or token in seen:
@@ -1529,10 +1612,46 @@ class Worker:
                 raise LocalActivityExecutionAborted("workflow history page was not acknowledged")
             if any(not isinstance(event, dict) for event in page["history_events"]):
                 raise LocalActivityExecutionAborted("workflow history page contains a malformed event")
+            task["_sticky_resume_token"] = token
+            task["_sticky_resume_offset"] = len(history)
             history.extend(page["history_events"])
             token = page.get("next_history_page_token")
         return history
 
+    def _sticky_cache_key(self, task: dict[str, Any]) -> CacheKey | None:
+        if not self._sticky_cache.enabled:
+            return None
+        workflow_id, run_id = task.get("workflow_id"), task.get("run_id")
+        if not isinstance(workflow_id, str) or not workflow_id or not isinstance(run_id, str) or not run_id:
+            return None
+        return workflow_id, run_id, self.build_id.strip() if self.build_id is not None else DEFAULT_SDK_VERSION
+
+    def _sticky_cache_claim(
+        self, task: dict[str, Any], history: list[dict[str, Any]], commands: list[dict[str, Any]],
+    ) -> dict[str, Any] | None:
+        key = self._sticky_cache_key(task)
+        if key is None:
+            return None
+        if any(command.get("type") in {"complete_workflow", "fail_workflow", "continue_as_new"}
+               for command in commands):
+            self._sticky_cache.discard(key)
+            return None
+        if not self._sticky_cache.remember(
+            key, history, resume_token=task.get("_sticky_resume_token"),
+            resume_offset=task.get("_sticky_resume_offset", 0),
+        ):
+            return None
+        return {
+            "worker_id": self.worker_id, "workflow_id": key[0], "run_id": key[1], "build_id": key[2],
+            "ttl_seconds": self._sticky_cache.ttl_seconds,
+            "metrics": {name: self._sticky_cache.metrics()[name]
+                        for name in ("hit", "miss", "eviction", "forced_cold_replay")},
+        }
+
+    def sticky_cache_metrics(self) -> dict[str, int]:
+        """Return replay counters, retained entries and encoded history bytes."""
+        return self._sticky_cache.metrics()
+
     async def _refresh_cancellation_history(
         self, task: dict[str, Any], observed: CancellationRequest,
     ) -> list[dict[str, Any]]:
@@ -2416,6 +2535,7 @@ class Worker:
                     task_id=task_id,
                     attempt=attempt,
                     commands=[command],
+                    sticky_cache=self._sticky_cache_claim(task, history, [command]),
                 )
             except Exception as e:
                 log.warning("failed to complete workflow update task %s: %s", task_id, e)
@@ -2622,6 +2742,7 @@ class Worker:
                 commands=commands,
                 message_stream_cursors=outcome.message_stream_cursors,
                 message_stream_waits=outcome.message_stream_waits,
+                sticky_cache=self._sticky_cache_claim(task, history, commands),
             )
         except Exception as e:
             log.warning("failed to complete workflow task %s: %s", task_id, e)
@@ -2640,9 +2761,11 @@ class Worker:
         commands: list[dict[str, Any]],
         message_stream_cursors: list[dict[str, Any]] | None = None,
         message_stream_waits: list[dict[str, Any]] | None = None,
+        sticky_cache: dict[str, Any] | None = None,
     ) -> None:
         for completion_attempt in range(1, _WORKFLOW_TASK_COMPLETION_MAX_ATTEMPTS + 1):
             try:
+                completion_kwargs = {"sticky_cache": sticky_cache} if sticky_cache is not None else {}
                 await self.client.complete_workflow_task(
                     task_id=task_id,
                     lease_owner=self.worker_id,
@@ -2650,6 +2773,7 @@ class Worker:
                     commands=commands,
                     message_stream_cursors=message_stream_cursors,
                     message_stream_waits=message_stream_waits,
+                    **completion_kwargs,
                 )
                 return
             except Exception as error:
@@ -4353,6 +4477,7 @@ class Worker:
             await self.client.deregister_worker_registration(self.worker_id)
             self._registered = False
             log.info("worker %s deregistered", self.worker_id)
+        self._sticky_cache.clear()
 
     @staticmethod
     def _remaining_shutdown_time(deadline: float) -> float:

{durable_workflow-2.4.3 → durable_workflow-2.5.0/src/durable_workflow.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durable-workflow
-Version: 2.4.3
+Version: 2.5.0
 Summary: Python client and worker SDK for Durable Workflow Cloud and self-hosted Server
 Author: Durable Workflow Contributors
 License-Expression: MIT

{durable_workflow-2.4.3 → durable_workflow-2.5.0}/src/durable_workflow.egg-info/SOURCES.txt

@@ -8,6 +8,7 @@ src/durable_workflow/_cancellation_scope.py
 src/durable_workflow/_cancellation_scope_history.py
 src/durable_workflow/_cooperative_cancellation.py
 src/durable_workflow/_prepared_local_activity.py
+src/durable_workflow/_sticky_workflow_cache.py
 src/durable_workflow/activity.py
 src/durable_workflow/auth_composition.py
 src/durable_workflow/cancellation.py
@@ -106,6 +107,7 @@ tests/test_serializer.py
 tests/test_signals.py
 tests/test_sleep.py
 tests/test_standalone_activity_client.py
+tests/test_sticky_execution.py
 tests/test_storage_admission.py
 tests/test_sync.py
 tests/test_testing_harness.py

{durable_workflow-2.4.3 → durable_workflow-2.5.0}/tests/test_replay_regression_corpus.py

@@ -489,6 +489,27 @@ async def test_worker_replays_recorded_side_effects_across_history_pages() -> No
     assert serializer.decode_envelope(commands[0]["result"]) == 198
 
 
+@pytest.mark.asyncio
+async def test_replacement_with_sticky_suffix_fetches_cold_history_before_replay() -> None:
+    fixture = json.loads((FIXTURE_DIR / "sticky-cold-prefix-side-effects-avro.json").read_text(encoding="utf-8"))
+    history = fixture["history"]
+    client = AsyncMock(spec=Client)
+    client.workflow_task_history.return_value = {"history_events": history, "next_history_page_token": None}
+    worker = Worker(client, task_queue="cold-sticky", workflows=[PagedRecordedSideEffectsWorkflow])
+    commands = await worker._run_workflow_task_core({
+        "task_id": "replacement-task", "workflow_id": "cold-sticky", "run_id": "original-run",
+        "workflow_type": fixture["workflow"]["type"], "workflow_task_attempt": 2,
+        "history_events": history[-1:], "sticky_replay_mode": "sticky_hit_expected",
+        "arguments": serializer.encode([], codec="avro"), "payload_codec": "avro",
+    })
+    client.workflow_task_history.assert_awaited_once_with(
+        task_id="replacement-task", next_history_page_token="MA==", lease_owner=worker.worker_id,
+        workflow_task_attempt=2,
+    )
+    assert commands is not None
+    assert serializer.decode_envelope(commands[0]["result"]) == 198
+
+
 @pytest.mark.parametrize(
     "fixture",
     [

durable_workflow-2.5.0/tests/test_sticky_execution.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from durable_workflow import Client, Worker
+from durable_workflow._sticky_workflow_cache import StickyWorkflowCache
+from durable_workflow.errors import ServerError
+from durable_workflow.workflow import LocalActivityExecutionAborted
+
+
+def history(count: int = 6) -> list[dict[str, Any]]:
+    return [{"id": f"event-{sequence}", "sequence": sequence,
+             "event_type": "WorkflowStarted" if sequence == 1 else "SideEffectRecorded",
+             "payload": {"value": {"number": sequence}}} for sequence in range(1, count + 1)]
+
+
+def task(events: list[dict[str, Any]], *, mode: str = "sticky_hit_expected",
+         last: int = 6, token: str | None = None) -> dict[str, Any]:
+    return {"task_id": "current-task", "workflow_id": "workflow", "run_id": "run",
+            "workflow_task_attempt": 3, "history_events": events,
+            "sticky_replay_mode": mode, "last_history_sequence": last,
+            "next_history_page_token": token}
+
+
+@pytest.mark.parametrize(("capacity", "max_bytes", "ttl"), [
+    (-1, 100, 1), (True, 100, 1), (1, 0, 1), (1, True, 1),
+    (1, 100, 0), (1, 100, 3601), (1, 100, 1.5),
+])
+def test_cache_limits_reject_invalid_values(capacity: int, max_bytes: int, ttl: int) -> None:
+    with pytest.raises(ValueError, match="sticky_cache"):
+        StickyWorkflowCache(capacity, max_bytes, ttl)
+
+
+def test_lru_entry_limit_and_build_run_isolation() -> None:
+    cache = StickyWorkflowCache(2, 10_000, 300)
+    a, b, c = ("wf", "run-a", "build-a"), ("wf", "run-b", "build-a"), ("wf", "run-a", "build-b")
+    assert cache.remember(a, history())
+    assert cache.remember(b, history())
+    assert cache.lookup(a) is not None  # Keep A, evict B.
+    assert cache.remember(c, history())
+    assert cache.lookup(b) is None
+    assert cache.lookup(a) is not None
+    assert cache.lookup(c) is not None
+    assert cache.metrics()["eviction"] == 1
+    assert cache.metrics()["entries"] == 2
+
+
+def test_byte_limit_evicts_and_rejects_oversized_history() -> None:
+    sizing = StickyWorkflowCache(10, 100_000, 300)
+    assert sizing.remember(("a", "run", "build"), history(2))
+    one_size = sizing.metrics()["history_bytes"]
+    cache = StickyWorkflowCache(10, one_size * 2 - 1, 300)
+    assert cache.remember(("a", "run", "build"), history(2))
+    assert cache.remember(("b", "run", "build"), history(2))
+    assert cache.metrics()["entries"] == 1
+    assert cache.metrics()["eviction"] == 1
+    assert not cache.remember(("b", "run", "build"), history(10))
+    assert cache.metrics()["entries"] == 0
+    assert cache.metrics()["history_bytes"] == 0
+
+
+def test_cache_snapshot_is_independent_of_input_and_replayed_mutation() -> None:
+    cache = StickyWorkflowCache(1, 10_000, 300)
+    original = history()
+    key = ("wf", "run", "build")
+    assert cache.remember(key, original)
+    original[0]["payload"]["value"]["number"] = 999
+    first = cache.lookup(key)
+    assert first is not None
+    first[0][0]["payload"]["value"]["number"] = 888
+    second = cache.lookup(key)
+    assert second is not None
+    assert second[0] == history()
+
+
+def test_expiry_disabled_cache_and_malformed_history_are_not_retained() -> None:
+    cache = StickyWorkflowCache(1, 10_000, 1)
+    key = ("wf", "run", "build")
+    with patch("durable_workflow._sticky_workflow_cache.time.monotonic", return_value=10):
+        assert cache.remember(key, history())
+    with patch("durable_workflow._sticky_workflow_cache.time.monotonic", return_value=11):
+        assert cache.lookup(key) is None
+        assert cache.metrics()["history_bytes"] == 0
+    assert not StickyWorkflowCache(0, 10_000, 300).remember(key, history())
+    assert not cache.remember(key, history()[1:])
+    broken = history()
+    broken[2]["sequence"] = 9
+    assert not cache.remember(key, broken)
+    broken = history()
+    broken[0]["payload"] = object()
+    assert not cache.remember(key, broken)
+
+
+@pytest.mark.asyncio
+async def test_warm_history_reuses_server_cursor_on_current_lease_and_advances_boundary() -> None:
+    client = AsyncMock(spec=Client)
+    client.workflow_task_history.side_effect = [
+        {"history_events": history()[2:4], "next_history_page_token": "server-cursor-four"},
+        {"history_events": history()[4:], "next_history_page_token": None},
+    ]
+    worker = Worker(client, task_queue="q", worker_id="holder", build_id="build", sticky_cache_capacity=2)
+    first = task(history()[:2], mode="cold_replay", token="server-cursor-two")
+    loaded = await worker._load_workflow_claim_history(first)
+    claim = worker._sticky_cache_claim(first, loaded, [{"type": "start_timer"}])
+    assert claim is not None
+    assert claim["build_id"] == "build"
+    assert claim["metrics"]["miss"] == 1
+    client.workflow_task_history.reset_mock(side_effect=True)
+    client.workflow_task_history.side_effect = [
+        {"history_events": history(8)[4:6], "next_history_page_token": "server-cursor-six"},
+        {"history_events": history(8)[6:], "next_history_page_token": None},
+    ]
+    current = task(history(8)[:2], last=8, token="server-cursor-two")
+    current["total_history_events"] = current.pop("last_history_sequence")
+    assert await worker._load_workflow_claim_history(current) == history(8)
+    assert [call.kwargs["next_history_page_token"] for call in client.workflow_task_history.await_args_list] == [
+        "server-cursor-four", "server-cursor-six",
+    ]
+    assert all(call.kwargs["task_id"] == "current-task" and call.kwargs["lease_owner"] == "holder"
+               and call.kwargs["workflow_task_attempt"] == 3 for call in client.workflow_task_history.await_args_list)
+    worker._sticky_cache_claim(current, history(8), [{"type": "start_timer"}])
+    retained = worker._sticky_cache.lookup(("workflow", "run", "build"))
+    assert retained is not None
+    assert retained[1:] == ("server-cursor-six", 6)
+    assert worker.sticky_cache_metrics()["hit"] == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("change", ["prefix", "tail", "shorter", "cursor"])
+async def test_changed_history_or_rejected_cursor_fetches_complete_authoritative_history(change: str) -> None:
+    client = AsyncMock(spec=Client)
+    worker = Worker(client, task_queue="q", build_id="build", sticky_cache_capacity=1)
+    worker._sticky_cache.remember(("workflow", "run", "build"), history(),
+                                 resume_token="old-cursor", resume_offset=4)
+    authoritative = history()
+    inline = authoritative[:2]
+    last = 6
+    if change == "prefix":
+        authoritative[1]["payload"] = {"changed": True}
+    elif change == "shorter":
+        authoritative = history(4)
+        inline, last = authoritative[:2], 4
+    replies: list[Any] = [{"history_events": authoritative[2:], "next_history_page_token": None}]
+    if change == "tail":
+        authoritative[5]["payload"] = {"changed": True}
+        replies.insert(0, {"history_events": authoritative[4:], "next_history_page_token": None})
+    if change == "cursor":
+        replies.insert(0, ServerError(400, {"reason": "invalid_page_token"}))
+    client.workflow_task_history.side_effect = replies
+    loaded = await worker._load_workflow_claim_history(task(inline, last=last, token="current-cursor"))
+    assert loaded == authoritative
+    assert client.workflow_task_history.await_args_list[-1].kwargs["next_history_page_token"] == "current-cursor"
+    assert worker.sticky_cache_metrics()["hit"] == 0
+    assert worker.sticky_cache_metrics()["forced_cold_replay"] == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["sticky_hit_expected", "forced_cold_replay"])
+async def test_suffix_without_matching_prefix_uses_lease_scoped_cold_fetch(mode: str) -> None:
+    client = AsyncMock(spec=Client)
+    client.workflow_task_history.return_value = {"history_events": history(), "next_history_page_token": None}
+    worker = Worker(client, task_queue="q", worker_id="replacement", sticky_cache_capacity=1)
+    assert await worker._load_workflow_claim_history(task(history()[4:], mode=mode)) == history()
+    client.workflow_task_history.assert_awaited_once_with(
+        task_id="current-task", next_history_page_token="MA==", lease_owner="replacement", workflow_task_attempt=3,
+    )
+    assert worker.sticky_cache_metrics()["forced_cold_replay"] == 1
+
+
+@pytest.mark.asyncio
+async def test_build_mismatch_forced_replay_and_replacement_do_not_use_retained_history() -> None:
+    client = AsyncMock(spec=Client)
+    client.workflow_task_history.return_value = {"history_events": history()[2:], "next_history_page_token": None}
+    worker = Worker(client, task_queue="q", build_id="new-build", sticky_cache_capacity=1)
+    worker._sticky_cache.remember(("workflow", "run", "old-build"), history(),
+                                 resume_token="old-cursor", resume_offset=4)
+    assert await worker._load_workflow_claim_history(task(history()[:2], token="current-cursor")) == history()
+    assert client.workflow_task_history.await_args.kwargs["next_history_page_token"] == "current-cursor"
+    assert worker.sticky_cache_metrics()["hit"] == 0
+    await worker.stop()
+    assert worker.sticky_cache_metrics()["entries"] == 0
+
+
+@pytest.mark.asyncio
+async def test_incomplete_cold_history_and_repeated_tokens_never_replay_partial_data() -> None:
+    client = AsyncMock(spec=Client)
+    worker = Worker(client, task_queue="q", sticky_cache_capacity=1)
+    client.workflow_task_history.return_value = {"history_events": history()[2:], "next_history_page_token": None}
+    with pytest.raises(LocalActivityExecutionAborted, match="complete canonical"):
+        await worker._load_workflow_claim_history(task(history()[4:]))
+    client.workflow_task_history.return_value = {"history_events": history()[:2], "next_history_page_token": "repeat"}
+    with pytest.raises(LocalActivityExecutionAborted, match="paging did not advance"):
+        await worker._load_workflow_claim_history(task(history()[:2], mode="cold_replay", token="repeat"))
+    client.complete_workflow_task.assert_not_awaited()
+
+
+def test_oversized_and_terminal_histories_issue_no_affinity_claim() -> None:
+    worker = Worker(AsyncMock(spec=Client), task_queue="q", build_id="build",
+                    sticky_cache_capacity=2, sticky_cache_max_bytes=100)
+    assert worker._sticky_cache_claim(task(history()), history(), [{"type": "start_timer"}]) is None
+    worker = Worker(AsyncMock(spec=Client), task_queue="q", build_id="build", sticky_cache_capacity=2)
+    assert worker._sticky_cache_claim(task(history()), history(), [{"type": "start_timer"}]) is not None
+    assert worker._sticky_cache_claim(task(history()), history(), [{"type": "complete_workflow"}]) is None
+    assert worker.sticky_cache_metrics()["entries"] == 0
+
+
+@pytest.mark.asyncio
+async def test_completion_transport_retries_identical_claim_and_omits_it_when_disabled() -> None:
+    client = AsyncMock(spec=Client)
+    worker = Worker(client, task_queue="q", worker_id="holder", build_id="build", sticky_cache_capacity=1)
+    claim = worker._sticky_cache_claim(task(history()), history(), [{"type": "start_timer"}])
+    client.complete_workflow_task.side_effect = [ServerError(503, {}), {}]
+    with patch("durable_workflow.worker.asyncio.sleep", new_callable=AsyncMock):
+        await worker._complete_workflow_task_with_retry(
+            task_id="current-task", attempt=3, commands=[{"type": "start_timer"}], sticky_cache=claim,
+        )
+    assert client.complete_workflow_task.await_args_list[0] == client.complete_workflow_task.await_args_list[1]
+    assert client.complete_workflow_task.await_args.kwargs["sticky_cache"] == claim
+    client.complete_workflow_task.reset_mock(side_effect=True)
+    await worker._complete_workflow_task_with_retry(task_id="next-task", attempt=1, commands=[{"type": "start_timer"}])
+    assert "sticky_cache" not in client.complete_workflow_task.await_args.kwargs
+
+
+@pytest.mark.asyncio
+async def test_client_completion_serializes_claim_on_worker_plane() -> None:
+    async with Client("https://server.example", worker_token="test") as client:
+        with patch.object(client, "_request", new_callable=AsyncMock) as request:
+            await client.complete_workflow_task(
+                task_id="task", lease_owner="worker", workflow_task_attempt=1,
+                commands=[{"type": "start_timer"}], sticky_cache={"worker_id": "worker"},
+            )
+        assert request.await_args.kwargs["worker"] is True
+        assert request.await_args.kwargs["json"]["sticky_cache"] == {"worker_id": "worker"}

{durable_workflow-2.4.3 → durable_workflow-2.5.0}/tests/test_worker.py

@@ -447,7 +447,7 @@ class TestWorkerRegistration:
                 "sticky_execution": {
                     "supported": False,
                     "minimum_protocol_version": "1.18",
-                    "reason": "python_worker_uses_complete_durable_history_replay",
+                    "reason": "sticky_cache_disabled",
                 },
             },
             "task_slots": {
@@ -1196,6 +1196,33 @@ class TestWorkerRegistration:
         mock_client.register_worker.assert_not_called()
 
 
+class TestStickyWorkerRegistration:
+    @pytest.mark.asyncio
+    async def test_only_enabled_worker_advertises_sticky_support(self, mock_client: AsyncMock) -> None:
+        info = compatible_cluster_info()
+        info["worker_protocol"]["server_capabilities"]["sticky_execution"] = {"supported": True}  # type: ignore[index]
+        mock_client.get_cluster_info.return_value = info
+        worker = Worker(mock_client, task_queue="q", sticky_cache_capacity=2)
+        await worker._register()
+        registration = mock_client.register_worker.await_args.kwargs
+        assert registration["capabilities"].count("sticky_execution") == 1
+        assert registration["capability_manifest"]["sticky_execution"] == {
+            "supported": True, "minimum_protocol_version": "1.18", "implementation": "bounded_durable_history_cache",
+        }
+        await worker.stop()
+
+    @pytest.mark.asyncio
+    async def test_enabled_cache_requires_server_support(self, mock_client: AsyncMock) -> None:
+        worker = Worker(mock_client, task_queue="q", sticky_cache_capacity=2)
+        with pytest.raises(RuntimeError, match="advertised sticky_execution support"):
+            await worker._register()
+        mock_client.register_worker.assert_not_awaited()
+
+    def test_manual_capability_cannot_advertise_disabled_cache(self, mock_client: AsyncMock) -> None:
+        with pytest.raises(ValueError, match="positive sticky_cache_capacity"):
+            Worker(mock_client, task_queue="q", capabilities=["sticky_execution"])
+
+
 class TestWorkflowTaskExecution:
     @pytest.mark.asyncio
     async def test_history_page_failure_does_not_replay_partial_history(