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
8 changes: 7 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 Down Expand Up @@ -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,
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
from typing import Any
from datetime import timedelta

from temporalio.common import WorkflowIDConflictPolicy

from agentex.types.task import Task
from agentex.types.agent import Agent
from agentex.types.event import Event
Expand Down Expand Up @@ -42,6 +44,8 @@ async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | N
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
return await self._temporal_client.start_workflow(
workflow=self._env_vars.WORKFLOW_NAME,
arg=CreateTaskParams(
Expand All @@ -52,6 +56,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:
Expand Down
110 changes: 110 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,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"]))
Loading