feat: support end_user_id for agentex agents & automatically include in spans

src/agentex/lib/core/temporal/interceptors/__init__.py

@@ -0,0 +1,13 @@
+from agentex.lib.core.temporal.interceptors.baggage_interceptor import (
+    END_USER_ID_HEADER,
+    DISABLE_TRACE_BAGGAGE_ENV,
+    TraceBaggageInterceptor,
+    trace_baggage_disabled,
+)
+
+__all__ = [
+    "END_USER_ID_HEADER",
+    "DISABLE_TRACE_BAGGAGE_ENV",
+    "TraceBaggageInterceptor",
+    "trace_baggage_disabled",
+]

src/agentex/lib/core/temporal/interceptors/baggage_interceptor.py

@@ -0,0 +1,156 @@
+"""Carries the caller-supplied ``end_user_id`` from a workflow into its activities.
+
+For Temporal agents the ACP server only *starts* the workflow; the work runs in a
+separate worker process, so a contextvar set at the ACP boundary never reaches the
+code that creates spans. The value does survive that hop inside the workflow start
+args and signal payloads, which is where the workflow-inbound half harvests it
+from. Spans created by workflow code are covered too, since those dispatch through
+a ``START_SPAN`` activity and so pass through the same activity-inbound hop.
+
+Replay-safe: headers are derived only from start/signal args, with no clock or I/O.
+
+Registered by default in ``AgentexWorker``; ``AGENTEX_DISABLE_TRACE_BAGGAGE``
+(``1``/``true``/``yes``/``on``) disables it.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, override
+
+from temporalio import workflow
+from temporalio.worker import (
+    Interceptor,
+    HandleSignalInput,
+    StartActivityInput,
+    ExecuteActivityInput,
+    ExecuteWorkflowInput,
+    StartLocalActivityInput,
+    ActivityInboundInterceptor,
+    WorkflowInboundInterceptor,
+    WorkflowOutboundInterceptor,
+)
+from temporalio.converter import default
+
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.core.tracing.baggage import set_end_user_id, reset_end_user_id
+
+logger = make_logger(__name__)
+
+END_USER_ID_HEADER = "agentex-end-user-id"
+_WORKFLOW_ATTR = "_agentex_end_user_id"
+DISABLE_TRACE_BAGGAGE_ENV = "AGENTEX_DISABLE_TRACE_BAGGAGE"
+
+
+def trace_baggage_disabled() -> bool:
+    raw = os.environ.get(DISABLE_TRACE_BAGGAGE_ENV, "").strip().lower()
+    return raw in ("1", "true", "yes", "on")
+
+
+def _extract_end_user_id(arg: Any) -> str | None:
+    # Args arrive as the params model or as its dumped dict, depending on how the
+    # caller serialized them.
+    if arg is None:
+        return None
+    value = arg.get("end_user_id") if isinstance(arg, dict) else getattr(arg, "end_user_id", None)
+    return value if isinstance(value, str) and value else None
+
+
+class TraceBaggageInterceptor(Interceptor):
+    """Threads the caller-supplied end user from workflow start args to activities."""
+
+    def __init__(self) -> None:
+        self._payload_converter = default().payload_converter
+
+    @override
+    def intercept_activity(self, next: ActivityInboundInterceptor) -> ActivityInboundInterceptor:
+        return _TraceBaggageActivityInboundInterceptor(next, self._payload_converter)
+
+    @override
+    def workflow_interceptor_class(self, input: Any) -> type[WorkflowInboundInterceptor] | None:  # noqa: ARG002
+        return _TraceBaggageWorkflowInboundInterceptor
+
+
+class _TraceBaggageWorkflowInboundInterceptor(WorkflowInboundInterceptor):
+    """Harvests the end user off inbound workflow args and signal args."""
+
+    @override
+    async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any:
+        self._stash(input.args[0] if input.args else None)
+        return await self.next.execute_workflow(input)
+
+    @override
+    async def handle_signal(self, input: HandleSignalInput) -> None:
+        # Signals (e.g. event/send) carry their own end user, which may differ
+        # from the one that created the task. Last one in wins so the activities
+        # a signal triggers are attributed to the caller that triggered them.
+        self._stash(input.args[0] if input.args else None)
+        return await self.next.handle_signal(input)
+
+    def _stash(self, arg: Any) -> None:
+        end_user_id = _extract_end_user_id(arg)
+        if end_user_id is None:
+            return
+        try:
+            setattr(workflow.instance(), _WORKFLOW_ATTR, end_user_id)
+        except Exception as exc:
+            logger.debug(f"Could not stash end_user_id on the workflow instance: {exc}")
+
+    @override
+    def init(self, outbound: WorkflowOutboundInterceptor) -> None:
+        self.next.init(_TraceBaggageWorkflowOutboundInterceptor(outbound, default().payload_converter))
+
+
+class _TraceBaggageWorkflowOutboundInterceptor(WorkflowOutboundInterceptor):
+    """Copies the stashed end user into activity headers."""
+
+    def __init__(self, next: WorkflowOutboundInterceptor, payload_converter: Any) -> None:
+        super().__init__(next)
+        self._payload_converter = payload_converter
+
+    @override
+    def start_activity(self, input: StartActivityInput) -> workflow.ActivityHandle[Any]:
+        self._add_header(input)
+        return self.next.start_activity(input)
+
+    @override
+    def start_local_activity(self, input: StartLocalActivityInput) -> workflow.ActivityHandle[Any]:
+        self._add_header(input)
+        return self.next.start_local_activity(input)
+
+    def _add_header(self, input: StartActivityInput | StartLocalActivityInput) -> None:
+        try:
+            end_user_id = getattr(workflow.instance(), _WORKFLOW_ATTR, None)
+            if not end_user_id:
+                return
+            headers = dict(input.headers or {})
+            headers[END_USER_ID_HEADER] = self._payload_converter.to_payload(end_user_id)
+            input.headers = headers
+        except Exception as exc:
+            logger.debug(f"Could not add end_user_id to activity headers: {exc}")
+
+
+class _TraceBaggageActivityInboundInterceptor(ActivityInboundInterceptor):
+    """Decodes the header into the tracing baggage contextvar for the activity."""
+
+    def __init__(self, next: ActivityInboundInterceptor, payload_converter: Any) -> None:
+        super().__init__(next)
+        self._payload_converter = payload_converter
+
+    @override
+    async def execute_activity(self, input: ExecuteActivityInput) -> Any:
+        end_user_id: str | None = None
+        if input.headers and END_USER_ID_HEADER in input.headers:
+            try:
+                end_user_id = self._payload_converter.from_payload(input.headers[END_USER_ID_HEADER], str)
+            except Exception as exc:
+                logger.debug(f"Could not decode the end_user_id activity header: {exc}")
+
+        if end_user_id is None:
+            return await self.next.execute_activity(input)
+
+        token = set_end_user_id(end_user_id)
+        try:
+            return await self.next.execute_activity(input)
+        finally:
+            reset_end_user_id(token)

src/agentex/lib/core/temporal/services/temporal_task_service.py

@@ -27,7 +27,13 @@ def __init__(
         self._env_vars = env_vars
 
 
-    async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | None) -> str:
+    async def submit_task(
+        self,
+        agent: Agent,
+        task: Task,
+        params: dict[str, Any] | None,
+        end_user_id: str | None = None,
+    ) -> str:
         """
         Submit a task to the async runtime for execution.
 
