mageflow 0.3.2__tar.gz → 0.3.3__tar.gz

{mageflow-0.3.2 → mageflow-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mageflow
-Version: 0.3.2
+Version: 0.3.3
 Summary: Manage Graph Execution Flow - A unified interface for task orchestration across different task managers
 Project-URL: Homepage, https://imaginary-cherry.github.io/mageflow/
 Project-URL: Documentation, https://imaginary-cherry.github.io/mageflow/
@@ -26,7 +26,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: System :: Distributed Computing
 Classifier: Typing :: Typed
 Requires-Python: <3.14,>=3.10
-Requires-Dist: thirdmagic<0.1.0,>=0.0.3
+Requires-Dist: thirdmagic<0.1.0,>=0.0.4
 Provides-Extra: dev
 Requires-Dist: black>=26.1.0; extra == 'dev'
 Requires-Dist: coverage[toml]<8.0.0,>=7.0.0; extra == 'dev'

{mageflow-0.3.2 → mageflow-0.3.3}/mageflow/callbacks.py

@@ -7,11 +7,16 @@ from typing import Any
 from hatchet_sdk import Context
 from hatchet_sdk.runnables.types import EmptyModel
 from pydantic import BaseModel
-
-from mageflow.utils.pythonic import flexible_call
+from thirdmagic.signature.retry_cache import (
+    retry_cache_ctx,
+    setup_retry_cache,
+    teardown_retry_cache,
+)
 from thirdmagic.task.model import TaskSignature
 from thirdmagic.task_def import MageflowTaskDefinition
 
+from mageflow.utils.pythonic import flexible_call
+
 
 class AcceptParams(Enum):
     JUST_MESSAGE = 1
@@ -27,6 +32,7 @@ def handle_task_callback(
     expected_params: AcceptParams = AcceptParams.NO_CTX,
     wrap_res: bool = True,
     send_signature: bool = False,
+    is_idempotent: bool = False,
 ):
     def task_decorator(func):
         @functools.wraps(func)
@@ -40,8 +46,18 @@ def handle_task_callback(
                 # NOTE: This should not run, the task should cancel, but just in case
                 return {"Error": "Task should have been canceled"}
             is_normal_run = lifecycle.is_vanilla_run()
+            is_task_finish = False
             signature = await lifecycle.start_task()
 
+            # Setup retry cache for signature idempotency on retries (durable tasks only)
+            cache_token = None
+            cache_state = None
+            if is_idempotent:
+                cache_state = await setup_retry_cache(
+                    ctx.workflow_id, ctx.attempt_number
+                )
+                cache_token = retry_cache_ctx.set(cache_state)
+
             # Add params if user requires
             if send_signature:
                 kwargs["signature"] = signature
@@ -54,19 +70,22 @@ def handle_task_callback(
                 else:
                     result = await flexible_call(func, message, ctx, *args, **kwargs)
             except asyncio.CancelledError as e:
+                is_task_finish = True
                 if not is_normal_run:
                     await lifecycle.task_failed(msg_data, e)
                 raise
             except Exception as e:
-                if is_normal_run:
-                    raise
-                if not TaskSignature.ClientAdapter.should_task_retry(
+                will_retry = TaskSignature.ClientAdapter.should_task_retry(
                     task_model, ctx.attempt_number, e
-                ):
-                    await lifecycle.task_failed(msg_data, e)
+                )
+                if not will_retry:
+                    is_task_finish = True
+                    if not is_normal_run:
+                        await lifecycle.task_failed(msg_data, e)
                 raise
             else:
                 # If this is a simple task, no signature, then we dont do any manipulation
+                is_task_finish = True
                 if is_normal_run:
                     return result
                 task_results = HatchetResult(hatchet_results=result)
@@ -76,6 +95,11 @@ def handle_task_callback(
                     return task_results
                 else:
                     return result
+            finally:
+                if cache_token is not None:
+                    retry_cache_ctx.reset(cache_token)
+                if is_task_finish and cache_state:
+                    await teardown_retry_cache(cache_state)
 
         wrapper.__signature__ = inspect.signature(func)
         return wrapper

{mageflow-0.3.2 → mageflow-0.3.3}/mageflow/chain/workflows.py

@@ -2,7 +2,6 @@ from logging import Logger
 from typing import Any
 
 from rapyer.fields import RapyerKey
-
 from thirdmagic.clients.lifecycle import BaseLifecycle
 
 

{mageflow-0.3.2 → mageflow-0.3.3}/mageflow/clients/hatchet/adapter.py

@@ -8,6 +8,13 @@ from hatchet_sdk.runnables.types import EmptyModel
 from hatchet_sdk.runnables.workflow import BaseWorkflow
 from pydantic import BaseModel, TypeAdapter
 from rapyer.fields import RapyerKey
+from thirdmagic.chain import ChainTaskSignature
+from thirdmagic.clients.base import BaseClientAdapter
+from thirdmagic.consts import TASK_ID_PARAM_NAME
+from thirdmagic.signature import Signature
+from thirdmagic.swarm import SwarmTaskSignature
+from thirdmagic.task import TaskSignature
+from thirdmagic.task_def import MageflowTaskDefinition
 
 from mageflow.chain.messages import ChainCallbackMessage, ChainErrorMessage
 from mageflow.clients.hatchet.workflow import MageflowWorkflow
@@ -25,13 +32,6 @@ from mageflow.swarm.messages import (
     SwarmErrorMessage,
     SwarmResultsMessage,
 )
-from thirdmagic.chain import ChainTaskSignature
-from thirdmagic.clients.base import BaseClientAdapter
-from thirdmagic.consts import TASK_ID_PARAM_NAME
-from thirdmagic.signature import Signature
-from thirdmagic.swarm import SwarmTaskSignature
-from thirdmagic.task import TaskSignature
-from thirdmagic.task_def import MageflowTaskDefinition
 
 
 class HatchetClientAdapter(BaseClientAdapter):
@@ -123,6 +123,15 @@ class HatchetClientAdapter(BaseClientAdapter):
     def extract_retries(self, client_task: BaseWorkflow) -> int:
         return client_task.tasks[0].retries
 
+    def _prepare_wf(self, signature: TaskSignature, set_return_field: bool, **kwargs):
+        total_kwargs = signature.kwargs | kwargs
+        workflow = self.hatchet.workflow(
+            name=signature.task_name, input_validator=signature.model_validators
+        )
+        return_field_name = signature.return_field_name if set_return_field else None
+        mageflow_wf = MageflowWorkflow(workflow, total_kwargs, return_field_name)
+        return mageflow_wf
+
     async def acall_signature(
         self,
         signature: TaskSignature,
@@ -134,17 +143,23 @@ class HatchetClientAdapter(BaseClientAdapter):
         if msg is None:
             msg = EmptyModel()
         options = self._update_options(signature, options)
-        total_kwargs = signature.kwargs | kwargs
-        workflow = self.hatchet.workflow(
-            name=signature.task_name, input_validator=signature.model_validators
-        )
-        mageflow_wf = MageflowWorkflow(
-            workflow,
-            total_kwargs,
-            signature.return_field_name if set_return_field else None,
-        )
+        mageflow_wf = self._prepare_wf(signature, set_return_field, **kwargs)
         return await mageflow_wf.aio_run_no_wait(msg, options)
 
+    async def await_signature(
+        self,
+        signature: "TaskSignature",
+        msg: Any,
+        set_return_field: bool,
+        options: TriggerWorkflowOptions = None,
+        **kwargs,
+    ):
+        if msg is None:
+            msg = EmptyModel()
+        options = self._update_options(signature, options)
+        mageflow_wf = self._prepare_wf(signature, set_return_field, **kwargs)
+        return await mageflow_wf.aio_run(msg, options)
+
     def should_task_retry(
         self,
         task_definition: MageflowTaskDefinition,

{mageflow-0.3.2 → mageflow-0.3.3}/mageflow/clients/hatchet/mageflow.py

@@ -17,6 +17,14 @@ from hatchet_sdk.runnables.types import (
 from hatchet_sdk.runnables.workflow import BaseWorkflow, Standalone
 from hatchet_sdk.worker.worker import LifespanFn
 from redis.asyncio import Redis
+from thirdmagic import chain, sign
+from thirdmagic.chain import ChainTaskSignature
+from thirdmagic.signature import Signature
+from thirdmagic.swarm import SwarmTaskSignature
+from thirdmagic.swarm.creator import SignatureOptions, swarm
+from thirdmagic.task import TaskInputType, TaskSignature, TaskSignatureConvertible
+from thirdmagic.task_def import MageflowTaskDefinition
+from thirdmagic.utils import HatchetTaskType
 from typing_extensions import override
 
 from mageflow.callbacks import AcceptParams, handle_task_callback
@@ -47,14 +55,6 @@ from mageflow.swarm.workflows import (
     swarm_item_failed,
 )
 from mageflow.utils.mageflow import does_task_wants_ctx
-from thirdmagic import chain, sign
-from thirdmagic.chain import ChainTaskSignature
-from thirdmagic.signature import Signature
-from thirdmagic.swarm import SwarmTaskSignature
-from thirdmagic.swarm.creator import SignatureOptions, swarm
-from thirdmagic.task import TaskInputType, TaskSignature, TaskSignatureConvertible
-from thirdmagic.task_def import MageflowTaskDefinition
-from thirdmagic.utils import HatchetTaskType
 
 Duration = timedelta | str
 
@@ -125,14 +125,16 @@ class HatchetMageflow(Hatchet):
             )
         )
 
-    def task_decorator(self, func: Callable, hatchet_task):
+    def task_decorator(self, func: Callable, hatchet_task, is_idempotent: bool = False):
         param_config = (
             AcceptParams.ALL
             if does_task_wants_ctx(func)
             else self.mageflow_config.param_config
         )
         send_signature = getattr(func, "__send_signature__", False)
-        handler_dec = handle_task_callback(param_config, send_signature=send_signature)
+        handler_dec = handle_task_callback(
+            param_config, send_signature=send_signature, is_idempotent=is_idempotent
+        )
         func = handler_dec(func)
         wf = hatchet_task(func)
         self._add_task_def(wf)
@@ -146,7 +148,9 @@ class HatchetMageflow(Hatchet):
         """
         hatchet_task = super().task(name=name, **kwargs)
 
-        decorator = functools.partial(self.task_decorator, hatchet_task=hatchet_task)
+        decorator = functools.partial(
+            self.task_decorator, hatchet_task=hatchet_task, is_idempotent=False
+        )
         return decorator
 
     @override
@@ -156,7 +160,9 @@ class HatchetMageflow(Hatchet):
         """
         hatchet_task = super().durable_task(name=name, **kwargs)
 
-        decorator = functools.partial(self.task_decorator, hatchet_task=hatchet_task)
+        decorator = functools.partial(
+            self.task_decorator, hatchet_task=hatchet_task, is_idempotent=True
+        )
 
         return decorator
 

{mageflow-0.3.2 → mageflow-0.3.3}/mageflow/clients/hatchet/workflow.py

@@ -3,7 +3,6 @@ from typing import Any
 from hatchet_sdk.runnables.workflow import Workflow
 from hatchet_sdk.utils.typing import JSONSerializableMapping
 from pydantic import BaseModel
-
 from thirdmagic.utils import deep_merge
 
 

{mageflow-0.3.2 → mageflow-0.3.3}/mageflow/config.py

@@ -1,26 +1,29 @@
 import dataclasses
-from dataclasses import dataclass, field
+from dataclasses import field
 from typing import Optional
 
-from mageflow.callbacks import AcceptParams
+from pydantic import Field
+from pydantic.dataclasses import dataclass
 from thirdmagic.chain import ChainTaskSignature
 from thirdmagic.consts import REMOVED_TASK_TTL
-from thirdmagic.signature import SignatureConfig
+from thirdmagic.signature import Signature, SignatureConfig
 from thirdmagic.swarm import SwarmTaskSignature
 from thirdmagic.swarm.state import PublishState
 from thirdmagic.task import TaskSignature
 
+from mageflow.callbacks import AcceptParams
+
 
 @dataclass
 class SignatureTTLConfig:
     active_ttl: Optional[int] = None  # seconds, None = use general
-    ttl_when_sign_done: Optional[int] = None
+    ttl_when_sign_done: Optional[int] = Field(default=None, ge=REMOVED_TASK_TTL)
 
 
 @dataclass
 class TTLConfig:
     active_ttl: int = 24 * 60 * 60  # general active TTL (default 24h)
-    ttl_when_sign_done: int = REMOVED_TASK_TTL  # general done TTL (default 5min)
+    ttl_when_sign_done: int = Field(default=REMOVED_TASK_TTL, ge=REMOVED_TASK_TTL)
     task: SignatureTTLConfig = field(default_factory=SignatureTTLConfig)
     chain: SignatureTTLConfig = field(default_factory=SignatureTTLConfig)
     swarm: SignatureTTLConfig = field(default_factory=SignatureTTLConfig)
@@ -44,4 +47,5 @@ def apply_ttl_config(ttl_config: TTLConfig):
         done_ttl = sig_config.ttl_when_sign_done or ttl_config.ttl_when_sign_done
 
         sig_type.Meta = dataclasses.replace(sig_type.Meta, ttl=active_ttl)
-        sig_type.SignatureSettings = SignatureConfig(ttl_when_sign_done=done_ttl)
+        if issubclass(sig_type, Signature):
+            sig_type.SignatureSettings = SignatureConfig(ttl_when_sign_done=done_ttl)

{mageflow-0.3.2 → mageflow-0.3.3}/mageflow/lifecycle/signature.py

@@ -2,7 +2,6 @@ import asyncio
 from typing import Any, Optional, cast
 
 import rapyer
-
 from thirdmagic.clients.lifecycle import BaseLifecycle
 from thirdmagic.container import ContainerTaskSignature
 from thirdmagic.signature import Signature

{mageflow-0.3.2 → mageflow-0.3.3}/mageflow/lifecycle/task.py

@@ -1,7 +1,6 @@
 from typing import Any
 
 from pydantic import BaseModel
-
 from thirdmagic.clients.lifecycle import BaseLifecycle
 
 

{mageflow-0.3.2 → mageflow-0.3.3}/mageflow/swarm/workflows.py

@@ -3,7 +3,6 @@ from typing import Any, Optional, cast
 
 import rapyer
 from rapyer.fields import RapyerKey
-
 from thirdmagic.clients.lifecycle import BaseLifecycle
 from thirdmagic.signature import Signature
 from thirdmagic.swarm import PublishState

{mageflow-0.3.2 → mageflow-0.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mageflow"
-version = "0.3.2"
+version = "0.3.3"
 description = "Manage Graph Execution Flow - A unified interface for task orchestration across different task managers"
 authors = [
     {name = "imaginary-cherry", email = "yedidyakfir@gmail.com"}
@@ -41,7 +41,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "thirdmagic>=0.0.3, <0.1.0",
+    "thirdmagic>=0.0.4, <0.1.0",
 ]
 
 [project.optional-dependencies]
@@ -103,3 +103,7 @@ exclude_lines = [
 
 [tool.coverage.html]
 directory = "htmlcov"
+
+[tool.ruff]
+select = ["I", "F401"]
+

{mageflow-0.3.2 → mageflow-0.3.3}/tests/integration/hatchet/models.py

@@ -1,7 +1,6 @@
 from typing import Any
 
 from pydantic import BaseModel, Field
-
 from thirdmagic.message import ReturnValue
 
 
@@ -50,6 +49,25 @@ class SignatureKeysResult(BaseModel):
     swarm_sub_task_keys: list[str]
     publish_state_key: str
 
+    def is_key_in_keys(self, key: str) -> bool:
+        if key in self.task_keys:
+            return True
+        if key in self.chain_sub_task_keys:
+            return True
+        if key in self.swarm_sub_task_keys:
+            return True
+        if key == self.publish_state_key:
+            return True
+        if key == self.chain_key:
+            return True
+        if key == self.swarm_key:
+            return True
+        return False
+
+
+class SignatureKeyWithWF(SignatureKeysResult):
+    workflow_id: str
+
 
 class MageflowTestError(Exception):
     pass

mageflow-0.3.3/tests/integration/hatchet/test_retry_cache.py

@@ -0,0 +1,65 @@
+import pytest
+from thirdmagic.task import TaskSignature
+
+from tests.integration.hatchet.conftest import HatchetInitData
+from tests.integration.hatchet.models import ContextMessage, SignatureKeyWithWF
+from tests.integration.hatchet.worker import retry_cache_durable_task
+
+
+async def _assert_retry_cache_idempotency(results: SignatureKeyWithWF):
+    workflow_id = results.workflow_id
+
+    # Assert - get the keys stored by each attempt from Redis
+    attempt_1_raw = await TaskSignature.Meta.redis.json().get(  # type: ignore[misc]
+        f"retry-cache-test:{workflow_id}:attempt-1"
+    )
+    attempt_2_raw = await TaskSignature.Meta.redis.json().get(  # type: ignore[misc]
+        f"retry-cache-test:{workflow_id}:attempt-2"
+    )
+
+    assert attempt_1_raw is not None, "Attempt 1 keys not found in Redis"
+    assert attempt_2_raw is not None, "Attempt 2 keys not found in Redis"
+    assert (
+        attempt_1_raw == attempt_2_raw
+    ), "Keys shouldn't be different between attempts"
+
+    last_sign = results.task_keys[-1]
+    results.task_keys = results.task_keys[:-1]
+    assert not results.is_key_in_keys(last_sign), "Last signature should not be in keys"
+    assert results.model_dump(mode="json") == attempt_1_raw
+
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test__retry_cache__durable_task_retry__no_duplicate_signatures(
+    hatchet_client_init: HatchetInitData,
+    test_ctx,
+    ctx_metadata,
+    trigger_options,
+):
+    # Arrange
+    hatchet = hatchet_client_init.hatchet
+
+    signature = await hatchet.asign(retry_cache_durable_task)
+    msg = ContextMessage(base_data=test_ctx)
+
+    # Act - trigger the durable task which will fail on attempt 1 and succeed on attempt 2
+    wf_res = await signature.aio_run(msg, options=trigger_options)
+
+    results = SignatureKeyWithWF(
+        **wf_res["retry_cache_durable_task"]["hatchet_results"]
+    )
+
+    await _assert_retry_cache_idempotency(results)
+
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test__retry_cache__task_retry__no_duplicate_signatures(
+    test_ctx, ctx_metadata, trigger_options
+):
+    # Arrange
+    msg = ContextMessage(base_data=test_ctx)
+
+    # Act - trigger directly via hatchet aio_run (not through a signature)
+    results = await retry_cache_durable_task.aio_run(msg, options=trigger_options)
+
+    await _assert_retry_cache_idempotency(results)

{mageflow-0.3.2 → mageflow-0.3.3}/tests/integration/hatchet/test_ttl.py

@@ -1,7 +1,7 @@
 import pytest
 
 from tests.integration.hatchet.conftest import HatchetInitData
-from tests.integration.hatchet.models import ContextMessage, SignatureKeysResult
+from tests.integration.hatchet.models import ContextMessage
 from tests.integration.hatchet.worker import (
     CHAIN_ACTIVE_TTL,
     SWARM_ACTIVE_TTL,

{mageflow-0.3.2 → mageflow-0.3.3}/tests/integration/hatchet/worker.py

@@ -4,7 +4,7 @@ import logging
 import os
 from datetime import datetime, timedelta
 
-from thirdmagic.consts import TASK_ID_PARAM_NAME
+from thirdmagic.consts import REMOVED_TASK_TTL, TASK_ID_PARAM_NAME
 from thirdmagic.task import TaskSignature
 
 # Start coverage if COVERAGE_PROCESS_START is set
@@ -32,6 +32,7 @@ from tests.integration.hatchet.models import (
     MessageWithMsgResults,
     MessageWithResult,
     SignatureKeysResult,
+    SignatureKeyWithWF,
     SleepTaskMessage,
 )
 
@@ -53,11 +54,11 @@ hatchet = Hatchet(debug=True, config=config_obj)
 
 # Per-type TTL configuration for tests
 TASK_ACTIVE_TTL = 600  # 10 minutes
-TASK_DONE_TTL = 60  # 1 minute
+TASK_DONE_TTL = REMOVED_TASK_TTL + 60  # 1 minute
 CHAIN_ACTIVE_TTL = 900  # 15 minutes
-CHAIN_DONE_TTL = 90  # 1.5 minutes
+CHAIN_DONE_TTL = REMOVED_TASK_TTL + 90  # 1.5 minutes
 SWARM_ACTIVE_TTL = 1200  # 20 minutes
-SWARM_DONE_TTL = 120  # 2 minutes
+SWARM_DONE_TTL = REMOVED_TASK_TTL + 120  # 2 minutes
 MAX_DONE_TTL = max(TASK_DONE_TTL, CHAIN_DONE_TTL, SWARM_DONE_TTL)
 
 TEST_MAGEFLOW_CONFIG = MageflowConfig(
@@ -225,6 +226,55 @@ async def create_signatures_for_ttl_test(msg: ContextMessage) -> SignatureKeysRe
     )
 
 
+@hatchet.durable_task(
+    name="retry_cache_durable_task",
+    input_validator=ContextMessage,
+    retries=3,
+    execution_timeout=timedelta(seconds=60),
+)
+@hatchet.with_ctx
+async def retry_cache_durable_task(
+    msg: ContextMessage, ctx: Context
+) -> SignatureKeyWithWF:
+    # Create standalone task signatures
+    sig1 = await hatchet.asign(task1)
+    sig2 = await hatchet.asign(task2)
+
+    # Create a chain
+    chain_sub1 = await hatchet.asign(task1)
+    chain_sub2 = await hatchet.asign(task2)
+    chain_sig = await hatchet.achain([chain_sub1, chain_sub2])
+
+    # Create a swarm
+    swarm_sub1 = await hatchet.asign(task1)
+    swarm_sub2 = await hatchet.asign(task2)
+    swarm_sig = await hatchet.aswarm([swarm_sub1, swarm_sub2], is_swarm_closed=True)
+
+    # Collect all created signature keys
+    results = SignatureKeyWithWF(
+        task_keys=[sig1.key, sig2.key],
+        chain_key=chain_sig.key,
+        chain_sub_task_keys=[chain_sub1.key, chain_sub2.key],
+        swarm_key=swarm_sig.key,
+        swarm_sub_task_keys=[swarm_sub1.key, swarm_sub2.key],
+        publish_state_key=swarm_sig.publishing_state_id,
+        workflow_id=ctx.workflow_id,
+    )
+    all_keys = results.model_dump(mode="json")
+
+    # Store keys in Redis for test verification, keyed by attempt number
+    attempt_key = f"retry-cache-test:{ctx.workflow_id}:attempt-{ctx.attempt_number}"
+    await TaskSignature.Meta.redis.json().set(attempt_key, "$", all_keys)  # type: ignore[misc]
+
+    if ctx.attempt_number == 1:
+        raise ValueError("Intentional first attempt failure for retry cache test")
+
+    new_task = await hatchet.asign(task2)
+    results.task_keys.append(new_task.key)
+
+    return results
+
+
 workflows = [
     task1,
     task2,
@@ -245,6 +295,7 @@ workflows = [
     cancel_retry,
     accept_msg_results,
     create_signatures_for_ttl_test,
+    retry_cache_durable_task,
 ]
 
 

{mageflow-0.3.2 → mageflow-0.3.3}/tests/unit/assertions.py

@@ -2,7 +2,6 @@ from typing import Literal, TypeVar
 
 from rapyer.fields import RapyerKey
 from redis.asyncio import Redis
-
 from thirdmagic.consts import REMOVED_TASK_TTL
 from thirdmagic.task import TaskSignature
 

{mageflow-0.3.2 → mageflow-0.3.3}/tests/unit/callbacks/conftest.py

@@ -5,15 +5,15 @@ from unittest.mock import AsyncMock, MagicMock
 
 import pytest_asyncio
 from hatchet_sdk import Context
-
-import mageflow
-from mageflow.callbacks import AcceptParams, handle_task_callback
-from tests.integration.hatchet.models import ContextMessage
 from thirdmagic.consts import TASK_ID_PARAM_NAME
 from thirdmagic.signature.status import SignatureStatus
 from thirdmagic.task import TaskSignature
 from thirdmagic.task_def import MageflowTaskDefinition
 
+import mageflow
+from mageflow.callbacks import AcceptParams, handle_task_callback
+from tests.integration.hatchet.models import ContextMessage
+
 
 @dataclass
 class MockContextConfig:
@@ -86,6 +86,7 @@ def handler_factory(
     send_signature: bool = False,
     return_value: Any = "success_result",
     raises: BaseException | None = None,
+    durable: bool = False,
 ):
     tracked_calls: list[CallTracker] = []
 
@@ -93,6 +94,7 @@ def handler_factory(
         expected_params=expected_params,
         wrap_res=wrap_res,
         send_signature=send_signature,
+        is_idempotent=durable,
     )
     async def handler(*args, **kwargs):
         tracked_calls.append(CallTracker(args=args, kwargs=kwargs))