handle same workflow task/create gracefully (WorkflowAlreadyStartedError)

Author: alvinkam2001

src/agentex/lib/core/clients/temporal/temporal_client.py

@@ -5,7 +5,11 @@
 from collections.abc import Callable
 
 from temporalio.client import Client, WorkflowExecutionStatus
-from temporalio.common import RetryPolicy as TemporalRetryPolicy, WorkflowIDReusePolicy
+from temporalio.common import (
+    RetryPolicy as TemporalRetryPolicy,
+    WorkflowIDReusePolicy,
+    WorkflowIDConflictPolicy,
+)
 from temporalio.service import RPCError, RPCStatusCode
 from temporalio.converter import PayloadCodec, DataConverter
 
@@ -151,6 +155,7 @@ async def start_workflow(
         self,
         *args: Any,
         duplicate_policy: DuplicateWorkflowPolicy = DuplicateWorkflowPolicy.ALLOW_DUPLICATE,
+        id_conflict_policy: WorkflowIDConflictPolicy = WorkflowIDConflictPolicy.UNSPECIFIED,
         retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
         task_timeout: timedelta = timedelta(seconds=10),
         execution_timeout: timedelta | None = None,
@@ -163,6 +168,7 @@ async def start_workflow(
             task_timeout=task_timeout,
             execution_timeout=execution_timeout,
             id_reuse_policy=DUPLICATE_POLICY_TO_ID_REUSE_POLICY[duplicate_policy],
+            id_conflict_policy=id_conflict_policy,
             **kwargs,
         )
         return workflow_handle.id

src/agentex/lib/core/temporal/services/temporal_task_service.py

@@ -6,6 +6,8 @@
 from contextlib import contextmanager
 from collections.abc import Iterator
 
+from temporalio.common import WorkflowIDConflictPolicy
+
 from agentex.types.task import Task
 from agentex.types.agent import Agent
 from agentex.types.event import Event
@@ -89,6 +91,8 @@ async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | N
         # value bounds the whole continue-as-new chain's wall-clock lifetime.
         timeout_seconds = self._env_vars.WORKFLOW_EXECUTION_TIMEOUT_SECONDS
         execution_timeout = timedelta(seconds=timeout_seconds) if timeout_seconds and timeout_seconds > 0 else None
+        # USE_EXISTING makes task/create idempotent
+        # If same task ID is already running Temporal returns a handle to the existing run instead of raising WorkflowAlreadyStarted
         with _acp_dispatch_span("acp.task_create", task_id=task.id):
             return await self._temporal_client.start_workflow(
                 workflow=self._env_vars.WORKFLOW_NAME,
@@ -100,6 +104,7 @@ async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | N
                 id=task.id,
                 task_queue=self._env_vars.WORKFLOW_TASK_QUEUE,
                 execution_timeout=execution_timeout,
+                id_conflict_policy=WorkflowIDConflictPolicy.USE_EXISTING,
             )
 
     async def get_state(self, task_id: str) -> WorkflowState:

tests/lib/core/services/test_temporal_task_service.py

@@ -0,0 +1,110 @@
+"""Unit tests for TemporalTaskService idempotency behavior.
+
+Covers the ``task/create`` idempotency guarantee: duplicate submits for the
+same task ID must not raise ``WorkflowAlreadyStartedError``. The service
+achieves this by passing ``WorkflowIDConflictPolicy.USE_EXISTING`` to Temporal,
+which returns a handle to the existing run instead of erroring.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import Mock, AsyncMock
+
+import pytest
+from temporalio.common import WorkflowIDConflictPolicy
+
+from agentex.types.task import Task
+from agentex.types.agent import Agent
+from agentex.lib.core.clients.temporal.temporal_client import TemporalClient
+from agentex.lib.core.temporal.services.temporal_task_service import TemporalTaskService
+
+
+def _agent() -> Agent:
+    return Agent(
+        id="test-agent-456",
+        name="test-agent",
+        description="test-agent",
+        acp_type="async",
+        created_at="2023-01-01T00:00:00Z",
+        updated_at="2023-01-01T00:00:00Z",
+    )
+
+
+def _task() -> Task:
+    return Task(id="test-task-123", status="RUNNING")
+
+
+def _env_vars() -> Mock:
+    env_vars = Mock()
+    env_vars.WORKFLOW_NAME = "test-workflow"
+    env_vars.WORKFLOW_TASK_QUEUE = "test-queue"
+    env_vars.WORKFLOW_EXECUTION_TIMEOUT_SECONDS = 0
+    return env_vars
+
+
+class TestSubmitTaskIdempotency:
+    async def test_submit_task_uses_use_existing_conflict_policy(self) -> None:
+        """Duplicate task/create must be idempotent.
+
+        Passing ``WorkflowIDConflictPolicy.USE_EXISTING`` tells Temporal to
+        return the existing workflow handle instead of raising
+        ``WorkflowAlreadyStartedError`` when a run with that ID is already
+        active. Without this, load-balanced agentex-agent replicas racing on
+        the same task ID surface Temporal's start conflict as an error log.
+        """
+        temporal_client = Mock()
+        temporal_client.start_workflow = AsyncMock(return_value="test-task-123")
+
+        service = TemporalTaskService(temporal_client=temporal_client, env_vars=_env_vars())
+
+        result = await service.submit_task(agent=_agent(), task=_task(), params=None)
+
+        temporal_client.start_workflow.assert_awaited_once()
+        kwargs = temporal_client.start_workflow.await_args.kwargs
+        assert kwargs["id_conflict_policy"] == WorkflowIDConflictPolicy.USE_EXISTING
+        assert kwargs["id"] == "test-task-123"
+        assert result == "test-task-123"
+
+
+class TestTemporalClientConflictPolicyPlumbing:
+    """Boundary tests: TemporalClient.start_workflow must forward
+    ``id_conflict_policy`` to the underlying temporalio client, and default
+    to ``UNSPECIFIED`` so callers that don't opt in keep their current
+    behavior (Temporal server treats UNSPECIFIED as FAIL on start).
+    """
+
+    async def test_forwards_id_conflict_policy_when_set(self) -> None:
+        inner_client = Mock()
+        inner_handle = Mock()
+        inner_handle.id = "wf-1"
+        inner_client.start_workflow = AsyncMock(return_value=inner_handle)
+
+        tc = TemporalClient(temporal_client=inner_client)
+
+        await tc.start_workflow(
+            workflow="w",
+            arg={},
+            id="id-1",
+            task_queue="q",
+            id_conflict_policy=WorkflowIDConflictPolicy.USE_EXISTING,
+        )
+
+        kwargs = inner_client.start_workflow.await_args.kwargs
+        assert kwargs["id_conflict_policy"] == WorkflowIDConflictPolicy.USE_EXISTING
+
+    async def test_default_conflict_policy_is_unspecified(self) -> None:
+        inner_client = Mock()
+        inner_handle = Mock()
+        inner_handle.id = "wf-1"
+        inner_client.start_workflow = AsyncMock(return_value=inner_handle)
+
+        tc = TemporalClient(temporal_client=inner_client)
+
+        await tc.start_workflow(workflow="w", arg={}, id="id-1", task_queue="q")
+
+        kwargs = inner_client.start_workflow.await_args.kwargs
+        assert kwargs["id_conflict_policy"] == WorkflowIDConflictPolicy.UNSPECIFIED
+
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(pytest.main([__file__, "-v"]))