diff --git a/src/agentex/lib/core/temporal/interceptors/__init__.py b/src/agentex/lib/core/temporal/interceptors/__init__.py new file mode 100644 index 000000000..80d547564 --- /dev/null +++ b/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", +] diff --git a/src/agentex/lib/core/temporal/interceptors/baggage_interceptor.py b/src/agentex/lib/core/temporal/interceptors/baggage_interceptor.py new file mode 100644 index 000000000..8c56a0858 --- /dev/null +++ b/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) diff --git a/src/agentex/lib/core/temporal/services/temporal_task_service.py b/src/agentex/lib/core/temporal/services/temporal_task_service.py index 5f6c0c381..6228bc691 100644 --- a/src/agentex/lib/core/temporal/services/temporal_task_service.py +++ b/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(), ) diff --git a/src/agentex/lib/core/temporal/workers/worker.py b/src/agentex/lib/core/temporal/workers/worker.py index 2b4958b1f..ec099046a 100644 --- a/src/agentex/lib/core/temporal/workers/worker.py +++ b/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}") diff --git a/src/agentex/lib/core/tracing/baggage.py b/src/agentex/lib/core/tracing/baggage.py new file mode 100644 index 000000000..51a4db044 --- /dev/null +++ b/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} diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index a22bfd658..dcda9c79a 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/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) 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 1ea8e82e6..6443ba37c 100644 --- a/src/agentex/lib/sdk/fastacp/base/base_acp_server.py +++ b/src/agentex/lib/sdk/fastacp/base/base_acp_server.py @@ -28,6 +28,7 @@ from agentex.protocol.json_rpc import JSONRPCError, JSONRPCRequest, JSONRPCResponse from agentex.lib.utils.model_utils import BaseModel from agentex.lib.utils.registration import register_agent +from agentex.lib.core.tracing.baggage import set_end_user_id # from agentex.lib.sdk.fastacp.types import BaseACPConfig from agentex.lib.environment_variables import EnvironmentVariables, refreshed_environment_variables @@ -197,6 +198,11 @@ async def _handle_jsonrpc(self, request: Request): params_data["request"] = {"headers": custom_headers} params = params_model.model_validate(params_data) + # Not reset: each request runs in its own asyncio task, so contexts are + # already isolated, and a reset here would fire before a streaming + # generator had finished producing spans. + set_end_user_id(getattr(params, "end_user_id", None)) + if method in RPC_SYNC_METHODS: handler = self._handlers[method] result = await handler(params) diff --git a/src/agentex/lib/sdk/fastacp/impl/temporal_acp.py b/src/agentex/lib/sdk/fastacp/impl/temporal_acp.py index 1a9cce7a8..7eacf201e 100644 --- a/src/agentex/lib/sdk/fastacp/impl/temporal_acp.py +++ b/src/agentex/lib/sdk/fastacp/impl/temporal_acp.py @@ -105,7 +105,10 @@ async def handle_task_create(params: CreateTaskParams) -> None: logger.info(f"TemporalACP received task create rpc call for task {params.task.id}") if self._temporal_task_service is not None: await self._temporal_task_service.submit_task( - agent=params.agent, task=params.task, params=params.params + agent=params.agent, + task=params.task, + params=params.params, + end_user_id=params.end_user_id, ) @self.on_task_event_send @@ -118,6 +121,7 @@ async def handle_event_send(params: SendEventParams) -> None: task=params.task, event=params.event, request=params.request, + end_user_id=params.end_user_id, ) except Exception as e: diff --git a/src/agentex/protocol/acp.py b/src/agentex/protocol/acp.py index 7e310cd89..7e7fe67fd 100644 --- a/src/agentex/protocol/acp.py +++ b/src/agentex/protocol/acp.py @@ -21,6 +21,14 @@ class RPCMethod(str, Enum): TASK_INTERRUPT = "task/interrupt" +END_USER_ID_DESCRIPTION = ( + "Opaque identifier for the end user this call is made on behalf of. Not " + "persisted on the task, so it is only present on the calls whose caller " + "supplied it. The framework stamps it onto trace spans as __end_user_id__ so " + "agents do not have to thread it through their own code." +) + + class CreateTaskParams(BaseModel): """Parameters for task/create method. @@ -29,18 +37,20 @@ class CreateTaskParams(BaseModel): task: The task to be created. params: The parameters for the task as inputted by the user. request: Additional request context including headers forwarded to this agent. + end_user_id: The end user this call was made on behalf of. """ agent: Agent = Field(..., description="The agent that the task was sent to") task: Task = Field(..., description="The task to be created") params: dict[str, Any] | None = Field( - None, + default=None, description="The parameters for the task as inputted by the user", ) request: dict[str, Any] | None = Field( default=None, description="Additional request context including headers forwarded to this agent", ) + end_user_id: str | None = Field(default=None, description=END_USER_ID_DESCRIPTION) class SendMessageParams(BaseModel): @@ -52,6 +62,7 @@ class SendMessageParams(BaseModel): content: The message that was sent to the agent. stream: Whether to stream the message back to the agentex server from the agent. request: Additional request context including headers forwarded to this agent. + end_user_id: The end user this call was made on behalf of. """ agent: Agent = Field(..., description="The agent that the message was sent to") @@ -67,6 +78,7 @@ class SendMessageParams(BaseModel): default=None, description="Additional request context including headers forwarded to this agent", ) + end_user_id: str | None = Field(default=None, description=END_USER_ID_DESCRIPTION) class SendEventParams(BaseModel): @@ -77,6 +89,7 @@ class SendEventParams(BaseModel): task: The task that the message was sent to. event: The event that was sent to the agent. request: Additional request context including headers forwarded to this agent. + end_user_id: The end user this call was made on behalf of. """ agent: Agent = Field(..., description="The agent that the event was sent to") @@ -86,6 +99,7 @@ class SendEventParams(BaseModel): default=None, description="Additional request context including headers forwarded to this agent", ) + end_user_id: str | None = Field(default=None, description=END_USER_ID_DESCRIPTION) class CancelTaskParams(BaseModel): @@ -95,6 +109,7 @@ class CancelTaskParams(BaseModel): agent: The agent that the task was sent to. task: The task that was cancelled. request: Additional request context including headers forwarded to this agent. + end_user_id: The end user this call was made on behalf of. """ agent: Agent = Field(..., description="The agent that the task was sent to") @@ -103,6 +118,7 @@ class CancelTaskParams(BaseModel): default=None, description="Additional request context including headers forwarded to this agent", ) + end_user_id: str | None = Field(default=None, description=END_USER_ID_DESCRIPTION) class InterruptTaskParams(BaseModel): diff --git a/tests/lib/core/temporal/interceptors/__init__.py b/tests/lib/core/temporal/interceptors/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/lib/core/temporal/interceptors/test_baggage_interceptor.py b/tests/lib/core/temporal/interceptors/test_baggage_interceptor.py new file mode 100644 index 000000000..63851899c --- /dev/null +++ b/tests/lib/core/temporal/interceptors/test_baggage_interceptor.py @@ -0,0 +1,303 @@ +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from temporalio.converter import default + +from agentex.types.task import Task +from agentex.types.agent import Agent +from agentex.protocol.acp import CreateTaskParams +from agentex.lib.core.tracing.baggage import get_end_user_id, set_end_user_id, reset_end_user_id +from agentex.lib.core.temporal.workers.worker import _with_default_interceptors +from agentex.lib.core.temporal.interceptors.baggage_interceptor import ( + END_USER_ID_HEADER, + DISABLE_TRACE_BAGGAGE_ENV, + TraceBaggageInterceptor, + _extract_end_user_id, + trace_baggage_disabled, + _TraceBaggageActivityInboundInterceptor, + _TraceBaggageWorkflowInboundInterceptor, + _TraceBaggageWorkflowOutboundInterceptor, +) + +MODULE = "agentex.lib.core.temporal.interceptors.baggage_interceptor" + +PAYLOAD_CONVERTER = default().payload_converter + + +@pytest.fixture(autouse=True) +def _clear_baggage(): + token = set_end_user_id(None) + yield + reset_end_user_id(token) + + +def _create_task_params(end_user_id: str | None) -> CreateTaskParams: + return CreateTaskParams( + agent=Agent( + id="agent-1", + name="agent", + acp_type="async", + description="test agent", + created_at="2026-01-01T00:00:00Z", + updated_at="2026-01-01T00:00:00Z", + ), + task=Task(id="task-1", status="RUNNING"), + end_user_id=end_user_id, + ) + + +class TestExtractEndUserId: + def test_reads_from_a_params_model(self): + assert _extract_end_user_id(_create_task_params("user-1")) == "user-1" + + def test_reads_from_a_dumped_dict(self): + assert _extract_end_user_id({"end_user_id": "user-1"}) == "user-1" + + def test_none_arg(self): + assert _extract_end_user_id(None) is None + + def test_missing_field(self): + assert _extract_end_user_id({"task": {}}) is None + + def test_null_field(self): + assert _extract_end_user_id(_create_task_params(None)) is None + + def test_empty_string_is_treated_as_absent(self): + assert _extract_end_user_id({"end_user_id": ""}) is None + + def test_non_string_is_treated_as_absent(self): + assert _extract_end_user_id({"end_user_id": 123}) is None + + +class TestWorkflowInbound: + def _interceptor(self) -> tuple[_TraceBaggageWorkflowInboundInterceptor, Any, SimpleNamespace]: + next_inbound = AsyncMock() + instance = SimpleNamespace() + return _TraceBaggageWorkflowInboundInterceptor(next_inbound), next_inbound, instance + + async def test_execute_workflow_stashes_the_end_user_on_the_instance(self): + interceptor, next_inbound, instance = self._interceptor() + input = SimpleNamespace(args=[_create_task_params("user-1")]) + + with patch(f"{MODULE}.workflow.instance", return_value=instance): + await interceptor.execute_workflow(input) # type: ignore[arg-type] + + assert instance._agentex_end_user_id == "user-1" + next_inbound.execute_workflow.assert_called_once_with(input) + + async def test_execute_workflow_stashes_nothing_when_absent(self): + interceptor, _next, instance = self._interceptor() + input = SimpleNamespace(args=[_create_task_params(None)]) + + with patch(f"{MODULE}.workflow.instance", return_value=instance): + await interceptor.execute_workflow(input) # type: ignore[arg-type] + + assert not hasattr(instance, "_agentex_end_user_id") + + async def test_execute_workflow_tolerates_empty_args(self): + interceptor, next_inbound, instance = self._interceptor() + input = SimpleNamespace(args=[]) + + with patch(f"{MODULE}.workflow.instance", return_value=instance): + await interceptor.execute_workflow(input) # type: ignore[arg-type] + + next_inbound.execute_workflow.assert_called_once_with(input) + + async def test_signal_overwrites_the_stashed_end_user(self): + interceptor, _next, instance = self._interceptor() + + with patch(f"{MODULE}.workflow.instance", return_value=instance): + await interceptor.execute_workflow(SimpleNamespace(args=[_create_task_params("creator")])) # type: ignore[arg-type] + await interceptor.handle_signal(SimpleNamespace(args=[{"end_user_id": "signaller"}])) # type: ignore[arg-type] + + assert instance._agentex_end_user_id == "signaller" + + async def test_signal_without_an_end_user_keeps_the_existing_one(self): + interceptor, _next, instance = self._interceptor() + + with patch(f"{MODULE}.workflow.instance", return_value=instance): + await interceptor.execute_workflow(SimpleNamespace(args=[_create_task_params("creator")])) # type: ignore[arg-type] + await interceptor.handle_signal(SimpleNamespace(args=[{"end_user_id": None}])) # type: ignore[arg-type] + + assert instance._agentex_end_user_id == "creator" + + async def test_a_failure_to_stash_does_not_break_the_workflow(self): + interceptor, next_inbound, _instance = self._interceptor() + input = SimpleNamespace(args=[_create_task_params("user-1")]) + + with patch(f"{MODULE}.workflow.instance", side_effect=RuntimeError("not in a workflow")): + await interceptor.execute_workflow(input) # type: ignore[arg-type] + + next_inbound.execute_workflow.assert_called_once_with(input) + + +class TestWorkflowOutbound: + def _interceptor(self) -> tuple[_TraceBaggageWorkflowOutboundInterceptor, Any]: + next_outbound = MagicMock() + return _TraceBaggageWorkflowOutboundInterceptor(next_outbound, PAYLOAD_CONVERTER), next_outbound + + def _decode(self, input: Any) -> str: + return PAYLOAD_CONVERTER.from_payload(input.headers[END_USER_ID_HEADER], str) + + def test_start_activity_adds_the_header(self): + interceptor, next_outbound = self._interceptor() + input = SimpleNamespace(headers=None) + + with patch(f"{MODULE}.workflow.instance", return_value=SimpleNamespace(_agentex_end_user_id="user-1")): + interceptor.start_activity(input) # type: ignore[arg-type] + + assert self._decode(input) == "user-1" + next_outbound.start_activity.assert_called_once_with(input) + + def test_start_local_activity_adds_the_header(self): + interceptor, next_outbound = self._interceptor() + input = SimpleNamespace(headers=None) + + with patch(f"{MODULE}.workflow.instance", return_value=SimpleNamespace(_agentex_end_user_id="user-1")): + interceptor.start_local_activity(input) # type: ignore[arg-type] + + assert self._decode(input) == "user-1" + next_outbound.start_local_activity.assert_called_once_with(input) + + def test_existing_headers_are_preserved(self): + interceptor, _next = self._interceptor() + other = PAYLOAD_CONVERTER.to_payload("keep-me") + input = SimpleNamespace(headers={"other": other}) + + with patch(f"{MODULE}.workflow.instance", return_value=SimpleNamespace(_agentex_end_user_id="user-1")): + interceptor.start_activity(input) # type: ignore[arg-type] + + assert input.headers["other"] is other + assert self._decode(input) == "user-1" + + def test_no_header_when_nothing_is_stashed(self): + interceptor, next_outbound = self._interceptor() + input = SimpleNamespace(headers=None) + + with patch(f"{MODULE}.workflow.instance", return_value=SimpleNamespace()): + interceptor.start_activity(input) # type: ignore[arg-type] + + assert input.headers is None + next_outbound.start_activity.assert_called_once_with(input) + + def test_a_failure_to_stamp_does_not_break_the_activity(self): + interceptor, next_outbound = self._interceptor() + input = SimpleNamespace(headers=None) + + with patch(f"{MODULE}.workflow.instance", side_effect=RuntimeError("boom")): + interceptor.start_activity(input) # type: ignore[arg-type] + + next_outbound.start_activity.assert_called_once_with(input) + + def test_inbound_init_installs_the_outbound_interceptor(self): + next_inbound = MagicMock() + interceptor = _TraceBaggageWorkflowInboundInterceptor(next_inbound) + + interceptor.init(MagicMock()) + + installed = next_inbound.init.call_args.args[0] + assert isinstance(installed, _TraceBaggageWorkflowOutboundInterceptor) + + +class TestActivityInbound: + def _interceptor(self, observed: list[str | None]) -> _TraceBaggageActivityInboundInterceptor: + next_inbound = MagicMock() + + async def execute_activity(_input: Any) -> str: + observed.append(get_end_user_id()) + return "result" + + next_inbound.execute_activity = execute_activity + return _TraceBaggageActivityInboundInterceptor(next_inbound, PAYLOAD_CONVERTER) + + async def test_header_is_visible_to_the_activity(self): + observed: list[str | None] = [] + interceptor = self._interceptor(observed) + headers = {END_USER_ID_HEADER: PAYLOAD_CONVERTER.to_payload("user-1")} + + result = await interceptor.execute_activity(SimpleNamespace(headers=headers)) # type: ignore[arg-type] + + assert result == "result" + assert observed == ["user-1"] + + async def test_context_is_reset_after_the_activity(self): + interceptor = self._interceptor([]) + headers = {END_USER_ID_HEADER: PAYLOAD_CONVERTER.to_payload("user-1")} + + await interceptor.execute_activity(SimpleNamespace(headers=headers)) # type: ignore[arg-type] + + assert get_end_user_id() is None + + async def test_context_is_reset_even_when_the_activity_raises(self): + next_inbound = MagicMock() + + async def boom(_input: Any) -> None: + raise ValueError("activity failed") + + next_inbound.execute_activity = boom + interceptor = _TraceBaggageActivityInboundInterceptor(next_inbound, PAYLOAD_CONVERTER) + headers = {END_USER_ID_HEADER: PAYLOAD_CONVERTER.to_payload("user-1")} + + with pytest.raises(ValueError, match="activity failed"): + await interceptor.execute_activity(SimpleNamespace(headers=headers)) # type: ignore[arg-type] + + assert get_end_user_id() is None + + async def test_no_header_leaves_the_context_alone(self): + observed: list[str | None] = [] + interceptor = self._interceptor(observed) + + await interceptor.execute_activity(SimpleNamespace(headers=None)) # type: ignore[arg-type] + + assert observed == [None] + + async def test_an_undecodable_header_does_not_break_the_activity(self): + observed: list[str | None] = [] + interceptor = self._interceptor(observed) + headers = {END_USER_ID_HEADER: "not-a-payload"} + + result = await interceptor.execute_activity(SimpleNamespace(headers=headers)) # type: ignore[arg-type] + + assert result == "result" + assert observed == [None] + + +class TestDefaultRegistration: + def test_registered_by_default(self, monkeypatch): + monkeypatch.delenv(DISABLE_TRACE_BAGGAGE_ENV, raising=False) + assert any(isinstance(i, TraceBaggageInterceptor) for i in _with_default_interceptors([])) + + def test_placed_before_user_interceptors(self, monkeypatch): + monkeypatch.delenv(DISABLE_TRACE_BAGGAGE_ENV, raising=False) + user = MagicMock() + result = _with_default_interceptors([user]) + assert isinstance(result[0], TraceBaggageInterceptor) + assert result[1] is user + + def test_not_duplicated_when_already_supplied(self, monkeypatch): + monkeypatch.delenv(DISABLE_TRACE_BAGGAGE_ENV, raising=False) + mine = TraceBaggageInterceptor() + result = _with_default_interceptors([mine]) + assert result == [mine] + + def test_the_input_list_is_not_mutated(self, monkeypatch): + monkeypatch.delenv(DISABLE_TRACE_BAGGAGE_ENV, raising=False) + supplied: list[Any] = [] + _with_default_interceptors(supplied) + assert supplied == [] + + @pytest.mark.parametrize("value", ["1", "true", "yes", "on", "TRUE", " On "]) + def test_kill_switch_removes_it(self, monkeypatch, value): + monkeypatch.setenv(DISABLE_TRACE_BAGGAGE_ENV, value) + assert trace_baggage_disabled() is True + assert _with_default_interceptors([]) == [] + + @pytest.mark.parametrize("value", ["0", "false", "no", "off", ""]) + def test_other_values_keep_it_enabled(self, monkeypatch, value): + monkeypatch.setenv(DISABLE_TRACE_BAGGAGE_ENV, value) + assert trace_baggage_disabled() is False + assert any(isinstance(i, TraceBaggageInterceptor) for i in _with_default_interceptors([])) diff --git a/tests/lib/core/tracing/test_baggage.py b/tests/lib/core/tracing/test_baggage.py new file mode 100644 index 000000000..3976ef94b --- /dev/null +++ b/tests/lib/core/tracing/test_baggage.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock + +import pytest + +from agentex.lib.core.tracing.trace import Trace, AsyncTrace +from agentex.lib.core.tracing.baggage import ( + END_USER_ID_SPAN_KEY, + END_USER_ID_MAX_LENGTH, + get_baggage, + get_end_user_id, + set_end_user_id, + reset_end_user_id, + merge_baggage_into_span_data, +) + + +@pytest.fixture(autouse=True) +def _clear_baggage(): + """Keep contextvar state from leaking between tests in the same task.""" + token = set_end_user_id(None) + yield + reset_end_user_id(token) + + +class TestSetAndGet: + def test_round_trips_a_value(self): + set_end_user_id("user-1") + assert get_end_user_id() == "user-1" + + def test_defaults_to_none(self): + assert get_end_user_id() is None + + def test_reset_restores_the_previous_value(self): + set_end_user_id("outer") + token = set_end_user_id("inner") + assert get_end_user_id() == "inner" + reset_end_user_id(token) + assert get_end_user_id() == "outer" + + def test_strips_surrounding_whitespace(self): + set_end_user_id(" user-1\n") + assert get_end_user_id() == "user-1" + + def test_blank_becomes_none(self): + set_end_user_id(" ") + assert get_end_user_id() is None + + def test_truncates_to_the_backend_cap(self): + set_end_user_id("u" * (END_USER_ID_MAX_LENGTH + 50)) + value = get_end_user_id() + assert value is not None + assert len(value) == END_USER_ID_MAX_LENGTH + + def test_non_string_is_ignored(self): + """Defensive: values ultimately originate from a deserialized payload.""" + set_end_user_id(1234) # type: ignore[arg-type] + assert get_end_user_id() is None + + +class TestGetBaggage: + def test_empty_when_unset(self): + assert get_baggage() == {} + + def test_carries_the_reserved_key(self): + set_end_user_id("user-1") + assert get_baggage() == {END_USER_ID_SPAN_KEY: "user-1"} + + +class TestMergeBaggageIntoSpanData: + def test_returns_data_unchanged_when_no_baggage(self): + data = {"a": 1} + assert merge_baggage_into_span_data(data) == {"a": 1} + + def test_creates_a_dict_from_none(self): + set_end_user_id("user-1") + assert merge_baggage_into_span_data(None) == {END_USER_ID_SPAN_KEY: "user-1"} + + def test_none_stays_none_when_no_baggage(self): + assert merge_baggage_into_span_data(None) is None + + def test_merges_alongside_existing_keys(self): + set_end_user_id("user-1") + assert merge_baggage_into_span_data({"a": 1}) == {END_USER_ID_SPAN_KEY: "user-1", "a": 1} + + def test_explicit_call_site_data_wins(self): + set_end_user_id("from-baggage") + merged = merge_baggage_into_span_data({END_USER_ID_SPAN_KEY: "from-call-site"}) + assert merged == {END_USER_ID_SPAN_KEY: "from-call-site"} + + def test_list_data_is_left_alone(self): + set_end_user_id("user-1") + data = [{"a": 1}] + assert merge_baggage_into_span_data(data) == [{"a": 1}] + + +class TestSpanStamping: + """The merge has to happen before the span is handed to the processors.""" + + def _sync_trace(self) -> Trace: + return Trace(processors=[], client=MagicMock(), trace_id="trace-1") + + def _async_trace(self) -> AsyncTrace: + return AsyncTrace(processors=[], client=MagicMock(), trace_id="trace-1") + + def test_sync_start_span_stamps_the_end_user(self): + set_end_user_id("user-1") + span = self._sync_trace().start_span(name="foo") + assert span.data == {END_USER_ID_SPAN_KEY: "user-1"} + + def test_sync_start_span_preserves_call_site_data(self): + set_end_user_id("user-1") + span = self._sync_trace().start_span(name="foo", data={"a": 1}) + assert span.data == {END_USER_ID_SPAN_KEY: "user-1", "a": 1} + + def test_sync_start_span_leaves_data_none_without_baggage(self): + span = self._sync_trace().start_span(name="foo") + assert span.data is None + + def test_sync_end_span_restamps_data_replaced_wholesale_after_start(self): + """Agent code may assign span.data outright, dropping what start merged in.""" + trace = self._sync_trace() + set_end_user_id("user-1") + span = trace.start_span(name="foo") + span.data = {"a": 1} + trace.end_span(span) + assert span.data == {END_USER_ID_SPAN_KEY: "user-1", "a": 1} + + def test_sync_end_span_keeps_the_value_merged_at_start(self): + """The ordinary path: the same Span object carries data through to end.""" + trace = self._sync_trace() + set_end_user_id("user-1") + span = trace.start_span(name="foo") + trace.end_span(span) + assert span.data == {END_USER_ID_SPAN_KEY: "user-1"} + + def test_sync_span_context_manager_stamps_the_end_user(self): + trace = self._sync_trace() + set_end_user_id("user-1") + with trace.span(name="foo") as span: + assert span is not None + assert span.data == {END_USER_ID_SPAN_KEY: "user-1"} + + async def test_async_start_span_stamps_the_end_user(self): + set_end_user_id("user-1") + span = await self._async_trace().start_span(name="foo") + assert span.data == {END_USER_ID_SPAN_KEY: "user-1"} + + async def test_async_start_span_preserves_call_site_data(self): + set_end_user_id("user-1") + span = await self._async_trace().start_span(name="foo", data={"a": 1}) + assert span.data == {END_USER_ID_SPAN_KEY: "user-1", "a": 1} + + async def test_async_end_span_restamps_data_replaced_wholesale_after_start(self): + trace = self._async_trace() + set_end_user_id("user-1") + span = await trace.start_span(name="foo") + span.data = {"a": 1} + await trace.end_span(span) + assert span.data == {END_USER_ID_SPAN_KEY: "user-1", "a": 1} + + async def test_async_end_span_keeps_the_value_merged_at_start(self): + trace = self._async_trace() + set_end_user_id("user-1") + span = await trace.start_span(name="foo") + await trace.end_span(span) + assert span.data == {END_USER_ID_SPAN_KEY: "user-1"} + + +class TestContextIsolation: + """The whole point of the contextvar: no cross-request bleed.""" + + async def test_concurrent_tasks_do_not_see_each_others_values(self): + observed: dict[str, str | None] = {} + + async def request(name: str, end_user_id: str) -> None: + set_end_user_id(end_user_id) + await asyncio.sleep(0) # let the other task run and set its own value + observed[name] = get_end_user_id() + + await asyncio.gather(request("a", "user-a"), request("b", "user-b")) + + assert observed == {"a": "user-a", "b": "user-b"} + + async def test_a_child_tasks_value_does_not_leak_to_the_parent(self): + set_end_user_id("parent") + + async def child() -> None: + set_end_user_id("child") + + await asyncio.create_task(child()) + + assert get_end_user_id() == "parent" + + async def test_a_task_created_after_set_inherits_the_value(self): + """Why the ACP server's background request tasks are safe.""" + set_end_user_id("user-1") + observed: list[str | None] = [] + + async def child() -> None: + observed.append(get_end_user_id()) + + await asyncio.create_task(child()) + + assert observed == ["user-1"] diff --git a/tests/test_acp_interrupt.py b/tests/test_acp_interrupt.py index f53cfa6b4..6e138b6b6 100644 --- a/tests/test_acp_interrupt.py +++ b/tests/test_acp_interrupt.py @@ -60,8 +60,15 @@ def test_params_model_registered(self) -> None: assert PARAMS_MODEL_BY_METHOD[RPCMethod.TASK_INTERRUPT] is InterruptTaskParams def test_params_mirror_cancel_shape(self) -> None: - """InterruptTaskParams mirrors CancelTaskParams field-for-field.""" - assert set(InterruptTaskParams.model_fields) == set(CancelTaskParams.model_fields) + """InterruptTaskParams mirrors CancelTaskParams, minus ``end_user_id``. + + ``end_user_id`` is the one intentional divergence: the control plane only + accepts it on the four RPCs an end user's request can originate + (task/create, message/send, event/send, task/cancel). task/interrupt is + driven by the agent itself via the REST API, so there is no end user to + attribute. + """ + assert set(InterruptTaskParams.model_fields) == set(CancelTaskParams.model_fields) - {"end_user_id"} assert set(InterruptTaskParams.model_fields) == {"agent", "task", "request"} def test_params_validate_round_trip(self) -> None: diff --git a/tests/test_end_user_id_propagation.py b/tests/test_end_user_id_propagation.py new file mode 100644 index 000000000..c8ac7e28c --- /dev/null +++ b/tests/test_end_user_id_propagation.py @@ -0,0 +1,252 @@ +"""End-to-end propagation of the caller-supplied ``end_user_id`` onto spans. + +Three legs: the protocol field itself, the sync ACP server binding it at the +JSON-RPC choke point, and the Temporal ACP server forwarding it into the workflow +start args and event signal payload for the baggage interceptor to pick up. +""" + +from __future__ import annotations + +from typing import Any, override +from unittest.mock import Mock, AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient + +from agentex.types.task import Task +from agentex.types.agent import Agent +from agentex.types.event import Event +from agentex.protocol.acp import ( + RPCMethod, + SendEventParams, + CancelTaskParams, + CreateTaskParams, + SendMessageParams, + InterruptTaskParams, +) +from agentex.lib.core.tracing.trace import Trace +from agentex.lib.core.tracing.baggage import END_USER_ID_SPAN_KEY, get_end_user_id +from agentex.lib.environment_variables import EnvironmentVariables +from agentex.types.task_message_content import TextContent +from agentex.lib.sdk.fastacp.impl.temporal_acp import TemporalACP +from agentex.lib.sdk.fastacp.base.base_acp_server import BaseACPServer +from agentex.lib.core.temporal.services.temporal_task_service import TemporalTaskService + + +def _agent() -> Agent: + return Agent( + id="agent-1", + name="test-agent", + description="test agent", + acp_type="async", + created_at="2026-01-01T00:00:00Z", + updated_at="2026-01-01T00:00:00Z", + ) + + +def _task() -> Task: + return Task(id="task-1", status="RUNNING") + + +def _event() -> Event: + return Event( + id="event-1", + agent_id="agent-1", + task_id="task-1", + sequence_id=1, + content=TextContent(author="user", content="hi"), + ) + + +# --------------------------------------------------------------------------- +# Protocol +# --------------------------------------------------------------------------- + + +class TestProtocol: + @pytest.mark.parametrize( + "model,kwargs", + [ + (CreateTaskParams, {}), + (SendMessageParams, {"content": TextContent(author="user", content="hi")}), + (SendEventParams, {"event": _event()}), + (CancelTaskParams, {}), + ], + ) + def test_field_is_accepted(self, model: type, kwargs: dict[str, Any]) -> None: + params = model(agent=_agent(), task=_task(), end_user_id="user-1", **kwargs) + assert params.end_user_id == "user-1" + + @pytest.mark.parametrize( + "model,kwargs", + [ + (CreateTaskParams, {}), + (SendMessageParams, {"content": TextContent(author="user", content="hi")}), + (SendEventParams, {"event": _event()}), + (CancelTaskParams, {}), + ], + ) + def test_field_is_optional(self, model: type, kwargs: dict[str, Any]) -> None: + """Older control planes omit it entirely.""" + params = model(agent=_agent(), task=_task(), **kwargs) + assert params.end_user_id is None + + def test_interrupt_does_not_carry_it(self) -> None: + """task/interrupt has no end_user_id on the control plane either.""" + assert "end_user_id" not in InterruptTaskParams.model_fields + + +# --------------------------------------------------------------------------- +# Sync agents +# --------------------------------------------------------------------------- + + +_observed_span_data: list[Any] = [] + + +class _SpanCreatingServer(BaseACPServer): + """Creates a span inside the handler, exactly as an agent's own code would.""" + + __test__ = False + + @override + def _setup_handlers(self) -> None: + @self.on_message_send + async def handler(params: SendMessageParams) -> TextContent: # type: ignore[reportUnusedFunction] + trace = Trace(processors=[], client=MagicMock(), trace_id="trace-1") + span = trace.start_span(name="work") + _observed_span_data.append(span.data) + return TextContent(author="agent", content=str(get_end_user_id())) + + +def _message_send_request(end_user_id: str | None) -> dict[str, Any]: + params: dict[str, Any] = { + "agent": { + "id": "agent-1", + "name": "test-agent", + "description": "d", + "acp_type": "sync", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", + }, + "task": {"id": "task-1"}, + "content": {"type": "text", "author": "user", "content": "hi"}, + "stream": False, + } + if end_user_id is not None: + params["end_user_id"] = end_user_id + return {"jsonrpc": "2.0", "method": RPCMethod.MESSAGE_SEND.value, "params": params, "id": 1} + + +class TestSyncAgentPropagation: + @pytest.fixture(autouse=True) + def _reset_observed(self): + _observed_span_data.clear() + yield + _observed_span_data.clear() + + def test_span_created_in_the_handler_is_stamped(self) -> None: + client = TestClient(_SpanCreatingServer.create()) + + response = client.post("/api", json=_message_send_request("user-1")) + + assert response.status_code == 200 + assert _observed_span_data == [{END_USER_ID_SPAN_KEY: "user-1"}] + + def test_handler_sees_the_value_in_context(self) -> None: + client = TestClient(_SpanCreatingServer.create()) + + response = client.post("/api", json=_message_send_request("user-1")) + + assert response.json()["result"]["content"]["content"] == "user-1" + + def test_nothing_is_stamped_when_the_caller_omits_it(self) -> None: + client = TestClient(_SpanCreatingServer.create()) + + response = client.post("/api", json=_message_send_request(None)) + + assert response.status_code == 200 + assert _observed_span_data == [None] + + def test_one_request_does_not_bleed_into_the_next(self) -> None: + """The failure mode this whole design exists to avoid.""" + client = TestClient(_SpanCreatingServer.create()) + + client.post("/api", json=_message_send_request("user-a")) + client.post("/api", json=_message_send_request(None)) + + assert _observed_span_data == [{END_USER_ID_SPAN_KEY: "user-a"}, None] + + +# --------------------------------------------------------------------------- +# Temporal agents +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_temporal_client() -> AsyncMock: + client = AsyncMock() + client.start_workflow = AsyncMock(return_value="workflow-1") + client.send_signal = AsyncMock(return_value=None) + return client + + +@pytest.fixture +def task_service(mock_temporal_client: AsyncMock) -> TemporalTaskService: + env_vars = Mock(spec=EnvironmentVariables) + env_vars.WORKFLOW_NAME = "test-workflow" + env_vars.WORKFLOW_TASK_QUEUE = "test-queue" + env_vars.WORKFLOW_EXECUTION_TIMEOUT_SECONDS = None + return TemporalTaskService(temporal_client=mock_temporal_client, env_vars=env_vars) + + +class TestTemporalPropagation: + async def test_submit_task_puts_it_in_the_workflow_start_args( + self, task_service: TemporalTaskService, mock_temporal_client: AsyncMock + ) -> None: + await task_service.submit_task(agent=_agent(), task=_task(), params=None, end_user_id="user-1") + + arg = mock_temporal_client.start_workflow.call_args.kwargs["arg"] + assert arg.end_user_id == "user-1" + + async def test_submit_task_defaults_to_none( + self, task_service: TemporalTaskService, mock_temporal_client: AsyncMock + ) -> None: + await task_service.submit_task(agent=_agent(), task=_task(), params=None) + + arg = mock_temporal_client.start_workflow.call_args.kwargs["arg"] + assert arg.end_user_id is None + + async def test_send_event_puts_it_in_the_signal_payload( + self, task_service: TemporalTaskService, mock_temporal_client: AsyncMock + ) -> None: + await task_service.send_event(agent=_agent(), task=_task(), event=_event(), end_user_id="user-1") + + payload = mock_temporal_client.send_signal.call_args.kwargs["payload"] + assert payload["end_user_id"] == "user-1" + + async def test_acp_task_create_handler_forwards_it( + self, task_service: TemporalTaskService, mock_temporal_client: AsyncMock + ) -> None: + acp = TemporalACP(temporal_address="localhost:7233", temporal_task_service=task_service) + acp._setup_handlers() + + handler = acp._handlers[RPCMethod.TASK_CREATE] + assert handler is not None + await handler(CreateTaskParams(agent=_agent(), task=_task(), end_user_id="user-1")) + + arg = mock_temporal_client.start_workflow.call_args.kwargs["arg"] + assert arg.end_user_id == "user-1" + + async def test_acp_event_send_handler_forwards_it( + self, task_service: TemporalTaskService, mock_temporal_client: AsyncMock + ) -> None: + acp = TemporalACP(temporal_address="localhost:7233", temporal_task_service=task_service) + acp._setup_handlers() + + handler = acp._handlers[RPCMethod.EVENT_SEND] + assert handler is not None + await handler(SendEventParams(agent=_agent(), task=_task(), event=_event(), end_user_id="user-1")) + + payload = mock_temporal_client.send_signal.call_args.kwargs["payload"] + assert payload["end_user_id"] == "user-1"