fix(runtime): require explicit resume identity capability (#6571)

* fix(runtime): require explicit resume identity capability

Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com>

* fix(runtime): guard key-only idempotency admission

---------

Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com>
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

README.md

@@ -578,6 +578,12 @@ it does not make old workers safe for resume admission. An upgraded worker retur
 inspect that run and the current thread state before deciding to submit a new
 action. See the [run API contract](backend/docs/API.md#create-run).
 
+Custom run stores must explicitly support every supplied idempotency field.
+Keyed input and resume requests return HTTP 503 before run admission when the
+store cannot guarantee that support. The built-in Memory and SQL stores support
+both fields; older custom stores with an explicit `idempotency_key` parameter
+continue to accept keyed input requests.
+
 #### Option 2: Local Development
 
 If you prefer running services locally:

backend/packages/harness/deerflow/runtime/runs/manager.py

@@ -185,19 +185,28 @@ class PersistenceRetryPolicy:
 
 
 class RunIdempotencyUnsupported(RuntimeError):
-    """The configured store cannot safely persist keyed-resume identity."""
+    """The configured store cannot safely persist the requested idempotency fields."""
 
 
-def _store_accepts_idempotency_request(store: RunStore) -> bool:
-    """Best-effort capability check for old explicit RunStore overrides."""
+def _store_accepts_idempotency_parameter(store: RunStore, parameter_name: str) -> bool:
+    """Require an explicit keyword parameter or an opt-in owned by the method."""
     callable_ = store.create_thread_operation_atomic
     if getattr(callable_, "__func__", None) is RunStore.create_thread_operation_atomic:
         return False
     try:
         parameters = inspect.signature(callable_).parameters
     except (TypeError, ValueError):
+        return False
+    parameter = parameters.get(parameter_name)
+    if parameter is not None and parameter.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY):
         return True
-    return "idempotency_request" in parameters or any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values())
+    if not any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()):
+        return False
+    implementation_owner = next(
+        (candidate for candidate in type(store).__mro__ if "create_thread_operation_atomic" in candidate.__dict__),
+        None,
+    )
+    return bool(implementation_owner and implementation_owner.__dict__.get("supports_idempotency_request_kwargs", False))
 
 
 @dataclass
@@ -1728,8 +1737,10 @@ class RunManager:
             # 2) Persist to store while still holding the local lock. The
             #    store is the source of truth for cross-process atomicity.
             if self._store is not None:
-                if idempotency_request is not None and not _store_accepts_idempotency_request(self._store):
+                if idempotency_request is not None and not _store_accepts_idempotency_parameter(self._store, "idempotency_request"):
                     raise RunIdempotencyUnsupported("The configured RunStore does not support keyed resume idempotency")
