fix(task-create): handle WorkflowAlreadyStartedError gracefully (#489)

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
 
@@ -15,6 +19,7 @@
     TaskStatus,
     RetryPolicy,
     WorkflowState,
+    ConflictWorkflowPolicy,
     DuplicateWorkflowPolicy,
 )
 from agentex.lib.core.clients.temporal.utils import get_temporal_client
@@ -75,6 +80,13 @@
     DuplicateWorkflowPolicy.TERMINATE_IF_RUNNING: WorkflowIDReusePolicy.TERMINATE_IF_RUNNING,
 }
 
+CONFLICT_POLICY_TO_ID_CONFLICT_POLICY = {
+    ConflictWorkflowPolicy.UNSPECIFIED: WorkflowIDConflictPolicy.UNSPECIFIED,
+    ConflictWorkflowPolicy.FAIL: WorkflowIDConflictPolicy.FAIL,
+    ConflictWorkflowPolicy.USE_EXISTING: WorkflowIDConflictPolicy.USE_EXISTING,
+    ConflictWorkflowPolicy.TERMINATE_EXISTING: WorkflowIDConflictPolicy.TERMINATE_EXISTING,
+}
+
 
 class TemporalClient:
     def __init__(
@@ -151,18 +163,28 @@ async def start_workflow(
         self,
         *args: Any,
         duplicate_policy: DuplicateWorkflowPolicy = DuplicateWorkflowPolicy.ALLOW_DUPLICATE,
+        conflict_policy: ConflictWorkflowPolicy = ConflictWorkflowPolicy.UNSPECIFIED,
         retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
         task_timeout: timedelta = timedelta(seconds=10),
         execution_timeout: timedelta | None = None,
         **kwargs: Any,
     ) -> str:
+        if (
+            duplicate_policy == DuplicateWorkflowPolicy.TERMINATE_IF_RUNNING
+            and conflict_policy != ConflictWorkflowPolicy.UNSPECIFIED
+        ):
+            raise ValueError(
+                "conflict_policy cannot be set when duplicate_policy is TERMINATE_IF_RUNNING; "
+                "use ConflictWorkflowPolicy.TERMINATE_EXISTING instead"
+            )
         temporal_retry_policy = TemporalRetryPolicy(**retry_policy.model_dump(exclude_unset=True))
         workflow_handle = await self.client.start_workflow(
             *args,
             retry_policy=temporal_retry_policy,
             task_timeout=task_timeout,
             execution_timeout=execution_timeout,
             id_reuse_policy=DUPLICATE_POLICY_TO_ID_REUSE_POLICY[duplicate_policy],
+            id_conflict_policy=CONFLICT_POLICY_TO_ID_CONFLICT_POLICY[conflict_policy],
             **kwargs,
         )
         return workflow_handle.id

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

@@ -40,6 +40,13 @@ class DuplicateWorkflowPolicy(str, Enum):
     TERMINATE_IF_RUNNING = "TERMINATE_IF_RUNNING"
 
 
+class ConflictWorkflowPolicy(str, Enum):
+    UNSPECIFIED = "UNSPECIFIED"
+    FAIL = "FAIL"
+    USE_EXISTING = "USE_EXISTING"
+    TERMINATE_EXISTING = "TERMINATE_EXISTING"
+
+
 class TaskStatus(str, Enum):
     CANCELED = "CANCELED"
     COMPLETED = "COMPLETED"

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

@@ -11,7 +11,7 @@
 from agentex.types.event import Event
 from agentex.protocol.acp import SendEventParams, CreateTaskParams, InterruptTaskParams
 from agentex.lib.environment_variables import EnvironmentVariables
-from agentex.lib.core.clients.temporal.types import WorkflowState
+from agentex.lib.core.clients.temporal.types import WorkflowState, ConflictWorkflowPolicy
 from agentex.lib.core.temporal.types.workflow import SignalName
 from agentex.lib.core.clients.temporal.temporal_client import TemporalClient
 
@@ -89,6 +89,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 +102,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,
+                conflict_policy=ConflictWorkflowPolicy.USE_EXISTING,
             )
 
     async def get_state(self, task_id: str) -> WorkflowState:

tests/lib/core/services/test_temporal_task_service.py

@@ -0,0 +1,140 @@
+"""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 ``ConflictWorkflowPolicy.USE_EXISTING`` through the
+``TemporalClient`` wrapper, which maps to ``WorkflowIDConflictPolicy.USE_EXISTING``
+on the underlying temporalio client so Temporal 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.types import (
+    ConflictWorkflowPolicy,
+    DuplicateWorkflowPolicy,
+)
+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 ``ConflictWorkflowPolicy.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["conflict_policy"] == ConflictWorkflowPolicy.USE_EXISTING
+        assert kwargs["id"] == "test-task-123"
+        assert result == "test-task-123"
+
+
+class TestTemporalClientConflictPolicyPlumbing:
+    """Boundary tests: ``TemporalClient.start_workflow`` wraps
+    ``WorkflowIDConflictPolicy`` in a local ``ConflictWorkflowPolicy`` enum
+    (mirroring the existing ``DuplicateWorkflowPolicy`` pattern) so SDK users
+    don't have to import ``temporalio.common`` to opt into non-default behavior.
+    Also guards the incompatible ``TERMINATE_IF_RUNNING`` + explicit-conflict
+    combo client-side rather than letting it round-trip to the frontend as
+    ``InvalidArgument``.
+    """
+
+    async def test_forwards_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",
+            conflict_policy=ConflictWorkflowPolicy.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
+
+    async def test_terminate_if_running_with_explicit_conflict_policy_raises(self) -> None:
+        """temporalio rejects this combo at the frontend as InvalidArgument;
+        we fail fast client-side with a clearer message.
+        """
+        inner_client = Mock()
+        inner_client.start_workflow = AsyncMock()
+
+        tc = TemporalClient(temporal_client=inner_client)
+
+        with pytest.raises(ValueError, match="TERMINATE_EXISTING"):
+            await tc.start_workflow(
+                workflow="w",
+                arg={},
+                id="id-1",
+                task_queue="q",
+                duplicate_policy=DuplicateWorkflowPolicy.TERMINATE_IF_RUNNING,
+                conflict_policy=ConflictWorkflowPolicy.USE_EXISTING,
+            )
+
+        inner_client.start_workflow.assert_not_awaited()
+
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(pytest.main([__file__, "-v"]))