fix: limit swarm callback concurrency

.github/workflows/ci.yml

@@ -14,7 +14,7 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
       - name: Install tools
         run: pip install ruff black
@@ -54,10 +54,10 @@ jobs:
     needs: [validate-branch-target]
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
       - name: Set up Python 3.13
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.13'
           cache: 'pip'
@@ -78,10 +78,10 @@ jobs:
     needs: [validate-branch-target]
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
       - name: Set up Python 3.13
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.13'
           cache: 'pip'
@@ -102,10 +102,10 @@ jobs:
     needs: [validate-branch-target]
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
       - name: Set up Python 3.13
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.13'
           cache: 'pip'
@@ -153,10 +153,10 @@ jobs:
         hatchet-version: ['1.20', '1.21', '1.22', '1.23']
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
           cache: 'pip'
@@ -188,10 +188,10 @@ jobs:
         hatchet-version: ['1.22', '1.23']
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
           cache: 'pip'
@@ -237,10 +237,10 @@ jobs:
     needs: [validate-branch-target]
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
       - name: Set up Python 3.13
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.13'
           cache: 'pip'
@@ -262,10 +262,10 @@ jobs:
     needs: [validate-branch-target]
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
       - name: Set up Python 3.13
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.13'
           cache: 'pip'
@@ -278,7 +278,7 @@ jobs:
           pip install testcontainers click
 
       - name: Setup Node.js