@@ -48,6 +54,7 @@ async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | N
                 agent=agent,
                 task=task,
                 params=params,
+                end_user_id=end_user_id,
             ),
             id=task.id,
             task_queue=self._env_vars.WORKFLOW_TASK_QUEUE,
@@ -62,7 +69,14 @@ async def get_state(self, task_id: str) -> WorkflowState:
             workflow_id=task_id,
         )
 
-    async def send_event(self, agent: Agent, task: Task, event: Event, request: dict | None = None) -> None:
+    async def send_event(
+        self,
+        agent: Agent,
+        task: Task,
+        event: Event,
+        request: dict | None = None,
+        end_user_id: str | None = None,
+    ) -> None:
         return await self._temporal_client.send_signal(
             workflow_id=task.id,
             signal=SignalName.RECEIVE_EVENT.value,
@@ -71,6 +85,7 @@ async def send_event(self, agent: Agent, task: Task, event: Event, request: dict
                 task=task,
                 event=event,
                 request=request,
+                end_user_id=end_user_id,
             ).model_dump(),
         )
 

src/agentex/lib/core/temporal/workers/worker.py

@@ -31,6 +31,11 @@
 from agentex.lib.utils.registration import register_agent
 from agentex.lib.environment_variables import EnvironmentVariables
 from agentex.lib.core.compat.version_guard import assert_backend_compatible
