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 <noreply@anthropic.com>

Author: mohammadatallah-scale

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:

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}")

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

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()

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)

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)