From d718f788899243775f2bd0d07adc149025b9f621 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Wed, 7 Oct 2026 08:04:28 -0400 Subject: [PATCH 01/13] feat(tracing): stamp sgp evals run and row ids on SGP spans Spans of a task whose task_metadata carries sgp_evals now get flat sgp_evals_generation_run_id, sgp_evals_row_id and sgp_evals_attempt_idx keys on the SGP copy. Sync and async agents read the task in the ACP server. Temporal agents carry it in the workflow memo and activity headers to the worker. Co-Authored-By: Claude Sonnet 5.5 --- .../services/temporal_task_service.py | 5 + .../lib/core/temporal/workers/worker.py | 3 +- .../processors/sgp_tracing_processor.py | 11 +- src/agentex/lib/core/tracing/sgp_evals.py | 72 ++++++ .../lib/core/tracing/sgp_evals_interceptor.py | 70 ++++++ .../lib/sdk/fastacp/base/base_acp_server.py | 5 + tests/lib/core/tracing/test_sgp_evals.py | 219 ++++++++++++++++++ 7 files changed, 380 insertions(+), 5 deletions(-) create mode 100644 src/agentex/lib/core/tracing/sgp_evals.py create mode 100644 src/agentex/lib/core/tracing/sgp_evals_interceptor.py create mode 100644 tests/lib/core/tracing/test_sgp_evals.py 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 774de4d9e..f13bc3eca 100644 --- a/src/agentex/lib/core/temporal/services/temporal_task_service.py +++ b/src/agentex/lib/core/temporal/services/temporal_task_service.py @@ -10,6 +10,7 @@ from agentex.types.agent import Agent from agentex.types.event import Event from agentex.protocol.acp import SendEventParams, CreateTaskParams, InterruptTaskParams +from agentex.lib.core.tracing import sgp_evals from agentex.lib.environment_variables import EnvironmentVariables from agentex.lib.core.clients.temporal.types import WorkflowState, ConflictWorkflowPolicy from agentex.lib.core.temporal.types.workflow import SignalName @@ -91,6 +92,9 @@ async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | N 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 + # Eval tasks carry their span attrs in the memo so the worker can stamp them. + eval_attrs = sgp_evals.span_attrs_from_task_metadata(task.task_metadata) + memo_kwargs: dict[str, Any] = {"memo": {sgp_evals.MEMO_KEY: eval_attrs}} if eval_attrs else {} with _acp_dispatch_span("acp.task_create", task_id=task.id): return await self._temporal_client.start_workflow( workflow=self._env_vars.WORKFLOW_NAME, @@ -103,6 +107,7 @@ async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | N task_queue=self._env_vars.WORKFLOW_TASK_QUEUE, execution_timeout=execution_timeout, conflict_policy=ConflictWorkflowPolicy.USE_EXISTING, + **memo_kwargs, ) 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 ba8f87de5..de2904c41 100644 --- a/src/agentex/lib/core/temporal/workers/worker.py +++ b/src/agentex/lib/core/temporal/workers/worker.py @@ -34,6 +34,7 @@ from agentex.lib.core.tracing.span_queue import shutdown_default_span_queue from agentex.lib.core.compat.version_guard import assert_backend_compatible from agentex.lib.core.observability.sgp_obs_setup import init_sgp_obs, shutdown_sgp_obs +from agentex.lib.core.tracing.sgp_evals_interceptor import SGPEvalsInterceptor from agentex.lib.core.tracing.tracing_processor_manager import shutdown_sync_tracing_processors logger = make_logger(__name__) @@ -274,7 +275,7 @@ async def run( build_id=str(uuid.uuid4()), debug_mode=debug_enabled, # Disable deadlock detection in debug mode # Temporal inherits client tracing before these business interceptors. - interceptors=self.interceptors, + interceptors=[*self.interceptors, SGPEvalsInterceptor()], ) logger.info(f"Starting workers for task queue: {self.task_queue}") diff --git a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index 9ee269231..7d47f84fd 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -11,7 +11,7 @@ from scale_gp_beta.lib.tracing.span import Span as SGPSpan from agentex.types.span import Span -from agentex.lib.core.tracing import code_revision +from agentex.lib.core.tracing import sgp_evals, code_revision from agentex.lib.types.tracing import SGPTracingProcessorConfig from agentex.lib.utils.logging import make_logger from agentex.lib.core.observability import tracing_metrics_recording as _metrics @@ -71,7 +71,7 @@ def _add_source_to_span(span: Span, env_vars: EnvironmentVariables) -> None: def _sgp_metadata(span: Span) -> Any: - """Metadata for the SGP write: ``span.data`` plus the opt-in commit SHA. + """Metadata for the SGP write: ``span.data`` plus the opt-in commit SHA and evals run/row ids. Returns a COPY rather than mutating ``span``. ``trace.py`` hands the same Span instance to every registered processor, so anything written onto @@ -83,11 +83,14 @@ def _sgp_metadata(span: Span) -> Any: leak like that today. Left as-is: changing five long-shipped fields is not this change's business.) """ + extra: dict[str, Any] = dict(sgp_evals.attrs_for_span(span)) commit_sha = code_revision.commit_sha() - if commit_sha is None: + if commit_sha is not None: + extra[code_revision.COMMIT_SHA_KEY] = commit_sha + if not extra: return span.data if isinstance(span.data, dict): - return {**span.data, code_revision.COMMIT_SHA_KEY: commit_sha} + return {**span.data, **extra} # List-shaped data is an accepted `data` shape and has nowhere to put a # metadata key; leave it untouched rather than dropping the caller's data. return span.data diff --git a/src/agentex/lib/core/tracing/sgp_evals.py b/src/agentex/lib/core/tracing/sgp_evals.py new file mode 100644 index 000000000..e77ab99d1 --- /dev/null +++ b/src/agentex/lib/core/tracing/sgp_evals.py @@ -0,0 +1,72 @@ +"""Run/row attribution for spans of SGP evals generation-unit tasks. + +The evals service tags each unit's task with ``task_metadata`` carrying +``sgp_evals`` plus the run, row and attempt it belongs to. The ids are copied +onto the SGP copy of every span as flat ``sgp_evals_*`` keys so a spans search +on ``extra_metadata`` finds a unit's spans. Flat because sgp-traces treats +dotted keys differently on ClickHouse and Postgres. +""" + +from __future__ import annotations + +import threading +from typing import Any +from collections import OrderedDict + +from agentex.types.span import Span + +__all__ = ("MEMO_KEY", "register_task", "register_task_metadata", "span_attrs_from_task_metadata", "attrs_for_span") + +TASK_METADATA_MARKER = "sgp_evals" +SPAN_KEY_PREFIX = "sgp_evals_" +_TASK_METADATA_KEYS = ("generation_run_id", "row_id", "attempt_idx") +# Temporal workflow memo key the ACP server sets so the worker process can stamp the same attrs. +MEMO_KEY = "sgp_evals_span_attrs" + +# Only eval tasks are stored, so this stays tiny. The LRU bound caps a long-lived agent process. +_MAX_TASKS = 10_000 +_attrs_by_task: OrderedDict[str, dict[str, Any]] = OrderedDict() +_lock = threading.Lock() + + +def span_attrs_from_task_metadata(task_metadata: Any) -> dict[str, Any] | None: + """The flat span attrs for an evals generation-unit task, or None for any other task.""" + if not isinstance(task_metadata, dict) or task_metadata.get(TASK_METADATA_MARKER) is None: + return None + attrs = { + f"{SPAN_KEY_PREFIX}{key}": task_metadata[key] for key in _TASK_METADATA_KEYS if task_metadata.get(key) is not None + } + return attrs or None + + +def register_task(task_id: str, attrs: dict[str, Any]) -> None: + with _lock: + _attrs_by_task[task_id] = attrs + _attrs_by_task.move_to_end(task_id) + while len(_attrs_by_task) > _MAX_TASKS: + _attrs_by_task.popitem(last=False) + + +def register_task_metadata(task_id: str, task_metadata: Any) -> dict[str, Any] | None: + """Remember an eval task's span attrs. No-op (and no lookup cost) for every other task.""" + attrs = span_attrs_from_task_metadata(task_metadata) + if attrs is not None: + register_task(task_id, attrs) + return attrs + + +def attrs_for_span(span: Span) -> dict[str, Any]: + """Attrs for the task a span belongs to, found by ``span.task_id`` then ``span.trace_id``.""" + if not _attrs_by_task: + return {} + with _lock: + for key in (span.task_id, span.trace_id): + if key and key in _attrs_by_task: + return dict(_attrs_by_task[key]) + return {} + + +def clear() -> None: + """Reset the registry (test isolation).""" + with _lock: + _attrs_by_task.clear() diff --git a/src/agentex/lib/core/tracing/sgp_evals_interceptor.py b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py new file mode 100644 index 000000000..49b8a4e2b --- /dev/null +++ b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py @@ -0,0 +1,70 @@ +"""Carry an eval task's span attrs from its Temporal workflow memo into the worker's activities. + +Spans are emitted by activities in the worker process, which never sees the ACP +server's task. The ACP server puts the attrs in the workflow memo, the outbound +interceptor copies them onto each activity's headers, and the activity +interceptor registers them for the task so the SGP processor can stamp spans. +Non-eval workflows have no memo entry and get no header. +""" + +from __future__ import annotations + +from typing import Any, override + +from temporalio import activity, workflow +from temporalio.worker import ( + Interceptor, + StartActivityInput, + ExecuteActivityInput, + ActivityInboundInterceptor, + WorkflowInboundInterceptor, + WorkflowOutboundInterceptor, +) +from temporalio.converter import default + +from agentex.lib.utils.logging import make_logger +from agentex.lib.core.tracing.sgp_evals import MEMO_KEY, register_task + +logger = make_logger(__name__) + +ATTRS_HEADER = "sgp-evals-span-attrs" +_converter = default().payload_converter + + +class SGPEvalsInterceptor(Interceptor): + @override + def intercept_activity(self, next: ActivityInboundInterceptor) -> ActivityInboundInterceptor: + return _ActivityInbound(next) + + @override + def workflow_interceptor_class(self, input: Any) -> type[WorkflowInboundInterceptor] | None: + return _WorkflowInbound + + +class _WorkflowInbound(WorkflowInboundInterceptor): + @override + def init(self, outbound: WorkflowOutboundInterceptor) -> None: + super().init(_WorkflowOutbound(outbound)) + + +class _WorkflowOutbound(WorkflowOutboundInterceptor): + @override + def start_activity(self, input: StartActivityInput) -> workflow.ActivityHandle[Any]: + attrs = workflow.memo_value(MEMO_KEY, default=None) + if isinstance(attrs, dict) and attrs: + input.headers = {**input.headers, ATTRS_HEADER: _converter.to_payload(attrs)} + return super().start_activity(input) + + +class _ActivityInbound(ActivityInboundInterceptor): + @override + async def execute_activity(self, input: ExecuteActivityInput) -> Any: + payload = input.headers.get(ATTRS_HEADER) + # The workflow id is the task id (see TemporalTaskService.submit_task). + task_id = activity.info().workflow_id + if payload is not None and task_id: + try: + register_task(task_id, _converter.from_payload(payload, dict)) + except Exception: + logger.warning("failed to read sgp evals span attrs from activity headers", exc_info=True) + return await super().execute_activity(input) diff --git a/src/agentex/lib/sdk/fastacp/base/base_acp_server.py b/src/agentex/lib/sdk/fastacp/base/base_acp_server.py index 50c304c92..bec8ca787 100644 --- a/src/agentex/lib/sdk/fastacp/base/base_acp_server.py +++ b/src/agentex/lib/sdk/fastacp/base/base_acp_server.py @@ -24,6 +24,7 @@ SendMessageParams, InterruptTaskParams, ) +from agentex.lib.core.tracing import sgp_evals from agentex.lib.utils.logging import make_logger, ctx_var_request_id from agentex.protocol.json_rpc import JSONRPCError, JSONRPCRequest, JSONRPCResponse from agentex.lib.utils.model_utils import BaseModel @@ -351,6 +352,10 @@ async def _handle_jsonrpc(self, request: Request): params_data["request"] = {"headers": custom_headers} params = params_model.model_validate(params_data) + task = getattr(params, "task", None) + if task is not None: + sgp_evals.register_task_metadata(task.id, task.task_metadata) + if method in RPC_SYNC_METHODS: handler = self._handlers[method] result = await handler(params) diff --git a/tests/lib/core/tracing/test_sgp_evals.py b/tests/lib/core/tracing/test_sgp_evals.py new file mode 100644 index 000000000..7abeb9589 --- /dev/null +++ b/tests/lib/core/tracing/test_sgp_evals.py @@ -0,0 +1,219 @@ +"""Spans of an SGP evals generation-unit task carry flat ``sgp_evals_*`` keys on the SGP copy. + +Expectations come from the contract with sgp-evaluations: the task is created with +``task_metadata = {"sgp_evals": "generation-unit", "generation_run_id", "row_id", "attempt_idx"}`` +and its spans must be findable by ``sgp_evals_generation_run_id``, ``sgp_evals_row_id`` and +``sgp_evals_attempt_idx``. +""" + +from __future__ import annotations + +import uuid +from typing import Any +from datetime import UTC, datetime +from unittest.mock import Mock, AsyncMock, patch + +import pytest +from temporalio.worker import StartActivityInput + +from agentex.types.span import Span +from agentex.types.task import Task +from agentex.types.agent import Agent +from agentex.protocol.acp import RPCMethod, CreateTaskParams, SendMessageParams +from agentex.lib.core.tracing import sgp_evals, sgp_evals_interceptor as interceptor +from agentex.types.task_message_content import TextContent +from agentex.lib.sdk.fastacp.impl.sync_acp import SyncACP +from agentex.lib.core.clients.temporal.types import ConflictWorkflowPolicy +from agentex.lib.core.temporal.services.temporal_task_service import TemporalTaskService +from agentex.lib.core.tracing.processors.sgp_tracing_processor import _sgp_metadata + +EVAL_METADATA = { + "sgp_evals": "generation-unit", + "generation_run_id": "run-42", + "row_id": "row-7", + "attempt_idx": 2, +} +EXPECTED_ATTRS = { + "sgp_evals_generation_run_id": "run-42", + "sgp_evals_row_id": "row-7", + "sgp_evals_attempt_idx": 2, +} + + +@pytest.fixture(autouse=True) +def _clean_registry(): + sgp_evals.clear() + yield + sgp_evals.clear() + + +def _span(trace_id: str = "task-1", task_id: str | None = None, data: Any = None) -> Span: + return Span( + id=str(uuid.uuid4()), + name="s", + start_time=datetime.now(UTC), + trace_id=trace_id, + task_id=task_id, + data=data, + ) + + +def _agent() -> Agent: + return Agent( + id="a1", + name="a", + description="a", + acp_type="async", + created_at="2023-01-01T00:00:00Z", + updated_at="2023-01-01T00:00:00Z", + ) + + +class TestAttrsFromTaskMetadata: + def test_eval_task_yields_flat_keys(self) -> None: + assert sgp_evals.span_attrs_from_task_metadata(EVAL_METADATA) == EXPECTED_ATTRS + + @pytest.mark.parametrize("metadata", [None, {}, {"generation_run_id": "run-42"}, {"other": 1}, "x"]) + def test_non_eval_task_yields_nothing(self, metadata: Any) -> None: + assert sgp_evals.span_attrs_from_task_metadata(metadata) is None + + +class TestSGPMetadata: + def test_registered_task_stamps_sgp_copy_not_span_data(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + span = _span(data={"caller": "kept"}) + + metadata = _sgp_metadata(span) + + assert metadata == {"caller": "kept", **EXPECTED_ATTRS} + assert span.data == {"caller": "kept"} + + def test_matches_on_span_task_id_when_trace_id_differs(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + assert _sgp_metadata(_span(trace_id="other", task_id="task-1", data={}))["sgp_evals_row_id"] == "row-7" + + def test_other_tasks_and_non_eval_tasks_are_untouched(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + sgp_evals.register_task_metadata("task-2", {"unrelated": True}) + assert _sgp_metadata(_span(trace_id="task-2", data={"k": 1})) == {"k": 1} + assert _sgp_metadata(_span(trace_id="task-3", data={"k": 1})) == {"k": 1} + + def test_registry_is_bounded(self) -> None: + with patch.object(sgp_evals, "_MAX_TASKS", 2): + for i in range(3): + sgp_evals.register_task_metadata(f"t{i}", EVAL_METADATA) + assert sgp_evals.attrs_for_span(_span(trace_id="t0")) == {} + assert sgp_evals.attrs_for_span(_span(trace_id="t2")) == EXPECTED_ATTRS + + +class TestAcpServer: + """Sync and async (base) agents: the ACP server sees the task and spans are emitted in-process.""" + + async def _dispatch(self, acp: SyncACP, method: RPCMethod, params: Any) -> Any: + class _Req: + headers: dict[str, str] = {} + + async def json(self) -> dict[str, Any]: + return {"jsonrpc": "2.0", "method": method.value, "params": params.model_dump(mode="json"), "id": "1"} + + return await acp._handle_jsonrpc(_Req()) # pyright: ignore[reportArgumentType] + + async def test_task_create_registers_eval_task(self) -> None: + acp = SyncACP() + seen: list[Any] = [] + + @acp.on_task_create + async def handler(params: CreateTaskParams) -> None: + seen.append(_sgp_metadata(_span(trace_id=params.task.id, data={}))) + + task = Task(id="task-1", status="RUNNING", task_metadata=EVAL_METADATA) + await self._dispatch(acp, RPCMethod.TASK_CREATE, CreateTaskParams(agent=_agent(), task=task, params=None)) + + assert seen == [EXPECTED_ATTRS] + + async def test_message_send_registers_eval_task(self) -> None: + acp = SyncACP() + seen: list[Any] = [] + + @acp.on_message_send + async def handler(params: SendMessageParams) -> Any: + seen.append(_sgp_metadata(_span(trace_id=params.task.id, data={}))) + return None + + task = Task(id="task-1", status="RUNNING", task_metadata=EVAL_METADATA) + params = SendMessageParams(agent=_agent(), task=task, content=TextContent(author="user", content="hi"), stream=False) + await self._dispatch(acp, RPCMethod.MESSAGE_SEND, params) + + assert seen == [EXPECTED_ATTRS] + + async def test_non_eval_task_is_not_registered(self) -> None: + acp = SyncACP() + + @acp.on_task_create + async def handler(params: CreateTaskParams) -> None: + return None + + task = Task(id="task-1", status="RUNNING", task_metadata={"team": "x"}) + await self._dispatch(acp, RPCMethod.TASK_CREATE, CreateTaskParams(agent=_agent(), task=task, params=None)) + + assert sgp_evals.attrs_for_span(_span(trace_id="task-1")) == {} + + +def _env_vars() -> Mock: + env_vars = Mock() + env_vars.WORKFLOW_NAME = "wf" + env_vars.WORKFLOW_TASK_QUEUE = "q" + env_vars.WORKFLOW_EXECUTION_TIMEOUT_SECONDS = 0 + return env_vars + + +class TestTemporal: + """Temporal agents: attrs ride the workflow memo, then activity headers, into the worker process.""" + + async def _submit(self, task_metadata: dict[str, Any] | None) -> dict[str, Any]: + client = Mock() + client.start_workflow = AsyncMock(return_value="task-1") + service = TemporalTaskService(temporal_client=client, env_vars=_env_vars()) + await service.submit_task( + agent=_agent(), task=Task(id="task-1", task_metadata=task_metadata), params=None + ) + return client.start_workflow.await_args.kwargs + + async def test_eval_task_workflow_starts_with_attrs_in_memo(self) -> None: + kwargs = await self._submit(EVAL_METADATA) + assert kwargs["memo"] == {"sgp_evals_span_attrs": EXPECTED_ATTRS} + assert kwargs["conflict_policy"] == ConflictWorkflowPolicy.USE_EXISTING + + async def test_non_eval_task_workflow_has_no_memo(self) -> None: + assert "memo" not in await self._submit({"team": "x"}) + assert "memo" not in await self._submit(None) + + async def test_memo_flows_through_activity_headers_to_worker_registry(self) -> None: + sent: dict[str, Any] = {} + next_outbound = Mock() + next_outbound.start_activity = lambda input: sent.update(headers=input.headers) + outbound = interceptor._WorkflowOutbound(next_outbound) + start_input = Mock(spec=StartActivityInput, headers={}) + + with patch.object(interceptor.workflow, "memo_value", return_value=EXPECTED_ATTRS): + outbound.start_activity(start_input) + + inbound_next = Mock() + inbound_next.execute_activity = AsyncMock(return_value="ok") + inbound = interceptor._ActivityInbound(inbound_next) + info = Mock(workflow_id="task-1") + with patch.object(interceptor.activity, "info", return_value=info): + await inbound.execute_activity(Mock(headers=sent["headers"])) + + assert sgp_evals.attrs_for_span(_span(trace_id="task-1")) == EXPECTED_ATTRS + + def test_workflow_without_memo_adds_no_header(self) -> None: + sent: dict[str, Any] = {} + next_outbound = Mock() + next_outbound.start_activity = lambda input: sent.update(headers=input.headers) + start_input = Mock(spec=StartActivityInput, headers={}) + + with patch.object(interceptor.workflow, "memo_value", return_value=None): + interceptor._WorkflowOutbound(next_outbound).start_activity(start_input) + + assert sent["headers"] == {} From 43211ba5d35223f44a9eeb78711d982a88fd8250 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Wed, 7 Oct 2026 08:10:05 -0400 Subject: [PATCH 02/13] fix(tracing): keep list-shaped eval spans searchable and drop stale eval attrs Co-Authored-By: Claude Sonnet 5.5 --- src/agentex/lib/core/temporal/workers/worker.py | 2 +- .../lib/core/tracing/processors/sgp_tracing_processor.py | 9 ++++++--- src/agentex/lib/core/tracing/sgp_evals.py | 3 +++ tests/lib/core/tracing/test_sgp_evals.py | 9 +++++++++ 4 files changed, 19 insertions(+), 4 deletions(-) diff --git a/src/agentex/lib/core/temporal/workers/worker.py b/src/agentex/lib/core/temporal/workers/worker.py index de2904c41..6c2a9afbe 100644 --- a/src/agentex/lib/core/temporal/workers/worker.py +++ b/src/agentex/lib/core/temporal/workers/worker.py @@ -275,7 +275,7 @@ async def run( build_id=str(uuid.uuid4()), debug_mode=debug_enabled, # Disable deadlock detection in debug mode # Temporal inherits client tracing before these business interceptors. - interceptors=[*self.interceptors, SGPEvalsInterceptor()], + interceptors=[SGPEvalsInterceptor(), *self.interceptors], ) logger.info(f"Starting workers for task queue: {self.task_queue}") diff --git a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index 7d47f84fd..e30b97e6f 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -83,7 +83,8 @@ def _sgp_metadata(span: Span) -> Any: leak like that today. Left as-is: changing five long-shipped fields is not this change's business.) """ - extra: dict[str, Any] = dict(sgp_evals.attrs_for_span(span)) + eval_attrs = sgp_evals.attrs_for_span(span) + extra: dict[str, Any] = dict(eval_attrs) commit_sha = code_revision.commit_sha() if commit_sha is not None: extra[code_revision.COMMIT_SHA_KEY] = commit_sha @@ -91,8 +92,10 @@ def _sgp_metadata(span: Span) -> Any: return span.data if isinstance(span.data, dict): return {**span.data, **extra} - # List-shaped data is an accepted `data` shape and has nowhere to put a - # metadata key; leave it untouched rather than dropping the caller's data. + # List-shaped data has nowhere to put a key. Eval spans must stay searchable by run and row, + # so their list moves under "data". Otherwise it is left untouched rather than reshaped. + if eval_attrs and isinstance(span.data, list): + return {**extra, "data": span.data} return span.data diff --git a/src/agentex/lib/core/tracing/sgp_evals.py b/src/agentex/lib/core/tracing/sgp_evals.py index e77ab99d1..cf63e3da3 100644 --- a/src/agentex/lib/core/tracing/sgp_evals.py +++ b/src/agentex/lib/core/tracing/sgp_evals.py @@ -52,6 +52,9 @@ def register_task_metadata(task_id: str, task_metadata: Any) -> dict[str, Any] | attrs = span_attrs_from_task_metadata(task_metadata) if attrs is not None: register_task(task_id, attrs) + elif _attrs_by_task: + with _lock: + _attrs_by_task.pop(task_id, None) return attrs diff --git a/tests/lib/core/tracing/test_sgp_evals.py b/tests/lib/core/tracing/test_sgp_evals.py index 7abeb9589..75ff723a1 100644 --- a/tests/lib/core/tracing/test_sgp_evals.py +++ b/tests/lib/core/tracing/test_sgp_evals.py @@ -98,6 +98,15 @@ def test_other_tasks_and_non_eval_tasks_are_untouched(self) -> None: assert _sgp_metadata(_span(trace_id="task-2", data={"k": 1})) == {"k": 1} assert _sgp_metadata(_span(trace_id="task-3", data={"k": 1})) == {"k": 1} + def test_eval_span_with_list_data_keeps_list_and_gains_the_ids(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + assert _sgp_metadata(_span(data=[{"a": 1}])) == {**EXPECTED_ATTRS, "data": [{"a": 1}]} + + def test_task_that_stops_being_an_eval_task_stops_being_stamped(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + sgp_evals.register_task_metadata("task-1", {"team": "x"}) + assert _sgp_metadata(_span(data={"k": 1})) == {"k": 1} + def test_registry_is_bounded(self) -> None: with patch.object(sgp_evals, "_MAX_TASKS", 2): for i in range(3): From a3482979766d411683e3ea0167e23d61a1a46b18 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Wed, 7 Oct 2026 08:19:36 -0400 Subject: [PATCH 03/13] fix(tracing): stamp sgp evals ids on local activity spans too Local activities go through start_local_activity, which the outbound interceptor did not override, so their spans carried no run or row id. Also pin the worker's interceptor wiring. Co-Authored-By: Claude Opus 5.5 --- .../lib/core/tracing/sgp_evals_interceptor.py | 16 +++++++-- tests/lib/core/tracing/test_sgp_evals.py | 34 ++++++++++++++++++- 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/src/agentex/lib/core/tracing/sgp_evals_interceptor.py b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py index 49b8a4e2b..8b20ea4e9 100644 --- a/src/agentex/lib/core/tracing/sgp_evals_interceptor.py +++ b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py @@ -16,6 +16,7 @@ Interceptor, StartActivityInput, ExecuteActivityInput, + StartLocalActivityInput, ActivityInboundInterceptor, WorkflowInboundInterceptor, WorkflowOutboundInterceptor, @@ -47,14 +48,23 @@ def init(self, outbound: WorkflowOutboundInterceptor) -> None: super().init(_WorkflowOutbound(outbound)) +def _add_attrs_header(input: StartActivityInput | StartLocalActivityInput) -> None: + attrs = workflow.memo_value(MEMO_KEY, default=None) + if isinstance(attrs, dict) and attrs: + input.headers = {**input.headers, ATTRS_HEADER: _converter.to_payload(attrs)} + + class _WorkflowOutbound(WorkflowOutboundInterceptor): @override def start_activity(self, input: StartActivityInput) -> workflow.ActivityHandle[Any]: - attrs = workflow.memo_value(MEMO_KEY, default=None) - if isinstance(attrs, dict) and attrs: - input.headers = {**input.headers, ATTRS_HEADER: _converter.to_payload(attrs)} + _add_attrs_header(input) return super().start_activity(input) + @override + def start_local_activity(self, input: StartLocalActivityInput) -> workflow.ActivityHandle[Any]: + _add_attrs_header(input) + return super().start_local_activity(input) + class _ActivityInbound(ActivityInboundInterceptor): @override diff --git a/tests/lib/core/tracing/test_sgp_evals.py b/tests/lib/core/tracing/test_sgp_evals.py index 75ff723a1..5d564f65e 100644 --- a/tests/lib/core/tracing/test_sgp_evals.py +++ b/tests/lib/core/tracing/test_sgp_evals.py @@ -11,10 +11,11 @@ import uuid from typing import Any from datetime import UTC, datetime +from contextlib import ExitStack from unittest.mock import Mock, AsyncMock, patch import pytest -from temporalio.worker import StartActivityInput +from temporalio.worker import StartActivityInput, StartLocalActivityInput from agentex.types.span import Span from agentex.types.task import Task @@ -24,6 +25,7 @@ from agentex.types.task_message_content import TextContent from agentex.lib.sdk.fastacp.impl.sync_acp import SyncACP from agentex.lib.core.clients.temporal.types import ConflictWorkflowPolicy +from agentex.lib.core.temporal.workers.worker import AgentexWorker from agentex.lib.core.temporal.services.temporal_task_service import TemporalTaskService from agentex.lib.core.tracing.processors.sgp_tracing_processor import _sgp_metadata @@ -226,3 +228,33 @@ def test_workflow_without_memo_adds_no_header(self) -> None: interceptor._WorkflowOutbound(next_outbound).start_activity(start_input) assert sent["headers"] == {} + + def test_local_activities_get_the_header_too(self) -> None: + sent: dict[str, Any] = {} + next_outbound = Mock() + next_outbound.start_local_activity = lambda input: sent.update(headers=input.headers) + start_input = Mock(spec=StartLocalActivityInput, headers={}) + + with patch.object(interceptor.workflow, "memo_value", return_value=EXPECTED_ATTRS): + interceptor._WorkflowOutbound(next_outbound).start_local_activity(start_input) + + assert interceptor.ATTRS_HEADER in sent["headers"] + + async def test_worker_runs_with_the_interceptor_ahead_of_agent_interceptors(self) -> None: + agent_interceptor = interceptor.SGPEvalsInterceptor() + module = "agentex.lib.core.temporal.workers.worker" + with ExitStack() as stack: + stack.enter_context(patch(f"{module}.EnvironmentVariables")) + stack.enter_context(patch(f"{module}.init_sgp_obs")) + for name in ("shutdown_sgp_obs", "shutdown_default_span_queue", "shutdown_sync_tracing_processors", "get_temporal_client"): + stack.enter_context(patch(f"{module}.{name}", new=AsyncMock())) + worker_cls = stack.enter_context(patch(f"{module}.Worker")) + worker_cls.return_value.run = AsyncMock() + worker = AgentexWorker(task_queue="q", interceptors=[agent_interceptor]) + stack.enter_context(patch.object(worker, "start_health_check_server", new=AsyncMock())) + stack.enter_context(patch.object(worker, "_register_agent", new=AsyncMock())) + await worker.run(activities=[], workflow=object) + + interceptors = worker_cls.call_args.kwargs["interceptors"] + assert isinstance(interceptors[0], interceptor.SGPEvalsInterceptor) + assert interceptors[1:] == [agent_interceptor] From 618f226bd367e961b293d672aadf747a588def41 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Wed, 7 Oct 2026 08:21:30 -0400 Subject: [PATCH 04/13] fix(tracing): drop a worker's eval attrs when a task's activities carry none Co-Authored-By: Claude Sonnet 5.5 --- src/agentex/lib/core/tracing/sgp_evals.py | 9 +++++++-- src/agentex/lib/core/tracing/sgp_evals_interceptor.py | 7 +++++-- tests/lib/core/tracing/test_sgp_evals.py | 9 +++++++++ 3 files changed, 21 insertions(+), 4 deletions(-) diff --git a/src/agentex/lib/core/tracing/sgp_evals.py b/src/agentex/lib/core/tracing/sgp_evals.py index cf63e3da3..d05d4a25c 100644 --- a/src/agentex/lib/core/tracing/sgp_evals.py +++ b/src/agentex/lib/core/tracing/sgp_evals.py @@ -52,10 +52,15 @@ def register_task_metadata(task_id: str, task_metadata: Any) -> dict[str, Any] | attrs = span_attrs_from_task_metadata(task_metadata) if attrs is not None: register_task(task_id, attrs) - elif _attrs_by_task: + else: + unregister_task(task_id) + return attrs + + +def unregister_task(task_id: str) -> None: + if _attrs_by_task: with _lock: _attrs_by_task.pop(task_id, None) - return attrs def attrs_for_span(span: Span) -> dict[str, Any]: diff --git a/src/agentex/lib/core/tracing/sgp_evals_interceptor.py b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py index 8b20ea4e9..5463f4328 100644 --- a/src/agentex/lib/core/tracing/sgp_evals_interceptor.py +++ b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py @@ -24,7 +24,7 @@ from temporalio.converter import default from agentex.lib.utils.logging import make_logger -from agentex.lib.core.tracing.sgp_evals import MEMO_KEY, register_task +from agentex.lib.core.tracing.sgp_evals import MEMO_KEY, register_task, unregister_task logger = make_logger(__name__) @@ -72,7 +72,10 @@ async def execute_activity(self, input: ExecuteActivityInput) -> Any: payload = input.headers.get(ATTRS_HEADER) # The workflow id is the task id (see TemporalTaskService.submit_task). task_id = activity.info().workflow_id - if payload is not None and task_id: + if task_id and payload is None: + # A reused workflow id must not inherit a previous eval run's ids. + unregister_task(task_id) + elif task_id and payload is not None: try: register_task(task_id, _converter.from_payload(payload, dict)) except Exception: diff --git a/tests/lib/core/tracing/test_sgp_evals.py b/tests/lib/core/tracing/test_sgp_evals.py index 5d564f65e..7e2464f67 100644 --- a/tests/lib/core/tracing/test_sgp_evals.py +++ b/tests/lib/core/tracing/test_sgp_evals.py @@ -218,6 +218,15 @@ async def test_memo_flows_through_activity_headers_to_worker_registry(self) -> N assert sgp_evals.attrs_for_span(_span(trace_id="task-1")) == EXPECTED_ATTRS + async def test_plain_activity_does_not_inherit_a_reused_workflow_ids_eval_attrs(self) -> None: + sgp_evals.register_task("task-1", EXPECTED_ATTRS) + inbound_next = Mock() + inbound_next.execute_activity = AsyncMock(return_value="ok") + with patch.object(interceptor.activity, "info", return_value=Mock(workflow_id="task-1")): + await interceptor._ActivityInbound(inbound_next).execute_activity(Mock(headers={})) + + assert sgp_evals.attrs_for_span(_span(trace_id="task-1")) == {} + def test_workflow_without_memo_adds_no_header(self) -> None: sent: dict[str, Any] = {} next_outbound = Mock() From ce0dcfa726d253c83edf4312de9b06db3cad797a Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Wed, 7 Oct 2026 08:43:47 -0400 Subject: [PATCH 05/13] fix(tracing): pin eval ids on a span when it starts so a reused task id cannot strip a queued span Co-Authored-By: Claude Sonnet 5.5 --- .../processors/sgp_tracing_processor.py | 2 + src/agentex/lib/core/tracing/sgp_evals.py | 47 ++++++++-- src/agentex/lib/core/tracing/trace.py | 4 + tests/lib/core/tracing/test_sgp_evals.py | 87 ++++++++++++++++++- 4 files changed, 132 insertions(+), 8 deletions(-) diff --git a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index e30b97e6f..32b74a773 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -154,6 +154,7 @@ def on_span_start(self, span: Span) -> None: def on_span_end(self, span: Span) -> None: sgp_span = _build_sgp_span(span, self.env_vars) sgp_span.end_time = span.end_time.isoformat() # type: ignore[union-attr] + sgp_evals.release_span(span.id) sgp_span.flush(blocking=False) @override @@ -255,6 +256,7 @@ async def on_spans_end(self, spans: list[Span]) -> None: for span in spans: sgp_span = _build_sgp_span(span, self.env_vars) sgp_span.end_time = span.end_time.isoformat() # type: ignore[union-attr] + sgp_evals.release_span(span.id) sgp_spans.append(sgp_span) await client.spans.upsert_batch(items=[s.to_request_params() for s in sgp_spans]) _metrics.record_export_success( diff --git a/src/agentex/lib/core/tracing/sgp_evals.py b/src/agentex/lib/core/tracing/sgp_evals.py index d05d4a25c..0a64e2c90 100644 --- a/src/agentex/lib/core/tracing/sgp_evals.py +++ b/src/agentex/lib/core/tracing/sgp_evals.py @@ -15,7 +15,15 @@ from agentex.types.span import Span -__all__ = ("MEMO_KEY", "register_task", "register_task_metadata", "span_attrs_from_task_metadata", "attrs_for_span") +__all__ = ( + "MEMO_KEY", + "register_task", + "register_task_metadata", + "span_attrs_from_task_metadata", + "attrs_for_span", + "capture_for_span", + "release_span", +) TASK_METADATA_MARKER = "sgp_evals" SPAN_KEY_PREFIX = "sgp_evals_" @@ -26,6 +34,8 @@ # Only eval tasks are stored, so this stays tiny. The LRU bound caps a long-lived agent process. _MAX_TASKS = 10_000 _attrs_by_task: OrderedDict[str, dict[str, Any]] = OrderedDict() +# Attrs captured when a span starts, so a later registry change cannot alter a span still queued for export. +_attrs_by_span: OrderedDict[str, dict[str, Any]] = OrderedDict() _lock = threading.Lock() @@ -63,18 +73,41 @@ def unregister_task(task_id: str) -> None: _attrs_by_task.pop(task_id, None) -def attrs_for_span(span: Span) -> dict[str, Any]: - """Attrs for the task a span belongs to, found by ``span.task_id`` then ``span.trace_id``.""" +def _lookup(span: Span) -> dict[str, Any]: + for key in (span.task_id, span.trace_id): + if key and key in _attrs_by_task: + return dict(_attrs_by_task[key]) + return {} + + +def capture_for_span(span: Span) -> None: + """Pin the span's attrs at start. No-op while no eval task is registered in this process.""" if not _attrs_by_task: + return + with _lock: + _attrs_by_span[span.id] = _lookup(span) + while len(_attrs_by_span) > _MAX_TASKS: + _attrs_by_span.popitem(last=False) + + +def release_span(span_id: str) -> None: + if _attrs_by_span: + with _lock: + _attrs_by_span.pop(span_id, None) + + +def attrs_for_span(span: Span) -> dict[str, Any]: + """Attrs captured at span start, else those of the task found by ``span.task_id`` then ``span.trace_id``.""" + if not _attrs_by_task and not _attrs_by_span: return {} with _lock: - for key in (span.task_id, span.trace_id): - if key and key in _attrs_by_task: - return dict(_attrs_by_task[key]) - return {} + if span.id in _attrs_by_span: + return dict(_attrs_by_span[span.id]) + return _lookup(span) def clear() -> None: """Reset the registry (test isolation).""" with _lock: _attrs_by_task.clear() + _attrs_by_span.clear() diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index 8f5260913..0871905e6 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -10,6 +10,7 @@ from agentex import Agentex, AsyncAgentex from agentex.types.span import Span +from agentex.lib.core.tracing import sgp_evals from agentex.lib.utils.logging import make_logger from agentex.lib.utils.model_utils import recursive_model_dump from agentex.lib.core.tracing.obs_ids import obs_correlation, warn_on_backend_drift @@ -293,6 +294,7 @@ def start_span( if obs_handle is not None: _register_obs_handle(span.id, obs_handle) + sgp_evals.capture_for_span(span) for processor in self.processors: _run_on_span_start(processor, span) @@ -457,6 +459,8 @@ async def start_span( if obs_handle is not None: _register_obs_handle(span.id, obs_handle) + sgp_evals.capture_for_span(span) + # Enqueueing the START event must not crash the app path either (same # principle as _run_on_span_start): swallow so start_span still returns # and end_span cleans up the handle. The processors' on_span_start runs diff --git a/tests/lib/core/tracing/test_sgp_evals.py b/tests/lib/core/tracing/test_sgp_evals.py index 7e2464f67..37f43ec51 100644 --- a/tests/lib/core/tracing/test_sgp_evals.py +++ b/tests/lib/core/tracing/test_sgp_evals.py @@ -22,12 +22,14 @@ from agentex.types.agent import Agent from agentex.protocol.acp import RPCMethod, CreateTaskParams, SendMessageParams from agentex.lib.core.tracing import sgp_evals, sgp_evals_interceptor as interceptor +from agentex.lib.types.tracing import SGPTracingProcessorConfig +from agentex.lib.core.tracing.trace import Trace, AsyncTrace from agentex.types.task_message_content import TextContent from agentex.lib.sdk.fastacp.impl.sync_acp import SyncACP from agentex.lib.core.clients.temporal.types import ConflictWorkflowPolicy from agentex.lib.core.temporal.workers.worker import AgentexWorker from agentex.lib.core.temporal.services.temporal_task_service import TemporalTaskService -from agentex.lib.core.tracing.processors.sgp_tracing_processor import _sgp_metadata +from agentex.lib.core.tracing.processors.sgp_tracing_processor import SGPSyncTracingProcessor, _sgp_metadata EVAL_METADATA = { "sgp_evals": "generation-unit", @@ -117,6 +119,89 @@ def test_registry_is_bounded(self) -> None: assert sgp_evals.attrs_for_span(_span(trace_id="t2")) == EXPECTED_ATTRS +class TestQueuedSpans: + """A span keeps the ids its task had when it started, even if the task id is reused before export.""" + + def test_captured_ids_survive_the_task_being_unregistered(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + span = _span(data={}) + sgp_evals.capture_for_span(span) + + sgp_evals.unregister_task("task-1") + + assert _sgp_metadata(span) == EXPECTED_ATTRS + + def test_captured_ids_survive_a_new_run_taking_the_task_id(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + span = _span(data={}) + sgp_evals.capture_for_span(span) + + sgp_evals.register_task("task-1", {"sgp_evals_row_id": "row-99"}) + + assert _sgp_metadata(span) == EXPECTED_ATTRS + + def test_plain_span_is_not_stamped_by_a_run_registered_after_it_started(self) -> None: + sgp_evals.register_task_metadata("other", EVAL_METADATA) + span = _span(trace_id="task-1", data={"k": 1}) + sgp_evals.capture_for_span(span) + + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + + assert _sgp_metadata(span) == {"k": 1} + + def test_release_drops_the_capture(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + span = _span(data={}) + sgp_evals.capture_for_span(span) + sgp_evals.release_span(span.id) + sgp_evals.unregister_task("task-1") + + assert sgp_evals.attrs_for_span(span) == {} + + def test_captures_are_bounded(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + spans = [_span(data={}) for _ in range(3)] + with patch.object(sgp_evals, "_MAX_TASKS", 2): + for span in spans: + sgp_evals.capture_for_span(span) + sgp_evals.unregister_task("task-1") + + assert sgp_evals.attrs_for_span(spans[0]) == {} + assert sgp_evals.attrs_for_span(spans[2]) == EXPECTED_ATTRS + + def test_sync_trace_start_span_captures_ids(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + span = Trace(processors=[], client=Mock(), trace_id="task-1").start_span(name="s") + + sgp_evals.unregister_task("task-1") + span.data = {} + + assert _sgp_metadata(span) == EXPECTED_ATTRS + + async def test_async_trace_start_span_captures_ids(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + span = await AsyncTrace(processors=[], client=Mock(), trace_id="task-1", span_queue=Mock()).start_span(name="s") + + sgp_evals.unregister_task("task-1") + span.data = {} + + assert _sgp_metadata(span) == EXPECTED_ATTRS + + def test_sgp_processor_releases_the_capture_when_the_span_ends(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + span = _span(data={}) + sgp_evals.capture_for_span(span) + span.end_time = datetime.now(UTC) + module = "agentex.lib.core.tracing.processors.sgp_tracing_processor" + with patch(f"{module}.tracing.init"), patch(f"{module}.EnvironmentVariables"): + processor = SGPSyncTracingProcessor(SGPTracingProcessorConfig(sgp_api_key="", sgp_account_id="")) + with patch(f"{module}._build_sgp_span", return_value=Mock()): + processor.on_span_end(span) + sgp_evals.unregister_task("task-1") + + assert sgp_evals.attrs_for_span(span) == {} + + class TestAcpServer: """Sync and async (base) agents: the ACP server sees the task and spans are emitted in-process.""" From f1d9eb45e9f7c31bd537ff0934033e2ed23e3b80 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Wed, 7 Oct 2026 08:56:22 -0400 Subject: [PATCH 06/13] fix(tracing): pin eval ids for plain spans and keep them until the upload succeeds Co-Authored-By: Claude Sonnet 5.5 --- .../processors/sgp_tracing_processor.py | 4 ++- src/agentex/lib/core/tracing/sgp_evals.py | 4 +-- tests/lib/core/tracing/test_sgp_evals.py | 34 ++++++++++++++++++- uv.lock | 4 +-- 4 files changed, 39 insertions(+), 7 deletions(-) diff --git a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index 32b74a773..2ec8a625c 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -256,9 +256,11 @@ async def on_spans_end(self, spans: list[Span]) -> None: for span in spans: sgp_span = _build_sgp_span(span, self.env_vars) sgp_span.end_time = span.end_time.isoformat() # type: ignore[union-attr] - sgp_evals.release_span(span.id) sgp_spans.append(sgp_span) await client.spans.upsert_batch(items=[s.to_request_params() for s in sgp_spans]) + # Released only after the upload so a queue retry rebuilds the span with the same ids. + for span in spans: + sgp_evals.release_span(span.id) _metrics.record_export_success( event_type="end", span_count=len(spans), processor="sgp" ) diff --git a/src/agentex/lib/core/tracing/sgp_evals.py b/src/agentex/lib/core/tracing/sgp_evals.py index 0a64e2c90..7fbeb5d24 100644 --- a/src/agentex/lib/core/tracing/sgp_evals.py +++ b/src/agentex/lib/core/tracing/sgp_evals.py @@ -81,9 +81,7 @@ def _lookup(span: Span) -> dict[str, Any]: def capture_for_span(span: Span) -> None: - """Pin the span's attrs at start. No-op while no eval task is registered in this process.""" - if not _attrs_by_task: - return + """Pin the span's attrs at start, empty included, so an eval task registered later cannot claim it.""" with _lock: _attrs_by_span[span.id] = _lookup(span) while len(_attrs_by_span) > _MAX_TASKS: diff --git a/tests/lib/core/tracing/test_sgp_evals.py b/tests/lib/core/tracing/test_sgp_evals.py index 37f43ec51..3fb5ce1f3 100644 --- a/tests/lib/core/tracing/test_sgp_evals.py +++ b/tests/lib/core/tracing/test_sgp_evals.py @@ -29,7 +29,11 @@ from agentex.lib.core.clients.temporal.types import ConflictWorkflowPolicy from agentex.lib.core.temporal.workers.worker import AgentexWorker from agentex.lib.core.temporal.services.temporal_task_service import TemporalTaskService -from agentex.lib.core.tracing.processors.sgp_tracing_processor import SGPSyncTracingProcessor, _sgp_metadata +from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( + SGPSyncTracingProcessor, + SGPAsyncTracingProcessor, + _sgp_metadata, +) EVAL_METADATA = { "sgp_evals": "generation-unit", @@ -149,6 +153,34 @@ def test_plain_span_is_not_stamped_by_a_run_registered_after_it_started(self) -> assert _sgp_metadata(span) == {"k": 1} + def test_plain_span_started_with_an_empty_registry_is_not_stamped_later(self) -> None: + span = _span(trace_id="task-1", data={"k": 1}) + sgp_evals.capture_for_span(span) + + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + + assert _sgp_metadata(span) == {"k": 1} + + async def test_async_processor_keeps_the_capture_when_the_upload_fails(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + span = _span(data={}) + sgp_evals.capture_for_span(span) + span.end_time = datetime.now(UTC) + module = "agentex.lib.core.tracing.processors.sgp_tracing_processor" + with patch(f"{module}.tracing.init"), patch(f"{module}.EnvironmentVariables"): + processor = SGPAsyncTracingProcessor(SGPTracingProcessorConfig(sgp_api_key="", sgp_account_id="")) + client = Mock() + client.spans.upsert_batch = AsyncMock(side_effect=[RuntimeError("boom"), None]) + with patch.object(processor, "_get_client", return_value=client): + with pytest.raises(RuntimeError): + await processor.on_spans_end([span]) + sgp_evals.unregister_task("task-1") + assert sgp_evals.attrs_for_span(span) == EXPECTED_ATTRS + + await processor.on_spans_end([span]) + + assert sgp_evals.attrs_for_span(span) == {} + def test_release_drops_the_capture(self) -> None: sgp_evals.register_task_metadata("task-1", EVAL_METADATA) span = _span(data={}) diff --git a/uv.lock b/uv.lock index 6c79c4ace..e15318da6 100644 --- a/uv.lock +++ b/uv.lock @@ -15,7 +15,7 @@ members = [ [[package]] name = "agentex-client" -version = "0.21.0" +version = "0.28.0" source = { editable = "." } dependencies = [ { name = "anyio" }, @@ -91,7 +91,7 @@ dev = [ [[package]] name = "agentex-sdk" -version = "0.21.0" +version = "0.28.0" source = { editable = "adk" } dependencies = [ { name = "agentex-client" }, From 3e31b85ba8f24ff0f229568e2c29e12d47bd442a Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Wed, 7 Oct 2026 08:56:29 -0400 Subject: [PATCH 07/13] chore: restore uv.lock Co-Authored-By: Claude Sonnet 5.5 --- uv.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/uv.lock b/uv.lock index e15318da6..6c79c4ace 100644 --- a/uv.lock +++ b/uv.lock @@ -15,7 +15,7 @@ members = [ [[package]] name = "agentex-client" -version = "0.28.0" +version = "0.21.0" source = { editable = "." } dependencies = [ { name = "anyio" }, @@ -91,7 +91,7 @@ dev = [ [[package]] name = "agentex-sdk" -version = "0.28.0" +version = "0.21.0" source = { editable = "adk" } dependencies = [ { name = "agentex-client" }, From 848212ae97666bbe5ab35cdb252c731d8db44520 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Wed, 7 Oct 2026 09:08:33 -0400 Subject: [PATCH 08/13] fix(tracing): keep eval span id captures apart from plain span pins so a span flood cannot evict them Co-Authored-By: Claude Sonnet 5.5 --- src/agentex/lib/core/tracing/sgp_evals.py | 18 +++++++++++++----- tests/lib/core/tracing/test_sgp_evals.py | 11 +++++++++++ 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/src/agentex/lib/core/tracing/sgp_evals.py b/src/agentex/lib/core/tracing/sgp_evals.py index 7fbeb5d24..e624edcd1 100644 --- a/src/agentex/lib/core/tracing/sgp_evals.py +++ b/src/agentex/lib/core/tracing/sgp_evals.py @@ -36,6 +36,8 @@ _attrs_by_task: OrderedDict[str, dict[str, Any]] = OrderedDict() # Attrs captured when a span starts, so a later registry change cannot alter a span still queued for export. _attrs_by_span: OrderedDict[str, dict[str, Any]] = OrderedDict() +# Spans pinned with no attrs live apart so a flood of plain spans cannot evict an eval span's pinned ids. +_plain_spans: OrderedDict[str, None] = OrderedDict() _lock = threading.Lock() @@ -83,24 +85,29 @@ def _lookup(span: Span) -> dict[str, Any]: def capture_for_span(span: Span) -> None: """Pin the span's attrs at start, empty included, so an eval task registered later cannot claim it.""" with _lock: - _attrs_by_span[span.id] = _lookup(span) - while len(_attrs_by_span) > _MAX_TASKS: - _attrs_by_span.popitem(last=False) + attrs = _lookup(span) + store: OrderedDict[str, Any] = _attrs_by_span if attrs else _plain_spans + store[span.id] = attrs or None + while len(store) > _MAX_TASKS: + store.popitem(last=False) def release_span(span_id: str) -> None: - if _attrs_by_span: + if _attrs_by_span or _plain_spans: with _lock: _attrs_by_span.pop(span_id, None) + _plain_spans.pop(span_id, None) def attrs_for_span(span: Span) -> dict[str, Any]: """Attrs captured at span start, else those of the task found by ``span.task_id`` then ``span.trace_id``.""" - if not _attrs_by_task and not _attrs_by_span: + if not _attrs_by_task and not _attrs_by_span and not _plain_spans: return {} with _lock: if span.id in _attrs_by_span: return dict(_attrs_by_span[span.id]) + if span.id in _plain_spans: + return {} return _lookup(span) @@ -109,3 +116,4 @@ def clear() -> None: with _lock: _attrs_by_task.clear() _attrs_by_span.clear() + _plain_spans.clear() diff --git a/tests/lib/core/tracing/test_sgp_evals.py b/tests/lib/core/tracing/test_sgp_evals.py index 3fb5ce1f3..85337b754 100644 --- a/tests/lib/core/tracing/test_sgp_evals.py +++ b/tests/lib/core/tracing/test_sgp_evals.py @@ -201,6 +201,17 @@ def test_captures_are_bounded(self) -> None: assert sgp_evals.attrs_for_span(spans[0]) == {} assert sgp_evals.attrs_for_span(spans[2]) == EXPECTED_ATTRS + def test_plain_spans_cannot_evict_an_eval_span_capture(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + eval_span = _span(data={}) + sgp_evals.capture_for_span(eval_span) + sgp_evals.unregister_task("task-1") + with patch.object(sgp_evals, "_MAX_TASKS", 2): + for _ in range(5): + sgp_evals.capture_for_span(_span(data={})) + + assert sgp_evals.attrs_for_span(eval_span) == EXPECTED_ATTRS + def test_sync_trace_start_span_captures_ids(self) -> None: sgp_evals.register_task_metadata("task-1", EVAL_METADATA) span = Trace(processors=[], client=Mock(), trace_id="task-1").start_span(name="s") From a789175375474d8819180220ae59e8ccb797d2ad Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Mon, 12 Oct 2026 11:55:08 -0400 Subject: [PATCH 09/13] fix(tracing): stamp eval ids on child workflow spans and skip span pins without an SGP processor Child workflows inherit no memo and their workflow id is not the task id, so the parent hands both down in the child's memo and activities carry the task id in a header. Span pinning now starts only once an SGP processor exists. Co-Authored-By: Claude Opus 5.5 --- .../processors/sgp_tracing_processor.py | 2 + src/agentex/lib/core/tracing/sgp_evals.py | 10 ++ .../lib/core/tracing/sgp_evals_interceptor.py | 41 ++++++-- tests/lib/core/tracing/test_sgp_evals.py | 98 ++++++++++++++++--- 4 files changed, 134 insertions(+), 17 deletions(-) diff --git a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index 2ec8a625c..75a1f576b 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -135,6 +135,7 @@ def __init__(self, config: SGPTracingProcessorConfig): disabled=disabled, ) self.env_vars = EnvironmentVariables.refresh() + sgp_evals.enable_capture() logger.info( "SGP tracing span-start upsert %s (%s)", "disabled — end-only ingest" if _skip_span_start_enabled() else "enabled", @@ -177,6 +178,7 @@ def __init__(self, config: SGPTracingProcessorConfig): asyncio.AbstractEventLoop, AsyncSGPClient ] = weakref.WeakKeyDictionary() self.env_vars = EnvironmentVariables.refresh() + sgp_evals.enable_capture() logger.info( "SGP tracing span-start upsert %s (%s)", "disabled — end-only ingest" if _skip_span_start_enabled() else "enabled", diff --git a/src/agentex/lib/core/tracing/sgp_evals.py b/src/agentex/lib/core/tracing/sgp_evals.py index e624edcd1..b3a65f8cc 100644 --- a/src/agentex/lib/core/tracing/sgp_evals.py +++ b/src/agentex/lib/core/tracing/sgp_evals.py @@ -21,6 +21,7 @@ "register_task_metadata", "span_attrs_from_task_metadata", "attrs_for_span", + "enable_capture", "capture_for_span", "release_span", ) @@ -39,6 +40,8 @@ # Spans pinned with no attrs live apart so a flood of plain spans cannot evict an eval span's pinned ids. _plain_spans: OrderedDict[str, None] = OrderedDict() _lock = threading.Lock() +# Only the SGP processor reads the pins, so a process without one never takes the lock per span. +_capture_enabled = False def span_attrs_from_task_metadata(task_metadata: Any) -> dict[str, Any] | None: @@ -82,8 +85,15 @@ def _lookup(span: Span) -> dict[str, Any]: return {} +def enable_capture() -> None: + global _capture_enabled + _capture_enabled = True + + def capture_for_span(span: Span) -> None: """Pin the span's attrs at start, empty included, so an eval task registered later cannot claim it.""" + if not _capture_enabled: + return with _lock: attrs = _lookup(span) store: OrderedDict[str, Any] = _attrs_by_span if attrs else _plain_spans diff --git a/src/agentex/lib/core/tracing/sgp_evals_interceptor.py b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py index 5463f4328..655d88093 100644 --- a/src/agentex/lib/core/tracing/sgp_evals_interceptor.py +++ b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py @@ -4,6 +4,7 @@ server's task. The ACP server puts the attrs in the workflow memo, the outbound interceptor copies them onto each activity's headers, and the activity interceptor registers them for the task so the SGP processor can stamp spans. +A child workflow gets the attrs and the task id in its own memo, since it inherits neither. Non-eval workflows have no memo entry and get no header. """ @@ -17,6 +18,7 @@ StartActivityInput, ExecuteActivityInput, StartLocalActivityInput, + StartChildWorkflowInput, ActivityInboundInterceptor, WorkflowInboundInterceptor, WorkflowOutboundInterceptor, @@ -29,6 +31,8 @@ logger = make_logger(__name__) ATTRS_HEADER = "sgp-evals-span-attrs" +TASK_ID_HEADER = "sgp-evals-task-id" +TASK_ID_MEMO_KEY = "sgp_evals_task_id" _converter = default().payload_converter @@ -48,10 +52,24 @@ def init(self, outbound: WorkflowOutboundInterceptor) -> None: super().init(_WorkflowOutbound(outbound)) -def _add_attrs_header(input: StartActivityInput | StartLocalActivityInput) -> None: +def _memo_attrs() -> dict[str, Any] | None: attrs = workflow.memo_value(MEMO_KEY, default=None) - if isinstance(attrs, dict) and attrs: - input.headers = {**input.headers, ATTRS_HEADER: _converter.to_payload(attrs)} + return attrs if isinstance(attrs, dict) and attrs else None + + +def _task_id() -> str: + # The workflow id is the task id (see TemporalTaskService.submit_task), except in a child workflow. + return workflow.memo_value(TASK_ID_MEMO_KEY, default=None) or workflow.info().workflow_id + + +def _add_attrs_header(input: StartActivityInput | StartLocalActivityInput) -> None: + attrs = _memo_attrs() + if attrs: + input.headers = { + **input.headers, + ATTRS_HEADER: _converter.to_payload(attrs), + TASK_ID_HEADER: _converter.to_payload(_task_id()), + } class _WorkflowOutbound(WorkflowOutboundInterceptor): @@ -65,19 +83,30 @@ def start_local_activity(self, input: StartLocalActivityInput) -> workflow.Activ _add_attrs_header(input) return super().start_local_activity(input) + @override + async def start_child_workflow(self, input: StartChildWorkflowInput) -> workflow.ChildWorkflowHandle[Any, Any]: + attrs = _memo_attrs() + if attrs: + input.memo = {**(input.memo or {}), MEMO_KEY: attrs, TASK_ID_MEMO_KEY: _task_id()} + return await super().start_child_workflow(input) + class _ActivityInbound(ActivityInboundInterceptor): @override async def execute_activity(self, input: ExecuteActivityInput) -> Any: payload = input.headers.get(ATTRS_HEADER) - # The workflow id is the task id (see TemporalTaskService.submit_task). + task_id_payload = input.headers.get(TASK_ID_HEADER) + # Older workers send no task id header, and outside a child workflow the workflow id is the task id. task_id = activity.info().workflow_id if task_id and payload is None: # A reused workflow id must not inherit a previous eval run's ids. unregister_task(task_id) - elif task_id and payload is not None: + elif payload is not None: try: - register_task(task_id, _converter.from_payload(payload, dict)) + if task_id_payload is not None: + task_id = _converter.from_payload(task_id_payload, str) + if task_id: + register_task(task_id, _converter.from_payload(payload, dict)) except Exception: logger.warning("failed to read sgp evals span attrs from activity headers", exc_info=True) return await super().execute_activity(input) diff --git a/tests/lib/core/tracing/test_sgp_evals.py b/tests/lib/core/tracing/test_sgp_evals.py index 85337b754..5d8afeb2b 100644 --- a/tests/lib/core/tracing/test_sgp_evals.py +++ b/tests/lib/core/tracing/test_sgp_evals.py @@ -15,7 +15,7 @@ from unittest.mock import Mock, AsyncMock, patch import pytest -from temporalio.worker import StartActivityInput, StartLocalActivityInput +from temporalio.worker import StartActivityInput, StartChildWorkflowInput, StartLocalActivityInput from agentex.types.span import Span from agentex.types.task import Task @@ -51,10 +51,27 @@ @pytest.fixture(autouse=True) def _clean_registry(): sgp_evals.clear() - yield + with patch.object(sgp_evals, "_capture_enabled", True): + yield sgp_evals.clear() +def _in_workflow(memo: dict[str, Any], workflow_id: str = "task-1") -> ExitStack: + stack = ExitStack() + stack.enter_context( + patch.object(interceptor.workflow, "memo_value", side_effect=lambda key, default=None: memo.get(key, default)) + ) + stack.enter_context(patch.object(interceptor.workflow, "info", return_value=Mock(workflow_id=workflow_id))) + return stack + + +async def _run_activity(headers: dict[str, Any], workflow_id: str) -> None: + inbound_next = Mock() + inbound_next.execute_activity = AsyncMock(return_value="ok") + with patch.object(interceptor.activity, "info", return_value=Mock(workflow_id=workflow_id)): + await interceptor._ActivityInbound(inbound_next).execute_activity(Mock(headers=headers)) + + def _span(trace_id: str = "task-1", task_id: str | None = None, data: Any = None) -> Span: return Span( id=str(uuid.uuid4()), @@ -244,6 +261,24 @@ def test_sgp_processor_releases_the_capture_when_the_span_ends(self) -> None: assert sgp_evals.attrs_for_span(span) == {} + def test_spans_are_not_pinned_until_an_sgp_processor_exists(self) -> None: + module = "agentex.lib.core.tracing.processors.sgp_tracing_processor" + with patch.object(sgp_evals, "_capture_enabled", False): + Trace(processors=[], client=Mock(), trace_id="task-1").start_span(name="s") + assert not sgp_evals._plain_spans + + with patch(f"{module}.tracing.init"), patch(f"{module}.EnvironmentVariables"): + SGPSyncTracingProcessor(SGPTracingProcessorConfig(sgp_api_key="", sgp_account_id="")) + Trace(processors=[], client=Mock(), trace_id="task-1").start_span(name="s") + assert len(sgp_evals._plain_spans) == 1 + + def test_async_sgp_processor_enables_pinning(self) -> None: + module = "agentex.lib.core.tracing.processors.sgp_tracing_processor" + with patch.object(sgp_evals, "_capture_enabled", False): + with patch(f"{module}.EnvironmentVariables"): + SGPAsyncTracingProcessor(SGPTracingProcessorConfig(sgp_api_key="", sgp_account_id="")) + assert sgp_evals._capture_enabled + class TestAcpServer: """Sync and async (base) agents: the ACP server sees the task and spans are emitted in-process.""" @@ -334,15 +369,10 @@ async def test_memo_flows_through_activity_headers_to_worker_registry(self) -> N outbound = interceptor._WorkflowOutbound(next_outbound) start_input = Mock(spec=StartActivityInput, headers={}) - with patch.object(interceptor.workflow, "memo_value", return_value=EXPECTED_ATTRS): + with _in_workflow({sgp_evals.MEMO_KEY: EXPECTED_ATTRS}): outbound.start_activity(start_input) - inbound_next = Mock() - inbound_next.execute_activity = AsyncMock(return_value="ok") - inbound = interceptor._ActivityInbound(inbound_next) - info = Mock(workflow_id="task-1") - with patch.object(interceptor.activity, "info", return_value=info): - await inbound.execute_activity(Mock(headers=sent["headers"])) + await _run_activity(sent["headers"], workflow_id="task-1") assert sgp_evals.attrs_for_span(_span(trace_id="task-1")) == EXPECTED_ATTRS @@ -361,7 +391,7 @@ def test_workflow_without_memo_adds_no_header(self) -> None: next_outbound.start_activity = lambda input: sent.update(headers=input.headers) start_input = Mock(spec=StartActivityInput, headers={}) - with patch.object(interceptor.workflow, "memo_value", return_value=None): + with _in_workflow({}): interceptor._WorkflowOutbound(next_outbound).start_activity(start_input) assert sent["headers"] == {} @@ -372,11 +402,57 @@ def test_local_activities_get_the_header_too(self) -> None: next_outbound.start_local_activity = lambda input: sent.update(headers=input.headers) start_input = Mock(spec=StartLocalActivityInput, headers={}) - with patch.object(interceptor.workflow, "memo_value", return_value=EXPECTED_ATTRS): + with _in_workflow({sgp_evals.MEMO_KEY: EXPECTED_ATTRS}): interceptor._WorkflowOutbound(next_outbound).start_local_activity(start_input) assert interceptor.ATTRS_HEADER in sent["headers"] + async def test_child_workflow_gets_the_attrs_and_task_id_in_its_memo(self) -> None: + sent: dict[str, Any] = {} + next_outbound = Mock() + next_outbound.start_child_workflow = AsyncMock(side_effect=lambda input: sent.update(memo=input.memo)) + start_input = Mock(spec=StartChildWorkflowInput, memo={"caller": "kept"}) + + with _in_workflow({sgp_evals.MEMO_KEY: EXPECTED_ATTRS}, workflow_id="task-1"): + await interceptor._WorkflowOutbound(next_outbound).start_child_workflow(start_input) + + assert sent["memo"] == { + "caller": "kept", + sgp_evals.MEMO_KEY: EXPECTED_ATTRS, + interceptor.TASK_ID_MEMO_KEY: "task-1", + } + + async def test_non_eval_child_workflow_memo_is_untouched(self) -> None: + sent: dict[str, Any] = {} + next_outbound = Mock() + next_outbound.start_child_workflow = AsyncMock(side_effect=lambda input: sent.update(memo=input.memo)) + start_input = Mock(spec=StartChildWorkflowInput, memo=None) + + with _in_workflow({}): + await interceptor._WorkflowOutbound(next_outbound).start_child_workflow(start_input) + + assert sent["memo"] is None + + async def test_child_workflow_activity_stamps_spans_of_the_parent_task(self) -> None: + sent: dict[str, Any] = {} + next_outbound = Mock() + next_outbound.start_activity = lambda input: sent.update(headers=input.headers) + child_memo = {sgp_evals.MEMO_KEY: EXPECTED_ATTRS, interceptor.TASK_ID_MEMO_KEY: "task-1"} + + with _in_workflow(child_memo, workflow_id="child-wf"): + interceptor._WorkflowOutbound(next_outbound).start_activity(Mock(spec=StartActivityInput, headers={})) + await _run_activity(sent["headers"], workflow_id="child-wf") + + assert sgp_evals.attrs_for_span(_span(trace_id="task-1")) == EXPECTED_ATTRS + assert sgp_evals.attrs_for_span(_span(trace_id="child-wf")) == {} + + async def test_header_without_task_id_registers_under_the_workflow_id(self) -> None: + headers = {interceptor.ATTRS_HEADER: interceptor._converter.to_payload(EXPECTED_ATTRS)} + + await _run_activity(headers, workflow_id="task-1") + + assert sgp_evals.attrs_for_span(_span(trace_id="task-1")) == EXPECTED_ATTRS + async def test_worker_runs_with_the_interceptor_ahead_of_agent_interceptors(self) -> None: agent_interceptor = interceptor.SGPEvalsInterceptor() module = "agentex.lib.core.temporal.workers.worker" From 8a170c1e705b2a7d525eb1172d889ced11c40766 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Mon, 12 Oct 2026 11:57:56 -0400 Subject: [PATCH 10/13] style: sort interceptor imports Co-Authored-By: Claude Opus 5.5 --- src/agentex/lib/core/tracing/sgp_evals_interceptor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/agentex/lib/core/tracing/sgp_evals_interceptor.py b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py index 655d88093..99f2196f1 100644 --- a/src/agentex/lib/core/tracing/sgp_evals_interceptor.py +++ b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py @@ -17,8 +17,8 @@ Interceptor, StartActivityInput, ExecuteActivityInput, - StartLocalActivityInput, StartChildWorkflowInput, + StartLocalActivityInput, ActivityInboundInterceptor, WorkflowInboundInterceptor, WorkflowOutboundInterceptor, From b45d3b30e685c6c2607fe2b50ef8867a06ee02c7 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Wed, 14 Oct 2026 08:45:58 -0400 Subject: [PATCH 11/13] fix(tracing): stamp an unpinned span only with ids its task had when it started A span that started before an SGP processor existed has no pin, so the export fell back to the live registry and could pick up a run registered later under the same task id. The registry now records when each task's ids took effect, and the fallback skips ids registered after the span started. Co-Authored-By: Claude Opus 5.5 --- src/agentex/lib/core/tracing/sgp_evals.py | 25 +++++++++++++------ tests/lib/core/tracing/test_sgp_evals.py | 29 ++++++++++++++++++++++- 2 files changed, 46 insertions(+), 8 deletions(-) diff --git a/src/agentex/lib/core/tracing/sgp_evals.py b/src/agentex/lib/core/tracing/sgp_evals.py index b3a65f8cc..19b71833c 100644 --- a/src/agentex/lib/core/tracing/sgp_evals.py +++ b/src/agentex/lib/core/tracing/sgp_evals.py @@ -11,6 +11,7 @@ import threading from typing import Any +from datetime import UTC, datetime from collections import OrderedDict from agentex.types.span import Span @@ -34,7 +35,8 @@ # Only eval tasks are stored, so this stays tiny. The LRU bound caps a long-lived agent process. _MAX_TASKS = 10_000 -_attrs_by_task: OrderedDict[str, dict[str, Any]] = OrderedDict() +# Each task's attrs with when they took effect, so an unpinned span only gets attrs its task had at its start. +_attrs_by_task: OrderedDict[str, tuple[dict[str, Any], datetime]] = OrderedDict() # Attrs captured when a span starts, so a later registry change cannot alter a span still queued for export. _attrs_by_span: OrderedDict[str, dict[str, Any]] = OrderedDict() # Spans pinned with no attrs live apart so a flood of plain spans cannot evict an eval span's pinned ids. @@ -56,7 +58,10 @@ def span_attrs_from_task_metadata(task_metadata: Any) -> dict[str, Any] | None: def register_task(task_id: str, attrs: dict[str, Any]) -> None: with _lock: - _attrs_by_task[task_id] = attrs + current = _attrs_by_task.get(task_id) + # Every event re-registers the same attrs, which must not move their start past the task's earlier spans. + if current is None or current[0] != attrs: + _attrs_by_task[task_id] = (attrs, datetime.now(UTC)) _attrs_by_task.move_to_end(task_id) while len(_attrs_by_task) > _MAX_TASKS: _attrs_by_task.popitem(last=False) @@ -78,13 +83,18 @@ def unregister_task(task_id: str) -> None: _attrs_by_task.pop(task_id, None) -def _lookup(span: Span) -> dict[str, Any]: +def _lookup(span: Span, started_at: datetime | None = None) -> dict[str, Any]: for key in (span.task_id, span.trace_id): - if key and key in _attrs_by_task: - return dict(_attrs_by_task[key]) + entry = _attrs_by_task.get(key) if key else None + if entry is not None and (started_at is None or entry[1] <= started_at): + return dict(entry[0]) return {} +def _aware(moment: datetime) -> datetime: + return moment if moment.tzinfo is not None else moment.replace(tzinfo=UTC) + + def enable_capture() -> None: global _capture_enabled _capture_enabled = True @@ -110,7 +120,7 @@ def release_span(span_id: str) -> None: def attrs_for_span(span: Span) -> dict[str, Any]: - """Attrs captured at span start, else those of the task found by ``span.task_id`` then ``span.trace_id``.""" + """Attrs captured at span start, else those its task (``span.task_id`` then ``span.trace_id``) had at its start.""" if not _attrs_by_task and not _attrs_by_span and not _plain_spans: return {} with _lock: @@ -118,7 +128,8 @@ def attrs_for_span(span: Span) -> dict[str, Any]: return dict(_attrs_by_span[span.id]) if span.id in _plain_spans: return {} - return _lookup(span) + # Unpinned (started before an SGP processor existed, or evicted): skip attrs registered after it started. + return _lookup(span, _aware(span.start_time)) def clear() -> None: diff --git a/tests/lib/core/tracing/test_sgp_evals.py b/tests/lib/core/tracing/test_sgp_evals.py index 5d8afeb2b..40ea52835 100644 --- a/tests/lib/core/tracing/test_sgp_evals.py +++ b/tests/lib/core/tracing/test_sgp_evals.py @@ -10,7 +10,7 @@ import uuid from typing import Any -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from contextlib import ExitStack from unittest.mock import Mock, AsyncMock, patch @@ -272,6 +272,33 @@ def test_spans_are_not_pinned_until_an_sgp_processor_exists(self) -> None: Trace(processors=[], client=Mock(), trace_id="task-1").start_span(name="s") assert len(sgp_evals._plain_spans) == 1 + def test_plain_span_started_before_pinning_is_not_stamped_by_a_later_run(self) -> None: + with patch.object(sgp_evals, "_capture_enabled", False): + span = Trace(processors=[], client=Mock(), trace_id="task-1").start_span(name="s") + span.start_time -= timedelta(seconds=1) + span.data = {"k": 1} + + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + + assert _sgp_metadata(span) == {"k": 1} + + def test_eval_span_started_before_pinning_keeps_its_ids_across_re_registration(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + with patch.object(sgp_evals, "_capture_enabled", False): + span = Trace(processors=[], client=Mock(), trace_id="task-1").start_span(name="s") + span.data = {} + + sgp_evals.register_task_metadata("task-1", dict(EVAL_METADATA)) + + assert _sgp_metadata(span) == EXPECTED_ATTRS + + def test_unpinned_span_with_a_naive_start_time_is_compared_as_utc(self) -> None: + sgp_evals.register_task_metadata("task-1", EVAL_METADATA) + span = _span(data={}) + span.start_time = datetime.now(UTC).replace(tzinfo=None) + timedelta(seconds=1) + + assert _sgp_metadata(span) == EXPECTED_ATTRS + def test_async_sgp_processor_enables_pinning(self) -> None: module = "agentex.lib.core.tracing.processors.sgp_tracing_processor" with patch.object(sgp_evals, "_capture_enabled", False): From 6a57f9ab5f30d74f19a159b5e991da4b79f3fd81 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Wed, 14 Oct 2026 09:09:43 -0400 Subject: [PATCH 12/13] fix(tracing): keep eval ids on spans that start on another Temporal worker start-span and end-span are separate activities, so a span can end on a worker that first registered its task after the span started, and the start-time gate dropped its ids. Ids read from a workflow memo are fixed for the workflow's life, so they now cover every span of the workflow. Only ids the ACP server registers keep the start-time gate. Co-Authored-By: Claude Opus 5.5 --- src/agentex/lib/core/tracing/sgp_evals.py | 17 +++++++++++++---- .../lib/core/tracing/sgp_evals_interceptor.py | 4 ++-- tests/lib/core/tracing/test_sgp_evals.py | 9 +++++++++ 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/src/agentex/lib/core/tracing/sgp_evals.py b/src/agentex/lib/core/tracing/sgp_evals.py index 19b71833c..a8ef10caa 100644 --- a/src/agentex/lib/core/tracing/sgp_evals.py +++ b/src/agentex/lib/core/tracing/sgp_evals.py @@ -19,6 +19,7 @@ __all__ = ( "MEMO_KEY", "register_task", + "register_workflow_task", "register_task_metadata", "span_attrs_from_task_metadata", "attrs_for_span", @@ -42,6 +43,7 @@ # Spans pinned with no attrs live apart so a flood of plain spans cannot evict an eval span's pinned ids. _plain_spans: OrderedDict[str, None] = OrderedDict() _lock = threading.Lock() +_WHOLE_WORKFLOW = datetime.min.replace(tzinfo=UTC) # Only the SGP processor reads the pins, so a process without one never takes the lock per span. _capture_enabled = False @@ -56,17 +58,24 @@ def span_attrs_from_task_metadata(task_metadata: Any) -> dict[str, Any] | None: return attrs or None -def register_task(task_id: str, attrs: dict[str, Any]) -> None: +def register_task(task_id: str, attrs: dict[str, Any], effective_from: datetime | None = None) -> None: with _lock: + since = effective_from or datetime.now(UTC) current = _attrs_by_task.get(task_id) # Every event re-registers the same attrs, which must not move their start past the task's earlier spans. - if current is None or current[0] != attrs: - _attrs_by_task[task_id] = (attrs, datetime.now(UTC)) + if current is not None and current[0] == attrs: + since = min(since, current[1]) + _attrs_by_task[task_id] = (attrs, since) _attrs_by_task.move_to_end(task_id) while len(_attrs_by_task) > _MAX_TASKS: _attrs_by_task.popitem(last=False) +def register_workflow_task(task_id: str, attrs: dict[str, Any]) -> None: + """Register ids read from a workflow memo, which also cover spans the workflow started before this worker saw it.""" + register_task(task_id, attrs, effective_from=_WHOLE_WORKFLOW) + + def register_task_metadata(task_id: str, task_metadata: Any) -> dict[str, Any] | None: """Remember an eval task's span attrs. No-op (and no lookup cost) for every other task.""" attrs = span_attrs_from_task_metadata(task_metadata) @@ -128,7 +137,7 @@ def attrs_for_span(span: Span) -> dict[str, Any]: return dict(_attrs_by_span[span.id]) if span.id in _plain_spans: return {} - # Unpinned (started before an SGP processor existed, or evicted): skip attrs registered after it started. + # Unpinned (started on another worker, before an SGP processor existed, or evicted). return _lookup(span, _aware(span.start_time)) diff --git a/src/agentex/lib/core/tracing/sgp_evals_interceptor.py b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py index 99f2196f1..f5f2df69d 100644 --- a/src/agentex/lib/core/tracing/sgp_evals_interceptor.py +++ b/src/agentex/lib/core/tracing/sgp_evals_interceptor.py @@ -26,7 +26,7 @@ from temporalio.converter import default from agentex.lib.utils.logging import make_logger -from agentex.lib.core.tracing.sgp_evals import MEMO_KEY, register_task, unregister_task +from agentex.lib.core.tracing.sgp_evals import MEMO_KEY, unregister_task, register_workflow_task logger = make_logger(__name__) @@ -106,7 +106,7 @@ async def execute_activity(self, input: ExecuteActivityInput) -> Any: if task_id_payload is not None: task_id = _converter.from_payload(task_id_payload, str) if task_id: - register_task(task_id, _converter.from_payload(payload, dict)) + register_workflow_task(task_id, _converter.from_payload(payload, dict)) except Exception: logger.warning("failed to read sgp evals span attrs from activity headers", exc_info=True) return await super().execute_activity(input) diff --git a/tests/lib/core/tracing/test_sgp_evals.py b/tests/lib/core/tracing/test_sgp_evals.py index 40ea52835..75d116d33 100644 --- a/tests/lib/core/tracing/test_sgp_evals.py +++ b/tests/lib/core/tracing/test_sgp_evals.py @@ -480,6 +480,15 @@ async def test_header_without_task_id_registers_under_the_workflow_id(self) -> N assert sgp_evals.attrs_for_span(_span(trace_id="task-1")) == EXPECTED_ATTRS + async def test_span_started_on_another_worker_keeps_its_ids_at_end(self) -> None: + span = _span(data={}) + span.start_time -= timedelta(seconds=1) + headers = {interceptor.ATTRS_HEADER: interceptor._converter.to_payload(EXPECTED_ATTRS)} + + await _run_activity(headers, workflow_id="task-1") + + assert _sgp_metadata(span) == EXPECTED_ATTRS + async def test_worker_runs_with_the_interceptor_ahead_of_agent_interceptors(self) -> None: agent_interceptor = interceptor.SGPEvalsInterceptor() module = "agentex.lib.core.temporal.workers.worker" From 0a79a58de8bc64aa28fd3e0ad14e7854da72943e Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Wed, 14 Oct 2026 09:09:43 -0400 Subject: [PATCH 13/13] ci(tutorials): pin pydantic below 2.14 for the openai-websearch-mcp server openai-websearch-mcp pins mcp 1.13.1, which imports a pydantic internal that pydantic 2.14 removed, so the server died at import and the 010, 020 and 040 tutorials timed out waiting for tool calls on every branch. Co-Authored-By: Claude Opus 5.5 --- .../tutorials/10_async/00_base/040_other_sdks/project/acp.py | 5 ++++- .../10_async/10_temporal/010_agent_chat/project/workflow.py | 3 ++- .../workflows/deep_research/performing_deep_research.py | 4 +++- .../050_agent_chat_guardrails/project/workflow.py | 3 ++- 4 files changed, 11 insertions(+), 4 deletions(-) diff --git a/examples/tutorials/10_async/00_base/040_other_sdks/project/acp.py b/examples/tutorials/10_async/00_base/040_other_sdks/project/acp.py index d2ec84fcd..7eedd51dc 100644 --- a/examples/tutorials/10_async/00_base/040_other_sdks/project/acp.py +++ b/examples/tutorials/10_async/00_base/040_other_sdks/project/acp.py @@ -54,7 +54,10 @@ class StateModel(BaseModel): args=["-y", "@modelcontextprotocol/server-sequential-thinking"], ), StdioServerParameters( - command="uvx", args=["openai-websearch-mcp"], env={"OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", "")} + command="uvx", + # mcp 1.13.1, pinned by openai-websearch-mcp, imports a pydantic internal that 2.14 removed. + args=["--with", "pydantic<2.14", "openai-websearch-mcp"], + env={"OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", "")}, ), ] diff --git a/examples/tutorials/10_async/10_temporal/010_agent_chat/project/workflow.py b/examples/tutorials/10_async/10_temporal/010_agent_chat/project/workflow.py index 3e3ac5b27..91cdfb217 100644 --- a/examples/tutorials/10_async/10_temporal/010_agent_chat/project/workflow.py +++ b/examples/tutorials/10_async/10_temporal/010_agent_chat/project/workflow.py @@ -55,7 +55,8 @@ class StateModel(BaseModel): # ), StdioServerParameters( command="uvx", - args=["openai-websearch-mcp"], + # mcp 1.13.1, pinned by openai-websearch-mcp, imports a pydantic internal that 2.14 removed. + args=["--with", "pydantic<2.14", "openai-websearch-mcp"], env={"OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", "")}, ), ] diff --git a/examples/tutorials/10_async/10_temporal/020_state_machine/project/workflows/deep_research/performing_deep_research.py b/examples/tutorials/10_async/10_temporal/020_state_machine/project/workflows/deep_research/performing_deep_research.py index 7bcb5a6d5..35e0fc25c 100644 --- a/examples/tutorials/10_async/10_temporal/020_state_machine/project/workflows/deep_research/performing_deep_research.py +++ b/examples/tutorials/10_async/10_temporal/020_state_machine/project/workflows/deep_research/performing_deep_research.py @@ -19,6 +19,8 @@ # this constraint every server dies at import and the agent silently makes zero tool # calls. Drop the pin once the servers support mcp 2.x. _MCP_PIN = ["--with", "mcp<2"] +# mcp 1.13.1, pinned by openai-websearch-mcp, imports a pydantic internal that 2.14 removed. +_WEBSEARCH_PIN = ["--with", "pydantic<2.14"] MCP_SERVERS = [ StdioServerParameters( @@ -27,7 +29,7 @@ ), StdioServerParameters( command="uvx", - args=[*_MCP_PIN, "openai-websearch-mcp"], + args=[*_MCP_PIN, *_WEBSEARCH_PIN, "openai-websearch-mcp"], env={ "OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", "") } diff --git a/examples/tutorials/10_async/10_temporal/050_agent_chat_guardrails/project/workflow.py b/examples/tutorials/10_async/10_temporal/050_agent_chat_guardrails/project/workflow.py index b54c8fade..88b066845 100644 --- a/examples/tutorials/10_async/10_temporal/050_agent_chat_guardrails/project/workflow.py +++ b/examples/tutorials/10_async/10_temporal/050_agent_chat_guardrails/project/workflow.py @@ -74,7 +74,8 @@ class StateModel(BaseModel): ), StdioServerParameters( command="uvx", - args=["openai-websearch-mcp"], + # mcp 1.13.1, pinned by openai-websearch-mcp, imports a pydantic internal that 2.14 removed. + args=["--with", "pydantic<2.14", "openai-websearch-mcp"], env={"OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", "")}, ), ]