+from agentex.lib.core.temporal.interceptors.baggage_interceptor import (
+    DISABLE_TRACE_BAGGAGE_ENV,
+    TraceBaggageInterceptor,
+    trace_baggage_disabled,
+)
 
 logger = make_logger(__name__)
 
@@ -91,6 +96,18 @@ def _validate_interceptors(interceptors: list) -> None:
             )
 
 
+def _with_default_interceptors(interceptors: list) -> list:
+    """Prepend the framework's own interceptors, first so user ones see their effect."""
+    if trace_baggage_disabled():
+        logger.info(f"Trace baggage propagation disabled via {DISABLE_TRACE_BAGGAGE_ENV}")
+        return list(interceptors)
+
+    if any(isinstance(i, TraceBaggageInterceptor) for i in interceptors):
+        return list(interceptors)
+
+    return [TraceBaggageInterceptor(), *interceptors]
+
+
 async def get_temporal_client(
     temporal_address: str,
     metrics_url: str | None = None,
@@ -203,6 +220,8 @@ async def run(
         if self.interceptors:
             _validate_interceptors(self.interceptors)
 
+        interceptors = _with_default_interceptors(self.interceptors)
+
         temporal_client = await get_temporal_client(
             temporal_address=os.environ.get("TEMPORAL_ADDRESS", "localhost:7233"),
             plugins=self.plugins,
@@ -229,7 +248,7 @@ async def run(
             max_concurrent_activities=self.max_concurrent_activities,
             build_id=str(uuid.uuid4()),
             debug_mode=debug_enabled,  # Disable deadlock detection in debug mode
-            interceptors=self.interceptors,  # Pass interceptors to Worker
+            interceptors=interceptors,  # Pass interceptors to Worker
         )
 
         logger.info(f"Starting workers for task queue: {self.task_queue}")

src/agentex/lib/core/tracing/baggage.py

@@ -0,0 +1,80 @@
+"""Request-scoped identifiers that the framework stamps onto trace spans.
+
+The tracing processors cannot read per-request context: they run on a long-lived
+queue-drain task whose contextvars were frozen when the first span in the process
+was enqueued (see ``span_queue.py``). So per-request values have to be attached to
+the ``Span`` object *before* it is enqueued, travelling with it as data rather than
+as ambient context.
+
+Deliberately one reserved, size-capped scalar rather than an arbitrary dict, which
+keeps caller-controlled free-form data (and its PHI and cardinality risks) off the
+tracing backend.
+"""
+
+from __future__ import annotations
+
+from typing import Any, TypeVar
+from contextvars import Token, ContextVar
+
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+# Matches the ``__key__`` style the SGP processor stamps its static keys with.
+END_USER_ID_SPAN_KEY = "__end_user_id__"
+# Mirrors the cap the control plane enforces on the RPC field.
+END_USER_ID_MAX_LENGTH = 256
+
+_end_user_id: ContextVar[str | None] = ContextVar("agentex_end_user_id", default=None)
+
+
+def set_end_user_id(value: str | None) -> Token[str | None]:
+    """Bind the end user for the current context, returning a reset token.
+
+    Normalized here rather than at the merge point, so a span never carries
+    something the backend would reject.
+    """
+    normalized: str | None = None
+    if isinstance(value, str):
+        stripped = value.strip()
+        if stripped:
+            normalized = stripped[:END_USER_ID_MAX_LENGTH]
+            if len(stripped) > END_USER_ID_MAX_LENGTH:
+                logger.warning(f"Truncated end_user_id from {len(stripped)} to {END_USER_ID_MAX_LENGTH} characters")
+    return _end_user_id.set(normalized)
+
+
+def get_end_user_id() -> str | None:
+    return _end_user_id.get()
+
+
+def reset_end_user_id(token: Token[str | None]) -> None:
+    _end_user_id.reset(token)
+
+
+def get_baggage() -> dict[str, str]:
+    """The span data keys to stamp for the current context."""
+    end_user_id = _end_user_id.get()
+    if end_user_id is None:
+        return {}
+    return {END_USER_ID_SPAN_KEY: end_user_id}
+
+
+DataT = TypeVar("DataT")
+
+
+def merge_baggage_into_span_data(data: DataT) -> DataT | dict[str, Any]:
+    """Merge the current context's baggage into already-serialized span data.
+
+    Explicit call-site data wins on collision. Span data may also be a list, which
+    has nowhere to put a key, so those are returned untouched.
+    """
+    baggage = get_baggage()
+    if not baggage:
+        return data
+    if data is None:
+        return dict(baggage)
+    if not isinstance(data, dict):
+        logger.debug(f"Skipping trace baggage merge for span data of type {type(data).__name__}")
+        return data
+    return {**baggage, **data}

src/agentex/lib/core/tracing/trace.py

@@ -11,6 +11,7 @@
 from agentex.types.span import Span
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.model_utils import recursive_model_dump
+from agentex.lib.core.tracing.baggage import merge_baggage_into_span_data
 from agentex.lib.core.tracing.span_error import set_span_error
 from agentex.lib.core.tracing.span_queue import (
     SpanEventType,
@@ -79,6 +80,7 @@ def start_span(
 
         serialized_input = recursive_model_dump(input) if input else None
         serialized_data = recursive_model_dump(data) if data else None
+        serialized_data = merge_baggage_into_span_data(serialized_data)
         id = str(uuid.uuid4())
 
         span = Span(
@@ -116,6 +118,9 @@ def end_span(
         span.input = recursive_model_dump(span.input) if span.input else None
         span.output = recursive_model_dump(span.output) if span.output else None
         span.data = recursive_model_dump(span.data) if span.data else None
+        # Redundant unless the caller assigned ``span.data`` wholesale after start,
+        # which would otherwise drop the baggage merged in there.
+        span.data = merge_baggage_into_span_data(span.data)
 
         for processor in self.processors:
             processor.on_span_end(span)
@@ -229,6 +234,7 @@ async def start_span(
 
         serialized_input = recursive_model_dump(input) if input else None
         serialized_data = recursive_model_dump(data) if data else None
+        serialized_data = merge_baggage_into_span_data(serialized_data)
         id = str(uuid.uuid4())
 
         span = Span(
@@ -266,6 +272,9 @@ async def end_span(
         span.input = recursive_model_dump(span.input) if span.input else None
         span.output = recursive_model_dump(span.output) if span.output else None
         span.data = recursive_model_dump(span.data) if span.data else None
+        # Redundant unless the caller assigned ``span.data`` wholesale after start,
+        # which would otherwise drop the baggage merged in there.
+        span.data = merge_baggage_into_span_data(span.data)
 
         if self.processors:
             self._span_queue.enqueue(SpanEventType.END, span.model_copy(deep=True), self.processors)