-        uses: actions/setup-node@v4
+        uses: actions/setup-node@v6
         with:
           node-version: lts/*
           cache: npm
@@ -306,7 +306,7 @@ jobs:
     needs: [validate-branch-target]
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
       - name: Install system dependencies
         run: |
@@ -347,10 +347,10 @@ jobs:
         backend: ['fakeredis', 'testcontainers']
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
       - name: Set up Python 3.13
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.13'
           cache: 'pip'

.github/workflows/codeql.yml

@@ -61,7 +61,7 @@ jobs:
         # your codebase is analyzed, see https://docs.github.com/en/code-security/code-scanning/creating-an-advanced-setup-for-code-scanning/codeql-code-scanning-for-compiled-languages
     steps:
     - name: Checkout repository
-      uses: actions/checkout@v4
+      uses: actions/checkout@v7
 
     # Add any setup steps before running the `github/codeql-action/init` action.
     # This includes steps like installing compilers or runtimes (`actions/setup-node`

.github/workflows/security.yml

@@ -25,15 +25,16 @@ jobs:
       name: security
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
         with:
           fetch-depth: 0
 
       - name: Run Gitleaks
-        uses: gitleaks/gitleaks-action@v2
+        uses: gitleaks/gitleaks-action@v3
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           GITLEAKS_LICENSE: ${{ secrets.GITLEAKS_LICENSE }}
+          NODE_OPTIONS: --no-deprecation
 
   summary:
     name: Security Summary

CHANGELOG.md

@@ -9,6 +9,7 @@
 ### 🐛 Fixed
 
 - **Keychain Access Error Handling**: The app now distinguishes between missing credentials (first launch) and macOS blocking keychain access (e.g. after app rename/update). A dedicated error screen guides the user to re-enter credentials instead of silently dropping into onboarding.
+- **Swarm Callback Concurrency**: Duplicate swarm done/error callbacks are now serialized per swarm item using the swarm and item IDs.
 
 ### 🔄 Changed
 

libs/mageflow/mageflow/clients/hatchet/mageflow.py

@@ -45,7 +45,7 @@
     lifespan_initialize,
     teardown_mageflow,
 )
-from mageflow.swarm.consts import SWARM_TASK_ID_PARAM_NAME
+from mageflow.swarm.consts import SWARM_ITEM_ID_PARAM_NAME, SWARM_TASK_ID_PARAM_NAME
 from mageflow.swarm.messages import (
     FillSwarmMessage,
     SwarmErrorMessage,
@@ -223,17 +223,29 @@ def init_mageflow_hatchet_tasks(self):
         on_chain_error_task = hatchet_chain_error(self.chain_error_task)
 
         # Swarm tasks
+        swarm_callback_concurrency = [
+            ConcurrencyExpression(
+                expression=(
+                    f"input.{SWARM_TASK_ID_PARAM_NAME} + ':' + "
+                    f"input.{SWARM_ITEM_ID_PARAM_NAME}"
+                ),
+                max_runs=1,
+                limit_strategy=ConcurrencyLimitStrategy.GROUP_ROUND_ROBIN,
+            )
+        ]
         swarm_done = self.hatchet.durable_task(
             name=ON_SWARM_ITEM_DONE,
             input_validator=SwarmResultsMessage,
             retries=5,
             execution_timeout=timedelta(minutes=1),
+            concurrency=swarm_callback_concurrency,
         )
         swarm_error = self.hatchet.durable_task(
             name=ON_SWARM_ITEM_ERROR,
             input_validator=SwarmErrorMessage,
             retries=5,
             execution_timeout=timedelta(minutes=5),
+            concurrency=swarm_callback_concurrency,
         )
         swarm_done = swarm_done(self.swarm_item_done)
         swarm_error = swarm_error(self.swarm_item_failed)

libs/mageflow/mageflow/clients/hatchet/workflow.py

@@ -1,7 +1,7 @@
 import functools
 import inspect
 from datetime import timedelta
-from typing import Any, TypedDict, Unpack
+from typing import Any, TypedDict, Unpack, cast
 
 from hatchet_sdk import Context
 from hatchet_sdk.rate_limit import RateLimit
@@ -122,5 +122,4 @@ def _serialize_input(self, input: Any) -> JSONSerializableMapping:
         else:
             return_field = input
 
-        full_msg = deep_merge(return_field, kwargs)
-        return super(MageflowWorkflow, self)._serialize_input(full_msg)
+        return cast(JSONSerializableMapping, deep_merge(return_field, kwargs))

libs/mageflow/mageflow/swarm/consts.py

@@ -1,2 +1,3 @@
 # Params
 SWARM_TASK_ID_PARAM_NAME = "swarm_task_id"
+SWARM_ITEM_ID_PARAM_NAME = "swarm_item_id"

libs/mageflow/mageflow/testing/_adapter.py

@@ -119,6 +119,8 @@ async def aio_result(self) -> dict:
 
 
 class TestClientAdapter(BaseClientAdapter):
+    __test__ = False
+
     def __init__(
         self,
         task_defs: dict[str, MageflowTaskDefinition] | None = None,

libs/mageflow/mageflow/testing/plugin.py

@@ -1,3 +1,5 @@
+import warnings
+
 import pytest
 import pytest_asyncio
 from thirdmagic.signature import Signature
@@ -31,7 +33,13 @@ async def _mageflow_redis_client(_mageflow_testing_config):
         yield client
         await client.aclose()
     else:
-        from testcontainers.redis import AsyncRedisContainer
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore",
+                message="The @wait_container_is_ready decorator is deprecated.*",
+                category=DeprecationWarning,
+            )
+            from testcontainers.redis import AsyncRedisContainer
 
         with AsyncRedisContainer(
             image="redis/redis-stack-server:7.2.0-v13"

libs/mageflow/pyproject.toml

@@ -125,7 +125,9 @@ directory = "htmlcov"
 markers = [
     "hatchet: tests specific to the Hatchet client integration",
 ]
+filterwarnings = [
+    "ignore:The @wait_container_is_ready decorator is deprecated.*:DeprecationWarning:testcontainers.redis",
+]
 
-[tool.ruff]
+[tool.ruff.lint]
 select = ["I", "F401"]
-

libs/mageflow/tests/unit/clients/test_hatchet_mageflow_tasks.py

@@ -0,0 +1,112 @@
+from datetime import timedelta
+
+from hatchet_sdk.runnables.types import ConcurrencyLimitStrategy
+
+from mageflow.clients.hatchet.mageflow import HatchetMageflow
+from mageflow.clients.inner_task_names import (
+    ON_CHAIN_END,
+    ON_CHAIN_ERROR,
+    ON_SWARM_ITEM_DONE,
+    ON_SWARM_ITEM_ERROR,
+    SWARM_FILL_TASK,
+)
+from mageflow.swarm.consts import SWARM_ITEM_ID_PARAM_NAME, SWARM_TASK_ID_PARAM_NAME
+
+
+def test_internal_swarm_callback_tasks_are_serialized_per_swarm_item(
+    hatchet_mock, redis_client, monkeypatch
+):
+    expected_task_count = 5
+    expected_registered_task_names = {
+        ON_CHAIN_END,
+        ON_CHAIN_ERROR,
+        ON_SWARM_ITEM_DONE,
+        ON_SWARM_ITEM_ERROR,
+        SWARM_FILL_TASK,
+    }
+    expected_expression = (
+        f"input.{SWARM_TASK_ID_PARAM_NAME} + ':' + input.{SWARM_ITEM_ID_PARAM_NAME}"
+    )
+    expected_callback_task_names = (ON_SWARM_ITEM_DONE, ON_SWARM_ITEM_ERROR)
+    expected_concurrency_count = 1
+    expected_max_runs = 1
+    expected_limit_strategy = ConcurrencyLimitStrategy.GROUP_ROUND_ROBIN
+    expected_done_timeout = timedelta(minutes=1)
+    expected_error_timeout = timedelta(minutes=5)
+
+    # Arrange
+    registered_tasks = {}
+
+    def fake_durable_task(**kwargs):
+        registered_tasks[kwargs["name"]] = kwargs
+
+        def decorator(func):
+            return func
+
+        return decorator
+
+    monkeypatch.setattr(hatchet_mock, "durable_task", fake_durable_task)
+    client = HatchetMageflow(hatchet=hatchet_mock, redis_client=redis_client)
+
+    # Act
+    tasks = client.init_mageflow_hatchet_tasks()
+
+    # Assert
+    assert len(tasks) == expected_task_count
+    assert set(registered_tasks) == expected_registered_task_names
+
+    for task_name in expected_callback_task_names:
+        concurrency = registered_tasks[task_name]["concurrency"]
+
+        assert len(concurrency) == expected_concurrency_count
+        assert concurrency[0].expression == expected_expression
+        assert concurrency[0].max_runs == expected_max_runs
+        assert concurrency[0].limit_strategy == expected_limit_strategy
+
+    assert (
+        registered_tasks[ON_SWARM_ITEM_DONE]["execution_timeout"]
+        == expected_done_timeout
+    )
+    assert (
+        registered_tasks[ON_SWARM_ITEM_ERROR]["execution_timeout"]
+        == expected_error_timeout
+    )
+
+
+def test_internal_swarm_fill_task_keeps_existing_concurrency_gate(
+    hatchet_mock, redis_client, monkeypatch
+):
+    expected_concurrency_count = 2
+    expected_expression = f"input.{SWARM_TASK_ID_PARAM_NAME}"
+    expected_cancel_newest_max_runs = 2
+    expected_cancel_newest_strategy = ConcurrencyLimitStrategy.CANCEL_NEWEST
+    expected_round_robin_max_runs = 1
+    expected_round_robin_strategy = ConcurrencyLimitStrategy.GROUP_ROUND_ROBIN
+
+    # Arrange
+    registered_tasks = {}
+
+    def fake_durable_task(**kwargs):
+        registered_tasks[kwargs["name"]] = kwargs
+
+        def decorator(func):
+            return func
+
+        return decorator
+
+    monkeypatch.setattr(hatchet_mock, "durable_task", fake_durable_task)
+    client = HatchetMageflow(hatchet=hatchet_mock, redis_client=redis_client)
+
+    # Act
+    client.init_mageflow_hatchet_tasks()
+
+    # Assert
+    concurrency = registered_tasks[SWARM_FILL_TASK]["concurrency"]
+
+    assert len(concurrency) == expected_concurrency_count
+    assert concurrency[0].expression == expected_expression
+    assert concurrency[0].max_runs == expected_cancel_newest_max_runs
+    assert concurrency[0].limit_strategy == expected_cancel_newest_strategy
+    assert concurrency[1].expression == expected_expression
+    assert concurrency[1].max_runs == expected_round_robin_max_runs
+    assert concurrency[1].limit_strategy == expected_round_robin_strategy