Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 23 additions & 1 deletion src/agentex/lib/core/clients/temporal/temporal_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -15,6 +19,7 @@
TaskStatus,
RetryPolicy,
WorkflowState,
ConflictWorkflowPolicy,
DuplicateWorkflowPolicy,
)
from agentex.lib.core.clients.temporal.utils import get_temporal_client
Expand Down Expand Up @@ -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__(
Expand Down Expand Up @@ -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
Expand Down
7 changes: 7 additions & 0 deletions src/agentex/lib/core/clients/temporal/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand All @@ -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:
Expand Down
140 changes: 140 additions & 0 deletions tests/lib/core/services/test_temporal_task_service.py
Original file line number Diff line number Diff line change
@@ -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"]))
Loading