+                if idempotency_key is not None and not _store_accepts_idempotency_parameter(self._store, "idempotency_key"):
+                    raise RunIdempotencyUnsupported("The configured RunStore does not support keyed run idempotency")
                 if multitask_strategy == "reject":
                     create_kwargs = {
                         "run_id": run_id,

backend/packages/harness/deerflow/runtime/runs/store/base.py

@@ -118,6 +118,15 @@ def run_is_before_cursor(
 
 
 class RunStore(abc.ABC):
+    # Explicit idempotency_key / idempotency_request keyword parameters each
+    # declare support for that field. An override accepting them only through
+    # **kwargs must opt in below and persist both the key and resume identity.
+    # The flag must be declared on the class that literally defines
+    # create_thread_operation_atomic: a flag-only child or mixin cannot opt in
+    # another class's method, and a new override cannot inherit a parent's opt-in.
+    # Inheriting the opted-in method unchanged preserves its capability.
+    supports_idempotency_request_kwargs = False
+
     async def list_changed(
         self,
         *,

backend/tests/AGENTS.md

@@ -22,6 +22,15 @@ Assert both writer formats: old workers reuse conflicting decisions, upgraded
 workers reject conflicts, and identity-less legacy rows fail closed even for
 identical retries. This pins why rollout requires routing to upgraded workers.
 
+Run-store capability tests in that file check keyed input separately from resume
+identity and assert that rejected admissions create no durable or local run.
+An explicit `idempotency_key` parameter must keep supporting keyed input even
+without resume support. For kwargs-only stores, pin the opt-in on the class
+that defines `create_thread_operation_atomic`: replacing an opted-in parent's
+method requires a fresh declaration; a flag-only child or mixin cannot opt in
+another class's implementation. Also cover unchanged inheritance and durable
+retry reuse, so the rejection tests cannot pass by rejecting every custom store.
+
 Upload case-collision coverage uses separate HTTP requests and preserves both
 reported payloads. Observe real filename-claim inputs to pin the disk seed;
 case-insensitive hosts can otherwise mask a missing seed through link retries.

backend/tests/test_thread_run_idempotency.py

@@ -1146,7 +1146,9 @@ async def test_keyed_input_remains_compatible_with_legacy_explicit_atomic_signat
                 idempotency_key=idempotency_key,
             )
 
-    manager = RunManager(store=LegacyAtomicMemoryStore())
+    store = LegacyAtomicMemoryStore()
+    manager = RunManager(store=store, worker_id="worker-a")
+    peer = RunManager(store=store, worker_id="worker-b")
 
     record = await manager.create_or_reject(
         "thread-1",
@@ -1154,6 +1156,12 @@ async def test_keyed_input_remains_compatible_with_legacy_explicit_atomic_signat
     )
 
     assert record.idempotency_key == "http-run:key"
+    persisted = await store.get(record.run_id)
+    assert persisted is not None
+    assert persisted["idempotency_key"] == "http-run:key"
+    reused = await peer.create_or_reject("thread-1", idempotency_key="http-run:key")
+    assert reused.run_id == record.run_id
+    assert reused.idempotency_reused is True
 
     with pytest.raises(RunIdempotencyUnsupported, match="does not support keyed resume"):
         await manager.create_or_reject(
@@ -1167,6 +1175,211 @@ async def test_keyed_input_remains_compatible_with_legacy_explicit_atomic_signat
         )
 
 
+@pytest.mark.anyio
+async def test_keyed_resume_rejects_ambiguous_kwargs_store_before_admission():
+    from deerflow.runtime import RunIdempotencyUnsupported
+
+    class DroppingKwargsStore(MemoryRunStore):
+        async def create_thread_operation_atomic(self, run_id, **kwargs):
+            kwargs.pop("idempotency_request", None)
+            return await super().create_thread_operation_atomic(run_id, **kwargs)
+
+    store = DroppingKwargsStore()
+    manager = RunManager(store=store)
+
+    with pytest.raises(RunIdempotencyUnsupported, match="does not support keyed resume"):
+        await manager.create_or_reject(
+            "thread-1",
+            idempotency_key="http-run:resume",
+            idempotency_request={
+                "version": 1,
+                "kind": "resume",
+                "sha256": "a" * 64,
+            },
+        )
+
+    assert await store.list_by_thread("thread-1") == []
+
+
+@pytest.mark.anyio
+async def test_keyed_resume_allows_explicit_kwargs_store_capability():
+    class ForwardingKwargsStore(MemoryRunStore):
+        supports_idempotency_request_kwargs = True
+
+        async def create_thread_operation_atomic(self, run_id, **kwargs):
+            return await super().create_thread_operation_atomic(run_id, **kwargs)
+
+    store = ForwardingKwargsStore()
+    manager = RunManager(store=store)
+    identity = {
+        "version": 1,
+        "kind": "resume",
+        "sha256": "a" * 64,
+    }
+
+    record = await manager.create_or_reject(
+        "thread-1",
+        idempotency_key="http-run:resume",
+        idempotency_request=identity,
+    )
+
+    persisted = await store.get(record.run_id)
+    assert persisted is not None
+    assert persisted["idempotency_request"] == identity
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("multitask_strategy", ["reject", "interrupt", "rollback"])
+async def test_keyed_input_rejects_ambiguous_kwargs_store_before_admission(multitask_strategy):
+    from deerflow.runtime import RunIdempotencyUnsupported
+
+    class DroppingKwargsStore(MemoryRunStore):
+        atomic_called = False
+
+        async def create_thread_operation_atomic(self, run_id, **kwargs):
+            self.atomic_called = True
+            kwargs.pop("idempotency_key", None)
+            return await super().create_thread_operation_atomic(run_id, **kwargs)
+
+    store = DroppingKwargsStore()
+    manager = RunManager(store=store)
+
+    with pytest.raises(RunIdempotencyUnsupported, match="does not support keyed run"):
+        await manager.create_or_reject(
+            "thread-1",
+            idempotency_key="http-run:input",
+            multitask_strategy=multitask_strategy,
+        )
+
+    assert store.atomic_called is False
+    assert await store.list_by_thread("thread-1") == []
+    assert await manager.list_by_thread("thread-1") == []
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("with_resume_identity", [False, True])
+async def test_explicit_resume_parameter_does_not_imply_key_support(with_resume_identity):
+    from deerflow.runtime import RunIdempotencyUnsupported
+
+    class DroppingKeyStore(MemoryRunStore):
+        async def create_thread_operation_atomic(self, run_id, *, idempotency_request=None, **kwargs):
+            kwargs.pop("idempotency_key", None)
+            return await super().create_thread_operation_atomic(run_id, idempotency_request=idempotency_request, **kwargs)
+
+    store = DroppingKeyStore()
+    manager = RunManager(store=store)
+    identity = {"version": 1, "kind": "resume", "sha256": "a" * 64} if with_resume_identity else None
+
+    with pytest.raises(RunIdempotencyUnsupported, match="does not support keyed run"):
+        await manager.create_or_reject("thread-1", idempotency_key="http-run:key", idempotency_request=identity)
+
+    assert await store.list_by_thread("thread-1") == []
+    assert await manager.list_by_thread("thread-1") == []
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("with_resume_identity", [False, True])
+@pytest.mark.parametrize("declaration", ["replaced_parent_method", "child_flag", "mixin_flag"])
+async def test_kwargs_store_capability_must_belong_to_implementation_owner(declaration, with_resume_identity):
+    from deerflow.runtime import RunIdempotencyUnsupported
+
+    class OptedInParent(MemoryRunStore):
+        supports_idempotency_request_kwargs = True
+
+        async def create_thread_operation_atomic(self, run_id, **kwargs):
+            return await super().create_thread_operation_atomic(run_id, **kwargs)
+
+    class ReplacingChild(OptedInParent):
+        async def create_thread_operation_atomic(self, run_id, **kwargs):
+            kwargs.pop("idempotency_key", None)
+            kwargs.pop("idempotency_request", None)
+            return await super().create_thread_operation_atomic(run_id, **kwargs)
+
+    class UnoptedParent(MemoryRunStore):
+        async def create_thread_operation_atomic(self, run_id, **kwargs):
+            kwargs.pop("idempotency_key", None)
+            kwargs.pop("idempotency_request", None)
+            return await super().create_thread_operation_atomic(run_id, **kwargs)
+
+    class FlagOnlyChild(UnoptedParent):
+        supports_idempotency_request_kwargs = True
+
+    class FlagMixin:
+        supports_idempotency_request_kwargs = True
+
+    class MixinStore(FlagMixin, UnoptedParent):
+        async def create_thread_operation_atomic(self, run_id, **kwargs):
+            return await super().create_thread_operation_atomic(run_id, **kwargs)
+
+    stores = {"replaced_parent_method": ReplacingChild, "child_flag": FlagOnlyChild, "mixin_flag": MixinStore}
+    store = stores[declaration]()
+    manager = RunManager(store=store)
+    identity = {"version": 1, "kind": "resume", "sha256": "a" * 64} if with_resume_identity else None
+
+    with pytest.raises(RunIdempotencyUnsupported):
+        await manager.create_or_reject("thread-1", idempotency_key="http-run:key", idempotency_request=identity)
+
+    assert await store.list_by_thread("thread-1") == []
+    assert await manager.list_by_thread("thread-1") == []
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("with_resume_identity", [False, True])
+async def test_kwargs_store_inherits_capability_with_unchanged_implementation(with_resume_identity):
+    class OptedInParent(MemoryRunStore):
+        supports_idempotency_request_kwargs = True
+
+        async def create_thread_operation_atomic(self, run_id, **kwargs):
+            return await super().create_thread_operation_atomic(run_id, **kwargs)
+
+    class InheritingStore(OptedInParent):
+        pass
+
+    store = InheritingStore()
+    owner = RunManager(store=store, worker_id="worker-a")
+    peer = RunManager(store=store, worker_id="worker-b")
+    identity = {"version": 1, "kind": "resume", "sha256": "a" * 64} if with_resume_identity else None
+
+    record = await owner.create_or_reject("thread-1", idempotency_key="http-run:key", idempotency_request=identity)
+    reused = await peer.create_or_reject("thread-1", idempotency_key="http-run:key", idempotency_request=identity)
+
+    assert reused.run_id == record.run_id
+    assert reused.idempotency_reused is True
+    persisted = await store.get(record.run_id)
+    assert persisted is not None
+    assert persisted["idempotency_key"] == "http-run:key"
+    assert persisted["idempotency_request"] == identity
+    assert len(await store.list_by_thread("thread-1")) == 1
+
+
+@pytest.mark.anyio
+async def test_start_run_returns_503_for_ambiguous_key_only_store(_stub_app_config):
+    class DroppingKeyStore(MemoryRunStore):
+        async def create_thread_operation_atomic(self, run_id, **kwargs):
+            kwargs.pop("idempotency_key", None)
+            return await super().create_thread_operation_atomic(run_id, **kwargs)
+
+    store = DroppingKeyStore()
+    manager = RunManager(store=store)
+
+    with (
+        patch("app.gateway.services.resolve_agent_factory", return_value=object()),
+        patch("app.gateway.services.run_agent", side_effect=AssertionError("worker must not attach")),
+        pytest.raises(HTTPException) as excinfo,
+    ):
+        await services.start_run(
+            RunCreateRequest(input={"messages": [{"role": "user", "content": "hello"}]}),
+            "thread-1",
+            _make_start_run_request(manager),
+            idempotency_key="http-run:input",
+        )
+
+    assert excinfo.value.status_code == 503
+    assert "does not support keyed run" in str(excinfo.value.detail)
+    assert await store.list_by_thread("thread-1") == []
+    assert await manager.list_by_thread("thread-1") == []
+
+
 @pytest.mark.anyio
 async def test_wait_retry_after_later_run_does_not_return_later_checkpoint(monkeypatch):
     """Complete two runs, then retry the first key: /wait must not return run B."""