diff --git a/.stats.yml b/.stats.yml index 04c8e43b9..25881e3c6 100644 --- a/.stats.yml +++ b/.stats.yml @@ -1,4 +1,4 @@ configured_endpoints: 75 -openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/sgp/agentex-sdk-644a4ec06aa1f055c614cbef3379684819a4edd84eeb20d2fb29ae01663622a3.yml -openapi_spec_hash: a6a4dc0c09691ac9783bf38e9653a464 +openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/sgp/agentex-sdk-330ce4f0d8feed6caeb73d6b12277cfd89f6ad85535b8c8a6f509743b0b6f8cb.yml +openapi_spec_hash: ed6b33682c511df6de538714c0864aa3 config_hash: 593e89b291976a5e84e4c3c3f8324354 diff --git a/src/agentex/lib/core/clients/temporal/temporal_client.py b/src/agentex/lib/core/clients/temporal/temporal_client.py index 02c44d81e..f835768a1 100644 --- a/src/agentex/lib/core/clients/temporal/temporal_client.py +++ b/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 diff --git a/src/agentex/lib/core/clients/temporal/utils.py b/src/agentex/lib/core/clients/temporal/utils.py index 95319720a..15b08cec6 100644 --- a/src/agentex/lib/core/clients/temporal/utils.py +++ b/src/agentex/lib/core/clients/temporal/utils.py @@ -9,6 +9,8 @@ from temporalio.converter import PayloadCodec, DataConverter from temporalio.contrib.pydantic import pydantic_data_converter +from agentex.lib.core.tracing.temporal import temporal_tracing_interceptors + # class DateTimeJSONEncoder(AdvancedJSONEncoder): # def default(self, o: Any) -> Any: # if isinstance(o, datetime.datetime): @@ -136,6 +138,9 @@ async def get_temporal_client( connect_kwargs: dict[str, Any] = { "target_host": temporal_address, "plugins": plugins, + # Propagate OTel trace context on outbound start_workflow / execute_activity + # (enabled by default; AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED=false to disable). + "interceptors": temporal_tracing_interceptors(), } if data_converter is not None: diff --git a/src/agentex/lib/core/temporal/services/temporal_task_service.py b/src/agentex/lib/core/temporal/services/temporal_task_service.py index 5f6c0c381..32cde071b 100644 --- a/src/agentex/lib/core/temporal/services/temporal_task_service.py +++ b/src/agentex/lib/core/temporal/services/temporal_task_service.py @@ -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 @@ -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( @@ -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: diff --git a/src/agentex/lib/core/temporal/workers/worker.py b/src/agentex/lib/core/temporal/workers/worker.py index 2b4958b1f..0cfe01185 100644 --- a/src/agentex/lib/core/temporal/workers/worker.py +++ b/src/agentex/lib/core/temporal/workers/worker.py @@ -29,6 +29,7 @@ from agentex.lib.utils.logging import make_logger from agentex.lib.utils.registration import register_agent +from agentex.lib.core.tracing.temporal import temporal_tracing_interceptors from agentex.lib.environment_variables import EnvironmentVariables from agentex.lib.core.compat.version_guard import assert_backend_compatible @@ -126,6 +127,9 @@ async def get_temporal_client( connect_kwargs: dict[str, Any] = { "target_host": temporal_address, "plugins": plugins, + # Propagate OTel trace context on outbound start_workflow / execute_activity + # (enabled by default; AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED=false to disable). + "interceptors": temporal_tracing_interceptors(), } if data_converter is not None: @@ -229,7 +233,9 @@ async def run( max_concurrent_activities=self.max_concurrent_activities, build_id=str(uuid.uuid4()), debug_mode=debug_enabled, # Disable deadlock detection in debug mode - interceptors=self.interceptors, # Pass interceptors to Worker + # Tracing interceptor OUTERMOST so business interceptors (and the spans + # they create) nest under the propagated workflow/activity span. + interceptors=[*temporal_tracing_interceptors(), *self.interceptors], ) logger.info(f"Starting workers for task queue: {self.task_queue}") diff --git a/src/agentex/lib/core/tracing/temporal.py b/src/agentex/lib/core/tracing/temporal.py new file mode 100644 index 000000000..484abc26b --- /dev/null +++ b/src/agentex/lib/core/tracing/temporal.py @@ -0,0 +1,73 @@ +"""OpenTelemetry trace-context propagation across Temporal boundaries. + +Temporal serializes ``start_workflow`` / ``execute_activity`` across (potentially +cross-process) boundaries, and does NOT carry the active W3C ``traceparent`` by +default. So any span created inside a workflow or activity becomes a **new +detached root** -- the trace shatters at every Temporal hop. + +This bites agentex directly: ``adk.tracing.span`` runs span creation as a +Temporal activity when ``in_temporal_workflow()`` is true, so without propagation +those business spans detach from the turn's obs trace. + +Wiring temporalio's first-party ``TracingInterceptor`` onto the Temporal client +and worker injects the active span context into Temporal headers on the caller +side and extracts + continues it on the workflow/activity side, using the global +OpenTelemetry propagator -- so ``client -> workflow -> activity`` is one trace. + +Enabled by DEFAULT. Set ``AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED=false`` +(also accepts ``0`` / ``no`` / ``off``) to turn it off. It also degrades to a +no-op -- and never raises -- if temporalio's OpenTelemetry contrib isn't +importable, so enabling it by default can't break a worker. +""" + +from __future__ import annotations + +import os +from typing import Any + +from agentex.lib.utils.logging import make_logger + +logger = make_logger(__name__) + +_ENABLE_ENV = "AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED" +_FALSEY = {"0", "false", "no", "off"} + + +def temporal_trace_interceptor_enabled() -> bool: + """Whether the Temporal OTel trace interceptor should be installed. + + Defaults to True; disabled only when ``AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED`` + is set to a falsy value (``0`` / ``false`` / ``no`` / ``off``).""" + return os.environ.get(_ENABLE_ENV, "true").strip().lower() not in _FALSEY + + +def temporal_tracing_interceptors() -> list[Any]: + """Interceptors that propagate OpenTelemetry trace context across Temporal. + + Returns ``[TracingInterceptor()]`` (enabled by default) so callers can splat + it into a client's / worker's ``interceptors=`` list. Returns ``[]`` when + disabled via env, or when temporalio's OpenTelemetry contrib is not + importable. Never raises -- observability wiring must not break a worker. + + ``TracingInterceptor`` implements both the client and worker interceptor + interfaces, so the same call is used on both sides: + - on the **client**, it injects context on outbound ``start_workflow`` / + ``execute_activity`` calls; + - on the **worker**, it extracts context and roots the workflow / activity + execution spans under it. + """ + if not temporal_trace_interceptor_enabled(): + logger.info("Temporal OTel trace interceptor disabled via %s", _ENABLE_ENV) + return [] + try: + from temporalio.contrib.opentelemetry import TracingInterceptor + + # Construct inside the try so a constructor failure (not just a missing + # contrib) also falls back to a no-op instead of aborting worker startup. + return [TracingInterceptor()] + except Exception as exc: # contrib unavailable OR constructor failure -> no-op, never raise + logger.warning( + "Temporal OTel trace interceptor unavailable (%s); traces will not propagate across Temporal boundaries.", + exc, + ) + return [] diff --git a/tests/lib/core/services/test_temporal_task_service.py b/tests/lib/core/services/test_temporal_task_service.py new file mode 100644 index 000000000..c5d8bd8b2 --- /dev/null +++ b/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"])) diff --git a/tests/lib/core/tracing/test_temporal_interceptor.py b/tests/lib/core/tracing/test_temporal_interceptor.py new file mode 100644 index 000000000..83c0c3681 --- /dev/null +++ b/tests/lib/core/tracing/test_temporal_interceptor.py @@ -0,0 +1,40 @@ +"""Unit tests for the Temporal OTel trace-interceptor wiring. + +Verifies the interceptor is on by default, the opt-out env flag, and the safe +no-op fallback when temporalio's OpenTelemetry contrib isn't importable. +""" + +import sys + +import pytest + +from agentex.lib.core.tracing import temporal as temporal_tracing + + +class TestTemporalTraceInterceptor: + def test_enabled_by_default(self, monkeypatch): + monkeypatch.delenv("AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED", raising=False) + assert temporal_tracing.temporal_trace_interceptor_enabled() is True + + interceptors = temporal_tracing.temporal_tracing_interceptors() + assert len(interceptors) == 1 + # temporalio's first-party OTel interceptor + assert type(interceptors[0]).__name__ == "TracingInterceptor" + + @pytest.mark.parametrize("value", ["false", "0", "no", "off", "FALSE", "Off"]) + def test_disabled_via_env(self, monkeypatch, value): + monkeypatch.setenv("AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED", value) + assert temporal_tracing.temporal_trace_interceptor_enabled() is False + assert temporal_tracing.temporal_tracing_interceptors() == [] + + @pytest.mark.parametrize("value", ["true", "1", "yes", "TRUE", "anything"]) + def test_enabled_for_non_falsy_values(self, monkeypatch, value): + monkeypatch.setenv("AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED", value) + assert temporal_tracing.temporal_trace_interceptor_enabled() is True + + def test_no_op_when_contrib_unimportable(self, monkeypatch): + # Enabled, but temporalio's OTel contrib not importable -> [] (never raises), + # so default-on can't break a worker that lacks the contrib. + monkeypatch.delenv("AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED", raising=False) + monkeypatch.setitem(sys.modules, "temporalio.contrib.opentelemetry", None) + assert temporal_tracing.temporal_tracing_interceptors() == []