diff --git a/agentex/src/domain/services/agent_acp_service.py b/agentex/src/domain/services/agent_acp_service.py index acb8b1f7..9035b619 100644 --- a/agentex/src/domain/services/agent_acp_service.py +++ b/agentex/src/domain/services/agent_acp_service.py @@ -84,13 +84,24 @@ } ) +# W3C trace-context headers. These are NOT x-* prefixed, so the allowlist below +# would otherwise strip them on the forward to the downstream agent — which +# detaches the agent's observability trace (and the Temporal workflow/activity it +# signals) from the ingress trace, so you can't follow a request from the API +# call to the agent's work. Forwarding them lets the agent continue the same +# trace. `baggage` additionally carries the business-trace sampling key. None are +# sensitive or hop-by-hop, so passing them through is safe. +TRACE_CONTEXT_HEADERS = frozenset({"traceparent", "tracestate", "baggage"}) + def filter_request_headers(headers: dict[str, str] | None) -> dict[str, str]: """ Filter request headers to only include safe custom headers. Security filtering rules: - 1. Allow only x-* prefixed headers (allowlist approach) + 1. Allow x-* prefixed headers (allowlist approach), plus the W3C + trace-context headers (traceparent/tracestate/baggage) so the downstream + agent's trace continues the ingress trace instead of detaching 2. Block hop-by-hop headers (connection, keep-alive, etc.) 3. Block sensitive headers (credentials, acting delegation, x-agent-api-key, x-selected-account-id) @@ -110,12 +121,27 @@ def filter_request_headers(headers: dict[str, str] | None) -> dict[str, str]: return { k: v for k, v in headers.items() - if k.lower().startswith("x-") + if (k.lower().startswith("x-") or k.lower() in TRACE_CONTEXT_HEADERS) and k.lower() not in HOP_BY_HOP_HEADERS and k.lower() not in BLOCKED_HEADERS } +def extract_trace_context_headers(headers: dict[str, str] | None) -> dict[str, str]: + """Pull just the W3C trace-context headers (case-insensitive) from an inbound + request. + + Used so trace context is forwarded on EVERY downstream operation + (task/create, event/send, message, streaming, cancel) rather than only the + call sites that happen to thread ``request_headers`` through get_headers(). + Keeping the agent's trace continuous with the ingress trace must not depend + on each caller remembering to pass headers. + """ + if not headers: + return {} + return {k: v for k, v in headers.items() if k.lower() in TRACE_CONTEXT_HEADERS} + + class AgentACPService(TaskMessageMixin): """ Client service for communicating with downstream ACP servers. @@ -283,6 +309,13 @@ async def get_headers( if request_headers is None: request_headers = dict(self._request.headers) filtered_request_headers = filter_request_headers(request_headers) + # Always forward inbound W3C trace-context, independent of whether the + # caller threaded request_headers through — otherwise task/create, + # message, streaming and cancel (which call get_headers(agent) with no + # request_headers) would drop traceparent and the downstream agent would + # start a detached trace. The inbound headers are on self._request. + inbound_headers = dict(self._request.headers) if getattr(self, "_request", None) is not None else {} + trace_context_headers = extract_trace_context_headers(inbound_headers) delegation_headers = self.get_delegation_headers(agent) auth_headers = await self.get_agent_auth_headers(agent) request_id = ctx_var_request_id.get(uuid4().hex) @@ -290,6 +323,7 @@ async def get_headers( # Later keys win. Client passthrough and delegation first; agent auth last. return { **filtered_request_headers, + **trace_context_headers, **delegation_headers, **auth_headers, "x-request-id": request_id, diff --git a/agentex/tests/unit/services/test_agent_acp_service.py b/agentex/tests/unit/services/test_agent_acp_service.py index 85bd96e7..44e437b0 100644 --- a/agentex/tests/unit/services/test_agent_acp_service.py +++ b/agentex/tests/unit/services/test_agent_acp_service.py @@ -21,6 +21,7 @@ from src.domain.repositories.agent_repository import AgentRepository from src.domain.services.agent_acp_service import ( AgentACPService, + extract_trace_context_headers, filter_request_headers, ) @@ -1131,3 +1132,46 @@ def test_blocks_user_api_key_and_acting_headers(self): } ) assert result == {"x-trace-id": "trace-1"} + + def test_forwards_w3c_trace_context_headers(self): + # traceparent/tracestate/baggage are not x-* but must survive the forward + # so the downstream agent's trace continues the ingress trace. + result = filter_request_headers( + { + "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + "tracestate": "vendor=value", + "baggage": "agentex.business_trace_id=t1", + "authorization": "Bearer x", # still blocked + "host": "gateway", # still hop-by-hop + } + ) + assert result == { + "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + "tracestate": "vendor=value", + "baggage": "agentex.business_trace_id=t1", + } + + +class TestExtractTraceContextHeaders: + def test_extracts_only_trace_context_case_insensitive(self): + # get_headers forwards these on every op regardless of request_headers, + # so extraction must pull them from inbound headers (case-insensitive) + # and ignore everything else. + result = extract_trace_context_headers( + { + "Traceparent": "00-abc-def-01", + "tracestate": "vendor=value", + "BAGGAGE": "agentex.business_trace_id=t1", + "x-api-key": "secret", + "content-type": "application/json", + } + ) + assert result == { + "Traceparent": "00-abc-def-01", + "tracestate": "vendor=value", + "BAGGAGE": "agentex.business_trace_id=t1", + } + + def test_empty_or_none(self): + assert extract_trace_context_headers(None) == {} + assert extract_trace_context_headers({}) == {}